PRiMeFlow

PRiMeFlow is a flow matching generative model designed for single-cell perturbation response modeling. Operating end-to-end in the full gene expression space and equipped with specialized architectures (e.g., U-Net), PRiMeFlow achieves state-of-the-art performance in inferring cellular response to small-molecule and (single/combinatorial) gene KD perturbations, demonstrated through comprehensive benchmarking in PerturBench. PRiMeFlow is also the best performing model on the Generalist Rank leaderboard in the ARC Virtual Cell Challenge.

PRiMeFlow leverages the PerturBench framework as dependency. While the original PerturBench provides standard baseline models (e.g., CPA, Biolord, SAMS-VAE), this repository extends the framework by introducing our novel flow matching based approach, along with necessary infrastructural updates to support PRiMeFlow training and evaluation.

Table of Contents

• Installation
• Reproduce PRiMeflow results in PerturBench
  • Dataset Access
  • Training and Evaluation
  • Prediction
• Reproduce PRiMeFlow Results in the Arc Virtual Cell Challenge
  • Pretraining and Finetuning Data Access
  • Pretraining on Perturbation Data Altas
  • Finetuning for VCC private test
  • Generating Predictions and Computing Arc VCC Metrics
• Tips
• Citation

Installation

conda create -n primeflow python=3.12
conda activate primeflow
git clone https://github.com/altoslabs/primeflow.git
cd primeflow
pip install -e .

Reproduce PRiMeflow results in PerturBench

Dataset Access

To download preprocessed benchmarking datasets (e.g., Norman19, Srivatsan20, Jiang24 and etc.), use the following command:

download --data-cache-dir=./datasets/

Manual split files required by certain benchmarking datasets (e.g., Jiang24) will also be downloaded.

Training and evaluation

To train the PRiMeFlow model with U-Net parameterization of the velocity field for a specific dataset (e.g., Jiang24), use one of our predefined configurations:

primeflow.train experiment=primeflow/jiang24/flow_matching_gaussian_source_ddp_unet.yaml

This will train a PRiMeFlow model from scratch, and automatically evaluate the best checkpoint in the validation split at the end of training. To reproduce our results in the PRiMeFlow manuscript, re-evaluate the best checkpoint on the test split:

primeflow.train experiment=primeflow/jiang24/flow_matching_gaussian_source_ddp_unet.yaml train=False ckpt_path=[best_ckpt_path] data/evaluation=mmd data.evaluation.split_value_to_evaluate=test

data/evaluation=mmd overrides the evaluation metrics to the ones reported in our manuscript.

Key inference-time parameters for PRiMeFlow include classifier_guidance_weight, solver and euler_steps.

• classifier_guidance_weight defaults to the value in the training configuration (e.g., 3.0), and can be modified by appending model.classifier_guidance_weight=[new_cfg_value] to the evaluation command.
• solver and euler_steps default to euler and 100. Enable Dormand Prince 45 ODE integration method by specifying +model.solver=dopri5 +model.euler=2.

Other notable inference-time parameters:

• model.prediction_batch_size: the number of trajectories being sampled at the same time. Reduce this number if you run into OOM during evaluation/prediction. Default: 128
• data.evaluation.use_synthetic_controls: predict control cells from the model and use them to replace the ground truth control cells when computing the evaluation metrics. Default: false
• data.evaluation.test_data_iter_factory.max_control_cells_per_covariate: the number of cells sampled per condition (perturbation+covariate). Default: 1000

The default config flow_matching_gaussian_source_hpo_jiang24_ddp_unet.yaml requires distributed data parallel training on a single node equipped with four GPUs. To re-configure the training run on single-GPU machines, simply override the command-line option:

primeflow.train experiment=primeflow/jiang24/flow_matching_gaussian_source_ddp_unet.yaml trainer.devices=1

To resume from an interrupted (e.g., preempted) run, append the checkpoint path:

primeflow.train experiment=primeflow/jiang24/flow_matching_gaussian_source_ddp_unet.yaml trainer.devices=1 ckpt_path=[resumed_ckpt_path]

Note that in order to reproduce our results in the PRiMeFlow manuscript, we recommend users to use our pre-defined configurations. We encourage users to study our configuration files to identify and modify interesting model hyperparameters and training/evaluation options in their own experiments.

Similar training configurations for Norman19 and Srivatsan20 datasets can be found in their respective folders: primeflow/norman19 and primeflow/sciplex3.

Prediction

An instruction dataframe is required for the prediction functionality. For example, for a model trained on the Jiang24 dataset, save a dataframe formatted like the example below as prediction_dataframe_jiang24_example.csv in ${paths.data_dir} (default: ./datasets/)

,condition,cell_type,treatment
0,PTEN,ht29,INS
1,ATF3,bxpc3,IFNG
2,control,ht29,INS
3,control,bxpc3,INS

A number of cells (determined by data.prediction.predict_data_iter_factory.max_control_cells_per_covariate) will be generated under each condition. Note that cell_type and treatment are covariates that the PRiMeFlow model is conditioned on during training in the Jiang24 dataset (refer to jiang24_no_controls.yaml).

