Corrected Samplers for Discrete Flow Models

Discrete flow models (DFMs) have been proposed to learn the data distribution on finite state space, offering a flexible framework as an alternative to discrete diffusion models. A line of recent work has studied samplers for discrete diffusion models, such as tau-leaping and Euler solver. However, these samplers require a large number of iterations to control discretization error, since the transition rates are frozen in time and evaluated at the initial state within each time interval. Moreover, theoretical results for these samplers often require boundedness conditions of the transition rate or they focus on a specific type of source distributions. To address those limitations, we establish non-asymptotic discretization error bounds for those samplers without any restriction on transition rates and source distributions, under the framework of discrete flow models. Furthermore, by analyzing a one-step lower bound of the Euler sampler, we propose two corrected samplers: \textit{time-corrected sampler} and \textit{location-corrected sampler}, which can reduce the discretization error of tau-leaping and Euler solver with almost no additional computational cost. We rigorously show that the location-corrected sampler has a lower complexity than existing parallel samplers. We validate the effectiveness of the proposed method by achieving better generation quality with reduced inference time on simulations and text-to-image generation tasks. Code can be found in \url{https://github.com/WanZhengyan/Corrected-Samplers-for-Discrete-Flow-Models}.

Discrete flow models provide a flexible way to generate discrete data, such as text tokens or image-related discrete representations. However, generating samples from these models can be slow, because existing sampling methods often need many small steps to maintain accuracy. This work studies how to make sampling for discrete flow models faster while still controlling the error. We develop theoretical guarantees for common samplers and propose simple corrections that improve their accuracy with little extra computational cost. As a result, the model can generate high-quality samples using fewer sampling steps. Overall, our work provides both theoretical understanding and practical tools for accelerating generation with discrete flow models.