Designing observational transcriptomic pretraining objectives for cellular perturbation prediction (Part 1)

This blog is based on a few of our findings in EMBER that I thought were interesting. There is a lot more in the paper so please check that out!

Read the paper · Code · Video explainer

The success of text LLMs largely stemmed from how downstream performance could be predictably scaled with compute and data, known as the scaling laws (1,2). It remains an open question on how to find similar scaling behaviors in other domains, such as biological data.

One particular task of interest is perturbation prediction, predictions of how a cell responds to an external perturbation, such as a gene knockout or chemical application (3,4).

Due to the limited amount of available public perturbation data* we focus on asking how to design pretraining objectives for observational transcriptomics data (scRNA-seq) to improve downstream performance.

Question: How can we design pretraining objectives for observational transcriptomic data that yield predictable improvements in downstream perturbation-response prediction?

In EMBER, we found that the gene representation level was a strong interface for transferring representations from pretrained transcriptomics models to downstream tasks.

Here, we’ll talk about some pretraining insights we found with gene representations and how to design better gene representations.

---

What does it look like to pretrain on single cell data?

Figure 1. EMBER pretraining. Each cell contributes 1,024 expressed and 256 non-expressed genes. The transformer produces contextualized gene embeddings used by two training objectives. This post focuses on masked expression reconstruction. View full size.

For EMBER, we chose to sample 1280 genes per cell, where 1024 genes are sampled from expressed genes and 256 are sampled from non-expressed genes. Then, we mask the gene expression of 50% of our 1280 sampled genes. After passing this sequence of 1280 gene tokens into a transformer, the output tokens are the contextualized gene embeddings, where we then apply two training objectives, masked expression reconstruction and cell subgraph contrastive learning (CSCL).

Why sample genes rather than use the full transcriptome?

Single-cell data is a tabular datatype that contains counts for usually 20K or more genes. One challenge of single-cell data is that it is incredibly sparse, usually 90% of measured genes have no expression.

It is arguable that maybe training on the full 20K transcriptome may capture more information, and we did that before with AIDO.Cell, but applying the attention mechanism over all 20K genes made the computational trade-offs quite large.

Just taking the 1024 expressed genes as input has empirically worked quite well for capturing an estimate of cell state.

We intentionally chose to focus on pretraining objectives, as we felt that was the biggest opportunity for scalable improvement.

For this blog post we will focus specifically on the effect of the second loss, masked expression reconstruction, and save CSCL for another day.

What are we actually learning (with reconstruction)?

Intuitively, reconstruction should organize gene embeddings in the best way to support… reconstruction (insightful, I know right?). Specifically, we hypothesized that genes that are “more difficult to reconstruct” may be geometrically separated from genes with much more predictable expression patterns.

So what would that mean? Aside from genes that were highly expressed, we hypothesized that a gene that is regulated by many other genes may be more difficult to reconstruct.

We measured this with Perturb-seq data, we measured a gene’s perturbation responsiveness as measuring how many perturbations regulate a readout gene. We calculated this for 4 immortalized cell lines with perturb-seq data available (Replogle-Nadig dataset) trained on cells that have different lineages from the cell lines.

We found that genes that were more perturbation responsive (genes regulated by many other genes), tended to lie further away from the centroid of gene embeddings in that cell (adjusting for expression magnitude, variance, and tokenization sampling).

Figure 2. Left: perturbation responsiveness counts how many perturbations affect a readout gene. Center: more responsive genes tend to lie farther from the gene-embedding centroid. Right: correlations are stronger for embeddings from the matched cell line than for unmatched cell lines. View full size.

The effect of pretraining data

This behavior appears to be dependent on the pretraining dataset composition as well. We chose 100 cell types that had a different lineage to the 4 immortalized cell lines, kept the number of cells fixed at 176K cells, and increased the number of cell types available in pretraining. We found there to be a monotonic relationship between the number of cell types available in pretraining, and the amount the embedding space recapitulated perturbation responsiveness.

Figure 3. With the number of pretraining cells held fixed, the correlation between distance from the gene-embedding centroid and perturbation responsiveness increases with cell-type diversity. View full size.

How could this be useful for perturbation prediction?

