[KongchangAI]
· 3 min read· 1,524 words

ODE Samplers for Discrete Diffusion Models: Exponential Residual Flows Accelerate Generation

ODE Samplers for Discrete Diffusion Models: Exponential Residual Flows Accelerate Generation

Proposes Exponential Residual Flows (ERFs) for deterministic ODE sampling in discrete diffusion models with consistency distillation acceleration.

This article introduces a new study on deterministic sampling for discrete diffusion models. Continuous diffusion can achieve deterministic sampling via probability flow ODE and be accelerated by consistency distillation, but discrete state spaces make direct ODE definition difficult due to discrete token jumps. The paper proposes Exponential Residual Flows (ERFs), introducing a continuous residual variable per token and position to embed discrete states into continuous space; residuals evolve deterministically by a piecewise ODE, and when they hit zero, the token jumps, with theoretical proof that the construction exactly preserves the original marginal distribution. The authors derive practical discretization schemes for masked and uniform diffusion, achieving superior results on code and text generation, with deterministic trajectories supporting consistency distillation to further reduce sampling steps.

Diffusion models have achieved remarkable success in continuous state spaces (e.g., image generation), where the sampling process can be made deterministic via the probability flow ordinary differential equation (ODE), enabling accelerated sampling and consistency distillation. However, whether a similar deterministic ODE formulation exists for discrete diffusion models (e.g., text and code generation) has remained an open question. A recent paper proposes Exponential Residual Flows (ERFs), which directly answers this question and offers a new approach to fast sampling for discrete diffusion models.

Paper title screenshot

The Challenge of Deterministic Sampling in Discrete Diffusion

The probability flow ODE in continuous diffusion models can exactly reproduce the model's marginal distribution, and its deterministic trajectory allows the use of black-box ODE solvers for fast sampling, as well as consistency distillation to further compress the number of sampling steps. But in discrete state spaces (e.g., tokens in a vocabulary), directly defining an ODE is challenging due to the discrete nature of state jumps: how can we construct a deterministic continuous evolution process without breaking the marginal distribution?

Traditional discrete diffusion sampling typically relies on stochastic methods (such as ancestral sampling). While these guarantee correct distributions, they require many steps and produce stochastic trajectories, making it difficult to apply acceleration techniques like consistency distillation. Therefore, exploring an ODE formulation for discrete diffusion has significant practical value.

Probability Flow ODE is a core concept in continuous diffusion models. In score-matching-based diffusion models, the forward noising process can be described as a stochastic differential equation (SDE), and the reverse generation process is also an SDE. The probability flow ODE is a deterministic ODE with the same marginal distributions as this SDE: it constructs a vector field such that the marginal distribution of particle trajectories evolving along this vector field matches that of the SDE at any time, while the trajectories themselves are deterministic. This property brings two key benefits: first, high-order black-box ODE solvers (e.g., Runge-Kutta) can be used to achieve high-precision sampling with fewer steps; second, the deterministic mapping enables consistency distillation—training a student model to directly map any noise level to the final sample, compressing the number of sampling steps to 1-4. This successful paradigm in continuous diffusion is precisely what discrete diffusion researchers aim to replicate.

Ancestral Sampling is the most commonly used sampling method for discrete diffusion models. Starting from a fully noised state, it samples step by step according to the transition probabilities of the reverse diffusion process: at each step, the next token is randomly drawn from a categorical distribution based on the current token and the diffusion model's reverse rates. Ancestral sampling strictly guarantees that generated samples follow the model distribution, but it has two inherent drawbacks: first, randomness is introduced at every step, making the sampling trajectory non-reproducible and high-variance; second, it typically requires a number of iterations comparable to the training diffusion steps (hundreds to thousands), resulting in low sampling efficiency. More importantly, because the trajectory is stochastic, it cannot be compressed into a few-step deterministic mapping via consistency distillation like continuous diffusion can—severely limiting the competitiveness of discrete diffusion in practical deployment. The motivation behind ERFs is precisely to circumvent these limitations of ancestral sampling.

The Core Idea of Exponential Residual Flows (ERFs)

The core innovation of ERFs is introducing a continuous residual variable for each token and position, thereby embedding the discrete state into a continuous space. Specifically, the model state consists of a discrete token sequence and a corresponding residual vector. The residuals evolve deterministically according to a piecewise-defined ODE, whose vector field depends on the current discrete state and the forward and reverse rate matrices of the diffusion model.

When a residual variable reaches zero, the corresponding discrete state jumps (i.e., the token changes). This mechanism is similar to "event-triggered" updates in continuous space, but strictly guarantees that the marginal distribution matches that of the original discrete diffusion model. The authors theoretically prove that this construction exactly preserves the marginal distribution of the original diffusion model, laying the foundation for deterministic sampling.

