Graham Heimberg
← All projects

Rethinking what foundation models are for

Most people use single-cell foundation models for one of two things: classifying cell types (which usually doesn't justify the cost of building a foundation model in the first place) or predicting the effect of an arbitrary gene knockout in a dish (which is both extremely difficult and, arguably, not always biologically meaningful). This project didn't train a new model. It asked a different question about the models we already have: how do you quantify a gene's importance to a cell's specific identity and function?

Image classifier attribution example: a neural network classifies a photo as a dog and highlights the pixels that drove that prediction
Attribution methods already have a home in image classifiers: this network calls the photo 96% dog, and highlighting which pixels drove that call shows it's actually looking at the dog.

The same idea carries over to a cell. Feed a cell's gene expression into a single-cell foundation model, and instead of just reading off its embedding, calculate an attribution score for every gene: how much did this specific gene move the cell to where it landed?

Diagram showing single-cell gene expression fed into a foundation model, producing both a cell embedding and a per-gene attribution score
FOXP3 gets a high attribution score (5.2) despite modest expression (7); MALAT1 gets a near-zero score despite the highest expression (10) of the four. Expression and importance aren't the same thing.

The intuitive assumption, that highly expressed genes matter more, is often backwards. Some of the most functionally important genes, like cytokines, are expressed at very low levels; some of the highest-expressed genes, like ribosomal genes, get filtered out precisely because they say almost nothing about a cell's role in its tissue.

Embedding shift comparison: altering the low-attribution gene MALAT1 barely moves a cell's embedding, while altering the high-attribution gene FOXP3 shifts it sharply into a different cluster
Proof in the embedding itself: zeroing out MALAT1 (low attribution) barely moves the cell; zeroing out FOXP3 (high attribution) shifts it sharply into a different neighborhood.
Scatter plot of standardized mean attribution versus standardized mean expression, highlighting B cell marker genes BANK1, MS4A1, and CD52 as high-attribution despite modest expression, and MALAT1, RPS27, and MT-CO1 as high-expression despite low attribution
The same pattern shows up in real B cells: marker genes like BANK1 and MS4A1 score high on attribution at modest expression, while housekeeping genes like MALAT1 score low on attribution despite dominating the expression profile.

Working out how to quantify "importance" properly, using what the foundation model has actually learned rather than raw expression, opened up several new tools: scalable searches for shared gene signatures across diseases, a cleaner way to simplify complex gene-expression data down to its essential patterns (dimensionality reduction) with much less background noise, and better ways to identify the genes that define a given cell type (marker genes). Applied at scale, that search for shared signatures is what turned up the paper's main finding: a single inflammatory gene program showing up across multiple, seemingly unrelated diseases, one that stayed invisible if you only looked at which genes were most highly expressed in each disease individually. A signature that broad is exactly the kind of thing worth chasing as a drug target, since treating it could help more than one patient population at once.

Gold, M., et al. Foundation Model Attributions Reveal Shared Inflammatory Program Across Diseases. Nature Biotechnology, 2026.