PaperlayerView exact revision on arXiv

Source-linked AI summary

PerturbDiff: Functional Diffusion for Single-Cell Perturbation Modeling

Xinyu Yuan, Xixian Liu, Ya Shi Zhang, Zuobai Zhang, Hongyu Guo, Jian Tang

arXiv:2602.19685v1cs.LGcs.AI

TL;DR

Single-cell perturbation prediction must use unpaired populations, while existing approaches often collapse latent response variability into one distribution. PerturbDiff embeds distributions in an RKHS and diffuses over them, achieving state-of-the-art results across perturbation benchmarks and adapting effectively in low-data settings.

  • Problem

    Destructive sequencing prevents observing the same cell before and after perturbation, while existing models typically collapse latent-factor-driven response variability into a single distribution.

  • Method

    PerturbDiff treats cell populations as random variables, embeds their distributions in an RKHS, and defines diffusion with an MMD-based denoising objective.

  • Results

    PerturbDiff achieves state-of-the-art performance across signaling, drug, and genetic perturbation benchmarks, particularly for differential-expression recovery and low-data adaptation.

  • Takeaways & Limitations

    The framework provides a distribution-level paradigm for virtual-cell modeling that captures response variability induced by unobserved biological and technical factors.

  • Takeaways & Limitations

    Downstream use requires care because training-data biases and over-interpretation of predicted effects can produce misleading biological conclusions; experimental validation remains necessary.

Abstract

from arXiv · show

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).

1. Introduction

Single-cell perturbation assays support biological discovery but cannot exhaustively measure the combinatorial intervention space, motivating virtual-cell models for unseen conditions. PerturbDiff addresses response variability by modeling distributions over cell distributions rather than a single fixed response distribution.

  • Single-cell perturbation measurements connect interventions to phenotypic outcomes, including immune signaling, drug responses, and gene-function discovery.
  • The combinatorial space of genes, drugs, dosages, and cellular contexts makes exhaustive experimentation costly and motivates computational prediction for unseen perturbations.
  • Destructive sequencing prevents true cell-to-cell correspondence, so random pairing learns average responses and misses cellular heterogeneity.
  • Unobserved latent factors create families of distinct cell distributions, shifting the objective from mapping one distribution to learning a distribution over distributions.
  • PerturbDiff embeds cell distributions in an RKHS and applies diffusion directly to these representations, producing an MMD-based objective for population alignment.
  • A pretraining strategy combines perturbation data with large-scale unperturbed RNA-seq data to cover more cellular states and improve adaptation when perturbed data are scarce.

2. Preliminary

The preliminary formulation represents cells as gene-expression vectors, populations as finite collections of cells, and distributions over the cell space. Perturbation modeling uses unpaired control and perturbed populations, while diffusion models provide a conditional generative framework with optional classifier-free guidance.

  • A cell is represented by a gene-expression vector in R^|G|, and the cell space contains all possible profiles.
  • A cell population is a finite collection of cells, while P denotes a probability distribution over the cell space.
  • Perturbation datasets contain unpaired control and perturbed populations, and the modeling goal is to predict the response after perturbation from control cells.
  • Context labels describe observable conditions such as cell type, donor, and batch, while perturbation labels identify interventions and define conditional cell distributions.
  • Diffusion models learn data distributions by reversing a forward Markov chain that progressively corrupts samples according to a noise schedule.
  • Classifier-free guidance steers sampling toward conditioning information by extrapolating between conditional and unconditional predictions, with w controlling guidance strength.

3. Related Work

Earlier perturbation models use deterministic or probabilistic cell-level mappings, often relying on random pairing and therefore capturing average effects. PerturbDiff instead models variability across distributions induced by unobserved latent factors.

  • Linear models, GEARS, and scGPT differ in modeling design but use cell-wise perturbation prediction based on randomly paired cells.
  • Random pairing drives existing cell-level methods toward average perturbation effects rather than heterogeneous cellular responses.
  • Most existing methods assume one perturbed distribution for each observed context and perturbation, implicitly marginalizing variability from unobserved factors.
  • PerturbDiff learns a diffusion process over the space of cell distributions to represent distribution-level variability caused by latent factors.

4. Method

