The Diffusion Duality, Chapter II: $Ψ$-Samplers

Justin Deschenaux; Caglar Gulcehre; Subham Sekhar Sahoo

cs.LGFeb 24, 2026

The Diffusion Duality, Chapter II: $Ψ$ -Samplers

Authors: Justin Deschenaux, Caglar Gulcehre, Subham Sekhar Sahoo

Organizations: EPFL, Lausanne, Switzerland · Microsoft AI · Cornell Tech, NY

Abstract

Uniform-state discrete diffusion models excel at few-step generation and guidance due to their ability to self-correct, making them preferred over autoregressive or Masked diffusion models in these settings. However, their sampling quality plateaus with ancestral samplers as the number of steps increases. We introduce a family of Predictor-Corrector (PC) samplers for discrete diffusion that generalize prior methods and apply to arbitrary noise processes. When paired with uniform-state diffusion, our samplers outperform ancestral sampling on both language and image modeling, achieving lower generative perplexity at matched unigram entropy on OpenWebText and better FID/IS scores on CIFAR10. Crucially, unlike conventional samplers, our PC methods continue to improve with more sampling steps. Taken together, these findings call into question the assumption that Masked diffusion is the inevitable future of diffusion-based language modeling. Beyond sampling, we develop a memory-efficient curriculum for the Gaussian relaxation training phase, reducing training time by 25% and memory by 33% compared to Duo while maintaining comparable perplexity on OpenWebText and LM1B and strong downstream performance. We release code, checkpoints, and a video-tutorial on: https://s-sahoo.com/duo-ch2

Figures & tables

Panel of Figure 1 — Figure 1: Performance on Language Modeling and Image Modeling. $\Psi$ -samplers generalize ReMDM ( Wang et al., 2025 ) to arbitrary noise distributions. (Left): Generative perplexity (Gen. PPL; $\downarrow$ ) as a function of NFEs, with nucleus sampling $p=0.9$ . $\Psi$ -samplers consistently improve with more steps, unlike ancestral sampling which plateaus. Curves are annotated with the average unigram entropy per sequence as a proxy for diversity. (Right): On CIFAR-10, $\Psi$ -samplers achieve better FID ( $\downarrow$ ) than MDLM (with ReMDM).

Figure 2:

\Psi

-samplers combine predictor and corrector steps. The predictor transitions from

{\mathbf{z}}_{t}

{\mathbf{z}}_{s}

via

{q_{s|t}}

, but fails to remask tokens in MDMs. The corrector steps inject noise via

q_{s}

, to revise earlier predictions. For

\kappa_{t}<1

, noise injection enables error correction while preserving the forward process marginals. Our framework extends prior PC methods ( Campbell et al., 2022 ; Gat et al., 2024 ; Wang et al., 2025 ) to arbitrary priors

\bm{\pi}

Figure 3: Efficient Curriculum for USDMs. Duo ( Sahoo et al., 2025a ) replaces discrete lookups with linear combinations of all $K$ embeddings: (1) Gaussian diffusion on one-hot representations, (2) Low-temperature ${\color[rgb]{0.6055,0,0}\mathrm{softmax}}$ , (3) weighted sum. $\text{Duo}^{++}$ exploits the sparsity of the tempered softmax (most weights are effectively zero), and simulate the k largest entries (out of K) using ordered statistics. The approximate normalizer ${\color[rgb]{0.6055,0,0}\tilde{Z}}$ admits a closed form expression ( 14 ). $\text{Duo}^{++}$ has a 33% lower memory and 25% faster training than Duo.

	Arc-e	Arc-c	HSwag	WinoG	PIQA	MathQA	OQA
AR Transformer	38.55	23.04	30.55	53.19	63.60	22.41	32.4
MDLM †	34.26	24.66	31.54	51.93	57.89	20.70	28.60
Duo	28.11	25.43	26.46	47.20	51.14	20.00	23.40
$\text{Duo}^{++}$ ( $k=2$ )	27.32	26.11	26.26	49.64	52.12	20.40	27.80
$\text{Duo}^{++}$ ( $k=3$ )	28.28	25.00	25.89	47.36	50.65	21.01	23.00
$\text{Duo}^{++}$ ( $k=5$ )	28.03	25.77	26.90	50.12	51.25	20.20	25.40

