scDFM: Distributional Flow Matching for Robust Single-Cell Perturbation Prediction (ICLR 2026)

Official repo for the paper scDFM, ICLR 2026.
Chenglei Yu^∗1,2, Chuanrui Wang^∗1, Bangyan Liao^∗1,2 & Tailin Wu^†1.

¹School of Engineering, Westlake University; ²Zhejaing University;

*Equal contribution, †Corresponding authors

Overview

We propose a novel distributional flow matching framework (scDFM) for robust single-cell perturbation prediction, which models the full distribution of perturbed cellular expression profiles conditioned on control states, thereby overcoming limitations of existing methods that rely on cell-level correspondences and fail to capture population-level transcriptional shifts.

Framework of paper:

Install dependencies

conda env create -f environment.yml

⏬ Dataset download

Put dataset into data file:

Alternative Data Access

We also provide the datasets via Google Drive. This folder contains:

The Norman dataset and its corresponding data splits.
The ComboSciPlex dataset.

Example directory layout after download (relative to repo root):

scDFM/
├─ data/
│  ├─ norman.h5ad
│  └─ combosciplex.h5ad
├─ src/
│  └─ ...
└─ run.sh

📥 Training

An example on additive task.

bash run.sh

🫡 Citation

If you find our work and/or our code useful, please cite us via:

@inproceedings{yu2026scdfm,
  title={sc{DFM}: Distributional Flow Matching Model for Robust Single-Cell Perturbation Prediction},
  author={Chenglei Yu and Chuanrui Wang and Bangyan Liao and Tailin Wu},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026},
  url={https://openreview.net/forum?id=QSGanMEcUV}
}

📚 Related Resources

AI for Scientific Simulation and Discovery Lab: https://github.com/AI4Science-WestlakeU