Generalizable single-cell perturbation response prediction using energy-guided flow matching
Authors: Jianan Wei, Jiajun Hong, Guikun Chen, Ning Yang, Lifeng Fan, Wenguan Wang
Organizations: Zhejiang University, Hangzhou, Zhejiang, China · Beijing Institute for General Artificial Intelligence, Beijing, China · Peking University Chengdu Academy for Advanced Interdisciplinary Biotechnologies, Chengdu, Sichuan, China
Predicting phenotypic and transcriptional responses to perturbations at single-cell resolution provides a powerful tool for probing biological systems. However, existing methods typically rely on fixed mappings learned during training, making it challenging to calibrate distribution shifts or adapt to novel perturbation conditions during inference. Here, we present scEGFlow, an energy-guided flow matching framework that dynamically bridges control and perturbed cellular states. scEGFlow models continuous transitions from control cell populations to perturbed states using conditional flow matching. It then applies condition-specific energy gradients to correct and steer these predictions, enabling flexible adjustments without retraining the flow model. Evaluations across benchmarks spanning imaging phenotypes and transcriptomic profiles show that scEGFlow outperforms existing methods in reconstructing response distributions under both seen and unseen perturbation conditions, faithfully preserving cellular manifold geometry and population heterogeneity. This advantage is notable when adapting to new conditions with only a few measured cells, consistently improving prediction accuracy. Furthermore, scEGFlow accurately recapitulates perturbation-induced up- and down-regulation patterns across consensus gene expression signatures, where energy guidance improves the agreement between predicted and observed regulatory directions. Ultimately, these findings demonstrate that scEGFlow provides a modular, generalizable, and steerable solution for single-cell perturbation modeling. By grounding generative flows in learned biological landscapes, this architecture establishes a new computational paradigm for navigating and manipulating cellular behavior in silico.
Figures & tables
Figure 1: Overview of scEGFlow. (a) Optimal transport establishes soft correspondences between unpaired control and perturbed cells through the transport matrix T∗ , yielding soft pseudo-pairs for training. (b) The perturbation vector p and biological context c are encoded as conditional inputs to the flow model, which predicts the velocity field vθ(xt,t,ξ) from the cell state xt and time t . (c) During inference, cells are progressively transported toward the perturbed distribution by alternating flow-based velocity updates with energy-guided refinement. The former follows the learned conditional velocity field, whereas the latter adjusts cellular states along the energy gradient toward condition-consistent low-energy regions.
Figure 2: Reconstruction of single-cell perturbation responses across diverse data modalities. (a) Marginal distributions of normalized pERK, Sox9, and pAKT marker intensities on the 4i [ 11 ] benchmark. Measured treated cells are shown in dark blue, and predictions from scEGFlow, scDFM, and PRNet are shown in red, gray, and light blue, respectively. (b–e) IID reconstruction performance on the 4i benchmark, evaluated using MMD 2 (b) , Sinkhorn divergence (c) , Pearson correlation between condition-level mean profiles (d) , and energy distance (e) . Lower values are better for MMD 2 , Sinkhorn divergence, and energy distance, whereas higher values are better for Pearson correlation. (f) Drug-level comparison between scEGFlow and the MeanExp baseline for six selected 4i perturbations. The horizontal axis shows the log 2 ratio of the energy distances obtained by MeanExp and scEGFlow. (g) UMAP visualization of predicted and measured 4i cells aggregated across all evaluated drug treatments in a shared phenotypic state space. Cells are colored by normalized mean pERK intensity. AU, arbitrary units. (h) UMAP comparison of SciPlex3 [ 3 ] IID response reconstruction across aggregated drug-dose conditions. Measured target cells are shown in gray and predicted cells from each method are shown in blue.
Method
4i
SciPlex3
Energy ↓
Sinkhorn ↓
MMD 2 ↓
Pearson ↑
Energy ↓
Sinkhorn ↓
MMD 2 ↓
Pearson ↑
MeanExp
0.0510
1.4463
0.4839
0.9925
2.2626
12.7494
0.6887
0.9914
CellOT [ 11 ]
3.1399
2.0230
0.0199
0.5513
–
–
–
–
biolord [ 22 ]
0.1507
1.2904
0.2411
0.9669
1.1494
0.9318
0.2192
0.9949
CPA [ 6 ]
0.6125
0.3091
0.0257
0.7455
2.9388
1.5674
0.2509
0.9476
PRNet [ 1 ]
0.0208
0.0254
0.0145
0.9910
0.6419
0.3788
0.0145
0.9856
Table 1: IID perturbation prediction on 4i [ 11 ] and SciPlex3 [ 3 ] .
Method
OOD drug-dose
OOD cellular context
Energy ↓
Sinkhorn ↓
MMD 2 ↓
Pearson ↑
Energy ↓
Sinkhorn ↓
MMD 2 ↓
Pearson ↑
MeanExp
0.8735
5.7422
0.6824
0.9966
52.9832
27.5550
0.7647
0.8090
biolord [ 22 ]
0.7792
0.5939
0.2114
0.9966
51.4051
26.7684
0.7584
0.8154
CPA [ 6 ]
2.2859
1.2211
0.2385
0.9607
14.9031
7.5212
0.4386
0.6901
PRNet [ 1 ]
0.4077
0.2211
0.0051
0.9909
55.4931
22.1669
0.1777
0.8003
CellFlow [ 12 ]
0.5218
0.3692
0.0065
0.9977
22.3046
12.2336
0.1624
0.9254
Table 2: OOD perturbation prediction under unseen drug-dose conditions and in an unseen cellular context on SciPlex3 [ 3 ] .
Figure 3: Generalization of single-cell perturbation responses across unseen chemical and cellular contexts. (a) UMAP visualization of unperturbed cells colored by three cell types (A549, K562, and MCF7 cells). (b) UMAP visualization of measured MCF7 cells treated with 20 representative compounds at the highest profiled dose. (c) Distribution-level prediction performance for held-out SciPlex3 [ 3 ] drug-dose conditions. Boxplots compare the energy distance and MMD 2 between predicted and measured target-cell populations over the OOD test set. (d) Gene-expression distributions of the five most strongly perturbed genes for three held-out drug-dose conditions, comparing control, measured perturbed, and scEGFlow-predicted cells. (e) Covariance matrices over the first 10 principal components for the held-out Mesna_0.1 condition, comparing the measured target population with populations predicted by scEGFlow and the baseline methods; all six matrices are computed in the same PCA basis fitted on the training cells. (f) Heatmap of the measured treated-versus-control response magnitude across representative drug-dose conditions in three cell types. Response magnitude is computed as the root-mean-square of the gene-wise treated-minus-control mean expression shifts after z-score standardization over the retained genes. (g) Gene-wise MMD 2 for scEGFlow and CellFlow on the most strongly perturbed genes, selected by the absolute treated-minus-control mean shift in the held-out A549 context. For each gene, cells across all evaluated drug-dose conditions are pooled before computing MMD 2 between measured and predicted expression distributions. Lower values indicate better agreement.
Figure 4: Generalization to unseen drug combinations and sample-efficient prediction. (a) UMAP visualization of measured ComboSciPlex [ 6 ] cells, showing the five held-out drug-combination responses and their positions relative to the response manifold. (b) Predicted versus measured mean gene expression for the five held-out drug combinations, shown for all genes and the top 50 differentially expressed genes (DEGs). (c) Expression distributions of the top two DEGs for each held-out drug combination, comparing control, measured treated, and scEGFlow-predicted cells. (d) Few-shot prediction on the 4i [ 11 ] benchmark under different numbers of target-condition training cells, evaluated using Pearson correlation, MMD 2 , Sinkhorn divergence, and energy distance. (e) UMAP comparison of zero-shot predictions, 3-shot predictions, 5-shot predictions and measured target cells, colored by normalized pERK intensity.
Figure 5: Directional concordance with L1000 [ 20 ] perturbational signatures. (a) Mean directional concordance between predicted perturbation responses and the corresponding L1000 Level 5 consensus signatures. (b) Context-level distribution of directional concordance across all 111 evaluation contexts. Each point represents one specific cell line-drug-dose condition.
Dataset
Cells/profiles
Features
Contexts
Conditions
Readout
4i [ 11 ]
77,743
48
1
27
imaging phenotype
SciPlex3 [ 3 ]
290,888
2,000
3
749
scRNA-seq
ComboSciPlex [ 6 ]
63,378
2,000
1
32
scRNA-seq
L1000 Level 3 [ 20 ]
24,104
978
21
110
gene-expression profiles
Table S1: Overview of the datasets used in this study.
Dataset
Feature dim.
Training steps
Batch size
Learning rate
4i [ 11 ]
48
100,000
256
10−6
SciPlex3 [ 3 ]
200
50,000
256
10−7
ComboSciPlex [ 6 ]
200
50,000
256
10−7
LINCS L1000 [ 20 ]
768
30,000
128
10−7
Table S2: Training hyperparameters for the conditional flow-matching model.
Dataset
Feature dim.
Training steps
Batch size
Learning rate
Weight decay
4i [ 11 ]
48
5,000
256
10−4
10−4
SciPlex3 [ 3 ]
200
300
256
10−4
10−4
ComboSciPlex [ 6 ]
200
5,000
256
10−4
10−4
LINCS L1000 [ 20 ]
768
300
128
10−4
10−4
Table S3: Training hyperparameters for the energy model.
Condition
Model
Energy ↓
Sinkhorn ↓
MMD 2 ↓
Pearson ↑
ENMD-2076_0.1
w/o energy guidance
0.2380
0.1360
0.0040
0.9947
ENMD-2076_0.1
w/ energy guidance
0.2050
0.1180
0.0035
0.9955
WP1066_0.01
w/o energy guidance
0.2520
0.1450
0.0041
0.9942
WP1066_0.01
w/ energy guidance
0.2180
0.1260
0.0036
0.9951
Mesna_0.1
w/o energy guidance
0.3530
0.2035
0.0054
0.9907
Mesna_0.1
w/ energy guidance
0.3084
0.1733
0.0046
0.9929
Table S4: Prediction performance with and without energy guidance on three held-out OOD drug-dose conditions in SciPlex3 [ 3 ] . Both settings share the same trained flow model.
Condition
Model
Energy ↓
Sinkhorn ↓
MMD 2 ↓
Pearson ↑
Dacinostat + Danusertib
w/o energy guidance
25.8631
100.5147
0.0656
0.9588
Dacinostat + Danusertib
w/ energy guidance
21.7346
83.7241
0.0476
0.9720
Givinostat + Cediranib
w/o energy guidance
29.4728
112.7362
0.0498
0.9792
Givinostat + Cediranib
w/ energy guidance
24.1183
91.5836
0.0524
0.9760
Panobinostat + Alvespimycin
w/o energy guidance
64.7315
195.6841
0.1215
0.7904
Panobinostat + Alvespimycin
w/ energy guidance
55.8427
168.9142
0.0987
0.8210
Table S5: Prediction performance with and without energy guidance on five held-out ComboSciPlex [ 6 ] OOD drug combinations. Both settings share the same trained flow model. Average rows are obtained by pooling all test cells from the five held-out combinations and recomputing each metric on the pooled measured and predicted populations.
λ
Energy ↓
Sinkhorn ↓
MMD 2 ↓
Pearson ↑
0.0000
0.2810
0.1615
0.0045
0.9932
0.0125
0.2705
0.1554
0.0043
0.9936
0.0250
0.2588
0.1480
0.0041
0.9940
0.0500
0.2438
0.1391
0.0039
0.9945
0.1000
0.2479
0.1415
0.0040
0.9944
0.2000
0.2676
0.1532
0.0043
0.9937
Table S6: Sensitivity of SciPlex3 [ 3 ] OOD drug-dose prediction to the guidance coefficient λ .
λ
Energy ↓
Sinkhorn ↓
MMD 2 ↓
Pearson ↑
0.0000
29.9767
111.6462
0.0618
0.9597
0.0125
28.4163
105.7328
0.0589
0.9621
0.0250
26.6835
99.2846
0.0554
0.9646
0.0500
24.8578
93.0152
0.0527
0.9667
0.1000
25.1936
94.2671
0.0532
0.9662
0.2000
27.1849
101.4385
0.0569
0.9637
Table S7: Sensitivity of ComboSciPlex [ 6 ] OOD drug combination prediction to the guidance coefficient λ .
Predicting single-cell responses to genetic and chemical perturbations is important for understanding regulatory mechanisms and prioritizing therapeutic interventions. Existing methods have made progress by modeling perturbation-conditioned expression shifts with neural networks, graph priors, foundation-model representations, or distributional generative models. However, combinatorial perturbation prediction remains challenging because much of the response can be explained by additive single-perturbation effects, while the biologically informative signal often lies in a smaller non-additive residual induced by perturbation interactions. Directly generating the full post-perturbation expression profile may therefore allocate substantial capacity to the dominant additive component and underemphasize synergy-specific structure. To address this issue, we propose SynGeneFlow, a synergy-aware flow matching framework that anchors the source distribution at an additive prior and learns a residual flow toward the true perturbed state. SynGeneFlow further incorporates perturbation-conditioned gene interaction modeling, where a static gene-gene mask provides biological structure and a perturbation-dependent deformation bias dynamically modulates gene-gene coupling. To preserve population-level response structure, we align both endpoint distributions and biological residual distributions with MMD-based regularization. We evaluate SynGeneFlow on genetic and drug perturbation benchmarks under observed-single, held-out, and cross-modality settings. The results demonstrate that residualizing flow matching around additive perturbation effects provides a biologically grounded and robust framework for single-cell perturbation response prediction.
Jingrui Wu, Wangyun Dan
Department of Statistics, Rice University, Houston, TX, USA · Department of Biostatistics and Bioinformatics, Emory University, Atlanta, GA, USA
Predicting cellular responses to perturbations is a central problem in cellular biology, with broad applications in systems biology and drug discovery. This task is challenging because cellular responses can be complex and cell-state dependent, intrinsic cell-to-cell variability can be confounded with perturbation effects, and destructive single-cell RNA sequencing precludes paired measurements of the same cell before and after treatment. Flow matching transports control cells to perturbed states flexibly, but acting on the full cell state can confound perturbation effects with pre-existing cell-to-cell variability. Disentangled approaches separate responsive from invariant components, but model perturbations through prescribed mechanisms, such as latent shifts or graph edits, limiting their flexibility. We address both limitations in a unified framework. A variational encoder disentangles each cell into an invariant block, capturing state unaffected by the perturbation, and a responsive block, capturing state it changes, through conditional priors and an information-theoretic invariance constraint. Conditional flow matching transports only the responsive block, conditioned on the perturbation and invariant state, yielding a flexible, data-driven model of perturbation effects without confounding pre-existing variability. Across several benchmarks, our method outperforms the strongest published method in settings involving combinatorial and unseen perturbation prediction.
Mustapha Bounoua, Giulio Franzese, Pietro Michiardi
Predicting the effects of perturbations in-silico on cell state can identify drivers of cell behavior at scale and accelerate drug discovery. However, modeling challenges remain due to the inherent heterogeneity of single cell gene expression and the complex, latent gene dependencies. Here, we present PRiMeFlow, an end-to-end flow matching based approach to directly model the effects of genetic and small molecule perturbations in the gene expression space. The distribution-fitting approach taken by PRiMeFlow enables it to accurately approximate the empirical distribution of single-cell gene expression, which we demonstrate through extensive benchmarking inside PerturBench. Through ablation studies, we also validate important model design choices such as operating in gene expression space and parameterizing the velocity field with a U-Net architecture. Finally, by scaling PRiMeFlow to a broad perturbation data atlas spanning multiple datasets and employing a carefully designed pretraining-finetuning strategy, we demonstrate its outstanding performance on the H1 human embryonic stem cells from the ARC Virtual Cell Challenge benchmark.
Zichao Yan, Yan Wu, Mica Xu Ji +10
Altos Labs · Current affiliation: Verge Genomics · Current affiliation: Loka +1