Work done as final project for 6.8700/HST.507[J] (Graduate Computational Biology) at MIT.
Modern single-cell perturbation assays now enable systematic measurement of how genetic and chemical interventions change cellular transcriptional states at large scale. Despite this progress, accurately predicting transcriptional responses in new cell types, contexts, or perturbations remains a central challenge in functional genomics and drug discovery. Recent benchmarks have shown that deep learning models often fail to improve upon simple linear or mean expression baselines, which suggests that they do not fully capture the underlying biological mechanisms that govern cellular response.
A progression of methods, from lightweight models such as ScAPE to diffusion based, variational, and transformer architectures including CPA, scGPT, and STATE, has steadily increased improved performance on several perturbation datasets. However, their improvements remain modest. Moreover, most approaches that perform better treat perturbation prediction as a high dimensional regression problem , which is inherently not interpretable or generalizable. Thus, we ask whether encoding biological structure (specifically from gene regulatory networks) into the learning process can improve prediction of transcriptional responses while maintaining biological interpretability.
In single-cell perturbation modeling, the central task is to predict the transcriptional response of a cell to a specified genetic or chemical perturbation. Given a representation of the unperturbed cell state and a description of the perturbation, the model outputs gene-level changes, typically a post-perturbation gene expression profile, or log fold-changes relative to control. Performance is assessed in terms of generalization to new cell types, contexts, and perturbations, i.e. accuracy in predicting which genes respond and how strongly they respond.
To quantify performance, we use the metrics outlined by STATE and the Virtual Cell Challenge (VCC). DE Overlap@N measures how many of the top-k truly differentially expressed genes (DEGs) are recovered within the model's top-k predictions. DE Precision@N instead quantifies the fraction of the model's top-k predicted DEGs that are truly DE. MSE and MAE evaluate the magnitude error of predicted log fold-changes. DE Spearman computes the Spearman rank correlation between predicted and true log fold-changes restricted to significantly DE genes, capturing how well the model preserves the ordering of gene-level DE effects.
STATE is a large-scale foundation model for predicting transcriptional responses to genetic and chemical perturbations at single-cell resolution. It is composed of two components. The State Transition (ST) module learns how expression changes between perturbed and matched control cells, modeling perturbation-induced differences in gene expression. The State Embedding (SE) module is trained on large observational single-cell atlases and maps cells into context-rich embeddings that capture baseline cell identity and state. During perturbation prediction, STATE combines the perturbation-aware ST module with these SE embeddings to predict gene-level responses. Trained on hundreds of millions of cells, STATE achieves strong cross-cell-type generalization, and thus, we adopt STATE-142M as a baseline reference model for benchmarking perturbation prediction.
We obtain the working GRN by running inference with scPRINT, a single-cell foundation model trained to infer gene-gene regulatory relationships from large-scale transcriptomic data. Given single-cell expression profiles, scPRINT assigns an edge weight to each gene pair that reflects how strongly one gene is inferred to influence the other, thereby defining a weighted gene regulatory network. In our setting, we run scPRINT on exactly the set of genes for which we predict perturbation responses and use the resulting edge weights to construct a GRN whose nodes align with the output space of STATE.
We investigated two post-training methods, the first of which was post-hoc smoothing. Post-hoc smoothing applies a graph-based diffusion operator directly to the final output of the transition model, the predicted gene expression changes produced by the ST decoder head (the last linear projection from the latent representation to the gene expression space). In this paradigm, all model parameters, remain frozen. Only the decoded prediction vector is transformed using the structure of the GRN. Given a GRN with adjacency matrix $A$ and degree matrix $D$, we construct the symmetric normalized graph Laplacian, which turns the GRN topology into an explicit, differentiable prior on how gene responses should behave across the network: $$ L = I - D^{-1/2} A D^{-1/2}. $$ We obtain smoothed predictions by diffusing the decoder output $y_{\text{pred}} \in \mathbb{R}^G$ over the GRN via the resolvent operator: $$ y_{\text{smoothed}} = (I + \tau L)^{-1} y_{\text{pred}}, $$ Equivalently, $y_{\text{smoothed}}$ is the unique minimizer of the quadratic objective: $$ y_{\text{smoothed}} = \arg\min_{x \in \mathbb{R}^G} \left[ \frac{1}{2} \lVert x - y_{\text{pred}} \rVert_2^2 + \frac{\tau}{2} x^\top L x \right] $$ which encourages transformations to stay close to the original prediction $y_{\text{pred}}$ and enforces smoothness on the GRN, meaning connected genes should have more consistent perturbation responses unless the data strongly suggests otherwise. Unfortunately, post-hoc smoothing yields subpar results across all metrics. We hypothesize that this is because it applies a fixed linear transformation that enforces smoothness on the final predictions irrespective of context, damping sharp gene-specific perturbation effects and spreading them across neighbors in the GRN, thereby weakening the signals that distinguish responsive from non-responsive genes.
The second post-training method we investigated was fine-tuning with GRN guidance. The goal of GRN-guided fine-tuning is to refine the pretrained perturbation prediction model on new experimental conditions while preserving meaningful biological structures learned during pretraining. Therefore, this approach augments the base training objective with two additional regularization terms that encourages preservation of gene-gene structure and retains knowledge learned from pretraining. In this paradigm, we keep the entire ST module trainable. We formulate the fine-tuning objective: $$ L = L_{\text{base}} + \lambda L_{\text{GRN}} + \mu \lVert W - W_0 \rVert_2^2, $$ where $L_{\text{base}}$ is the original pre-training loss, and $$ L_{\text{GRN}} = \frac{1}{|E|} \sum_{(i,j) \in E} \lVert W_i - W_j \rVert_2^2 $$ encourages smoothness across genes connected in the GRN, with $W_i$ denoting the row of the ST decoder head corresponding to gene $i$. The proximity term $\lVert W - W_0 \rVert_2^2$ penalizes deviations of the current ST decoder weights $W$ at a given fine-tuning step from the pretrained weights $W_0$. The hyperparameters $\lambda$ and $\mu$ control the strength of GRN smoothing and proximity regularization, respectively, and $t$ denotes the GRN edge subset size. As before, genes linked in the GRN are encouraged to remain close in the readout space, while large shifts in the model's outputs are penalized. We find that incorporating GRN structure admits improvements over the pre-trained baseline, especially for DE-focused metrics. GRN-guided fine-tuning for 2000 steps yields the largest increase in Overlap@N (+80.8%), Precision@N (+52.4%), and DE Spearman (+12.4%). Shorter schedules (800-1200 steps) reduce MAE/MSE, with fine-tuning for 1200 steps achieving the lowest MSE relative to baseline (-22.5%). Overall, GRN guidance shapes the model toward lower reconstruction error at intermediate steps and better recovery of biologically meaningful DE patterns at longer fine-tuning schedules.
Perturbation prediction models are most useful for in silico screening, where they can forecast transcriptomic responses to large panels of CRISPR edits or small molecules across many cell types, enabling researchers to prioritize a small subset of high-value experiments. They also support target discovery and mechanism-of-action analysis by identifying the genes and pathways most strongly affected by a given intervention, which can refine hypotheses about causal drivers and downstream programs. In addition, by simulating combination treatments or multi-gene perturbations, these models make it feasible to explore large combinatorial spaces that would be expensive or impractical to cover systematically in wet lab experiments.