Attribution-based diagnostics for single-cell foundation models
Authors: Farnoush Rezaei Jafari, Tristan Lazard, Kenza Bouzid, Shruthi Bannur, Jeremiah Wander, Ashley Mae Conard, Stephanie Hyland, Lorin Crawford, Daniel Coelho de Castro
Organizations: BIFOLD - Berlin Institute for the Foundations of Learning and Data · Technical University of Berlin · BIFOLD – Berlin Institute for the Foundations of Learning and Data, Berlin, Germany · Microsoft Research
Single-cell RNA foundation models (scFMs) are increasingly used for in-silico biological discovery, yet predictive performance alone does not reveal whether their decisions rely on biologically meaningful signals or model-specific shortcuts. We systematically investigate post-hoc attribution of downstream predictions across two tasks: cell-type annotation using pretrained Cell2Sentence and scGPT models, and perturbation prediction using scGPT+STATE in a controlled setting. For cell annotation, we use attribution to assess whether predictions rely on established marker genes and how their relevance varies with model architecture, expression rank, and the surrounding cellular context. For perturbation prediction, we compare attribution patterns with known regulatory ground-truth in a synthetic setting to assess what they reveal about the model's learned computations. Across these tasks, attribution reveals biological feature use, contextual interactions, contributions from different input pathways, learning progression, associations with prediction quality, and potential pitfalls that are not fully captured by standard performance metrics. Experiments with complete and limited regulatory knowledge further demonstrate that conclusions drawn from attribution depend strongly on the choice of biological reference, baseline, and null model. Together, these findings highlight the value of attribution as a flexible post-hoc diagnostic for understanding model behavior, identifying potential issues, and informing the development of scFMs.
Figures & tables
Figure 1: Overview of our framework for explaining downstream predictions of single-cell foundation models. A.1. Layer-wise Relevance Propagation (LRP) decomposes a cell-type probe logit into per-gene relevance scores, which are compared with known markers of the predicted cell type. A.2. Symbolic XAI decomposes the probe logit into relevance scores for gene subsets, which can be combined into logical queries to quantify how the relevance of marker genes changes in different contextual settings. B.3. For each perturbation of gene Gi and output gene Gj , LRP attributes the predicted expression of Gj to the model inputs, quantifying the contribution of Gi . Repeating this across perturbations yields a gene–gene relevance matrix that can be interpreted as an attribution-based estimate of regulatory influence and compared with the ground-truth gene regulatory network (GRN). Red paths in A and B indicate relevance propagation through the models, with line width and color intensity encoding relevance magnitude.
Figure 2: Mean marker-recovery average precision (AP) across tasks and models. Class-stratified bootstrapping 95% CIs are also shown, though many are too narrow to visualise at this scale. Relevance scores recover marker genes better than expression values and random ranking. For blood cell subtypes, AP is higher for correct than incorrect predictions, indicating that prediction errors are associated with weaker marker alignment. Incorrect predictions are reported only for blood subtypes, since each model makes at most five errors on the broader cell-type tasks. Marker prevalence and expected AP under random finite rankings are shown for reference.
Figure 3: Distribution of relevance scores across gene-rank bins, ordered by decreasing expression. Box plots show the distribution across test cells. Relevance is concentrated among early, highly expressed genes and decreases with rank, most strongly for C2S (Gemma2-2B), although this pattern alone does not establish positional shortcut behavior.
Figure 4: Context-sensitivity of C2S (Gemma2-2B) with K=1000 and scGPT Whole Human using symbolic queries. The relevance of qminimal is consistently concentrated near zero, whereas the full-context query qfull shows broader and often shifted relevance. This indicates that shared markers derive much of their relevance from interactions with the surrounding cellular context rather than from contributions confined to the shared-marker set.
Figure 5: Evaluation of regulatory alignment in both in-distribution contexts (0 and 1). (A) Binary recovery of direct edge existence, measured as average precision (AP). Empirical DE is clipped for better visualisation. (B) Sign agreement with nonzero ground-truth direct responses. (C) Spearman correlation with stationary total effects. Error bars for the true measurements reflect 95% bootstrap CIs over the 128 perturbations. The grey control error bars are the central 95% range of the perturbation-average evaluated against the randomised-GRN controls.
Figure 6: Regulator recovery under sparse prior knowledge for top-6 differentially expressed genes (DEGs), measured by recall@20. (A) Comparing the final model checkpoint (epoch 20) vs an early checkpoint (epoch 5) and random untrained model baselines. Recall increases with training, and at epoch 5 LRP signatures exceed the random untrained models while predicted DE does not. Results under full knowledge (all direct edges) are also shown for reference as hollow markers, though the numeric scales are not directly comparable as the denominators differ. (B) Recovering known regulatory edges, alongside randomised controls (details in Section 4.4 ).
Table B.3: Pooling strategies selected for the final-layer probes used in attribution analyses. Each strategy is selected separately for each model, context length, and prediction task based on validation performance. A dash denotes an unevaluated combination. Context length K applies only to C2S.
Figure B.1: Probe accuracies across model depths and pooling strategies. Broad cell types are predicted more accurately than fine-grained cell subtypes. Single-token representations become increasingly predictive toward the final layers, while longer C2S input contexts generally improve accuracy. Here, K denotes the number of input genes, and [CLS] denotes the classification token in scGPT.
Half-rule ( Rezaei Jafari et al., 2024 ; Arras et al., 2019 )
✓
×
×
AH-rule ( Ali et al., 2022 )
✓
×
×
AttnLRP ( Achtibat et al., 2024 )
×
✓
✓
Appendix
Table C.2: Propagation rules used in our LRP implementation for the three single-cell architectures considered in this study. A green tick ( ✓ ) indicates that the rule was applied, while a red cross ( × ) indicates that it was not used.
Figure H.1: Evaluation of regulatory alignment in both in-distribution contexts (0 and 1). (A) Binary recovery of direct edge existence, measured as average precision (AP). Empirical DE is clipped for better visualisation. (B) Sign agreement with nonzero ground-truth direct responses. (C) Spearman correlation with stationary total effects. Error bars for the true measurements reflect 95% bootstrap CIs over the 128 perturbations. The grey control error bars are the central 95% range of the perturbation-average evaluated against the randomised-GRN controls. This is an extension of Fig. 5 , including all-cell gene LRP as well as interventional signatures: perturbation ablation and gene masking.
We present a systematic evaluation framework - thirty-seven analyses, 153 statistical tests, four cell types, two perturbation modalities - for assessing mechanistic interpretability in single-cell foundation models. Applying this framework to scGPT and Geneformer, we find that attention patterns encode structured biological information with layer-specific organisation - protein-protein interactions in early layers, transcriptional regulation in late layers - but this structure provides no incremental value for perturbation prediction: trivial gene-level baselines outperform both attention and correlation edges (AUROC 0.81-0.88 versus 0.70), pairwise edge scores add zero predictive contribution, and causal ablation of regulatory heads produces no degradation. These findings generalise from K562 to RPE1 cells; the attention-correlation relationship is context-dependent, but gene-level dominance is universal. Cell-State Stratified Interpretability (CSSI) addresses an attention-specific scaling failure, improving GRN recovery up to 1.85x. The framework establishes reusable quality-control standards for the field.
Ihor Kendiukhov
Department of Computer Science, University of Tübingen, Tübingen, Germany
Background: Single-cell foundation models such as Geneformer and scGPT encode rich biological information, but whether this includes causal regulatory logic rather than statistical co-expression remains unclear. Sparse autoencoders (SAEs) can resolve superposition in neural networks by decomposing dense activations into interpretable features, yet they have not been systematically applied to biological foundation models. Results: We trained TopK SAEs on residual stream activations from all layers of Geneformer V2-316M (18 layers, d=1152) and scGPT whole-human (12 layers, d=512), producing atlases of 82525 and 24527 features, respectively. Both atlases confirm massive superposition, with 99.8 percent of features invisible to SVD. Systematic characterization reveals rich biological organization: 29 to 59 percent of features annotate to Gene Ontology, KEGG, Reactome, STRING, or TRRUST, with U-shaped layer profiles reflecting hierarchical abstraction. Features organize into co-activation modules (141 in Geneformer, 76 in scGPT), exhibit causal specificity (median 2.36x), and form cross-layer information highways (63 to 99.8 percent). When tested against genome-scale CRISPRi perturbation data, only 3 of 48 transcription factors (6.2 percent) show regulatory-target-specific feature responses. A multi-tissue control yields marginal improvement (10.4 percent, 5 of 48 TFs), establishing model representations as the bottleneck. Conclusions: These models have internalized organized biological knowledge, including pathway membership, protein interactions, functional modules, and hierarchical abstraction, yet they encode minimal causal regulatory logic. We release both feature atlases as interactive web platforms enabling exploration of more than 107000 features across 30 layers of two leading single-cell foundation models.
Ihor Kendiukhov
Department of Computer Science, University of Tübingen, Tübingen, Germany
Single-cell RNA sequencing (scRNA-seq) data exhibit strong and reproducible statistical structure. This has motivated the development of large-scale foundation models, such as TranscriptFormer, that use transformer-based architectures to learn a generative model for gene expression by embedding genes into a latent vector space. These embeddings have been used to obtain state-of-the-art (SOTA) performance on downstream tasks such as cell-type classification, disease-state prediction, and cross-species learning. Here, we ask whether similar performance can be achieved without utilizing computationally intensive deep learning-based representations. Using simple, interpretable pipelines that rely on careful normalization and linear methods, we obtain SOTA or near SOTA performance across multiple benchmarks commonly used to evaluate single-cell foundation models, including outperforming foundation models on out-of-distribution tasks involving novel cell types and organisms absent from the training data. Our findings highlight the need for rigorous benchmarking and suggest that the biology of cell identity can be captured by simple linear representations of single cell gene expression data.
Huan Souza, Pankaj Mehta
Department of Physics, Boston University, Boston, MA, 02215, USA · Faculty of Computing and Data Science, Boston University, Boston, MA, 02215, USA