Table 1: Accuracy on multiple-choice question answering datasets. Abbreviations: Arc-e (ARC-Easy), Arc-c (ARC-Challenge), HSwag (HellaSwag), WinoG (Winogrande), PIQA (Physical Intelligence Question Answering), OQA (OpenBookQA). † Results from Deschenaux et al. (2025) .

\text{Duo}^{++}

(

k=2

) achieves slightly higher accuracy than Duo on 4 out of 6 tasks. Overall,

\text{Duo}^{++}

matches Duo’s performance while using 25% fewer flops . The highest accuracy among USDMs is bolded . The absolute best per column is underlined .

Figure 4: Illustration of the possible evolution of

t

and the associated

\kappa_{t}

. In practice, we use

\kappa_{t}

close to 1 during the PC phase.

	LM1B	OWT
Autoregressive
Transformer †	22.3	17.5
Masked Diffusion
SEDD Absorb †	32.7	24.1
MDLM †	27.0	23.2
Uniform-state Diffusion

Table 2: Test perplexity (PPL) on LM1B and OWT. Lower is better. † Results from Sahoo et al. (2025a) . Best Uniform-state diffusion numbers are bolded . Duo and

\text{Duo}^{++}

achieve comparable performance across both datasets while requiring 25% fewer GPU-hours (Table 4 ), demonstrating the effectiveness of our memory-efficient curriculum.

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Polynomial approximation and approximation error, compared to the series approximation, truncated at 150 terms. The degree- $9$ polynomial (left) achieves orders of magnitude lower error than the degree- $5$ polynomial (center) and sigmoid (right) approximations.

Component	Value
Vocab size	256
Number of ResNet blocks per scale	2
Base channels	128
Channel multiplier per scale	(1,2,2,2)
Attention resolutions	16
Conditional embedding dimension	128

Appendix

Table 3: Model architecture on CIFAR10

Figure 6:

\Psi

-samplers, which generalize ReMDM, significantly improve the Inception Score on CIFAR-10, compared to ancestral sampling.

Method	Throughput	Peak Memory
	(samples/s) $\uparrow$	(GiB) $\downarrow$
Duo	81.8	94.3
Duo (after the CL)	122.4	63.3
$\text{Duo}^{++}$ ( $k\in\{2,3,5\}$ )	121.9	63.4

Appendix

Table 4: Training efficiency comparison between Duo and

\text{Duo}^{++}

on 138M parameter models. All measurements are conducted on a training job on 8 NVIDIA GH200-120GB GPU with batch size 32. We report the average throughput in sequence per second. The row “Duo (after CL)” denotes the resources consumption of Duo after the Curriculum phase. The impact of

k

is minimal when

k\in\{2,3,5\}

, and

\text{Duo}^{++}

uses similar resources.

	PTB	Wiki	LM1B	LBD	AG News	PubMed	ArXiv
Autoregressive
Transformer †	82.05	25.75	51.25	51.28	52.09	49.01	41.73
Diffusion (138M)
SEDD Uniform †	105.51	41.10	82.62	57.29	82.64	55.89	50.86
UDLM †	112.82	39.42	77.59	53.57	80.96	50.98	44.08
Duo †	89.35	33.57	73.86	49.78	67.81	44.48	40.39

Appendix

Table 5: Zero-shot perplexity (PPL) on seven datasets. Lower is better. † Results taken from Sahoo et al. (2025a) .

\text{Duo}^{++}

(

k=2

) achieves a slightly lower zero-shot perplexity than Duo on 6 of 7 datasets.

Figure 7: Marginal distributions of the top-5 entries using a tokenizer with 100k tokens, inverse temperature 100, and log signal-to-noise ratio

-2