Start the prediction by running:

primeflow.predict experiment=primeflow/jiang24/flow_predict_jiang24_example

Modify crucial inference-time parameters such as classifier_guidance_weight, solver, euler_steps and max_control_cells_per_covariate, whenever necessary.

For massive dataframes (containing hundreds or thousands of perturbation+covariate conditions), distribute the workload using predict_chunk.py:

primeflow.predict_chunk experiment=primeflow/jiang24/flow_predict_jiang24_example +chunk_idx=0 +total_chunks=4

This divides the dataframe into 4 chunks, predicting only the subset defined by chunk_idx. This is highly recommended for array jobs in slurm-based HPC clusters.

Reproduce PRiMeFlow results in the Arc Virtual Cell Challenge

Pretraining and finetuning data access

Our pretraining data corpus is curated from a number of public sources (excluding any proprietary data):

• Roohani et al. (2025)
• Replogle et al. (2022)
• Nadig et al. (2025)
• McFaline-Figueroa et al. (2024)
• Replogle et al. (2020)
• Jiang et al. (2024)
• Feng et al. (2026)
• Zhu et al. (2025)

Refer to the PRiMeFlow manuscript for more detailed descriptions. The combined and preprocessed pretraining and finetuning dataset can be accessed at altoslabs/primeflow-vcc-datasets, along with gene feature, split and prediction dataframe csvs. We have provided a download script: download_vcc_datasets.sh.

Note on disk space. The .gz files total ~244 GB compressed, most are from the twelve CD4T datasets (~198 GB). If space is constrained, comment those out to download only the nine remaining files (~47 GB). Leave extra room for ungzipping.

Also consider pulling in additional public Perturb-Seq data, if storage/compute budget permits:

• X Atlas Orion
• Arc Virtual Cell Atlas

The ESM2 embeddings used for gene name representations are derived from the Arc Institute's competition support set and processed as described below. A preprocessed copy is also available in our HuggingFace repo.

1. Download competition_support_set.zip and unzip it to find ESM2_pert_features.pt.
2. Open ESM2_pert_features.pt with torch.load, convert it to parquet format with pandas.DataFrame.to_parquet.
3. Save it to the data folder as ESM2_pert_features.parquet.

Feel free to use other embedding sources (e.g., LLM-based, interactome-extracted):

model.dynamics_model.gene_embedding_parquet_filepath=[path_to_gene_emb_parquet]

Expression genes missing in the parquet file will use a randomly initialized nn.Parameter inside the U-Net velocity field. Set model.dynamics_model.gene_embedding_parquet_filepath=null if you do not wish to use any gene embeddings.

The Virtual Cell Challenge data can be accessed following the instruction here. adata_Training.h5ad, gene_names.csv and pert_counts_Test.csv will be needed.

Pretraining on Perturbation Data Altas

The pretraining experiment configuration is provided in flow_matching_gaussian_source_multifile_h5.yaml. By default, 12 GPUs are required and are distributed across two nodes (e.g., 2x6 setting). An example slurm job script is provided in vcc_pretrain.sh. To pretrain on a subset of the datasets, for example non-CD4T, add to the cmd: data=vcc_multifile_h5_nocd4t_no_controls.

Override the command line options if you wish to point the paths to different dataset, split or expression gene feature sets:

primeflow.train \
  experiment=primeflow/vcc/flow_matching_gaussian_source_annbacked_unet \
  data.data.filename='["path_to_h5ad_1", "path_to_h5ad_2", ...]' \  # the pretraining perturbation data atlas
  data.splitter.split_path=[path_to_split_csv] \  # a csv file assigning per cell in the datasets to either train or val split 
  data.data_iter_factory.feature_filter_path=[path_to_gene_feature_csv] \  # a csv file containing a list of expression genes that will be used for model training
  model.dynamics_model.gene_embedding_parquet_filepath=[path_to_gene_emb_parquet]

Important pretraining trick

After 1 million iterations of training, consider (1) re-initializing optimizer and LR scheduler (but not resetting global steps/epochs), (2) reducing the learning from the default 5e-5 to 1e-5, and (3) increasing the overall batch size by adding more devices. Continue for another 1M training iterations to further improve model performance on the validation set.

Re-initialize optimizer/scheduler, then use the cmd options below to adjust learning rates and world size:

model.lr_scheduler.conf.warmup_start_lr=1e-5 \
model.lr_scheduler.conf.eta_min=1e-5 \
model.lr=1e-5 \
trainer.devices=4 \
trainer.num_nodes=4 \

Finetuning for VCC private test

Slurm job script for finetuning is provided in vcc_finetune.sh. Several key command line options that need to be specified:

