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

TL;DR

scDFM introduces distribution-aware perturbation modeling for single-cell transcriptomes by combining conditional flow matching with an MMD objective and a perturbation-aware backbone (PAD-Transformer). It explicitly models population-level distribution shifts and gene regulatory structure to predict post-perturbation states under both seen and combinatorial perturbations. Across genetic and drug perturbation benchmarks, scDFM outperforms baselines in pointwise accuracy and distributional fidelity, including unseen perturbations, with notable improvements in combinatorial settings. This approach enables robust in silico perturbation screening and provides a foundation for scalable, distribution-aware cellular response modeling.

Abstract

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

Figures (8)

• Figure 1: Overview of scDFM, which models perturbation-specific cell state transitions as a flow matching process from noise to perturbed expression. The PAD-Transformer predicts time-dependent velocities conditioned on control cell context and perturbation embedding, while gene–gene masked attention and differential Transformer layers capture biological dependencies. Final distributional alignment is enforced via MMD regularization.
• Figure 2: The additive setting (Section \ref{['section:experiment-norman-additive']}) tests generalization to unseen doubles when all singles are observed, while the holdout setting (Section \ref{['section:experiment-norman-holdout']}) evaluates prediction of entirely unobserved singles and their combinations
• Figure 3: Double perturbation prediction error ($L_2$). Our method achieves the lowest error distribution, outperforming both additive and baseline models.
• Figure 4: UMAP visualizations of perturbed cell states. Removing MMD leads to clear distributional mismatches, where generated cells deviate from the ground truth manifold.
• Figure 5: Ablation study on the Norman holdout setting.