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

Ninghan Fan; Qi Liu; Xunuo Zhu; Yukai Sun; Luyuan Chen; Xuheng Zhou; Yuetian Du; Ming Kong

cs.AIAug 15, 2026

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

Authors: Ninghan Fan, Qi Liu, Xunuo Zhu, Yukai Sun, Luyuan Chen, Xuheng Zhou, Yuetian Du, Ming Kong,

Organizations: Zhejiang University · City University of Hong Kong

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}$ } (Discrete Diffusion with Regulation Reinforcement), 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.

Figures & tables

Figure 1: Paradigm comparison. Each row summarizes paradigm properties and traces gene-expression generation, with web nodes and aligned vector entries denoting genes and their expression values. We compare traditional single-pass prediction, DDPM-based global denoising, and two forms of $D^{2}R^{2}$ : random unmasking (MDDM) and regulation-guided unmasking.

Figure 2: Overview of the $D^{2}R^{2}$ framework. The left panel shows the iterative generation loop: after the MDDM predicts the selected expression tokens, RPM observes the updated partial state and recomputes which genes should be resolved next. The upper-right panel shows the biological prior used to initialize the policy, while the lower-right panel illustrates GRPO that optimizes the order to each perturbation.

		Differential Expression			Absolute Expression
Dataset	Method	Pearson $\Delta$ $\uparrow$	Cos. LogFC $\uparrow$	Cos. LogFC Rank $\downarrow$	Cos. PCA $\uparrow$	Sym. KL $\downarrow$
Norman19	SAMS-VAE (NeurIPS 2023)	0.553 $\pm$ 0.079	0.560 $\pm$ 0.043	0.028 $\pm$ 0.010	0.251 $\pm$ 0.042	50.888 $\pm$ 4.215
	Biolord (Nat. Biotechnol. 2024)	0.404 $\pm$ 0.069	0.388 $\pm$ 0.051	0.046 $\pm$ 0.010	0.035 $\pm$ 0.016	0.526 $\pm$ 0.040
	STATE (NeurIPS 2025)	0.467 $\pm$ 0.020	0.379 $\pm$ 0.020	0.525 $\pm$ 0.017	-0.009 $\pm$ 0.039	0.273 $\pm$ 0.014
	Cell Flow (bioRxiv 2025)	0.489 $\pm$ 0.077	0.373 $\pm$ 0.066	0.329 $\pm$ 0.013	0.084 $\pm$ 0.066	0.207 $\pm$ 0.024
	Squidiff (Nat. Methods 2025)	0.291 $\pm$ 0.001	0.226 $\pm$ 0.002	0.427 $\pm$ 0.007	0.110 $\pm$ 0.000	2.861 $\pm$ 0.019

Table 1: Main results on Norman19 and H1. All methods are evaluated over five independent training runs using seeds

\{0,1,2,3,42\}

, and results are reported as mean

\pm

standard deviation. For

D^{2}R^{2}

, MDDM checkpoint is fixed across runs, while RPM is independently trained for each seed. Metrics are grouped by differential and absolute expression. Best and second-best results are bolded and underlined, respectively.

Figure 3: Violin plots showing the distribution of evaluation metrics across perturbations in Norman19. Mann-Whitney

U

test, * P < 0.05, *** P < 0.001.

Model-Uncertainty Orders
Order	Pearson $\Delta$ $\uparrow$	Cos. LogFC $\uparrow$	Cos. LogFC Rank $\downarrow$	Cos. PCA $\uparrow$	Sym. KL $\downarrow$
Random	0.623 (ref.)	0.720 (ref.)	0.032 (ref.)	0.597 (ref.)	0.183 (ref.)
Confidence high	0.590 ( $-5.3$ %)	0.537 ( $-25.4$ %)	0.030 ( $+5.2$ %)	0.353 ( $-40.9$ %)	0.410 ( $-124.4$ %)
Confidence low	0.672 ( $+7.9$ %)	0.522 ( $-27.5$ %)	0.068 ( $-114.1$ %)	0.640 ( $+7.2$ %)	0.203 ( $-11.0$ %)
Entropy high	0.624 ( $+0.2$ %)	0.452 ( $-37.2$ %)	0.073 ( $-130.9$ %)	0.564 ( $-5.6$ %)	0.206 ( $-12.5$ %)
Entropy low	0.591 ( $-5.1$ %)	0.538 ( $-25.3$ %)	0.030 ( $+5.2$ %)	0.355 ( $-40.5$ %)	0.399 ( $-118.5$ %)

Table 2: Generation-order ablation on Norman19. All methods use the same frozen MDDM. Random averages multiple random orders. Cells show score and relative performance gain over Random, with signs aligned so that positive always indicates improvement; for Cos. LogFC Rank and Sym. KL, this corresponds to relative error reduction. Orange intensity indicates a positive gain. Percentages are computed before rounding.

Figure 4: Biological analysis. (a) Proportions of regulatory and regulated genes across generation steps under the RPM-refined generation order. (b) Boxplot comparing

\Delta\mathrm{Order}

