scDFM: Distributional Flow Matching Model for Robust Single-Cell Perturbation Prediction

Chenglei Yu; Chuanrui Wang; Bangyan Liao; Tailin Wu

Source-linked AI summary

scDFM: Distributional Flow Matching Model for Robust Single-Cell Perturbation Prediction

Chenglei Yu, Chuanrui Wang, Bangyan Liao, Tailin Wu

arXiv:2602.07103v1Feb 6, 2026q-bio.QMcs.AI

TL;DR

Single-cell perturbation prediction must handle unpaired, sparse, noisy data while capturing population-level shifts and generalizing to unseen combinations. scDFM combines conditional flow matching, MMD distributional alignment, and the PAD-Transformer to model perturbed cell distributions. Across genetic and drug benchmarks, it consistently outperforms prior methods, with particularly strong combinatorial generalization.

Problem
Unpaired single-cell measurements and sparse, noisy observations limit cell-level supervision and make population-level perturbation responses difficult to model.
Method
scDFM uses conditional flow matching to generate perturbed cell distributions from control states, with MMD regularization and a PAD-Transformer for structured, noise-robust modeling.
Results
scDFM consistently outperforms prior methods across genetic and drug perturbation benchmarks, demonstrating strong generalization in unseen and combinatorial settings.
Takeaways & Limitations
Distribution-level generative modeling, combined with biologically structured attention, supports robust in silico prediction of single-cell perturbation responses.
Takeaways & Limitations
Conditional flow matching primarily enforces local trajectory consistency, so generated terminal distributions may still deviate from ground-truth perturbed populations without distributional regularization.

Abstract

from arXiv · show

A central goal in systems biology and drug discovery is to predict the transcriptional response of cells to perturbations. This task is challenging due to the noisy and sparse nature of single-cell measurements, as well as the fact that perturbations often induce population-level shifts rather than changes in individual cells. Existing deep learning methods typically assume cell-level correspondences, limiting their ability to capture such global effects. We present scDFM, a generative framework based on conditional flow matching that models the full distribution of perturbed cells conditioned on control states. By incorporating a maximum mean discrepancy (MMD) objective, our method aligns perturbed and control populations beyond cell-level correspondences. To further improve robustness to sparsity and noise, we introduce the Perturbation-Aware Differential Transformer (PAD-Transformer), a backbone architecture that leverages gene interaction graphs and differential attention to capture context-specific expression changes. Across multiple genetic and drug perturbation benchmarks, scDFM consistently outperforms prior methods, demonstrating strong generalization in both unseen and combinatorial settings. In the combinatorial setting, it reduces mean squared error by 19.6% relative to the strongest baseline. These results highlight the importance of distribution-level generative modeling for robust in silico perturbation prediction. The code is available at https://github.com/AI4Science-WestlakeU/scDFM

1 INTRODUCTION

Single-cell perturbation prediction is hindered by unpaired, noisy measurements and by models that overlook population-level effects and gene dependencies. scDFM addresses these challenges with distribution-level training and a perturbation-aware backbone, evaluated across genetic and drug perturbation settings.

Research challenge: Unpaired measurements prevent observing the same cell before and after perturbation, making cell-level supervision impossible.Standard pointwise losses therefore inadequately capture population responses.
Research challenge: Sparse, zero-inflated, noisy data and context-dependent nonlinear effects complicate single-cell perturbation modeling.Dropout, batch effects, and uneven sequencing depth can distort biological signals and gene–gene correlations.
Research challenge: Gene-independent or shallow backbones struggle to generalize to new cell types and out-of-distribution perturbations.The passage reports that such models can overfit noise rather than extract meaningful biological signal.
Proposed approach: scDFM uses conditional flow matching to reconstruct the distribution of perturbed single-cell gene expression.The framework targets both population-level distributional fidelity and noisy, interdependent gene regulation.
Proposed approach: MMD directly aligns generated and true post-perturbation cell distributions beyond paired per-cell losses.Combined with flow matching, it guides the model toward reproducing population shifts.
Proposed approach: PAD-Transformer combines gene–gene masked attention with differential attention to model regulatory dependencies and perturbation-specific interactions.The architecture is designed to filter noise and preserve biologically structured signals.
Evaluation: scDFM is evaluated on Norman genetic and Combosiplex drug perturbation benchmarks, including additive and holdout generalization settings.The evaluations target accuracy and robustness in combinatorial single-cell perturbation prediction.

2 RELATED WORK

