alphaFeedback

WorksComputer Science

An ODE-based Sampler for Discrete Diffusion Models

AI contributions
  • A generative AI assistant
View details

Abstract

The marginals of a continuous diffusion model can be reproduced by a deterministic ordinary differential equation (ODE), known as the probability flow ODE, which provides a foundation for accelerated sampling and consistency distillation. In this paper we address the question whether an analogous formulation exists for discrete diffusion models and answer it affirmatively by introducing exponential residual flows (ERFs). Our framework augments the discrete state with a continuous residual variable for each token and position. These residuals evolve deterministically according to piecewise-defined ODEs whose vector fields depend on the current discrete state and the forward and backward rate matrices of the prescribed diffusion model. The discrete state changes whenever a residual reaches zero. We prove that this construction preserves the marginals of the original diffusion. We derive a practical discretization of our sampler and specialize it to masked and uniform diffusion. The deterministic trajectories of ERFs support consistency distillation and yield empirical improvements over previous approaches across a variety of benchmarks including code and text generation tasks.

Introduction

Discrete diffusion models have emerged as a compelling alternative to autoregressive models. By allowing multiple positions to be updated in parallel, they offer a path toward faster decoding [23, 26] and potentially more efficient steering and control [29]. In continuous time, these models are described by forward and reverse continuous-time Markov chains [3, 23]. Standard simulation of the reverse process is inherently stochastic: Gillespie’s algorithm draws fresh randomness at each jump [14], while ancestral sampling of discrete-time reverse chains likewise relies on stochastic transitions [30, 26, 19, 29].

This raises a natural question: can one design a discrete diffusion sampler whose only randomness is the initial condition, as in the probability-flow ODE for continuous diffusion models [32 a]? A direct analogue on a finite state space is impossible because a deterministic map cannot split probability mass. Continuous auxiliary randomness is therefore required at initialization, particularly for masked diffusion, whose terminal distribution is a Dirac mass. Existing uses of auxiliary randomness serve different purposes. The projected Gaussian trajectories of Sahoo et al. [27] act as more of a reparameterization and do not define a deterministic flow driven solely by the current discrete state. Gumbel Distillation (Zhang et al., 2026) instead uses Gumbel noise to construct deterministic training targets from an autoregressive teacher for parallel decoding, rather than a marginal-preserving deterministic flow for a discrete diffusion process.

Our contributions are as follows:

  • Our central contribution is a new deterministic sampler for discrete diffusion models, which we call the exponential residual flow (ERF). We augment the state with non-negative auxiliary variables, called residuals, initialized as independent exponential random variables, one for each token and position. Their evolution is then governed by an ODE determined by the forward and reverse rates of the discrete dynamics that governed jumps of the discrete state. A token change is triggered whenever a residual reaches zero. In contrast to Gillespie’s first-reaction construction, where residual clocks are refreshed after each jump, ERF continues to evolve all remaining residuals deterministically. Consequently, once the initial augmented state has been sampled, no additional randomness is required. We prove that these augmented dynamics preserve the marginal distribution of the original CTMC at every time. At each fixed time, the residuals remain independent standard exponential random variables and are independent of the discrete state. The resulting process generally differs from the reverse Markov chain and need not be Markovian.

  • To make this construction applicable to learned diffusion models, we introduce a novel numerical integration scheme that freezes the leave-one-out denoiser over each sampling interval while preserving the known time dependence of the transition rates. We derive explicit formulations for both masked and uniform diffusion models, and show that the resulting scheme is exact for factorized data distributions when the denoiser is exact.

  • We evaluate ERF against ancestral sampling under matched model-evaluation budgets and find comparable generation quality. The deterministic nature of ERF also makes it directly amenable to consistency distillation. We exploit this property to train few-step student models and evaluate the resulting speed–quality trade-offs on text generation, code generation, and Sudoku completion.

