Deterministic method speeds discrete diffusion sampling for text and code
An ODE-based Sampler for Discrete Diffusion Models
Machine Learning
Summary
This paper focuses on a type of machine learning model called discrete diffusion models, which help generate things like text and code. The authors develop a new way to simulate these models using a set of mathematical rules called ordinary differential equations (ODEs), making the process faster and more reliable. They do this by adding a small continuous value to each discrete part, which changes over time until a switch happens, preserving the model's behavior. Their approach improves performance in tasks like code and text generation compared to previous methods.
What this means in practice
- •For natural language processing engineers: Generate text with faster and more consistent discrete diffusion sampling using ODE-based deterministic trajectories.
- •For software development tool creators: Improve automatic code generation by integrating efficient sampling from discrete diffusion models based on exponential residual flows.
Authors
Samson Gourevitch, Yazid Janati, Eric Moulines, Eric P. Xing, Alain Durmus, Giacomo Zanella
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.