🤖 AI

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

The paper introduces $D^{2}R^{2}$ , a discrete diffusion model enhanced with regulation reinforcement that improves single-cell perturbation prediction by progressively generating gene expression responses in a biologically informed order rather than predicting the entire profile simultaneously.

Original authors: Ninghan Fan, Qi Liu, Xunuo Zhu, Yukai Sun, Luyuan Chen, Xuheng Zhou, Yuetian Du, Ming Kong, Xiaojun Zhu, Jie Liu, Zhan Zhou, Qiang Zhu

Published 2026-08-18

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ninghan Fan, Qi Liu, Xunuo Zhu, Yukai Sun, Luyuan Chen, Xuheng Zhou, Yuetian Du, Ming Kong, Xiaojun Zhu, Jie Liu, Zhan Zhou, Qiang Zhu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Cells are the fundamental units of life, each acting as a complex factory where thousands of genes work together to maintain function and respond to changes. When scientists want to understand how a cell reacts to a specific stressor, such as a drug or a genetic change, they often need to see how the activity of every single gene shifts. This field, known as functional genomics, aims to predict these shifts without having to run expensive and time-consuming experiments for every possible scenario. For years, researchers have built computer models to simulate these cellular responses, treating the cell's genetic activity as a single, unified picture that is generated all at once. However, this approach overlooks a crucial biological reality: genes do not react in isolation or simultaneously. Instead, they operate in a chain of command, where the response of one gene often triggers or influences the response of another, creating a specific sequence of events that unfolds over time.

A new study introduces a method called D2R2 that changes how these predictions are made by respecting this natural order. Rather than guessing the entire genetic profile in one go, the researchers designed a system that builds the prediction gene by gene, step by step. The core idea is that the order in which genes are predicted matters just as much as the prediction itself. To achieve this, the team combined two distinct tools. First, they used a model that generates the actual genetic values, treating them like a sequence of discrete steps rather than a continuous blur. Second, and more importantly, they added a "regulatory policy" module that acts as a guide, deciding which gene to predict next based on biological rules. This guide starts with a map of how genes normally control one another in an unperturbed cell, but it is smart enough to adapt that map when a specific change occurs, learning which genes become most important to predict first under new conditions.

The researchers tested this approach on two large datasets containing real-world data from human cells. One dataset involved over a thousand different genetic changes in a specific type of blood cell, while the other focused on hundreds of genetic interventions in embryonic stem cells. In these tests, the new method outperformed all existing models across every metric they measured. It was particularly successful at capturing the subtle differences between how cells respond to different types of changes, a task where previous models often struggled. The study showed that the method's success relied heavily on the order in which it generated the data. When the researchers forced the system to pick genes at random, the results were poor. When they tried to use the model's own uncertainty to decide the order, the results were inconsistent and often worse than random. However, when the system followed the biological map of gene regulation, the predictions improved significantly. Even more telling, when the researchers deliberately reversed the biological order—predicting the genes that should come last first—the performance dropped below even the random attempts, proving that following the natural direction of biological influence is essential.

Further analysis revealed that the system learned to prioritize the right genes at the right times. It consistently chose to predict the activity of regulatory genes, which act as the master switches for the cell, before predicting the genes they control. This mimics the way a real cell processes information, establishing the upstream context before filling in the downstream details. The system also learned to identify specific genes that become active only under certain conditions, bringing them forward in the prediction sequence when they were relevant. By treating the generation of a cell's genetic response as a guided, sequential process rather than a single snapshot, this work establishes a new way to model cellular behavior. It demonstrates that the path a prediction takes is just as important as the destination, offering a more accurate and biologically grounded way to understand how cells react to the world around them.

Technical Summary: D2R2 – Discrete Diffusion with Regulation Reinforcement for Single-Cell Perturbation Prediction

Problem Statement

Accurate prediction of single-cell transcriptomic responses to genetic perturbations is a central challenge in functional genomics and drug discovery. While high-throughput experimental data exists, generating it remains costly, motivating the development of in silico virtual cell models. Existing approaches, including transformer-based single-pass predictors, variational autoencoders, and diffusion models, predominantly follow a profile-level generation paradigm. Whether generating the entire expression profile at once or iteratively denoising the whole profile, these methods leave the order in which individual gene responses are generated unmodeled. This omission is critical because gene responses are interdependent; predicting them in a structured order could allow earlier predictions to provide informative context for subsequent ones. Furthermore, simple ordering strategies (e.g., random selection or model uncertainty) fail to exploit the underlying gene regulatory structure, while fixed regulatory priors derived from control cells cannot adapt to specific perturbations or the evolving generation state.

Methodology: D2R2 Framework

The authors propose D2R2 (Discrete Diffusion with Regulation Reinforcement), which reformulates perturbation prediction as a regulation-guided, gene-wise progressive generation task. The framework consists of two primary components:

1. Masked Discrete Diffusion Model (MDDM)

D2R2 represents continuous expression values as ordinal discrete tokens via equal-frequency binning.

