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

Justin Deschenaux; Caglar Gulcehre; Subham Sekhar Sahoo

Source-linked AI summary

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

Justin Deschenaux, Caglar Gulcehre, Subham Sekhar Sahoo

arXiv:2602.21185v2Feb 24, 2026cs.LG

TL;DR

Uniform-state diffusion can self-correct but ancestral sampling plateaus at high step counts, leaving its quality and likelihood below masked diffusion in some regimes. Duo++ introduces general Ψ-posteriors and predictor-corrector samplers, alongside a memory-efficient curriculum, improving high-step generation while reducing training cost. The resulting samplers match masked diffusion on language, achieve stronger CIFAR-10 FID and IS, and continue improving with more function evaluations.

Problem
Uniform-state diffusion supports self-correction but ancestral sampling plateaus at high step counts, while masked diffusion has stronger quality and likelihood in those regimes.
Method
Duo++ uses non-Markovian Ψ-posteriors to construct predictor-corrector samplers for arbitrary discrete-diffusion noise priors and reformulates the Gaussian-relaxation curriculum for efficient training.
Results
Ψ-samplers match masked diffusion on natural-language generation, achieve stronger FID and IS on CIFAR-10, continue improving with NFEs, and reduce training time by 25% and peak memory by 33%.
Takeaways & Limitations
Predictor-corrector sampling narrows uniform-state diffusion’s high-step gap with masked diffusion while extending the framework beyond uniform and masked noise.
Takeaways & Limitations
Masked diffusion cannot remask tokens after unmasking, and explicit Gaussian-latent materialization is memory-intensive during curriculum training.

Abstract

from arXiv · show

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

1 INTRODUCTION

Duo++ extends uniform-state discrete diffusion with Ψ-posteriors and predictor-corrector samplers to address ancestral sampling’s high-step plateau. It also introduces an efficient curriculum that improves the training cost of uniform-state models.

Motivation: Uniform-state diffusion permits repeated token revisions, enabling self-correction and making it effective for few-step and guided generation.Unlike masked diffusion, uniform-state models can revise tokens multiple times during generation.
Motivation: High-step generation quality and likelihood remain weaker for uniform-state diffusion than for masked diffusion, while prior predictor-corrector methods for uniform-state models remain underexplored.This motivates improving both sampling behavior and model capacity.
Duo++: Duo++ introduces non-Markovian Ψ-posteriors that expand uniform-state diffusion and produce predictor-corrector samplers for discrete diffusion.The proposed posteriors are designed to align with intermediate diffusion marginals.
Results: Ψ-samplers improve text and image generation, scale better than ancestral samplers at high NFE, close the masked-diffusion gap for text, and surpass masked diffusion on image tasks.These results are reported across language and image generation settings.
Results: The reformulated curriculum provides a 2× speedup, reduces peak memory by 33% and end-to-end training time by 25%, and maintains similar perplexity and downstream accuracy.The reported efficiency gains target the curriculum-learning stage while preserving comparable model quality.

2 BACKGROUND

Discrete diffusion interpolates data toward a categorical prior through progressively noisier intermediate distributions. Masked and uniform-state priors produce different generation behavior, while Gaussian relaxations and predictor–corrector constructions extend the modeling and sampling framework.

Discrete diffusion models: Discrete diffusion defines increasingly noisy distributions that interpolate from data to a factorized categorical prior through independent Markovian transitions.The prior is a product of independent Cat(.; π) distributions across sequence positions.
Discrete diffusion models: The forward process samples z_t from Cat(.; α_t x_ℓ + (1 − α_t)π), with a monotonically decreasing noise schedule α_t.The learned reverse process p_θ maps noisy samples back toward the data distribution.
Masked diffusion processes: Masked diffusion uses a [MASK] prior, and once tokens are unmasked they cannot be remasked, which can cause compounding inference errors.The reverse process inherits the forward process’s inability to remask tokens.
Uniform-state diffusion processes: Uniform-state diffusion uses π = 1/K, allowing tokens to change values multiple times and supporting few-step generation and guidance.This revisability enables self-correction during generation, unlike masked diffusion.
Predictor–corrector methods: Ψ-samplers combine a predictor posterior with corrector noise, generalizing predictor–corrector methods to arbitrary priors while preserving forward-process marginals.The corrector injects noise to revise earlier predictions, while the predictor performs the transition between times.
Gaussian relaxation and curriculum: A Gaussian diffusion on one-hot representations can produce discrete latents with matching marginals through an arg max transformation, although the resulting trajectory need not be Markovian.Softmax relaxation then supplies partially clean embeddings during early curriculum training, but materializing full latent vectors is memory-intensive.

3 THE Ψ-POSTERIORS

