# PRiMeFlow: Capturing Complex Expression Heterogeneity in Perturbation Response Modelling

arXiv: https://arxiv.org/abs/2604.13986 · HTML: https://mlsanity.com/paper/2604.13986/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 45 references that are in the map.
- First posted on arXiv 15 Apr 2026.

## Builds on (most cited first, 16 of 16)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://mlsanity.com/paper/2107.14795/) (2021)
- [Modelling Cellular Perturbations with the Sparse Additive Mechanism Shift Variational Autoencoder](https://mlsanity.com/paper/2311.02794/) (2023)
- [PerturBench: Benchmarking Machine Learning Models for Cellular Perturbation Analysis](https://mlsanity.com/paper/2408.10609/) (2024)
- [PerturbDiff: Functional Diffusion for Single-Cell Perturbation Modeling](https://mlsanity.com/paper/2602.19685/) (2026)
- [Toward AI-Driven Digital Organism: Multiscale Foundation Models for Predicting, Simulating and Programming Biology at All Levels](https://mlsanity.com/paper/2412.06993/) (2024)
- [Scalable Single-Cell Gene Expression Generation with Latent Diffusion Models](https://mlsanity.com/paper/2511.02986/) (2025)

## Built on by (most cited first, 1 of 1)

- [Scaling an Autoregressive Transformer for Single-Cell Generation](https://mlsanity.com/paper/2608.02961/) (2026)

## Abstract (arXiv)

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.