Paper extends probability-flow ODE to discrete diffusion models

Discrete diffusion models have a deterministic analogue to the probability flow ODE that powers accelerated sampling in continuous diffusion. That is the central claim of a new preprint from Samson Gourevitch, Yazid Janati, Eric Moulines, Eric P. Xing, Alain Durmus and Giacomo Zanella, posted on arXiv.

Their construction, which they call exponential residual flows, "augments the discrete state with a continuous residual variable for each token and position." The residuals move along piecewise-defined ODEs whose vector fields depend on the current discrete state and the forward and backward rate matrices of the prescribed diffusion. As the authors put it, "The discrete state changes whenever a residual reaches zero." They prove the construction preserves the marginals of the original diffusion, and specialize the sampler to the masked and uniform cases most of the field is actually using.

The payoff, per the abstract, is that the deterministic trajectories "support consistency distillation and yield empirical improvements over previous approaches across a variety of benchmarks including code and text generation tasks." The abstract names no specific benchmarks and reports no numerical gains.

Two of the researchers we track in our Who's Who directory shared the preprint shortly after it went up.