The paper introduces Ψ-posteriors as non-Markovian alternatives that preserve the marginals of discrete diffusion and yield predictor-corrector samplers for arbitrary noise priors. These samplers add controlled noise to ancestral reverse transitions, enabling error correction and subsuming earlier predictor-corrector methods.

Ψ-posteriors are non-Markovian generative processes that share the same marginals as the original discrete diffusion process.They apply to both Masked and Uniform-state diffusion.
The Ψ reverse posterior linearly combines standard reverse posteriors with the forward-process prior, using π = m for MDMs and π = 1/K for USDMs.The mixing coefficient κ_t controls the interpolation; κ_t = 1 recovers ancestral sampling.
For κ_t < 1, Ψ-sampling adds noise to the ancestral transition, with q_t acting as the corrector and q_s|t as the predictor.This construction is analogous to predictor-corrector methods in Gaussian diffusion.
Different κ_t choices recover prior predictor-corrector formulations, so the Ψ framework extends them to arbitrary noise priors.The framework subsumes methods from Campbell et al., Gat et al., and ReMDM.
The offset term enables error correction by allowing generated tokens to be replaced when the denoiser is imperfect.For MDMs, it permits previously denoised tokens to return to the masked state; for USDMs, the offset supports correction through the uniform prior.

4 SCALABLE CURRICULUM FOR FASTER TRAINING

The scalable curriculum addresses the cost of replacing discrete token inputs with low-temperature softmax relaxations over very large vocabularies. It exploits softmax sparsity to approximate the weighted embedding using only top-k entries without materializing the full K-dimensional vector.

K > 100,000 makes naive materialization of a K-dimensional softmax weight vector for every token infeasible during training.The curriculum uses a very low temperature, τ = 10^-3, which concentrates most softmax mass on a few entries.
Duo++ samples the top-k Gaussian values and their indices using ordered statistics, requiring only O(k) random variables instead of constructing the full vector.The top-k set accounts for whether the shifted clean-token coordinate is selected.
The approximate softmax-weighted embedding retains only the selected top-k embedding rows while approximating the contribution of unsampled terms in the normalizer.The unsampled terms are individually small but may have a non-negligible aggregate.
The curriculum objective computes the diffusion transformation operator on the fly using a Taylor expansion rather than relying on precomputed cached pairs.Caching many (α_t, T(α̃_t)) pairs is described as cumbersome.

5 EXPERIMENTS

Experiments show that Ψ-samplers improve discrete diffusion generation across language and image tasks, while Duo++'s efficient curriculum reduces training cost without sacrificing comparable performance.

Ψ-samplers improve text and image quality, enabling uniform-state diffusion models to outperform masked diffusion models in sample quality.
5.1.1 LANGUAGE MODELING: Beyond the sequence length, Ψ-samplers continue improving sample quality while ancestral sampling plateaus and the gap with masked diffusion closes.
5.1.1 LANGUAGE MODELING: Duo++ with Ψ-samplers outperforms MDLM with ReMDM and ancestral sampling across the full range of NFEs.
5.1.2 IMAGE MODELING: On CIFAR-10, Duo++ with Ψ-samplers achieves better FID and IS than MDLM with ancestral sampling and ReMDM.
5.2 FAST CURRICULUM: 33% lower peak memory and 25% faster training are achieved by the efficient curriculum while matching Duo on likelihood benchmarks and downstream tasks.
5.2 FAST CURRICULUM: Duo++ doubles curriculum-learning speed, and when curriculum learning covers half the training steps, it trains 25% faster at the 138M-parameter scale.

6 RELATED WORK AND DISCUSSION

The discussion situates Ψ-samplers as a general predictor-corrector framework for discrete diffusion, extending beyond uniform and masked noise while complementing other sampler designs.

The approach extends to discrete diffusion processes combining masked and uniform priors because it provides a predictor-corrector algorithm for arbitrary noise.
Ψ-samplers further generalize ReMDM and prior predictor-corrector methods to support arbitrary diffusion processes.
Ψ-samplers use uniform steps but support any step-size schedule and use first-order information, complementing noise-adaptive and higher-order samplers.

7 CONCLUSION

The paper presents Ψ-posteriors and Ψ-samplers as a unified predictor-corrector framework for discrete diffusion, extending masked methods to arbitrary noise priors. Duo++ improves generation quality and scaling while reducing curriculum-training cost.

Ψ-posteriors use non-Markovian superposition while preserving the marginals of the corresponding Markovian discrete diffusion process.
Ψ-samplers subsume prior masked-diffusion predictor-corrector samplers and extend naturally to uniform-prior discrete diffusion.
Duo++ matches Masked Diffusion Models on natural-language generation and achieves stronger FID and IS scores on CIFAR-10.
Performance continues improving with NFEs for Ψ-samplers, whereas ancestral samplers plateau as sampling steps increase.
The scalable curriculum reduces peak memory usage by 33% and training time by 25%.

