PRiMeFlow: Capturing Complex Expression Heterogeneity in Perturbation Response Modelling

Zichao Yan; Yan Wu; Mica Xu Ji; Chaitra Agrahar; Esther Wershof; Marcel Nassar; Mehrshad Sadria; Ridvan Eksi

cs.LGApr 15, 2026

PRiMeFlow: Capturing Complex Expression Heterogeneity in Perturbation Response Modelling

Authors: Zichao Yan, Yan Wu, Mica Xu Ji, Chaitra Agrahar, Esther Wershof, Marcel Nassar, Mehrshad Sadria, Ridvan Eksi,

Organizations: Altos Labs · Current affiliation: Verge Genomics · Current affiliation: Loka · Current affiliation: Independent Researcher

Abstract

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
Model	GEX	GEX rank	PCA	recall
PRiMeFlow CFG1	$\mathbf{0.13\pm 4\times 10^{-3}}$	$\mathbf{0.13\pm 8\times 10^{-3}}$	$\mathbf{0.14\pm 8\times 10^{-3}}$	$\mathbf{0.26\pm 1\times 10^{-2}}$
PRiMeFlow CFG3	$0.25\pm 2\times 10^{-3}$	$0.22\pm 7\times 10^{-3}$	$0.31\pm 1\times 10^{-2}$	$\mathbf{0.27\pm 6\times 10^{-3}}$
PRiMeFlow CFG5	$0.47\pm 1\times 10^{-2}$	$0.29\pm 2\times 10^{-2}$	$0.60\pm 2\times 10^{-2}$	$0.25\pm 5\times 10^{-3}$
FM PCA CFG1	$1.2\pm 7\times 10^{-3}$	$0.26\pm 6\times 10^{-3}$	$0.50\pm 2\times 10^{-2}$	$0.042\pm 2\times 10^{-3}$
FM PCA CFG3	$1.3\pm 5\times 10^{-3}$	$0.25\pm 7\times 10^{-3}$	$0.44\pm 1\times 10^{-2}$	$0.071\pm 3\times 10^{-3}$

Table 2: Distributional metric performance of PRiMeFlow on the Srivatsan20 dataset.

Model	MMD	MMD	MMD	DEG
Model	GEX	GEX rank	PCA	recall
PRiMeFlow CFG1	$\mathbf{0.26\pm 7\times 10^{-3}}$	$0.019\pm 4\times 10^{-3}$	$\mathbf{0.34\pm 2\times 10^{-2}}$	$0.54\pm 2\times 10^{-2}$
PRiMeFlow CFG3	$0.61\pm 3\times 10^{-2}$	$0.066\pm 9\times 10^{-3}$	$0.86\pm 5\times 10^{-2}$	$\mathbf{0.61\pm 8\times 10^{-3}}$
PRiMeFlow CFG5	$1.3\pm 4\times 10^{-2}$	$0.17\pm 1\times 10^{-2}$	$1.8\pm 5\times 10^{-2}$	$0.56\pm 1\times 10^{-2}$
FM PCA CFG1	$0.77\pm 4\times 10^{-3}$	$0.059\pm 3\times 10^{-3}$	$0.61\pm 5\times 10^{-3}$	$0.37\pm 2\times 10^{-3}$
FM PCA CFG3	$1.1\pm 8\times 10^{-3}$	$0.098\pm 2\times 10^{-3}$	$0.82\pm 9\times 10^{-3}$	$0.45\pm 2\times 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 $\uparrow$	PDS $\uparrow$	MAE $\downarrow$	SPEARMAN $\uparrow$	SPEARMAN_LFC $\uparrow$	AUPRC $\uparrow$	PEARSON $\uparrow$
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 $\uparrow$	PDS $\uparrow$	MAE $\downarrow$	SPEARMAN $\uparrow$	SPEARMAN_LFC $\uparrow$	AUPRC $\uparrow$	PEARSON $\uparrow$
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
Model	logFC	logFC rank	mean	mean rank
PRiMeFlow CFG1	$0.46\pm 1\times 10^{-2}$	$0.12\pm 1\times 10^{-2}$	$0.019\pm 5\times 10^{-4}$	$0.14\pm 6\times 10^{-3}$
PRiMeFlow CFG3	$0.43\pm 9\times 10^{-3}$	$\mathbf{0.10\pm 4\times 10^{-3}}$	$0.029\pm 3\times 10^{-4}$	$0.22\pm 5\times 10^{-3}$
PRiMeFlow CFG5	$0.39\pm 1\times 10^{-2}$	$\mathbf{0.10\pm 2\times 10^{-3}}$	$0.041\pm 4\times 10^{-4}$	$0.29\pm 1\times 10^{-2}$
PRiMeFlow MLP CFG1	$-0.012\pm 4\times 10^{-4}$	$0.46\pm 1\times 10^{-2}$	$0.31\pm 9\times 10^{-4}$	$0.42\pm 1\times 10^{-2}$
PRiMeFlow MLP CFG3	$-0.0071\pm 2\times 10^{-3}$	$0.45\pm 3\times 10^{-2}$	$0.33\pm 5\times 10^{-3}$	$0.44\pm 2\times 10^{-2}$