• data.data.filename=[path_to_fientune_dataset]: the finetuning dataset is a subset of the pretraining data corpus that contains the (1) 150 VCC training perturbations in H1 hESCs, (2) 100 VCC private test perturbations found in public datasets and (3) control cells in H1 hESCs and all public data sources.
• data.splitter.split_path=[path_to_finetune_split]: the split file for the finetuning dataset. Note that 38 out of 43 perturbations originally held-out in the validation split are reassigned to the training split.
• ckpt_path=[path_to_pretrained_ckpt]: checkpoint of a pretrained model

For more details on the finetuning dataset and split, refer to the PRiMeFlow manuscript, appendix section B.2.

Generating Predictions and Computing Arc VCC Metrics

The config file required for prediction is provided in flow_matching_predict.yaml. Add the following to the config:

• ckpt_path: point to a pretrained or finetuned model checkpoint
• data.prediction.prediction_dataframe_path: included in the download from our HuggingFace repo. A fixed amount of cells (default: 4200) will be generated per condition out of the 100 VCC private test perturbations in total.

An example slurm array job script can be found in vcc_predict.sh. Use the commented block to generate synthetic control cells (only need to run it once) with a lower classifier-free guidance weight (e.g., 5 compared to perturbed cells' cfg weight at 20), in order to maximize the empirical performance of perturbation discrimination and DEGs recall.

Once predictions have been fully generated, use prep_with_synthetic_controls_final_test.ipynb to combine all predicted perturbed cells, as well as synthetic control cells, into a single h5ad object. Finally, use the cell_eval_final_test.ipynb for scoring.

Refer to the PRiMeFlow manuscript appendix section B.3, for a complete description of our inference time strategies.

Tips

• Add time delays between slurm array jobs to prevent multiple jobs from writing to the same output_dir (e.g., use a SLEEP_TIME=$((SLURM_ARRAY_TASK_ID*10))).
• In our pre-defined training configurations, the learning rate is mostly constant with a short warm-up period. Explore alternative lr_scheduler that better suites your task.
• Modify paths.output_dir if you wish to designate a directory name for your experiment (default: datetime)

To train data coupling based PRiMeFlow (e.g., U-Net parameterized velocity to map empirical control cells to perturbed cells using optimal transport), use the following example for Srivatsan20 dataset:

primeflow.train \
experiment=primeflow/sciplex3/flow_matching_gaussian_source_ddp_unet.yaml \
data=sciplex3 \
data/transform=prime_fm_pipeline_controls \
data/transform/flow_matching_util@data.transform.conf.flow_matching_util=ot_conditional_flow_matching \

For mini-batch optimal transport, set data.loader.use_m_per_class_sampler=True, and adjust the per-condition number of samples: data.loader.num_samples_per_condition as needed.

To add a new dataset, follow the instruction in PerturBench for data preprocessing. A couple of useful tips:

• we recommend using the normalized and log1p'ed gene expression values for model training
• manually create a csv file that maps each cell id to either train, val or test split, like in jiang24_split.csv (refer to Dataset Access about downloading this file)
• create a new data config file following the example of jiang24_no_controls.yaml.
  • Replace data.filename and splitter.split_path accordingly, and pass a list of gene names in csv format to data_iter_factory.feature_filter_path, if you want to filter down to a subset of expression genes (otherwise, all expression genes in the dataset will be used for model training/evaluation). Modify key fields: data_iter_factory.perturbation_key, data_iter_factory.perturbation_control_value and data_iter_factory.covariate_keys, based on the information in your dataset's metadata.
  • Don't worry about data_iter_factory.perturbation_combination_delimiter if your data do not contain combinatorial perturbations. Otherwise, make sure the metadata column data_iter_factory.perturbation_key actually uses data_iter_factory.perturbation_combination_delimiter to separate different perturbations (e.g., KLF1+MAP2K6 in norman19_processed.h5ad.

Citation

If you have used PRiMeFlow or PerturBench in your work, please consider citing Yan, Wu, Ji, Agrahar, Wershof, Nassar, Sadria, and Eksi et al, 2026:

@article{yan2026primeflow,
  title={PRiMeFlow: Capturing Complex Expression Heterogeneity in Perturbation Response Modelling},
  author={Yan, Zichao and Wu, Yan and Ji, Mica Xu and Agrahar, Chaitra and Wershof, Esther and Nassar, Marcel and Sadria, Mehrshad and Eksi, Ridvan and Trifonov, Vladimir and Ibarra, Ignacio and others},
  journal={arXiv preprint arXiv:2604.13986},
  year={2026}
}

and Wu, Wershof, Shmon, Nassar, Osinski, Eksi, and Yan et al, 2025:

@inproceedings{wu2025perturbench,
      title={PerturBench: Benchmarking Machine Learning Models for Cellular Perturbation Analysis}, 
      author={Yan Wu and Esther Wershof and Sebastian M Schmon and Marcel Nassar and Błażej Osiński and Ridvan Eksi and Zichao Yan and Rory Stark and Kun Zhang and Thore Graepel},
      booktitle={The Thirty-ninth Conference on Neural Information Processing Systems (NeurIPS) Datasets and Benchmarks Track},
      year={2025},
      url={https://openreview.net/forum?id=PPPDuyiZaG},
}