$D^{2}R^{2}$: Discrete Diffusion with Regulation Reinforcement for Single-Cell Perturbation Prediction

2026-08-15

AI Summary

This study addresses the neglect of temporal gene response dynamics in existing single-cell perturbation prediction methods by proposing the D²R² framework. This approach reformulates prediction as a regulation-guided, progressive gene-by-generation process, integrating masked discrete diffusion models with reinforcement learning (GRPO). Notably, it is the first to treat generation order as an interpretable dimension subject to joint optimization. Experiments on the Norman19 dataset demonstrate state-of-the-art performance across five metrics with strong competitiveness. Furthermore, ablation studies confirm that biologically informed ordering significantly outperforms baseline strategies, and the optimized generation paths accurately capture key regulatory factors, thereby achieving a unification of high predictive accuracy and model interpretability.

Abstract

Predicting single-cell transcriptomic responses to genetic perturbations is central to functional genomics and virtual-cell modeling. Existing approaches, however, typically predict an entire expression profile as a whole, leaving the order in which individual gene responses are generated unmodeled. To address this problem, we introduce \textbf{$D^{2}R^{2}$} (\textbf{D}iscrete \textbf{D}iffusion with \textbf{R}egulation \textbf{R}einforcement), which reformulates perturbation prediction as regulation-guided gene-wise progressive generation. A Masked Discrete Diffusion Model represents expression as ordinal tokens and reconstructs a fully masked profile step by step, allowing generated gene responses to condition those that remain masked. A Regulatory Policy Module initializes the generation policy from a gene regulatory network inferred from control cells and adapts it to the perturbation and current partially generated state. Then, group-relative policy optimization refines only the ordering policy using final perturbation-effect agreement as reward. Across Norman19 and VCC-H1, $D^{2}R^{2}$ achieves the best performance on all five metrics on Norman19 and remains competitive on H1. Controlled ablations holding the generator and generation budget fixed show that biological-prior ordering improves over random ordering and is more reliable than uncertainty-based heuristics, whereas reversing the biological-prior ordering degrades every metric. Biological analyses further show that the refined policy prioritizes regulatory genes early while promoting perturbation-specific transcription factors and responsive genes. These results establish gene generation order as an effective, controllable, and biologically interpretable dimension of single-cell perturbation prediction.

Problem

• Single-Cell Perturbation Prediction
• Gene Expression Profile
• Generation Order
• Functional Genomics

Innovation

• Discrete Diffusion
• Progressive Generation
• Regulation Reinforcement
• Gene Regulatory Network
• Policy Optimization