Prior work spans foundation models, perturbation-specific architectures, and generative trajectory models for single-cell biology. scDFM connects these directions through attention-based representations and distributional generative modeling.

Foundation models: Foundation models learn general-purpose cell and gene embeddings from large collections of unperturbed expression profiles.These embeddings can be transferred to downstream tasks with minimal supervision.
Perturbation-specific models: Perturbation-specific models encode perturbation composition and biological priors to improve extrapolation to unseen conditions.Examples incorporate shared latent spaces, gene–gene co-expression, Gene Ontology, or probabilistic formulations.
scDFM positioning: scDFM draws on attention-based encoders while modeling perturbation-specific transitions through flow matching and distributional alignment.Its overview combines PAD-Transformer dynamics, gene–gene attention, and MMD regularization.
Generative approaches: Recent generative approaches use diffusion-based models to learn continuous trajectories between control and perturbed cellular states.Their performance is described as sensitive to the selected embedding space.

3 METHOD

scDFM learns a conditional generative process that transforms noisy expression states into perturbed-cell distributions, conditioned on control states and perturbation signals. Its training combines flow-matching trajectory learning with MMD-based endpoint alignment and a PAD-Transformer backbone that models gene dependencies and perturbation-specific dynamics.

Problem setup: Each training instance contains a control state, a multi-hot perturbation condition, and a post-perturbation state; scDFM learns pθ(x | cx, cp).The model targets the population-level distribution of perturbed states and novel perturbation combinations.
Distribution-level regularization: The model combines conditional flow-matching loss with multi-kernel MMD regularization to align generated terminal samples with the ground-truth perturbed distribution.MMD addresses the limitation that local trajectory consistency alone may not constrain global population-level outcomes.
Conditional flow matching: Flow matching learns a conditional velocity field that transforms noisy gene-expression states into perturbed states over t ∈ [0, 1].The source is a noisy expression distribution, while the target is the conditional perturbed-expression distribution.
Distribution-level regularization: The endpoint estimate is formed from an intermediate state and the learned velocity, producing terminal samples that are compared with empirical post-perturbation cells.The predicted endpoint is x̂1 = xt + (1 − t) · vθ(xt | t, cx, cp).
Backbone design: PAD-Transformer uses gene–gene relationships in attention and jointly encodes control and perturbed states to model perturbation-specific dynamics under noisy, sparse conditions.Its components include graph-informed attention, differential attention, and perturbation-aware latent refinement.

4 EXPERIMENT

scDFM is evaluated against statistical, deep-learning, graph-based, foundation-model, and flow-matching baselines across additive, holdout, and drug-perturbation settings. It achieves strong reconstruction, distributional, and differential-expression performance, with ablations showing that MMD and structured attention contribute to robustness.

Experimental setup: The study benchmarks scDFM against statistical, linear, autoencoder, graph-based, foundation-model, and flow-matching baselines.The evaluation uses complementary whole-transcriptome, distribution-level, and differential-expression metrics.
Additive setting: The additive split tests double perturbations whose constituent singles were observed during training.This setting evaluates generalization to unseen combinations of known single perturbations.
Additive setting: The additive evaluation finds scDFM best or near-best across global reconstruction, distributional discrimination, and DE-Spearman ρ metrics.The method maintains a balanced profile, whereas baselines trade off pointwise accuracy, correlation, and distributional alignment.
Additive setting: 19.6%: scDFM reduces MSE versus CellFlow, from 0.00392 to 0.00315, while achieving discrimination score 0.9737.It also reports lower MAE, 0.02155 versus 0.02207, than CellFlow.
Holdout setting: In the holdout split, scDFM combines low error with strong distributional fidelity when singles and their associated doubles are withheld from training.This setting tests generalization to unseen individual perturbations and their combinations.
Drug perturbations: On ComboSciPlex drug perturbations, scDFM achieves the lowest L2, MSE, and MAE and the highest Pearson ∆ and DE-Spearman ρ.Its discrimination scores remain competitive while perturbation-specific differential-expression signals are recovered consistently.
Ablation study: Removing MMD causes the sharpest ablation decline and produces generated cells that deviate from the ground-truth distribution.Removing the gene-gene mask or Differential Transformer also reduces correlation with ground truth.

A.1 THE USE OF LARGE LANGUAGE MODELS (LLMS)

The paper uses large language models only as general-purpose tools for text refinement and language polishing; its core research ideas, datasets, and scientific conclusions are claimed as original contributions.