PerturbDiff models perturbation responses as a stochastic process over cell distributions rather than individual cells, capturing variability from unobserved factors. It embeds distributions in an RKHS, applies diffusion to their kernel mean embeddings, and trains with population-level matching supplemented by marginal pretraining.

  • Motivation: Existing methods often model one perturbed cell distribution for each observed condition, overlooking variability caused by unobservable latent factors.
  • Distribution-valued modeling: PerturbDiff treats the perturbed population as a distribution-valued random variable, whose realizations represent different plausible distributions under the same condition.The commonly learned target corresponds to the expectation of this random variable.
  • Hilbert-space representation: Cell distributions are mapped to kernel mean embeddings in an RKHS, where they become points suitable for interpolation, discrepancy measurement, and diffusion.Empirical control and perturbed batches provide the embeddings used for training.
  • Diffusion modeling: The forward process progressively noises perturbed embeddings toward a Gaussian-like reference, while the learned reverse process predicts denoised target embeddings.Gaussian random elements and a variational reverse formulation make diffusion in the infinite-dimensional RKHS tractable.
  • Training objective: Training matches predicted and real cell populations with MMD, yielding a distribution-aware objective beyond pointwise cell reconstruction.The objective is implemented using empirical cell batches as finite-sample kernel mean embeddings.
  • Pretraining: Marginal pretraining uses large-scale unperturbed RNA-seq data before perturbation fine-tuning, improving data efficiency and generalization when perturbation data are scarce.The pretraining stage includes 61 million cells spanning hundreds of cell types.
  • Modeling perspective: PerturbDiff formulates population matching as learning a stochastic distribution over cell distributions, with MMD arising from squared RKHS distance between kernel mean embeddings.

5. Experiment

Across signaling, drug, and genetic benchmarks, PerturbDiff generally outperforms strong baselines, especially on differential-expression recovery and cross-perturbation consistency. Marginal pretraining improves zero-shot and low-data adaptation, while scaling and loss ablations reveal important design constraints.

  • Overall performance: Across 12 metrics, PerturbDiff (From Scratch) outperforms nearly all baselines on PBMC and Tahoe100M and ranks second on Replogle.The benchmark covers signaling, drug, and genetic perturbations.
  • Overall performance: PerturbDiff improves differential-expression recovery beyond average-expression matching, with substantial gains on AUPRC and AUROC.It more clearly separates true DE from non-DE genes than STATE, which assigns many non-DE genes large −log10(padj) values.
  • Pretraining and low-data adaptation: PerturbDiff (Finetuned) consistently improves over training from scratch and matches STATE on Replogle, although average-expression accuracy on Tahoe100M is slightly lower.Finetuning further improves DE-related metrics on PBMC and Tahoe100M.
  • Per-perturbation performance: PerturbDiff beats STATE on nearly all PBMC and Tahoe100M perturbation types for DE metrics, with win rates above 96–100%.These metrics include AUROC, AUPRC, and DEPrec.
  • Pretraining and low-data adaptation: Zero-shot pretrained models outperform random initialization on PBMC and Replogle, and finetuning substantially improves performance when perturbed cells are scarce.Downsampled PBMC experiments show finetuning outperforming training from scratch at both 1% and 5% data ratios.
  • Scaling and ablations: Scaling is non-monotonic: the 114M medium model is most stable, while the 239M model peaks at intermediate compute before degrading.The results favor moderate capacity and compute rather than aggressive scaling.
  • Scaling and ablations: Removing MMD consistently degrades performance, especially on DE metrics, while MSE is poorly suited to highly sparse single-cell expression data.Over 95% of PBMC expressions and over 60% of Replogle expressions are zero.

6. Conclusion

PerturbDiff models cell distributions as random variables by diffusing their RKHS embeddings, capturing response variability from unobserved factors. Across signaling, drug, and genetic benchmarks, it achieves state-of-the-art performance, particularly for differential-expression recovery and low-data adaptation.

  • Method: PerturbDiff applies functional diffusion to entire cell distributions, capturing variability from unobserved biological and technical latent factors.Empirical distributions are embedded in an RKHS, where diffusion induces an MMD-based denoising objective.
  • Results: Across signaling, drug, and genetic benchmarks, PerturbDiff achieves state-of-the-art performance, especially for perturbation-driven differential expression and low-data adaptation.These results support its use as a paradigm for virtual cell modeling.

