PRiMeFlow: Capturing Complex Expression Heterogeneity in Perturbation Response Modelling
For computational biologists and drug discovery researchers, PRiMeFlow improves perturbation response prediction by capturing expression heterogeneity, though it is an incremental advance over existing flow-based methods.
PRiMeFlow is a flow matching model that predicts perturbation effects on single-cell gene expression, achieving state-of-the-art performance by accurately approximating empirical distributions, as demonstrated in PerturBench and winning the Generalist Prize in the ARC Virtual Cell Challenge.
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. The PRiMeFlow architecture was used as the basis for the model that won the Generalist Prize in the first ARC Virtual Cell Challenge.