Genes that inherently organize by perturbation responsiveness can actually be a useful prior. Here is a simple toy example:

For a given gene knockout, can a logistic regression classifier distinguish differentially expressed genes (DE genes) from those that are unchanged?

We tracked logistic regression accuracy over many perturbations, and also calculated Spearman correlation over many perturbations while the reconstruction-only version of EMBER trained. We found that linear regression performance tracked closely with how well the model recapitulated perturbation responsiveness.

Figure 4. Left: held-out readout-gene classification and correlation with perturbation responsiveness improve during pretraining. Right: NUP54 knockdown in Jurkat, with five-fold cross-validation balanced accuracy of 0.80. The histogram shows signed distances from the logistic-regression decision plane. View full size.

As a specific case study, we can look at NUP54 knockdown in Jurkat, and see that the logistic regression hyperplane appears to separate the DE genes from the non-DE genes in this dataset.

Details on the logistic regression

We fit logistic regression to the normalized gene embeddings, $\hat{v} = v / \lVert v \rVert_2$, rather than the raw embeddings $v$, to predict whether a readout gene is differentially expressed under a given perturbation. It is the distance from the centroid that correlated with perturbation responsiveness. As a result, to set up a hyperplane to easier separate everything within a distance from the centroid, it would be easier to put the points on a sphere (just by renormalizing each gene by its norm). Hence the linear regression is actually on $v / \lVert v \rVert$ instead of just $v$.

We added a simple baseline where the model just took in expression magnitude and variance, to ensure that the gene embedding space was providing something more than just expression levels.

We use five-fold cross-validation over readout genes: each model is trained on four folds and evaluated on the remaining fold. We repeat this evaluation across over 500 different sampled perturbations.

For the NUP54 histogram, we pool the held-out results across all five folds. Each gene appears in the test set exactly once, and its plotted signed distance comes from the logistic-regression decision plane fitted without that gene.

So is reconstruction all we need?

In our paper, we found reconstruction alone to not be sufficient. Although reconstruction alone captures perturbation responsiveness, the gene embeddings do not organize by known gene ontology groupings (Figure 5). We also found that models trained with reconstruction alone do not perform well when holding out perturbations.

Figure 5. Left: gene-embedding clustering by Gene Ontology across models. Right: cell-type clustering for lineage markers, transcription factors, and housekeeping genes. Both panels report adjusted Rand index (ARI). NMI and more baselines can be found in the full paper. View full size.

We hypothesize that gene embeddings encoding molecular function is important for generalizing perturbations. Interestingly, our reconstruction model performed similarly to that of many existing scFMs that were also trained on only reconstruction-based objectives, suggesting that models trained on only reconstruction may learn similar patterns.

Next up

In short, we found that when trained with expression reconstruction, contextualized gene representations appear to organize in the latent space that is beneficial for downstream perturbation tasks. This pattern appears to scale with higher cell-type diversity in the pretraining corpus.

In part 2, we will further examine the cell subgraph contrastive learning training objective, which drives context-specific molecular function, but also works synergistically with reconstruction for better gene representations.

Limitations

Some key limitations are that these evaluations are on perturb-seq on immortalized cell lines, which may not generalize to other biological contexts.

Secondly, these experiments only go up to 100 cell types that are distinct from the 4 immortalized cell lines. It would be better to scale this analysis to 1000 or more cell types, but cell lines based on well studied cell types, such as Jurkat (T-cells) limit the number of other distinct cell types we could use in this analysis.

Lastly, I think there is a strong interplay between cell-type, donors, and experimental batch effects. We tried to minimize the number of donors and experimental batches in curating the pretraining datasets for this analysis, and did not investigate the isolated effects of these additional factors.

References

1. Hoffmann, J., et al. (2022). Training Compute-Optimal Large Language Models. arXiv:2203.15556.
2. Kaplan, J., et al. (2020). Scaling Laws for Neural Language Models. arXiv:2001.08361.
3. Roohani, Y. H., et al. (2025). Virtual Cell Challenge: Toward a Turing test for the virtual cell. Cell, 188(13), 3370–3374.
4. Xing, E. P., and Song, L. (2026). A world model of the virtual cell. Cell, 189(19), 5831–5844.