Notation. We consider integer-valued sequences in X=VLX = V^L, where V=⟦1,K⟧V = \llbracket1, K \rrbracket, and write ΔK={p∈RK:pk≥0, ∑k=1Kpk=1}\Delta_K = \{p \in\mathbb{R}^K : p_k \ge0,\ \sum_{k=1}^{K} p^k = 1\} for the probability simplex. For x∈Xx \in X, xℓx^\ell denotes its ℓ\ell-th token and x−ℓ∈VL−1x^{-\ell} \in V^{L-1} the sequence with position ℓ\ell removed. For masked diffusion, the mask token is m:=Km := K. For any vector x=(x1,…,xd)∈Vdx = (x^1,\ldots,x^d) \in V^d, we denote its one-hot matrix by x~∈{0,1}d×K\tilde{x} \in\{0,1\}^{d \times K}, with entries x~ℓk:=1xℓ=k\tilde{x}^{\ell k} := \mathbf{1}_{x^\ell=k}. For a single token x∈Vx \in V, x~∈RK\tilde{x} \in\mathbb{R}^K denotes its one-hot vector. For π∈ΔK\pi\in\Delta_K, Cat⁡(π)\operatorname{Cat}(\pi) denotes the categorical distribution, and Cat⁡(x;π):=⟨x~,π⟩\operatorname{Cat}(x;\pi) := \langle\tilde{x},\pi\rangle is the probability of xx under Cat⁡(π)\operatorname{Cat}(\pi). We use the convention sup⁡∅=−∞\sup\varnothing= -\infty.

Background on discrete diffusion models

We first briefly introduce discrete diffusion models [26, 30, 2, 19, 29]. This class of generative models constructs a stochastic process that transports a simple reference distribution, p1:=∏ℓ=1LCat⁡(πℓ)p_1 := \prod_{\ell=1}^{L} \operatorname{Cat}(\pi^\ell) with π∈(ΔK)L\pi\in(\Delta_K)^L, to a target data distribution p0p_0. To this end, we use a forward process that progressively corrupts the data distribution, ultimately transforming it into p1p_1. The generative model is then obtained by learning an approximation of the corresponding reverse process.

We consider as forward process a Continuous Time Markov Chain (CTMC) process (Yt)t∈[0,1](Y_t)_{t \in[0,1]} started at p0p_0 with transitions qt∣s(xt∣xs)=∏ℓ=1Lqt∣sℓ(xtℓ∣xsℓ)q_{t|s}(x_t|x_s) = \prod_{\ell=1}^{L} q_{t|s}^{\ell}(x_t^\ell|x_s^\ell), where xs,xt∈Xx_s,x_t \in X, 0≤s<t≤10 \le s < t \le1 and

qt∣sℓ(xtℓ∣xsℓ)=Cat⁡(xtℓ;αt∣sx~sℓ+(1−αt∣s)πℓ),0≤s<t≤1,(1)q_{t|s}^{\ell}(x_t^\ell|x_s^\ell) = \operatorname{Cat}(x_t^\ell; \alpha_{t|s}\tilde{x}_s^\ell+ (1-\alpha_{t|s})\pi^\ell), \qquad0 \le s < t \le1 , \tag*{(1)}

where x~s\tilde{x}_s denotes the one-hot encoding of xsx_s and (αt)t∈[0,1](\alpha_t)_{t\in[0,1]} is a monotone noise schedule with α0=1\alpha_0 = 1, α1≈0\alpha_1 \approx0, and αt∣s=αt/αs\alpha_{t|s} = \alpha_t/\alpha_s. The time-reversal of this process is obtained by reversing the direction of time and is defined as the process (Y1−t)t∈[0,1](Y_{1-t})_{t\in[0,1]} backward in time1 with Yt=lim⁡h↓0Yt−hY_t = \lim_{h\downarrow0} Y_{t-h}. For ease of notation, we denote YtY_t by XtX_t. Since YtY_t almost surely for any tt, the process XtX_t has marginals ptp_t, and allows by construction to transport the source distribution p1p_1 to the target p0p_0. In addition, the process (Xt)t∈[0,1](X_t)_{t\in[0,1]} is itself a CTMC with transitions

ps∣t(xs∣xt)=∑x0qs∣0,t(xs∣x0,xt) p0∣t(x0∣xt),(2)p_{s|t}(x_s|x_t) = \sum_{x_0} q_{s|0,t}(x_s|x_0,x_t)\,p_{0|t}(x_0|x_t) , \tag*{(2)}

with p0∣t(⋅∣xt)p_{0|t}(\cdot|x_t) the denoising posterior and qs∣0,tq_{s|0,t} the conditional distribution of YsY_s given (Y0,Yt)(Y_0,Y_t), whose explicit expression is given in Appendix § 11. During training, these reverse transitions are learned using a factorized model parameterized by θ\theta: ps∣tθ(xs∣xt)=∏ℓ=1Lps∣tθ(xsℓ∣xt)p^\theta_{s|t}(x_s|x_t) = \prod_{\ell=1}^{L} p^\theta_{s|t}(x_s^\ell|x_t). To obtain a practical model for inference, we then consider the following joint distribution:

