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

2026-02-06

AI Summary

Predicting transcriptional responses of cells to perturbations is challenging due to the high noise and sparsity of single-cell data, as well as the fact that perturbations typically induce shifts in population-level distributions rather than deterministic changes at the single-cell level. This work proposes the first distribution-level generative modeling framework for single-cell perturbation prediction, leveraging conditional flow matching to learn the full post-perturbation cell population distribution. By incorporating Maximum Mean Discrepancy (MMD) to align control and perturbed populations, the method eliminates reliance on one-to-one single-cell correspondences. Central to this approach is the Perturbation-Aware Differential Transformer (PAD-Transformer), which integrates gene interaction graph attention to effectively capture context-specific expression changes. The model significantly outperforms existing methods across diverse genetic and drug perturbation benchmarks, achieving a 19.6% reduction in mean squared error over the strongest baseline in combinatorial perturbation settings, demonstrating exceptional generalization capability.

Technology Category

• Machine Learning: Probabilistic Circuits and Graphical Models
• Natural Language Processing: Generation
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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

Problem

• single-cell perturbation prediction
• distributional shift
• transcriptional response
• noisy and sparse data
• population-level effects

Innovation

• distributional flow matching
• single-cell perturbation prediction
• maximum mean discrepancy
• differential attention
• generative modeling