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.
Figures & tables
Dataset
Single / Dual
Modality
Biological
Cells
Task
perturbations
states
Srivatsan20
188 / 0
small molecule
3
178,213
covariate transfer
Jiang24
219 / 0
genetic
30
1,628,476
covariate transfer
Norman19
155 / 131
genetic
1
91,168
combo prediction
Table 1: Summary of benchmarking datasets.
Figure 1: PRiMeFlow workflow overview. The velocity field is directly parameterized in the gene expression space, using a U-Net architecture, and is conditioned on information of the perturbation, covariates, as well as the time variable. The velocity field is learned with the standard conditional flow matching loss. At the inference time, perturbed cells are sampled by solving the initial value problems using the learned velocity field.
Model
MMD
MMD
MMD
DEG
GEX
GEX rank
PCA
recall
PRiMeFlow CFG1
0.13±4×10−3
0.13±8×10−3
0.14±8×10−3
0.26±1×10−2
PRiMeFlow CFG3
0.25±2×10−3
0.22±7×10−3
0.31±1×10−2
0.27±6×10−3
PRiMeFlow CFG5
0.47±1×10−2
0.29±2×10−2
0.60±2×10−2
0.25±5×10−3
FM PCA CFG1
1.2±7×10−3
0.26±6×10−3
0.50±2×10−2
0.042±2×10−3
FM PCA CFG3
1.3±5×10−3
0.25±7×10−3
0.44±1×10−2
0.071±3×10−3
Table 2: Distributional metric performance of PRiMeFlow on the Srivatsan20 dataset.
Model
MMD
MMD
MMD
DEG
GEX
GEX rank
PCA
recall
PRiMeFlow CFG1
0.26±7×10−3
0.019±4×10−3
0.34±2×10−2
0.54±2×10−2
PRiMeFlow CFG3
0.61±3×10−2
0.066±9×10−3
0.86±5×10−2
0.61±8×10−3
PRiMeFlow CFG5
1.3±4×10−2
0.17±1×10−2
1.8±5×10−2
0.56±1×10−2
FM PCA CFG1
0.77±4×10−3
0.059±3×10−3
0.61±5×10−3
0.37±2×10−3
FM PCA CFG3
1.1±8×10−3
0.098±2×10−3
0.82±9×10−3
0.45±2×10−3
Table 3: Distributional metric performance of PRiMeFlow on the Norman19 dataset.
Figure 2: Visualization of cells (ground truth and model predictions) in Srivatsan20 test split. Gray marks all the cells in the test split, used for computing PCA and illustrated as background. Green cells are the ground truth perturbed cells for one condition, repeated in each subplot. From each model, 100 cells are sampled, then projected to the same PCA space. A ): Ellagic acid perturbation in a549 cell line; B ) Pirarubicin perturbation in k562 cell line.
Model
DES ↑
PDS ↑
MAE ↓
SPEARMAN ↑
SPEARMAN_LFC ↑
AUPRC ↑
PEARSON ↑
PRiMeFlow pretrained
0.210
0.695
0.074
0.406
0.320
0.264
0.181
PRiMeFlow 300pts
0.223
0.725
0.084
0.431
0.381
0.264
0.193
PRiMeFlow 250pts
0.228
0.747
0.086
0.473
0.396
0.266
0.203
xBio
0.305
0.811
0.770
0.564
0.087
0.252
0.217
Mean Predictors
0.305
0.741
6.723
0.294
0.213
0.582
0.217
Shippers
0.354
0.699
0.231
0.000
0.227
0.576
0.123
Table 4: PRiMeFlow performance on H1 hESCs private test set using all seven ARC VCC metrics.
Appendix figures & tables20 assets
Supplementary material from the paper’s appendix.
Appendix
Model
DES ↑
PDS ↑
MAE ↓
SPEARMAN ↑
SPEARMAN_LFC ↑
AUPRC ↑
PEARSON ↑
PRiMeFlow pretrained
0.221
0.712
0.078
0.440
0.310
0.284
0.199
PRiMeFlow 200pts
0.224
0.817
0.092
0.603
0.386
0.289
0.216
x.Compass
0.344
0.852
4.215
-
-
-
-
Outlier
0.378
0.832
4.218
-
-
-
-
WIND
0.333
0.822
0.605
-
-
-
-
Appendix
Table 5: PRiMeFlow performance on H1 hESCs public test set using all seven ARC VCC metrics.
Model
Cosine
Cosine
RMSE
RMSE
logFC
logFC rank
mean
mean rank
PRiMeFlow CFG1
0.46±1×10−2
0.12±1×10−2
0.019±5×10−4
0.14±6×10−3
PRiMeFlow CFG3
0.43±9×10−3
0.10±4×10−3
0.029±3×10−4
0.22±5×10−3
PRiMeFlow CFG5
0.39±1×10−2
0.10±2×10−3
0.041±4×10−4
0.29±1×10−2
PRiMeFlow MLP CFG1
−0.012±4×10−4
0.46±1×10−2
0.31±9×10−4
0.42±1×10−2
PRiMeFlow MLP CFG3
−0.0071±2×10−3
0.45±3×10−2
0.33±5×10−3
0.44±2×10−2
Appendix
Table 6: Performance of PRiMeFlow on traditional pseudobulk-based metrics in Srivatsan20 dataset.
Model
Cosine
Cosine
RMSE
RMSE
logFC
logFC rank
mean
mean rank
PRiMeFlow CFG1
0.77±4×10−3
0.0057±8×10−4
0.046±1×10−3
0.022±6×10−3
PRiMeFlow CFG3
0.79±2×10−3
0.0072±3×10−3
0.074±2×10−3
0.077±8×10−3
PRiMeFlow CFG5
0.74±3×10−3
0.015±3×10−3
0.11±2×10−3
0.19±1×10−2
PRiMeFlow MLP CFG1
0.024±6×10−3
0.27±3×10−2
0.29±2×10−3
0.16±4×10−2
PRiMeFlow MLP CFG3
0.054±8×10−3
0.27±2×10−2
0.31±2×10−3
0.25±2×10−2
Appendix
Table 7: Performance of PRiMeFlow on traditional pseudobulk-based metrics in Norman19 dataset.
Model
Cosine
Cosine
RMSE
RMSE
logFC
logFC rank
mean
mean rank
PRiMeFlow CFG1
0.47±5×10−3
0.39±2×10−2
0.019±3×10−4
0.40±2×10−2
PRiMeFlow CFG3
0.41±4×10−3
0.38±1×10−2
0.034±5×10−4
0.43±1×10−2
PRiMeFlow CFG5
0.36±4×10−3
0.37±1×10−2
0.045±5×10−4
0.43±7×10−3
PRiMeFlow MLP CFG1
0.053±1×10−4
0.45±5×10−3
0.33±6×10−4
0.46±1×10−3
PRiMeFlow MLP CFG3
0.053±7×10−5
0.43±1×10−2
0.33±6×10−4
0.46±3×10−3
Appendix
Table 8: Performance of PRiMeFlow on traditional pseudobulk-based metrics in Jiang24 dataset.
Figure 3: Visualization of cells (ground truth and model predictions) in Srivatsan20 test split. Gray marks all the cells in the test split, used for computing PCA and illustrated as background. Green cells are the ground truth perturbed cells for one condition, repeated in each subplot. From each model, 100 cells are sampled, then projected to the same PCA space. A ): SB431542 perturbation in k562 cell line; B ) Obatoclax Mesylate (GX15-070) perturbation in mcf7 cell line.
Figure 4: Visualization of cells (ground truth and model predictions) in Srivatsan20 test split. Gray marks all the cells in the test split, used for computing PCA and illustrated as background. Green cells are the ground truth perturbed cells for one condition, repeated in each subplot. From each model, 100 cells are sampled, then projected to the same PCA space. A ): GSK J1 perturbation in k562 cell line; B ) Tacedinaline (C1994) perturbation in k562 cell line.
Figure 5: Visualization of cells (ground truth and model predictions) in Jiang24 test split. Gray marks all the cells in the test split, used for computing PCA and illustrated as background. Green cells are the ground truth perturbed cells for one condition, repeated in each subplot. From each model, 100 cells are sampled, then projected to the same PCA space. A ): DR1 perturbation in mcf7 cell line and treated with INS; B ) ZC3H3 (C1994) perturbation in hap1 cell line and treated with IFNG.
Figure 6: Visualization of cells (ground truth and model predictions) in Jiang24 test split. Gray marks all the cells in the test split, used for computing PCA and illustrated as background. Green cells are the ground truth perturbed cells for one condition, repeated in each subplot. From each model, 100 cells are sampled, then projected to the same PCA space. A ): SP1 perturbation in ht29 cell line and treated with INS; B ) PIK3CA perturbation in k562 cell line and treated with INS.
Hyperparameter
Value
Description
architecture
U-Net
Architecture choice for the velocity field
num_res_blocks
2
Number of residual blocks after each down/upsampling
hidden_dim
256
Base feature dim in U-Net and in input MLPs
dropout
0.1
Dropout rate
channel_mult
[1, 2, 2, 4, 4]
Hidden dim multipliers at each level
conv_resample
True
Learnable up/downsampling
Appendix
Table 12: Manually selected hyperparameters for PRiMeFlow in Srivatsan20 dataset.
Hyperparameter
Value
Description
architecture
U-Net
Architecture choice for the velocity field
num_res_blocks
2
Number of residual blocks after each down/upsampling
hidden_dim
256
Base feature dim in U-Net and in input MLPs
dropout
0.1
Dropout rate
channel_mult
[1, 1, 2, 2, 4]
Hidden dim multipliers at each level
conv_resample
True
Learnable up/downsampling
Appendix
Table 13: Manually selected hyperparameters for PRiMeFlow in Norman19 dataset.
Hyperparameter
Value
Description
architecture
U-Net
Architecture choice for the velocity field
num_res_blocks
2
Number of residual blocks after each down/upsampling
hidden_dim
256
Base feature dim in U-Net and in input MLPs
dropout
0.1
Dropout rate
channel_mult
[1, 1, 2, 2, 4]
Hidden dim multipliers at each level
conv_resample
True
Learnable up/downsampling
Appendix
Table 14: Manually selected hyperparameters for PRiMeFlow in Jiang24 dataset.
Hyperparameter
Value
Description
architecture
MLP
Architecture choice for the velocity field
hidden_dim
3072
Base feature dim in input MLPs and decoding MLP
dropout
0.0
Dropout rate
n_layers_gene_expression
1
MLP layers for gene expression input
n_layers_conditions
1
MLP layers for condition input
n_layers_time
3
MLP layers for time input
Appendix
Table 15: HPO selected hyperparameters for PRiMeFlow MLP in Srivatsan20 dataset.
Hyperparameter
Value
Description
architecture
MLP
Architecture choice for the velocity field
hidden_dim
2048
Base feature dim in input MLPs and decoding MLP
dropout
0.1
Dropout rate
n_layers_gene_expression
2
MLP layers for gene expression input
n_layers_conditions
3
MLP layers for condition input
n_layers_time
5
MLP layers for time input
Appendix
Table 16: HPO selected hyperparameters for PRiMeFlow MLP in Norman19 dataset.
Hyperparameter
Value
Description
architecture
MLP
Architecture choice for the velocity field
hidden_dim
3840
Base feature dim in input MLPs and decoding MLP
dropout
0.1
Dropout rate
n_layers_gene_expression
2
MLP layers for gene expression input
n_layers_conditions
2
MLP layers for condition input
n_layers_time
2
MLP layers for time input
Appendix
Table 17: HPO selected hyperparameters for PRiMeFlow MLP in Jiang24 dataset.
Hyperparameter
Value
Description
architecture
MLP
Architecture choice for the velocity field
hidden_dim
3072
Base feature dim in input MLPs and decoding MLP
dropout
0.2
Dropout rate
n_layers_gene_expression
2
MLP layers for gene expression input
n_layers_conditions
3
MLP layers for condition input
n_layers_time
4
MLP layers for time input
Appendix
Table 18: HPO selected hyperparameters for FM PCA in Srivatsan20 dataset.
Hyperparameter
Value
Description
architecture
MLP
Architecture choice for the velocity field
hidden_dim
1792
Base feature dim in input MLPs and decoding MLP
dropout
0.2
Dropout rate
n_layers_gene_expression
4
MLP layers for gene expression input
n_layers_conditions
2
MLP layers for condition input
n_layers_time
2
MLP layers for time input
Appendix
Table 19: HPO selected hyperparameters for FM PCA in Norman19 dataset.
Hyperparameter
Value
Description
architecture
MLP
Architecture choice for the velocity field
hidden_dim
3072
Base feature dim in input MLPs and decoding MLP
dropout
0.0
Dropout rate
n_layers_gene_expression
5
MLP layers for gene expression input
n_layers_conditions
3
MLP layers for condition input
n_layers_time
4
MLP layers for time input
Appendix
Table 20: HPO selected hyperparameters for FM PCA in Jiang24 dataset.
num_samples_per_condition
128
Number of cells sampled per condition to compute the mini-batch OT coupling
solver_type
unbalanced
OT solver algorithm
tau_a
1.0
Source marginal constraint weight
tau_b
1.0
Target marginal constraint weight
Appendix
Table 21: Additional hyperparameters for FM PCA OT in Srivatsan20 dataset.
num_samples_per_condition
256
Number of cells sampled per condition to compute the mini-batch OT coupling
solver_type
unbalanced
OT solver algorithm
tau_a
1.0
Source marginal constraint weight
tau_b
1.0
Target marginal constraint weight
Appendix
Table 22: Additional hyperparameters for FM PCA OT in Norman19 dataset.
num_samples_per_condition
256
Number of cells sampled per condition to compute the mini-batch OT coupling
solver_type
unbalanced
OT solver algorithm
tau_a
1.0
Source marginal constraint weight
tau_b
1.0
Target marginal constraint weight
Appendix
Table 23: Additional hyperparameters for FM PCA OT in Jiang24 dataset.
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 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.
Jianan Wei, Jiajun Hong, Guikun Chen +3
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 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