🤖 machine learning

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

This paper introduces a family of Predictor-Corrector samplers for discrete diffusion that overcome the performance plateau of ancestral sampling to improve generation quality with more steps, while also proposing a memory-efficient training curriculum that challenges the notion that Masked diffusion is the superior future for language modeling.

Original authors: Justin Deschenaux, Caglar Gulcehre, Subham Sekhar Sahoo

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Justin Deschenaux, Caglar Gulcehre, Subham Sekhar Sahoo

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

Imagine you are trying to write a story, but instead of starting with a blank page and writing word by word (like a traditional writer), you start with a page covered in static noise and gradually clean it up until a coherent story emerges. This is how Diffusion Models work. They are powerful AI tools used to generate images and text.

However, there's a catch. When these models try to generate text, they often get stuck in a "good enough" loop. If you ask them to write a long sentence, they might start strong but lose quality as they go, or they might get stuck repeating the same patterns.

This paper, titled "The Diffusion Duality, Chapter II: Ψ-Samplers," introduces a new set of tools (called Duo++ and Ψ-samplers) that fix these problems. Here is the breakdown in simple terms:

1. The Problem: The "One-and-Done" Mistake

Think of traditional text generation (like standard AI chatbots) as a train moving forward on a track. Once a train passes a station, it can't go back. If the AI makes a mistake early on, it has to keep writing from that mistake, which often ruins the whole sentence.

Some newer models (called Masked Diffusion Models) are better because they can "remask" (hide) a word and try again. But the authors found that the models they focused on (Uniform-State Diffusion Models) were even better at fixing mistakes early on, but they hit a "glass ceiling." No matter how much time you gave them to think (more steps), their quality stopped improving. They were like a student who studies hard but stops learning after a certain point.

2. The Solution: The "Editor and the Safety Net" (Ψ-Samplers)

The authors invented a new way to generate text called Ψ-samplers. Imagine a writer working on a draft:

The Predictor (The Draft): The AI guesses what the next word should be.
The Corrector (The Editor): This is the magic part. In the old methods, once a word was written, it was locked in. In this new method, the "Editor" is allowed to say, "Wait, that word doesn't fit. Let's hide it and try a different one," even if it was already written.

The paper calls this a Predictor-Corrector system. It's like having a safety net that catches the AI if it starts to fall. The authors proved that by adding this "safety net" (which they call a Ψ-posterior), the AI can keep improving its writing the longer you let it think. Unlike the old methods that hit a glass ceiling, these new samplers keep getting better and better the more steps you give them.

The Result:

For Text: The new method writes better sentences with less confusion (lower "perplexity") than the previous best methods, especially when given more time to think.
For Images: It also creates sharper, clearer pictures (better FID scores on CIFAR-10) than before.

3. The Efficiency Hack: The "Smart Menu" (Fast Curriculum)

Training these AI models is usually like trying to read a dictionary where every single word is a different flavor of ice cream. You have to taste every single one to find the right mix. This takes a huge amount of computer memory and time.

The authors realized that when the AI is "thinking" during training, it usually only cares about the top few "flavors" (words) and ignores the rest. The other flavors are so unlikely they might as well be zero.

So, they built a "Fast Curriculum" (a training schedule). Instead of tasting the whole dictionary, the AI now uses a smart menu that only looks at the top 2 or 3 most likely options.

The Analogy: Imagine you are ordering food. Instead of reading a 10,000-item menu, you only look at the top 3 items the chef recommends. You get the same great meal, but you save 33% of the time and memory.
The Result: They trained the model 25% faster and used 33% less memory, without losing any quality in the final product.

Summary of Claims

The paper claims three main things:

New Sampling Method: They created a general method (Ψ-samplers) that allows AI to "remask" and correct its own mistakes during text and image generation, leading to higher quality results that keep improving with more time.
Better Performance: Their new model (Duo++) beats previous state-of-the-art models on text generation (OpenWebText) and image generation (CIFAR-10).
Efficiency: They made the training process much cheaper and faster by ignoring the "noise" in the data, reducing memory use by a third and speeding up training by a quarter.

In short, they gave the AI a better editor and a smarter way to study, making it write better stories and draw better pictures without needing a supercomputer to do it.

Technical Summary: The Diffusion Duality, Chapter II: Ψ-Samplers

Problem Statement

Discrete diffusion models have emerged as a powerful alternative to autoregressive models for language and image generation, offering parallel decoding and bidirectional context. Among these, Uniform-State Diffusion Models (USDMs) excel in few-step generation and guided tasks due to their ability to self-correct by revising tokens multiple times. However, USDMs face two critical limitations:

Sampling Quality Plateau: Standard ancestral samplers for USDMs fail to improve generation quality as the number of sampling steps (NFEs) increases, unlike their performance in few-step regimes.
Training Inefficiency: Previous attempts to narrow the likelihood gap between USDMs and Masked Diffusion Models (MDMs) via curriculum learning (e.g., Duo) are computationally expensive, requiring high memory usage and long training times due to the need to materialize high-dimensional latent vectors during the training phase.

While Masked Diffusion Models (MDMs) generally achieve higher likelihoods, they suffer from compounding errors during inference because tokens, once unmasked, cannot be remasked. Recent work on Predictor-Corrector (PC) methods (e.g., ReMDM) has improved MDMs by allowing remasking, but similar advancements for USDMs have been underexplored, with existing PC methods for USDMs often underperforming ancestral samplers.

Methodology