p0:nθ(xt0:n)=p1(xtn)∏i=1npti−1∣tiθ(xti−1∣xti),(3)p^\theta_{0:n}(x_{t_0:n}) = p_1(x_{t_n}) \prod_{i=1}^{n} p^\theta_{t_{i-1}|t_i}(x_{t_{i-1}}|x_{t_i}) , \tag*{(3)}

for a choice of timesteps 0=t0<⋯<tn=10 = t_0 < \cdots< t_n = 1, setting si=ti−1s_i = t_{i-1} for notational convenience. A common parameterization used in the literature ([30, 26, 19, 29]) is ps∣tθ(xs∣xt)=qs∣0,t(xs∣x^0θ(xt,t),xt)p^\theta_{s|t}(x_s|x_t) = q_{s|0,t}(x_s|\hat{x}^\theta_0(x_t,t),x_t), where qs∣0,tq_{s|0,t} denotes the extension of the bridge transition in (17) to simplex-valued inputs, and x^0θ\hat{x}^\theta_0 denotes a neural network with values in (ΔK)L(\Delta_K)^L. In what follows we assume, without any loss of generality, that it approximates the leave-one-out denoiser, whose ℓ\ell-th component is (t,xt)↦E[X0ℓ∣Xt−ℓ=xt−ℓ](t,x_t) \mapsto\mathbb{E}[X^\ell_0 \mid X^{-\ell}_t = x^{-\ell}_t]; see Appendix § 11.

Rate matrices and discrete score. Another useful perspective on discrete diffusion is provided through rate matrices associated to the forward and backward processes ([3, 23] (Sun et al. 2023)). The rate matrix associated to the forward transitions (1) is given for any x,y∈Xx,y \in\mathcal{X} and t∈[0,1]t \in[0,1] by

Qt(y∣x)=∑ℓ=1L1y−ℓ=x−ℓQtℓ(yℓ∣xℓ),where Qtℓ(y∣x)=(−α˙tαt)⟨y~,πℓ⟩(x≠y).(4)Q_t(y|x) = \sum_{\ell=1}^{L} \mathbf{1}_{y^{-\ell}=x^{-\ell}}Q^\ell_t(y^\ell|x^\ell), \qquad\text{where } Q^\ell_t(y|x)=\left(-\frac{\dot{\alpha}_t}{\alpha_t}\right)\langle\tilde{y},\pi^\ell\rangle\quad(x\ne y). \tag*{(4)}