Impact Statement

The paper motivates distribution-level perturbation modeling for costly, heterogeneous single-cell experiments and describes datasets, preprocessing, splits, and pretraining resources used for evaluation.

  • Datasets: The study evaluates perturbation prediction on PBMC, Tahoe100M, and Replogle across different scales and heterogeneous experimental settings.These datasets range from hundreds of thousands to tens of millions of cells and vary in perturbations, cell types, and batches.
  • Preprocessing: PBMC and Tahoe100M undergo library-size normalization, log1p transformation, and rescaling, while Replogle uses the processed STATE protocol.The PBMC and Tahoe100M values are rescaled approximately into [0, 1), and Replogle data are filtered for on-target knockdown efficacy before normalization and log transformation.
  • Evaluation Design: The downstream datasets retain the top 2,000 highly variable genes, and dataset-specific holdouts test generalization under partial perturbation coverage.The split design includes donor, perturbation, and other dataset-specific holdout strategies.
  • Pretraining: Pretraining combines the perturbation datasets with single-cell RNA-seq data aggregated from 1,139 diverse CellxGene studies.The pretraining corpus is designed to broaden coverage of cellular states while avoiding information leakage through consistent splits.

A.4. Data Statistics Analysis for Perturbation Datasets

The datasets exhibit sparsity, imbalance, and structured distributional variation, supporting PerturbDiff’s choice to model perturbation responses at the population-distribution level.

  • Motivation: Perturbation datasets are heterogeneous and imbalanced across cells, perturbations, cell types, and experimental batches, motivating distribution-level modeling.These structures make individual-cell supervision difficult and unstable.
  • Perturbation Heterogeneity: Inter-perturbation MMD consistently exceeds intra-perturbation MMD, indicating structured distributional shifts rather than only stochastic cell-level noise.Figure 13 compares MMD within the same perturbation against MMD between different perturbations.
  • Perturbation Heterogeneity: Replogle shows larger and more variable inter-perturbation MMD than PBMC and Tahoe100M, while sparsity can compress observed distribution divergence.The passage cautions that smaller MMD in PBMC and Tahoe100M need not imply weaker biological effects.
  • Cell-Type Composition: Cell-type abundances vary sharply, yet inter-perturbation distributional shifts within cell types show no correlation with cell abundance.PBMC includes more than one million CD4 Naive cells but fewer than ten thousand Plasmablasts.
  • Distribution Representation: Kernel mean embeddings represent cell distributions as Hilbert-space points, providing a state space for diffusion over populations.The construction uses a positive-definite kernel and supports population-level similarities and discrepancies through RKHS inner products and norms.
  • Distribution Representation: Empirical cell distributions converge exponentially in supremum norm toward the underlying distribution as batch size grows.This result supports using empirical cell batches to represent the latent distributions.

B.1.7. TRACTABLE TRAINING AND SAMPLING.

PerturbDiff makes distribution-level diffusion computationally tractable by operating on empirical cell batches while preserving RKHS-based distributional objectives.

  • Empirical Cell-Space Computation: The method replaces probability distributions with empirical measures induced by batches of cells for practical training.Predicted and perturbed batches define empirical distributions used in the distribution-aware objective.
  • Distributional Objective: The resulting training signal penalizes density shifts, subpopulation reweighting, and other distributional effects beyond cell-wise reconstruction.Noise is injected by adding Gaussian noise to batched cell coordinates before embedding.
  • Noise Sampling: Gaussian noise in cell space approximates Gaussian random-element noise in the RKHS, enabling tractable forward diffusion.This avoids directly sampling function-valued Gaussian measures.
  • Model Implementation: The model uses a variance-preserving diffusion with an x0-predictor that reconstructs clean perturbed batches from noised batches conditioned on controls, context, and perturbation.Predicted expression values are constrained to be non-negative.
  • Training Objective: Training combines an MMD-based loss with cell-wise MSE, using the energy distance kernel to compare true and predicted empirical distributions.The default objective sets λMSE = 1.
  • Objective Interpretation: Locally, minimizing MMD is equivalent to minimizing a kernel-dependent matrix-weighted MSE that emphasizes informative population-statistic directions.The weighting matrix is determined by the kernel and the local geometry of the reference population.
  • Fine-Tuning: Pretrained input/output layers can be transferred, replaced randomly, or replaced using a 2k-gene vocabulary during fine-tuning.These strategies are compared empirically for downstream adaptation.