A.4 RECOVERING PREDICTOR-CORRECTOR METHODS FOR MASKED DIFFUSION

The Ψ-posterior framework recovers existing masked-diffusion predictor-corrector methods as special cases while supporting equivalent sampling schedules and efficient top-k Gaussian order-statistic sampling.

The remasking parameter σ_t controls how aggressively previously generated tokens are remasked, with σ_t = 0 reducing ReMDM to the standard MDM posterior.
Ψ-posteriors recover the ReMDM posterior when κ_t = 1 − σ_t/(1 − α_s), thereby generalizing ReMDM and earlier FB and DFM posteriors.
Cap, Rescale, and Loop schedules provide alternative ways to control remasking or alter the sampling-time evolution.
The efficient top-k routine samples Gaussian order statistics and assigns them to indices without materializing all K entries, using O(k) memory.
Inverse transform sampling generates variables with a target CDF by applying its inverse CDF to an independent uniform random variable.

B.3 EFFICIENT COMPUTATION OF T DURING TRAINING

The curriculum computes the diffusion transformation operator through a truncated series and a low-degree approximation rather than extensive precomputation. This reduces storage and accelerates training-time evaluation.

The curriculum replaces costly direct evaluation of T with an on-the-fly Taylor expansion and a low-degree polynomial approximation.
Caching fewer than 150 series terms replaces precomputing 100k evaluations of T and takes only a few seconds at training start.
The degree-9 polynomial has orders-of-magnitude lower approximation error than degree-5 and sigmoid alternatives.
The transformation operator’s sigmoid-like shape permits approximations with fewer than 10 parameters instead of 100k stored values or 300 series coefficients.
The curriculum’s top-k implementation avoids materializing the full Gaussian-diffused vector by sampling selected order statistics and approximating the softmax normalizer.

C.1.2 CIFAR10 (D3PM-LIKE ARCHITECTURE)

The CIFAR10 experiments use U-Net models and compare diffusion configurations, while the language experiments evaluate likelihood and downstream performance under established protocols.

CIFAR10 models use a U-Net backbone trained for 1.5M steps with batch size 128 and class conditioning with dropout rate 0.1.
MDLM and Duo experiments vary time conditioning and cosine versus log-linear noise schedules before ancestral-sampling evaluation.
Multiple-choice downstream evaluation selects the candidate answer with the highest estimated ELBO-based log-likelihood.
The language likelihood evaluation uses seven validation datasets and reports performance comparable to Duo.

D ADDITIONAL EXPERIMENTAL RESULTS

The efficient curriculum reproduces the reference method’s marginal distributions, while performance depends critically on the κ_t choice.

The efficient curriculum produces weights with the same marginal distributions as Sahoo et al. (2025a).This result is reported in supplementary analysis.
A poor κ_t choice can make Ψ-samplers underperform ancestral sampling.The authors identify κ_t selection as critical for strong performance.
The image experiments use CIFAR-10 with a U-Net and horizontal flipping, while language experiments use OpenWebText checkpoints from MDLM and Duo.

D.1.1 CIFAR-10

On CIFAR-10, Ψ-samplers improve image-generation quality over ancestral sampling and ReMDM, with performance depending on sampling and training choices.

Ψ-samplers significantly improve CIFAR-10 Inception Score compared with ancestral sampling.
The reported FID evaluation uses CIFAR-10 generated samples and compares training and sampling noise schedules, conditioning, nucleus sampling, and final-step decoding.
ReMDM improves the best FID from 24.73 to 23.71 using 4096 sampling steps, whereas Ψ-samplers reduce Duo’s FID from 23.71 to 20.71 with the log-linear schedule.
Duo improves from FID 25.63 with ancestral sampling to 15.05 with Ψ-samplers, while MDLM improves from 24.73 to 17.86.
Generative perplexity is evaluated with GPT-2 Large alongside unigram entropy as a proxy for repetitive-text effects.
The plotted Ψ-sampler schedules match ancestral-sampling unigram entropy when possible, then select the lowest generative perplexity.

D.3 TRAINING EFFICIENCY OF OUR FAST CURRICULUM

The paper reports extensive CIFAR-10 and language-modeling evaluations, while the fast curriculum reduces training memory and improves throughput relative to Duo.

Duo++ achieves slightly lower zero-shot perplexity than Duo on 6 of 7 datasets.
The efficient and naive implementations closely match in marginal distributions across the reported tokenizer, inverse-temperature, and signal-to-noise settings.
The reported CIFAR-10 FID comparisons cover ancestral sampling, ReMDM schedules, and Ψ-sampler schedules with varying noise injection.
Ψ-sampler CIFAR-10 evaluations include Inception Score and FID under activation intervals and constant κ_t settings.
OpenWebText evaluations report generative perplexity and unigram entropy for ancestral sampling and Ψ-samplers across non-distilled and distilled models.