Intermediate Layers Encode Optimal Biological Representations in Single-Cell Foundation Models

Abstract: Current single-cell foundation model benchmarks universally extract final layer embeddings, assuming these represent optimal feature spaces. We systematically evaluate layer-wise representations from scFoundation (100M parameters) and Tahoe-X1 (1.3B parameters) across trajectory inference and perturbation response prediction. Our analysis reveals that optimal layers are task-dependent (trajectory peaks at 60% depth, 31% above final layers) and context-dependent (perturbation optima shift 0-96% across T cell activation states). Notably, first-layer embeddings outperform all deeper layers in quiescent cells, challenging assumptions about hierarchical feature abstraction. These findings demonstrate that "where" to extract features matters as much as "what" the model learns, necessitating systematic layer evaluation tailored to biological task and cellular context rather than defaulting to final-layer embeddings.

• The paper shows intermediate layers yield up to 31% improved performance over final layers for trajectory inference in single-cell models.
• It employs layer-wise evaluations on scFoundation and Tahoe-X1 across tasks, revealing task- and context-dependent optimal layers.
• These findings challenge the universal use of final embeddings and advocate for systematic, layer-aware feature extraction in computational biology.

Layer-wise Feature Extraction in Single-Cell Foundation Models: Task- and Context-Dependent Optima

Introduction

The paper "Intermediate Layers Encode Optimal Biological Representations in Single-Cell Foundation Models" (2604.14838) interrogates the prevailing assumption that final-layer embeddings from single-cell foundation models universally provide optimal representations for downstream biological tasks. By conducting systematic, layer-wise evaluations across both trajectory inference and perturbation response prediction, using scFoundation (100M parameters) and Tahoe-X1 (1.3B parameters), the study demonstrates substantial task- and context-dependence in optimal layer selection. The findings reveal strong numerical advantages—intermediate layers outperform final layers by up to 31% in trajectory inference, while perturbation response prediction optima vary across the full depth of models depending on cellular activation state. These results challenge entrenched benchmarking habits and have substantial implications for the deployment and evaluation of foundation models in computational biology.

Methodological Framework

Two high-capacity, pretrained single-cell foundation models were analyzed: scFoundation, which integrates an asymmetric encoder-decoder architecture with read-depth aware pretraining, and Tahoe-X1, built on a scGPT-inspired transformer with masked expression prediction and dual gene/cell-aware decoders. Layer-wise embeddings were extracted post-feedforward, post-residual addition, and layer normalization, enabling normalized depth comparisons across architectures.

For trajectory inference, the study used the LT-scSeq dataset, featuring clonally-resolved hematopoietic differentiation, with diffusion pseudotime (DPT) providing ground truth temporal ordering. Performance was quantified via Spearman correlation between layer-derived DPT and reference pseudotime, leveraging robustness to monotonic transformations. Perturbation response prediction leveraged a genome-scale CRISPRi Perturb-seq dataset from primary human CD4+ T cells, stratified across Rest, Stim8hr, and Stim48hr activation states. Layer-wise correlation was measured between embedding-space perturbation effects and ground-truth differential expression via representational similarity analysis, with Spearman correlation capturing alignment in regulatory relationships.

Results: Trajectory Inference and Task-Dependent Layer Optima

Layer-wise analysis in trajectory inference revealed a pronounced task-dependent optimum for Tahoe-X1 at 60% model depth, where Spearman correlation with ground truth pseudotime ($\rho = 0.76$) exceeded the final layer performance by 31%. scFoundation showed a monotonic increase towards the penultimate layer, peaking at $\rho = 0.60$. Early layers encoded predominantly low-level or technical signals, with substantial semantic transitions emerging mid-depth.

Figure 1: Layer-wise pseudotime correlation; Tahoe-X1 peaks at 60% depth (Spearman $\rho = 0.76$), outperforming its final layer by 31%.

Tahoe-X1's performance profile exhibited U-shaped dynamics, with degradation at deeper layers, supporting the hypothesis that specialization for pretraining objectives in late layers reduces transferability to biological tasks. These findings are consistent with analogous studies in NLP, where intermediate layers manifest optimal representations for downstream cognitive tasks.

Results: Perturbation Response Prediction and Context Dependence

Analysis of perturbation modeling showed dramatic context-dependent variability in optimal layer extraction. For Tahoe-X1, the best performance shifted from the first layer in quiescent (Rest) T cells ($\rho = 0.37$) to the penultimate layer in the most activated state ($\rho = 0.54$), representing a 96 percentage point shift in normalized depth. scFoundation exhibited similar variability across activation states, with optima distributed between middle and final layers.

Figure 2: Layer-wise perturbation response correlation across T cell activation states, with optimal depths ranging from 0% to 96% within a single model.

Perturbation effect correlations improved systematically with T cell activation, indicating that activated states yield more coherent regulatory representations. Notably, in quiescent states, shallow representations outperformed deeper layers, suggesting that expression shifts associated with perturbations in resting cells are sufficiently captured at lower-level embeddings without hierarchical abstraction. These results contradict the assumption of monotonic improvement with layer depth and underscore that biological context fundamentally alters model representational dynamics.

Theoretical and Practical Implications

The paper asserts that "where" features are extracted is equally important as "what" the model learns. The bold claim—the universal reliance on final-layer embeddings is empirically invalidated—has significant ramifications for foundation model evaluation, benchmarking, and downstream pipeline design. Benchmarking frameworks and practitioners must support and scrutinize all layers, exposing layer-specific embeddings tailored to biological task semantics and cellular context.

The findings suggest that pretraining-induced specialization in upper layers may be detrimental for some downstream tasks, and that hierarchical abstraction does not always monotonically improve biological representation. The adoption of systematic layer evaluation will enable practitioners to harness substantial performance gains (up to 31%) and facilitate more faithful modeling of developmental trajectory and regulatory effects.

From a theoretical perspective, these results support the existence of task-relevant information bottlenecks, as observed in NLP, and highlight the necessity of context-aware representational analyses in biological modeling. As model capacities scale, investigation into inter-layer specialization, transferability, and semantic transitions will be critical for maximizing biological insight.

Future Directions

The study opens avenues for broader, architecture-independent evaluations across biological domains, including multi-omic integration, cell state annotation, and rare cell type prediction. Future research should interrogate layer-specific embedding generalizability under domain adaptation, assess robustness in low-data regimes, and explore protocol-guided layer fusion or ensemble extraction methods. Fine-grained ablation studies and contrastive analyses with architectural variants (e.g., attention patterns, block normalization, pretraining objectives) will further delineate the mechanistic underpinnings of layer specialization.

Conclusion

Optimal biological representations in single-cell foundation models are primarily determined by task and cellular context rather than model depth alone. Intermediate layers yield substantial performance improvements for trajectory inference, while perturbation response prediction displays strong context dependence—sometimes requiring shallow embeddings. The study challenges entrenched practices and advocates for systematic, layer-aware feature extraction to more fully exploit foundation model capacities and uncover complex biological phenomena.