PerturbDiff: Functional Diffusion for Single-Cell Perturbation Modeling
Authors: Xinyu Yuan, Xixian Liu, Ya Shi Zhang, Zuobai Zhang, Hongyu Guo, Jian Tang
Organizations: Mila - Québec AI Institute · University of Montréal · University of Ottawa · National Research Council of Canada · HEC Montréal · CIFAR AI Chair
Building Virtual Cells that can accurately simulate cellular responses to perturbations is a long-standing goal in systems biology. A fundamental challenge is that high-throughput single-cell sequencing is destructive: the same cell cannot be observed both before and after a perturbation. Thus, perturbation prediction requires mapping unpaired control and perturbed populations. Existing models address this by learning maps between distributions, but typically assume a single fixed response distribution when conditioned on observed cellular context (e.g., cell type) and the perturbation type. In reality, responses vary systematically due to unobservable latent factors such as microenvironmental fluctuations and complex batch effects, forming a manifold of possible distributions for the same observed conditions. To account for this variability, we introduce PerturbDiff, which shifts modeling from individual cells to entire distributions. By embedding distributions as points in a Hilbert space, we define a diffusion-based generative process operating directly over probability distributions. This allows PerturbDiff to capture population-level response shifts across hidden factors. Benchmarks on established datasets show that PerturbDiff achieves state-of-the-art performance in single-cell response prediction and generalizes substantially better to unseen perturbations. See our project page (https://katarinayuan.github.io/PerturbDiff-ProjectPage/), where code and data will be made publicly available (https://github.com/DeepGraphLearning/PerturbDiff).
Figures & tables
Figure 1: Distributional variability in single-cell perturbation data. (a) Traditional methods operate on unpaired control and perturbed cells, learning to map a control cell distribution to a perturbed one. (b) However, variations in cell distributions arise from unobserved latent factors, inducing a family of distinct cell distributions and shifting the objective to learning a distribution over cell distributions.
Figure 2: Overview of the PerturbDiff framework. (a) Distribution-valued random variables Dc,τ and Dc,τ in cell space are mapped to Hilbert-space elements μc,τ and μc∈Hk via kernel mean embedding. (b) Diffusion is defined on perturbed embeddings μ0:=μc,τ , with a denoising network predicting the target μθ . (c) Each MM-DiT block performs joint attention over control and perturbed token streams.
Figure 3: Perturbation prediction results across methods, metrics, and datasets. Each axis represents a performance metric, with higher values indicating better performance.
Figure 4Figure 5
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Table 2: Summary statistics of both perturbation and single-cell RNA-seq datasets used in this study.
Figure 10: The number of non-zero genes per cell across datasets under different gene sets: (a) downstream HVGs (b) pretraining gene set.
Figure 11: Distribution of cell counts per perturbation condition across datasets. Insets report the total number of perturbations and the average number of cells per perturbation. Lower within-perturbation MMD reflects regular cellular responses under the same perturbation, while higher between-perturbation MMD indicates distinct perturbation-specific effects.
Figure 12: Distribution of cell counts per (perturbation, cell type, experimental batch) combination across datasets. Insets report the total number of perturbations and the average number of cells per combination. Lower within-combination MMD reflects intrinsic cellular variability under identical perturbation, cell-type, and batch conditions, while higher between-combination MMD indicates perturbation-induced distributional differences beyond this baseline variability.
Figure 13: Distribution differences within and between perturbations. Violin plots show MMD values computed between cell distributions from the same perturbation (within) and from different perturbations (between) across perturbation datasets.
Figure 14: Cell counts per cell type (left panels (a), (c), and (e)) and between-perturbation MMD across example cell types (right panels (b), (d) and (f)), showing the heterogeneity in cell-type coverage and distributional differences across perturbations for specific cell types.
Figure 15: Per-metric scatter plot comparison between PerturbDiff (From Scratch) and STATE on PBMC dataset. Each panel corresponds to a different evaluation metric, plotting STATE (x-axis) against PerturbDiff (From Scratch) (y-axis) across held-out perturbations. The dashed diagonal indicates equal performance; points above the diagonal indicate higher performance for PerturbDiff (From Scratch) (reversed for MAE and MSE). Reported win rates denote the fraction of perturbations where PerturbDiff (From Scratch) achieves equal or better performance.
Figure 16: Per-metric scatter plot comparison between PerturbDiff (From Scratch) and STATE on Tahoe100M dataset. Each panel corresponds to a different evaluation metric, plotting STATE (x-axis) against PerturbDiff (From Scratch) (y-axis) across held-out perturbations. The dashed diagonal indicates equal performance; points above the diagonal indicate higher performance for PerturbDiff (From Scratch) (reversed for MAE and MSE). Reported win rates denote the fraction of perturbations where PerturbDiff (From Scratch) achieves equal or better performance.
Figure 17: Zero-shot performance on PBMC across pretraining steps, compared to a random baseline.
Figure 18: Zero-shot performance on Replogle across pretraining steps, compared to a random baseline.
Figure 19: Performance on downsampled PBMC (sample ratio 1% and 5%) across training steps for diverse metrics.
Figure 20: Compute and model-size scaling of PerturbDiff on PBMC for diverse metrics.
School of Computer Science and Engineering, South China University of Technology, Guangzhou, China · School of Future Technology, South China University of Technology, Guangzhou, China