LLMs are used for text refinement and language polishing, not as sources of the paper’s core scientific contributions.The authors state that the research ideas, datasets, and conclusions are their own original contributions.

A.2 MMD DETAILS

scDFM implements MMD with a dynamically scaled mixture of Gaussian RBF kernels. Bandwidths are adapted to mini-batch distances and combined across multiple scales for robustness to heterogeneous expression ranges.

MMD is implemented with a multi-kernel Gaussian RBF mixture rather than a single fixed kernel.
Bandwidths are recomputed each training step from pairwise squared distances within the current mini-batch.The median off-diagonal distance provides the reference scale for the bandwidths.
σℓ= √sℓ· m uses scales sℓ∈{0.5, 1.0, 2.0, 4.0} to adapt kernel bandwidths and improve robustness across heterogeneous gene-expression ranges.

A.3 TRAINING AND INFERENCE DETAILS

Training samples gene subsets, encodes gene identity with graph masks, and processes expression through PAD-Transformer blocks conditioned on perturbation and time. Inference evolves selected genes with an ODE solver, while the objective combines flow matching and dynamically scaled MMD.

Training: Training samples a subset I of s genes and restricts expression vectors and the gene graph to that subset.The sampled subset is redrawn during training unless otherwise specified.
Gene/context encoding: The gene encoder combines ordered gene identities with a masked subgraph, while the value embedder maps expression values into aligned gene tokens.
PAD-Transformer: PAD-Transformer injects perturbation and time embeddings, applies self- and cross-differential attention, and decodes the velocity field.
Objective: The objective combines conditional flow matching with an MMD term comparing predicted and ground-truth endpoint distributions on sampled genes.MMD uses a Gaussian RBF multi-kernel mixture with bandwidths dynamically estimated from batch distances.
Inference: Inference evolves the selected genes from t=0 to 1 with PAD-Transformer and an Euler or Heun ODE stepper.Full-vocabulary outputs require an optional post-hoc imputation head, which is not used in the experiments.
Complexity: Masked and differential attention scale as O(|I|2) per layer, while dynamic MMD adds O(B2) pairwise batch evaluations.The subset size balances accuracy and computational cost.

A.4.1 NORMAN DATASET

The Norman benchmark contains combinatorial CRISPRa perturbations in K562 cells and is processed into a gene-expression feature set for model training and evaluation. Experiments use reproducible train, validation, and test splits, including a combination-data split that holds out drug pairs.

Dataset: The Norman dataset profiles approximately 100 single-gene and 124 dual-gene CRISPRa activations in K562 cells.It is used to evaluate responses to combinatorial genetic perturbations.
Dataset: The study uses the publicly available scFoundation-reprocessed Norman dataset.
Preprocessing: Preprocessing normalizes library size to 10,000 counts per cell, applies log1p transformation, and selects the top 5,000 highly variable genes.
Feature selection: Including perturbed target genes yields 5,029 training genes, while evaluation focuses on the top 1,000 most highly expressed genes.
Reproducibility: Results are averaged across four independent random train/validation/test splits, with preprocessing configurations and split indices released in the public repository.
Combination split: Figure 6 marks training drug pairs with filled blue cells and held-out test pairs with red outlines.

A.4.2 COMBOSCIPLEX

The ComboSciPlex evaluation tests scDFM on held-out drug combinations, using distributional and pointwise metrics alongside robustness, case-study, and limitation analyses. Results indicate that scDFM better captures complex combinatorial responses while maintaining competitive stability, though broader contexts and richer biological priors remain future work.

Evaluation setting: ComboSciPlex evaluates held-out drug combinations, with single-agent profiles available for training and precise train/test pairs released with the code.
Metrics: The evaluation combines distributional and pointwise measures, including mean-level L2 distance, cell-level MSE and MAE, perturbation discrimination, and differential-expression Spearman correlation.
Limitations: The study uses a linear interpolant in log-normalized expression space, which may poorly approximate nonlinear, branching, or manifold-constrained biological transitions.
Limitations: The benchmarks exclude newer multi-context resources, and the Pearson correlation graph may miss nonlinear or causal gene-regulatory dependencies.
Case study: scDFM consistently places predictions within observed variance and recovers synergistic upregulation, whereas baselines often miss nonlinear responses or misestimate magnitudes.
Consistency: Across randomized folds, scDFM maintains low and competitive variance, while holdout difficulty increases variance for all methods and CPA is markedly less stable.