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.