between perturbation-specific TFs and non-specific TFs. Mann-Whitney

U

test,

P=1.6\times 10^{-15}

. (c) Boxplot showing the fraction of DE genes whose generation order was promoted by RPM or remained unpromoted. Each dot represents one perturbation, and paired fractions from the same perturbation are connected by lines. Paired Wilcoxon test,

P=1.7\times 10^{-9}

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Configuration
SAMS-VAE	Latent size 128; batch 256; dropout 0.4; learning rate $3.03\times 10^{-5}$
Biolord	Latent size 512; batch 1000; dropout 0.4; learning rate $1.67\times 10^{-4}$
STATE	Hidden size 768; eight Transformer layers; 12 attention heads; learning rate $10^{-4}$
Cell Flow	Batch 2000; hidden size 2048 with three layers; learning rate $10^{-5}$
Squidiff	Three layers; batch 64; learning rate $10^{-4}$ ; EMA 0.9999
scDFM	Hidden size 512; batch 32; learning rate $10^{-5}$

Appendix

Table 3: Model-specific configurations used for the baseline methods.

Parameter	Value
Batch size	128
Optimizer	AdamW
Weight decay	$10^{-2}$
Learning rate (Norman19 / H1)	$10^{-4}$ / $5\times 10^{-4}$
Warmup fraction	0.01
Gradient clipping	1.0

Appendix

Table 4: MDDM training configuration. Dataset-specific values are reported as Norman19 / H1.

Parameter	Value
Batch size	96
Group size	4
Optimizer	AdamW
Learning rate	$3\times 10^{-6}$
Weight decay	$10^{-4}$
Warmup fraction	0.01

Appendix

Table 5: GRPO training configuration for RPM optimization.

Figure 5: GSEA analysis using the $\Delta\mathrm{Order}$ values of all genes as the ranking metric. The GSEA analysis shows that RPM-promoted genes are enriched in cell cycle pathway (a) and G2-M transition pathway (b).

Single-Cell Perturbation Modeling

Single-cell perturbation modeling predicts how genetic or chemical interventions change cellular states across cells and contexts.

See all 61 papers

Differentiable Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation
Mar 7, 2026Andrea Giuseppe Di Francesco, Andrea Rubbi, Rishabh Jain +1Single-Cell Perturbation Modeling Retrieval-Augmented Generation
University of Cambridge Wellcome Sanger Institute
PerturbCellRL: Verifier-Guided Reinforcement Learning for Single-Cell Perturbation Prediction
Jun 26, 2026Dongxia Wu, Mingyu Li, Yuhui Zhang +4Single-Cell Perturbation Modeling Reinforcement Learning with Verifiable Rewards
Peking University Beijing, China · Stanford University Stanford, CA
CisTransCell: Single-Cell Perturbation Prediction via Gene Function, Regulatory Control, and Cellular Context
Jun 10, 2026Wei Zhang, Xun Jiang, Yuesi Xi +1Single-Cell Perturbation Modeling Genomics
1L3S Research Center, Leibniz Universität Hannover, Germany · School of Clinical Medicine & Laboratory Medicine, Jiangsu University, China · Institute for Information Processing (tnt), Leibniz Universität Hannover, Germany.

Mar 7, 2026J/K move · Enter open · S save

Differentiable Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

Andrea Giuseppe Di Francesco, Andrea Rubbi, Rishabh Jain, Pietro Liò

University of Cambridge Wellcome Sanger Institute

Single-Cell Perturbation Modeling Retrieval-Augmented Generation Adaptive Retrieval

Open paper AskPDF arXiv

Predicting transcriptional responses to genetic perturbations is fundamental to functional genomics and therapeutic discovery. Recent deep learning models have shown promise in single-cell perturbation response prediction, but they typically generate each response in isolation, without explicitly leveraging related perturbations. We introduce PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), a plug-in retrieval-and-conditioning module for generative cellular perturbation response. PT-RAG augments an existing perturbation-response backbone with learned access to related perturbation contexts. The key challenge is that relevance is not fixed in this setting: functionally related genes may elicit different effects across cell types. PT-RAG addresses this with a two-stage retrieval mechanism: GenePT-based semantic retrieval first identifies K candidate perturbations, after which a differentiable Gumbel-Softmax selector adaptively selects retrieved contexts conditioned on the control cell state, the query perturbation, and each candidate perturbation. We evaluate PT-RAG on two backbones, a STATE-style generator used as a frozen random reservoir and a fully trained scGPT, across cross-cell-type and cross-perturbation generalization tasks. PT-RAG consistently improves distributional similarity and often overall predictive quality; for example, on scGPT cross-cell-type results, the 2-Wasserstein distance drops by 5.9% relative to scGPT alone. The code to reproduce our experiments is available at https://github.com/difra100/PT-RAG_NIPS.

Abstract

Single-Cell Perturbation Modeling

Explore similar work

Differentiable Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

PerturbCellRL: Verifier-Guided Reinforcement Learning for Single-Cell Perturbation Prediction

CisTransCell: Single-Cell Perturbation Prediction via Gene Function, Regulatory Control, and Cellular Context