C.3. More Perturbation Prediction Results

PerturbDiff achieves low reconstruction errors across the three evaluation datasets, while mean baselines provide alternative average-expression references for comparison.

  • Reported Results: Figure 3 provides relative comparisons across methods and metrics, with detailed numerical results reported in Table 4.The radar plots emphasize higher-is-better metrics, while MSE and MAE are summarized separately.
  • MSE and MAE Results: PerturbDiff trained from scratch achieves low MSE and MAE reconstruction errors comparable to or better than strong baselines across all three datasets.STATE and Linear remain competitive on these metrics, particularly in lower-variance settings.
  • Mean Baselines: Mean baselines predict perturbed cells using averages computed per perturbation, cell type, batch, or across the full dataset.These variants establish average-expression references at different conditioning levels.

C.4. More Scatter Plot Comparison Results

PerturbDiff (From Scratch) shows systematic per-perturbation advantages over STATE on PBMC and Tahoe100M, especially for differential-expression and distributional metrics.

  • PerturbDiff (From Scratch) consistently outperforms STATE across PBMC perturbations on PRAUC, DEOver, and DEPrec.
  • Win rates exceed 88% on PBMC for PDSL1, PDSL2, and PDScos.
  • Win rates exceed 87% on Tahoe100M for PRAUC, DEPrec, MAE, MSE, and all PDS variants.
  • PerturbDiff (From Scratch) wins on more than 50% of Tahoe100M comparisons for every reported metric except LFCSpear.
  • The scatter comparisons indicate systematic gains across diverse perturbations, particularly on distribution-level response fidelity and perturbation discrimination.

C.5. More Zero-shot Results

Zero-shot performance varies by metric and dataset as marginal pretraining proceeds. Distribution- and perturbation-aware metrics show non-monotonic behavior, while stringent metrics distinguish pretrained models from random initialization.

  • MAE and MSE remain relatively stable across pretraining steps on PBMC and Replogle.
  • Distribution- and perturbation-aware metrics exhibit a U-shaped zero-shot trend as marginal pretraining increases.
  • On PBMC, random initialization produces seemingly non-trivial scores on several metrics despite no training, unlike Replogle.
  • Random initialization substantially underperforms pretrained counterparts on stringent metrics including DEOver, R2, PDCorr, DirAgr, MAE, and MSE.

C.6. More Limited Data Results

Marginal pretraining improves few-shot adaptation under limited PBMC data, producing faster, more stable optimization and stronger performance on several distribution- and perturbation-aware metrics.

  • Finetuning consistently converges faster and improves training stability across metrics on downsampled PBMC.
  • Under the 1% sample regime, models trained from scratch show pronounced training-step sensitivity, fluctuations, and early or mid-training degradation.
  • Finetuned models consistently outperform training from scratch on DEOver, DEPrec, DirAgr, LFCSpear, and all PDS variants.
  • The performance gaps between finetuning and training from scratch are larger at the 1% sample ratio.
  • Marginal pretraining improves data efficiency, optimization stability, and robustness under low-data regimes.

C.7. More Scaling Results

PerturbDiff’s scaling behavior is non-monotonic across model size and training compute, so stronger perturbation modeling depends on balancing capacity and compute rather than scaling either aggressively.

  • Increasing training compute does not uniformly improve PerturbDiff performance across model-size and compute scaling experiments.
  • Scaling is non-monotonic along both model-size and compute dimensions.
  • Effective perturbation modeling benefits from balanced model capacity and compute rather than aggressive scaling.
  • Zero-shot performance is evaluated across pretraining steps on PBMC and Replogle against a random baseline.
  • Downsampled PBMC experiments vary sample ratios of 1% and 5% across training steps, while scaling experiments assess compute and model size.
Loading 2602.19685v1…