Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Corrected Samplers for Discrete Flow Models

Zhengyan Wan, Yidong Ouyang, Liyan Xie, Fang Fang, Hongyuan Zha, Guang Cheng

TL;DR

This work addresses discretization errors in sampling for discrete flow models by establishing non-asymptotic total-variation bounds that do not require bounded transition rates or specific source distributions. It analyzes a one-step lower bound for the Euler sampler and introduces two corrections—time-corrected and location-corrected samplers—that reduce discretization error with almost no extra computation. The location-corrected approach also achieves lower iteration complexity than existing parallel samplers in certain settings, and experiments on simulations and text-to-image generation demonstrate improved sample quality and faster inference. The results offer a practical, theoretically-grounded boost to discrete-flow-based generative modeling with scalable parallelism.

Abstract

Discrete flow models (DFMs) have been proposed to learn the data distribution on a 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 iteration complexity than existing parallel samplers. We validate the effectiveness of the proposed method by demonstrating improved generation quality and reduced inference time on both simulation and text-to-image generation tasks. Code can be found in https://github.com/WanZhengyan/Corrected-Samplers-for-Discrete-Flow-Models.

Corrected Samplers for Discrete Flow Models

TL;DR

This work addresses discretization errors in sampling for discrete flow models by establishing non-asymptotic total-variation bounds that do not require bounded transition rates or specific source distributions. It analyzes a one-step lower bound for the Euler sampler and introduces two corrections—time-corrected and location-corrected samplers—that reduce discretization error with almost no extra computation. The location-corrected approach also achieves lower iteration complexity than existing parallel samplers in certain settings, and experiments on simulations and text-to-image generation demonstrate improved sample quality and faster inference. The results offer a practical, theoretically-grounded boost to discrete-flow-based generative modeling with scalable parallelism.

Abstract

Discrete flow models (DFMs) have been proposed to learn the data distribution on a 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 iteration complexity than existing parallel samplers. We validate the effectiveness of the proposed method by demonstrating improved generation quality and reduced inference time on both simulation and text-to-image generation tasks. Code can be found in https://github.com/WanZhengyan/Corrected-Samplers-for-Discrete-Flow-Models.
Paper Structure (42 sections, 14 theorems, 104 equations, 6 figures, 3 tables, 8 algorithms)

This paper contains 42 sections, 14 theorems, 104 equations, 6 figures, 3 tables, 8 algorithms.

Table of Contents

  1. Introduction
  2. Discrete Flow-based Models and Continuous-time Jump Processes
  3. Discrete Flow Models
  4. Two-stage Jump Processes
  5. Tau-leaping and Euler Sampler
  6. Non-asymptotic Error Bound
  7. One-step Analysis
  8. Time-corrected Sampler
  9. Algorithm
  10. Non-asymptotic Error Bound
  11. Location-corrected Sampler
  12. Algorithm
  13. Non-asymptotic Error Bound
  14. Choosing Time-discretization Scheme
  15. Experiments
  16. ...and 27 more sections

Key Result

Proposition 1

If the marginal transition rate $Q_t^d(x^d,z^d)$ can generate marginal probability path $p^d_t(x^d)$ for any $d\in{[ \mathcal{D} ]}$, then the joint transition rate $Q_t(x,z)=\sum_{d=1}^\mathcal{D}\delta_{x^{\backslash d}}(z^{\backslash d})Q_t^d(x^d,z^d)$ can generate $p_t(x_t)=\prod_{d=1}^\mathcal{

Figures (6)

  • Figure 1: Overview of the transition rates of samplers in the $k$-th time interval ($z^d\neq x^d$). Tau-leaping and Euler sampler freeze the time variables of both time schedule and the posterior in the transition rates. For the time-corrected sampler, we correct the time variable of the time schedule without additional computational cost. For the location-corrected sampler, we not only correct the time variable of the schedule but also correct the location after the first jump in each timestep by constructing a two-stage transition rate; this sampler requires at most 2 function calls in each timestep.
  • Figure 2: Performance comparison of different samplers with $3$-dimensional data distribution (vocabulary size is $8$). (Top) Uniform source distribution; (Bottom) masked source distribution.
  • Figure 3: GenEval scores of different samplers in text-to-image generation tasks.
  • Figure 4: Location-corrected sampler for few-step regime. (a) the transition structure in equation \ref{['eq:location-corrected rate']}; (b) the transition structure in equation \ref{['eq:location-corrected rate-general']}.
  • Figure 5: Performance comparison ($\mathcal{D}=6,9,12,15$, from top to bottom) with uniform source distribution. (Left) Total variation v.s. sampling time; (Mid) Total variation v.s. number of steps; (Right) Sampling time v.s. number of steps.
  • ...and 1 more figures

Theorems & Definitions (38)

  • Definition 1: CTMC
  • Proposition 1: Independent coupling
  • Remark 1: Parallel samplers
  • Proposition 2: Exit time
  • Remark 2: Masked source distribution
  • Remark 3: Target set mismatch for tau-leaping
  • Theorem 1: Tau-leaping and Euler sampler
  • Theorem 2: One-step lower bound for Euler sampler
  • Theorem 3: Time-corrected sampler
  • Remark 4: Randomized midpoint sampler
  • ...and 28 more