Rate Matrix is the mathematical foundation of discrete diffusion models. In the continuous-time Markov chain (CTMC) framework, the forward process of discrete diffusion is described by a rate matrix R_t: for any two distinct token categories i and j, R_t(i,j) represents the instantaneous rate of jumping from i to j, with diagonal elements determined by the condition that rows sum to negative values. The forward process gradually drives the data distribution toward a uniform or all-mask distribution; the reverse process requires the reverse rate matrix, which can be analytically computed from the forward rates and the model-learned scores (i.e., the denoising distribution). The vector field of ERFs is determined jointly by these forward and reverse rate matrices: the continuous evolution rate of the residual depends on the row of the rate matrix corresponding to the current discrete state, and when the residual hits zero, a jump is triggered with probabilities proportional to the rate matrix. Thus, understanding the rate matrix is a prerequisite for understanding the ERFs construction.

Discretization for Masked and Uniform Diffusion

To make ERFs practical, the authors derive a practical discretization scheme and apply it specifically to two common types of discrete diffusion: masked diffusion and uniform diffusion. In masked diffusion, tokens are replaced by masks with certain probabilities; in uniform diffusion, transitions occur uniformly across all token categories. For these two cases, the ERFs vector field can be expressed analytically, allowing efficient numerical integration.

The key to the discretization scheme is how to handle events where residuals hit zero. Since residual evolution is deterministic, it can be simulated via event-driven or fixed-step integration, and the paper provides a concrete algorithmic procedure. This deterministic trajectory not only supports standard ODE solvers but also paves the way for consistency distillation—distillation can train a few-step sampler, substantially reducing sampling costs.

Masked Diffusion and Uniform Diffusion are the two most commonly used concrete forms of discrete diffusion, differing in how the forward process corrupts data. Masked Diffusion replaces any token with a special [MASK] token at a certain rate during the forward process, with the final state being an all-mask sequence; its reverse process gradually predicts and fills in masked positions, which aligns closely with the training objective of BERT-style masked language models and is therefore particularly popular in text and code generation. Uniform Diffusion, on the other hand, transitions tokens to any other category in the vocabulary (including random tokens) at a uniform rate during the forward process, ultimately approaching a uniform distribution; it introduces no special tokens and is theoretically more symmetric, but reverse denoising is more difficult. The different rate matrix structures of the two diffusion types lead to different analytical expressions for the ERFs residual vector field; the paper derives efficient discretization schemes for each, enabling ERFs to be directly applied to existing mainstream discrete diffusion models.

Experimental Validation: Improvements in Code and Text Generation

The authors evaluate ERFs on a variety of benchmarks, including code generation and text generation tasks. Experimental results show that ERFs achieve empirical improvements in generation quality compared to previous discrete diffusion sampling methods. Because the deterministic trajectory supports consistency distillation, ERFs can further accelerate sampling, significantly reducing the number of sampling steps while maintaining generation quality.

Although the paper does not disclose specific quantitative metrics, the authors claim "superior performance to prior methods on multiple benchmarks," offering a new technical path for fast sampling of discrete diffusion models. Given the potential of discrete diffusion in natural language processing, code generation, and other domains, ERFs are expected to advance the practical deployment of these applications.

Significance and Outlook

The proposal of ERFs fills a gap in the theory of deterministic sampling for discrete diffusion models, extending the probability flow ODE idea from continuous diffusion to discrete spaces. Its theoretical guarantees and practical discretization scheme lay the foundation for future research. Possible future directions include: extending ERFs to more complex discrete structures (e.g., graphs, sequential decision-making), combining them with more efficient consistency distillation methods, and exploring applications on large language models. For practitioners working on generative model research, ERFs offer a new tool worth watching.

One-line Summary

Proposes Exponential Residual Flows (ERFs) to build deterministic ODE samplers for discrete diffusion models with consistency distillation acceleration.

Paragraph Summary

This article introduces a new study on deterministic sampling for discrete diffusion models. Continuous diffusion models can achieve deterministic sampling via the probability flow ODE and be substantially accelerated through consistency distillation, but discrete state spaces make it difficult to directly define an ODE due to the discrete nature of token jumps. The paper proposes Exponential Residual Flows (ERFs), which introduce a continuous residual variable for each token and position, embedding discrete states into continuous space; residuals evolve deterministically according to a piecewise ODE, and when they hit zero, the corresponding token jumps, with theoretical proof that this construction exactly preserves the marginal distribution of the original diffusion model. The authors derive practical discretization schemes for masked and uniform diffusion, achieving results superior to prior methods on code and text generation tasks; the deterministic trajectory also supports consistency distillation to further reduce sampling steps.

#discrete diffusion models#Exponential Residual Flows#probability flow ODE#consistency distillation#deterministic sampling#masked diffusion#uniform diffusion
Share:

Related articles