The paper introduces Duo++, a framework that addresses these challenges through two primary innovations: $\Psi$ -samplers for inference and a scalable curriculum for training.

1. $\Psi$ -Posteriors and $\Psi$ -Samplers

The authors propose a family of non-Markovian posteriors, termed $\Psi$ -posteriors, which generalize previous PC methods to arbitrary noise priors (both Masked and Uniform).

Definition: The $\Psi$ -posterior is defined as a linear superposition of the standard reverse posterior ( $q_{s|t}$ ) and the forward process marginal ( $q_s$ ):
$\Psi_{s|t}(\cdot|x^\ell, z^\ell_t) = \kappa_t q_{s|t}(\cdot|z^\ell_t, x^\ell) + (1-\kappa_t) q_s(\cdot|x^\ell)$
where $\kappa_t \in [0, 1]$ is a time-dependent coefficient.
Mechanism: This formulation allows the sampler to inject noise (via the $q_s$ $q_{s}$ term) during the generation process.
- For MDMs, this enables remasking: tokens that were previously decoded can be returned to the masked state if the denoiser makes an error, correcting compounding mistakes.
- For USDMs, this ensures every token retains a non-zero probability of being sampled, even if the denoiser assigns it near-zero probability, preventing the sampler from getting stuck in local optima.
Theoretical Guarantee: The $\Psi$ -posteriors are constructed to share the exact same marginals as the standard Markovian discrete diffusion process. This ensures that, given sufficient samples, the process converges to the correct data distribution.
Generalization: The framework subsumes existing PC methods (e.g., ReMDM, Campbell et al., Gat et al.) as special cases where specific schedules for $\kappa_t$ are chosen.

2. Scalable Curriculum for Faster Training

To address the computational bottleneck of the curriculum learning strategy proposed in Duo (Sahoo et al., 2025a), the authors introduce a memory-efficient implementation.

Sparsity Exploitation: The original curriculum replaces discrete token inputs with low-temperature softmax relaxations of Gaussian latents. For large vocabularies ( $K > 100,000$ ), materializing the full $K$ -dimensional weight vector is infeasible.
Top- $k$ Approximation: The authors observe that at low temperatures, the softmax distribution is sparse, concentrating mass on only a few entries. They propose approximating the full linear combination using only the top- $k$ entries ( $k \ll K$ ).
Efficient Sampling: Instead of generating the full vector and sorting, the method uses order statistics to sample the top- $k$ values and their corresponding indices directly from the underlying Gaussian distributions without materializing the full vector.
Normalization: A closed-form expression is derived to approximate the softmax normalizer ( $\tilde{Z}$ ) by accounting for the sampled top- $k$ terms and the expected contribution of the unsampled terms.
Efficiency: This approach reduces peak memory usage by 33% and accelerates training by 25% while maintaining the same likelihood performance.

Key Contributions

$\Psi$ -Posteriors: A unified framework for discrete diffusion that introduces non-Markovian superposition posteriors, enabling Predictor-Corrector sampling for arbitrary noise priors (Masked and Uniform) while preserving marginal consistency.
$\Psi$ -Samplers: A family of samplers that generalize prior PC methods (including ReMDM). Empirically, they allow USDMs to scale better with NFEs, closing the performance gap with MDMs in text generation and surpassing them in image generation.
Efficient Curriculum: A scalable training strategy that leverages the sparsity of low-temperature softmax relaxations to reduce memory and compute costs significantly without sacrificing model quality.

Experimental Results

Language Modeling (OpenWebText)

Performance: $\Psi$ -samplers significantly improve Generative Perplexity (Gen. PPL) for USDMs compared to ancestral sampling.
Scaling: Unlike ancestral sampling, which plateaus as NFEs increase, $\Psi$ -samplers continue to improve with more steps. At matched unigram entropy (diversity), Duo++ with $\Psi$ -samplers achieves lower Gen. PPL than MDLM with ReMDM.
Downstream Tasks: Duo++ matches the performance of the original Duo model on multiple-choice question answering tasks (e.g., ARC, HellaSwag) while using 25% fewer FLOPs.

Image Modeling (CIFAR-10)

Quality: On CIFAR-10, Duo++ with $\Psi$ -samplers achieves superior FID and Inception Score (IS) compared to both ancestral sampling and ReMDM.
Scaling: The improvement in image quality is consistent across varying numbers of sampling steps, with Duo++ outperforming MDLM variants.

Training Efficiency

Memory & Speed: The proposed efficient curriculum reduces peak memory usage by 33% and reduces total training time by 25% compared to the original Duo curriculum.
Likelihood: The models trained with the efficient curriculum achieve comparable perplexity on LM1B and OpenWebText to the original Duo models.

Significance and Claims

The paper argues that the assumption that Masked Diffusion is the inevitable future of diffusion-based language modeling is questionable. By introducing $\Psi$ -samplers, the authors demonstrate that Uniform-State Diffusion Models can achieve:

Superior Scaling: Performance that continues to improve with increased sampling steps, unlike ancestral samplers.
Competitive or Superior Quality: Matching or exceeding MDMs in text generation (when paired with PC sampling) and surpassing them in image generation.
Practical Efficiency: A training pipeline that is significantly more memory- and time-efficient than previous curriculum learning approaches.

The work suggests that the "duality" between Gaussian and discrete diffusion, combined with non-Markovian sampling strategies, offers a robust and scalable path forward for discrete generative modeling, challenging the dominance of masked approaches.