Generation Process: The model starts from a fully masked profile and reconstructs it step-by-step. At each step, a subset of unresolved (masked) genes is selected, and the MDDM predicts their expression tokens.
Contextual Conditioning: Once a token is predicted, it is fixed and serves as explicit context for subsequent steps. This allows generated gene responses to condition those that remain masked.
Training: The MDDM is trained with an $x_0$ -prediction objective on randomly masked positions, learning to reconstruct the original discrete tokens from partially masked profiles conditioned on the control profile, perturbation condition, and cellular covariates.

2. Regulatory Policy Module (RPM)

To determine which genes to generate next, D2R2 introduces the RPM, which learns an adaptive generation policy.

Biological Initialization: The policy is initialized using a biological prior derived from a Gene Regulatory Network (GRN) inferred from unperturbed control cells (using DeepSEM). Genes are ranked by PageRank scores on the GRN, prioritizing influential upstream regulators.
Adaptive Refinement via GRPO: The RPM transforms this static prior into a state- and perturbation-specific policy. Using Group-Relative Policy Optimization (GRPO), the RPM is refined while the MDDM generator is frozen.
- Reward Signal: The reward is based on the final agreement between the predicted perturbation effect and the ground truth (measured by correlation of differential expression).
- Mechanism: The policy samples subsets of genes to unmask. The optimization encourages the policy to select genes that, when resolved, provide the most informative context for the remaining genes, thereby improving the final prediction quality.

Training and Inference

The training proceeds in two stages:

MDDM Training: Learn to reconstruct expression tokens from masked inputs.
RPM Refinement: Initialize RPM to reproduce the biological prior, then freeze the MDDM and refine the ordering policy using GRPO with final perturbation-effect agreement as the reward.
During inference, the MDDM predicts tokens for the subset of genes selected by the RPM at each step, iterating until the full profile is generated.

Key Contributions

Gene-Wise Progressive Formulation: The paper formulates single-cell perturbation prediction as a gene-wise progressive generation task using a Masked Discrete Diffusion Model (MDDM), making the generation order explicit and allowing resolved genes to condition masked ones.
Regulatory Policy Module (RPM): The authors introduce a module that initializes an ordering policy from a control-derived biological prior and refines it using reinforcement learning (GRPO) while keeping the expression generator fixed. This produces generation orders specific to both the perturbation and the current generation state.
Controlled Analysis of Generation Order: Through controlled ablations, the study establishes generation order and its directionality as substantive design choices, demonstrating that biologically grounded ordering is superior to random or uncertainty-based strategies.

Experimental Results

The method was evaluated on two benchmarks: Norman19 (K562 cells, Perturb-seq) and VCC-H1 (H1 embryonic stem cells, CRISPRi).

Performance: D2R2 achieved the best performance across all five metrics on the Norman19 dataset and remained competitive (leading in most metrics) on the H1 dataset. It notably outperformed state-of-the-art baselines (including SAMS-VAE, Cell Flow, Squidiff, and scDFM) in differential expression metrics (Pearson $\Delta$ , Cos. LogFC, Cos. LogFC Rank).
Ablation Studies:
- Ordering Strategy: When the MDDM was fixed, the biological prior ordering significantly outperformed random ordering and uncertainty-based strategies (confidence/entropy).
- Directionality: Reversing the biological prior order caused performance to degrade below random ordering, indicating that following regulatory dependencies in the correct direction is essential.
- RPM Efficacy: Even with random initialization, the RPM improved metrics over random ordering. With biological initialization, the refined RPM improved all perturbation-effect metrics, demonstrating that reinforcement learning successfully adapts the static prior to specific perturbations.
Biological Analysis:
- The refined policy prioritized regulatory genes in early generation steps, establishing upstream context before generating downstream responses.
- RPM preferentially advanced perturbation-specific transcription factors (TFs) and responsive genes, capturing the specific regulatory rewiring induced by perturbations.
- Genes promoted by RPM were enriched for perturbation-responsive differentially expressed (DE) genes and cell-cycle-related pathways.

Significance and Claims

The paper claims that D2R2 establishes gene generation order as an effective, controllable, and biologically interpretable dimension of single-cell perturbation prediction. By decoupling expression modeling from generation-order learning, the framework demonstrates that:

Structure Matters: Explicitly modeling the order of gene resolution, guided by regulatory biology, yields superior predictive performance compared to profile-level generation or unstructured ordering.
Adaptability is Key: A static biological prior is insufficient; the ability to adapt the ordering policy to specific perturbations and the evolving generation state via reinforcement learning is crucial for capturing complex, perturbation-specific dependencies.
Interpretability: The generation order itself provides biological insight, revealing that the model naturally prioritizes upstream regulators and perturbation-specific drivers, aligning with known biological mechanisms of cellular response.

The authors conclude that their approach offers a new paradigm for virtual cell modeling that leverages regulatory structure not just for representation, but as a dynamic guide for the generative process.