PerturbDiff: Functional Diffusion for Single-Cell Perturbation Modeling
Abstract
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 capture 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, improving generalization. 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 (https://github.com/DeepGraphLearning/PerturbDiff) are publicly available.
Lay Summary
Scientists want to build “virtual cells”: computer models that predict how living cells will react to a drug, a gene edit, or another intervention. This is difficult because many single-cell experiments destroy the cell they measure, so researchers cannot observe the exact same cell before and after an intervention. Existing models usually compare groups of untreated and treated cells, but they often predict one average outcome for a given cell type and intervention. In reality, hidden factors, such as the local cell environment or experimental batch effects, can make the whole cell population respond in several plausible ways. We introduce PerturbDiff, a model that predicts changes in whole populations of cells rather than treating each cell separately. It learns how untreated cell populations can transform into realistic treated populations, while allowing the overall response to vary. Across benchmarks involving immune signals, drugs, and gene perturbations, PerturbDiff improves prediction accuracy, especially for identifying which genes are truly changed by an intervention. Better virtual-cell models could help researchers choose more promising experiments and study diseases and treatments more efficiently, while still using laboratory tests to validate the predictions.