References

  1. [1]David F. Anderson. A modified Next Reaction Method for simulating chemical systems with time dependent propensities and delays. The Journal of Chemical Physics, 127(21):214107, 2007. doi: 10.1063/1.2799998.
  2. [2]Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:17981–17993, 2021.
  3. [3]Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266–28279, 2022.
  4. [4]Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design, 2024. URL https://arxiv.org/abs/2402.04997.
  5. [5]Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https://arxiv.org/abs/2110.14168.
  6. [6]Giovanni Conforti, Alain Durmus, Le-Tuyet-Nhi Pham, and Gael Raoul. Non-asymptotic convergence of discrete diffusion models: Masked and random walk dynamics. arXiv preprint arXiv:2512.00580, 2025.
  7. [7]Justin Deschenaux and Caglar Gulcehre. Beyond autoregression: Fast LLMs via self-distillation through time. In International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=uZ5K4HeNwd.
  8. [8]Alain Durmus, Arnaud Guillin, and Pierre Monmarché. Piecewise deterministic markov processes and their invariant measures, 2021. URL https://arxiv.org/abs/1807.05421.
  9. [9]Luca Eyring, Vincent Pauline, Stefan Bauer, Alexey Dosovitskiy, and Zeynep Akata. Ddno: Discrete diffusion noise optimization. 2026.
  10. [10]Eugene A. Feinberg, Manasa Mandava, and Albert N. Shiryaev. On solutions of kolmogorov’s equations for jump markov processes, 2013. URL https://arxiv.org/abs/1301.6998.
  11. [11]Willy Feller. On the integro-differential equations of purely discontinuous markoff processes. Transactions of the American Mathematical Society, 48(3):488–515, 1940.
  12. [12]Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky T. Q. Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. In Advances in Neural Information Processing Systems, volume 37, 2024. doi: 10.52202/279917-4239. URL https://proceedings.neurips.cc/paper_files/paper/2024/hash/f0d629a734b56a642701bba7bc8bb3ed-Abstract-Conference.html.arxiv.org/abs/2407.15595
  13. [13]Michael A. Gibson and Jeshua Bruck. Efficient exact stochastic simulation of chemical systems with many species and many channels. The Journal of Physical Chemistry A, 104(9):1876–1889, 2000. doi: 10.1021/jp993732q.
  14. [14]Daniel T Gillespie. Exact stochastic simulation of coupled chemical reactions. The journal of physical chemistry, 81(25):2340–2361, 1977.
  15. [15]Daniel T Gillespie. Approximate accelerated stochastic simulation of chemically reacting systems. The Journal of chemical physics, 115(4):1716–1733, 2001.
  16. [16]Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http://Skylion007.github.io/OpenWebTextCorpus, 2019.
  17. [17]Google DeepMind. DiffusionGemma. https://deepmind.google/models/gemma/diffusiongemma/, June 2026. Accessed: 2026-06-10.
  18. [18]Samson Gourevitch, Yazid Janati, Dario Shariatian, Umut Simsekli, Eric Moulines, Eric P. Xing, and Alain Durmus. Uniform diffusion models revisited: Leave-one-out denoiser and absorbing state reformulation, 2026. URL https://arxiv.org/abs/2605.22765.
  19. [19]Emiel Hoogeboom, Didrik Nielsen, Priyank Jain, Patrick Forré, and Max Welling. Argmax flows: Learning categorical distributions with normalizing flows. In Third Symposium on Advances in Approximate Bayesian Inference, 2021.
  20. [20]Emiel Hoogeboom, David Ruhe, Jonathan Heek, Thomas Mensink, and Tim Salimans. Beyond single tokens: Distilling discrete diffusion models via discrete MMD, 2026. URL https://arxiv.org/abs/2603.20155.
  21. [21]David Li, Nikita Gushchin, Dmitry Abulhanov, Eric Moulines, Ivan Oseledets, Maxim Panov, and Alexander Korotin. Idlm: Inverse-distilled diffusion language models, 2026. URL https://arxiv.org/abs/2602.19066.
  22. [22]Bingbin Liu, Sebastien Bubeck, Ronen Eldan, Janardhan Kulkarni, Yuanzhi Li, Anh Nguyen, Rachel Ward, and Yi Zhang. TinyGSM: achieving ≥80% on GSM8K with small language models, 2023. URL https://arxiv.org/abs/2312.09241.
  23. [23]Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 32819–32848. PMLR, 2024. URL https://proceedings.mlr.press/v235/lou24a.html.arxiv.org/abs/2310.16834
  24. [24]Patrick Pynadath, Jiaxin Shi, and Ruqi Zhang. Generative frontiers: Why evaluation matters for diffusion language models, 2026. URL https://arxiv.org/abs/2604.02718.
  25. [25]Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
  26. [26]Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37:130136–130184, 2024.arxiv.org/abs/2406.07524
  27. [27]Subham Sekhar Sahoo, Justin Deschenaux, Aaron Gokaslan, Guanghan Wang, Justin Chiu, and Volodymyr Kuleshov. The diffusion quality. arXiv preprint arXiv:2506.10892, 2025.
  28. [28]Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=TIdIXTzphoI.arxiv.org/abs/2202.00512
  29. [29]Yair Schiff, Subham Sekhar Sahoo, Hao Phung, Guanghan Wang, Sam Boshar, Hugo Dalla-torre, Bernardo P de Almeida, Alexander Rush, Thomas Pierrot, and Volodymyr Kuleshov. Simple guidance mechanisms for discrete diffusion models. arXiv preprint arXiv:2412.10193, 2024.arxiv.org/abs/2412.10193
  30. [30]Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37:103131–103167, 2024.arxiv.org/abs/2406.04329
  31. [31]Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021a. URL https://openreview.net/forum?id=St1giarCHLP.arxiv.org/abs/2010.02502
  32. [32]Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b.
  33. [33]Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 32211–32252. PMLR, 23–29 Jul 2023. URL https://proceedings.mlr.press/v202/song23a.html.arxiv.org/abs/2303.01469

Paper details

vlnc:W_gjvptjgk
v1 · latestHistory (1)

published · Fresh — added to the Hub within the last 30 days

added ·

Usage

views · 0

downloads · 0

A view counts once per browser per day, a download each time the PDF is requested.

CC BY 4.0

Original source

  • Computer Science

A generative AI assistant · View statement →

Looking for related papers…
Contents