Appendix

Table 6: Performance of PRiMeFlow on traditional pseudobulk-based metrics in Srivatsan20 dataset.

Model	Cosine	Cosine	RMSE	RMSE
Model	logFC	logFC rank	mean	mean rank
PRiMeFlow CFG1	$0.77\pm 4\times 10^{-3}$	$0.0057\pm 8\times 10^{-4}$	$0.046\pm 1\times 10^{-3}$	$0.022\pm 6\times 10^{-3}$
PRiMeFlow CFG3	$\mathbf{0.79\pm 2\times 10^{-3}}$	$0.0072\pm 3\times 10^{-3}$	$0.074\pm 2\times 10^{-3}$	$0.077\pm 8\times 10^{-3}$
PRiMeFlow CFG5	$0.74\pm 3\times 10^{-3}$	$0.015\pm 3\times 10^{-3}$	$0.11\pm 2\times 10^{-3}$	$0.19\pm 1\times 10^{-2}$
PRiMeFlow MLP CFG1	$0.024\pm 6\times 10^{-3}$	$0.27\pm 3\times 10^{-2}$	$0.29\pm 2\times 10^{-3}$	$0.16\pm 4\times 10^{-2}$
PRiMeFlow MLP CFG3	$0.054\pm 8\times 10^{-3}$	$0.27\pm 2\times 10^{-2}$	$0.31\pm 2\times 10^{-3}$	$0.25\pm 2\times 10^{-2}$

Appendix

Table 7: Performance of PRiMeFlow on traditional pseudobulk-based metrics in Norman19 dataset.

Model	Cosine	Cosine	RMSE	RMSE
Model	logFC	logFC rank	mean	mean rank
PRiMeFlow CFG1	$0.47\pm 5\times 10^{-3}$	$0.39\pm 2\times 10^{-2}$	$0.019\pm 3\times 10^{-4}$	$0.40\pm 2\times 10^{-2}$
PRiMeFlow CFG3	$0.41\pm 4\times 10^{-3}$	$0.38\pm 1\times 10^{-2}$	$0.034\pm 5\times 10^{-4}$	$0.43\pm 1\times 10^{-2}$
PRiMeFlow CFG5	$0.36\pm 4\times 10^{-3}$	$0.37\pm 1\times 10^{-2}$	$0.045\pm 5\times 10^{-4}$	$0.43\pm 7\times 10^{-3}$
PRiMeFlow MLP CFG1	$0.053\pm 1\times 10^{-4}$	$0.45\pm 5\times 10^{-3}$	$0.33\pm 6\times 10^{-4}$	$0.46\pm 1\times 10^{-3}$
PRiMeFlow MLP CFG3	$0.053\pm 7\times 10^{-5}$	$0.43\pm 1\times 10^{-2}$	$0.33\pm 6\times 10^{-4}$	$0.46\pm 3\times 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.

Single-Cell Perturbation Modeling

Single-cell perturbation modeling predicts how genetic or chemical interventions change cellular states across cells and contexts.

See all 58 papers

Synergy-Aware Residual Flow Matching for Single-Cell Perturbation Prediction
Oct 18, 2026Jingrui Wu, Wangyun DanSingle-Cell Perturbation Modeling Flow-Based Generative Modeling
Department of Statistics, Rice University, Houston, TX, USA · Department of Biostatistics and Bioinformatics, Emory University, Atlanta, GA, USA
Generalizable single-cell perturbation response prediction using energy-guided flow matching
Sep 27, 2026Jianan Wei, Jiajun Hong, Guikun Chen +3Single-Cell Perturbation Modeling Single-Cell Transcriptomics
Zhejiang University, Hangzhou, Zhejiang, China · Beijing Institute for General Artificial Intelligence, Beijing, China · Peking University Chengdu Academy for Advanced Interdisciplinary Biotechnologies, Chengdu, Sichuan, China
DRIFT: Disentangled Responsive-Invariant Flow Transport for Single-Cell Perturbation Prediction
Sep 28, 2026Mustapha Bounoua, Giulio Franzese, Pietro MichiardiDisentangled Representation Learning Conditional Flow Matching
Department of Data Science EURECOM, France

Oct 18, 2026J/K move · Enter open · S save

Synergy-Aware Residual Flow Matching for Single-Cell Perturbation Prediction

Jingrui Wu, Wangyun Dan

Department of Statistics, Rice University, Houston, TX, USA · Department of Biostatistics and Bioinformatics, Emory University, Atlanta, GA, USA

Single-Cell Perturbation Modeling Flow-Based Generative Modeling Residual Learning

Open paper AskPDF arXiv

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.

Abstract

Single-Cell Perturbation Modeling

Explore similar work

Synergy-Aware Residual Flow Matching for Single-Cell Perturbation Prediction

Generalizable single-cell perturbation response prediction using energy-guided flow matching

DRIFT: Disentangled Responsive-Invariant Flow Transport for Single-Cell Perturbation Prediction