PaperlayerView exact revision on arXiv

Source-linked AI summary

Acceleration of Diffusion Language Model through Discrete Average Generator

Yidong Ouyang, Zhengyan Wan, Themis Haris, Tian Tan, Liqian Peng, Henry Li, Ziqian Lin, Jianhang Chen, Maryam Karimzadehgan, Alec Go, George Michailidis

arXiv:2609.38364v1stat.MLcs.AIcs.LG

TL;DR

Few-step generation in discrete diffusion remains challenging because sampling requires multiple sequential network evaluations. The paper introduces a CTMC average generator and a tractable training objective based on its self-consistency identity. It reports improved simulation and text-generation results, with comparable ImageNet performance.

  • Problem

    Discrete diffusion sampling remains computationally expensive because it requires sequential network evaluations, while few-step coordinate-wise updates lose dependencies between jointly updated coordinates.

    Source: §1 · p. 1

  • Method

    The paper defines a CTMC average generator as the normalized transition-kernel increment and uses its self-consistency identity with per-coordinate marginal projection for tractable training.

    Source: §1 · pp. 1–2

  • Results

    The method reduces simulation sampler total variation by up to 67%, achieves the lowest OpenWebText generative perplexity for 8–64 steps with 16× acceleration, and achieves comparable ImageNet performance.

    Source: §1 · p. 2, §6 · p. 11

  • Takeaways & Limitations

    The projected identity yields a closed equation whose transport term can be estimated without bias using one additional forward pass.

    Source: §6 · p. 11

    MeanFlow’s target depends on its unknown oracle average velocity, and approximating that target with the network prediction introduces bias and unstable training.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

Discrete diffusion models and flow matching have emerged as powerful frameworks for generative modeling over discrete state spaces, yet efficient few-step generation remains a fundamental challenge. In this work, we introduce the Discrete Average Generator, a principled extension of MeanFlow to Continuous-Time Markov Chains (CTMCs). Analogously to how MeanFlow defines an average velocity field over a time interval in continuous spaces, we define an average generator as the normalized increment of the transition kernel over a time interval. We show that this average generator satisfies a self-consistency identity, which provides the foundation for our training objective. We further develop training strategies that align with the standard training paradigm of diffusion language models while keeping the resulting objective tractable. When projected onto per-coordinate marginals, the self-consistency identity admits a closed-form expression, enabling efficient training and inference. In Potts model simulations, our objective reduces the total variation distance of the $K$-step sampler by up to 67%. On OpenWebText, our method achieves the lowest generative perplexity among the evaluated methods for 8 to 64 sampling steps while enabling a $16\times$ acceleration, and achieves comparable performance to existing methods on ImageNet.

1. Introduction

Few-step generation in discrete diffusion is costly because sampling requires sequential network evaluations and practical samplers lose dependencies between jointly updated coordinates. The paper extends MeanFlow to CTMCs with an average generator and a tractable training objective, reporting gains in simulations, text, and image generation.

1. Introduction

  • Discrete samplers require sequential network evaluations, and few-step coordinate-wise updates use instantaneous-rate laws that lose dependencies between coordinates updated together.These features make inference computationally expensive.Source: §1 · p. 1
  • The Discrete Average Generator extends MeanFlow to discrete state spaces governed by CTMCs, defining the average generator as the normalized transition-kernel increment over an interval.A single Euler step of size r−t with this generator reproduces the exact transition law.Source: §1 · pp. 1–2
  • Its average generator satisfies a self-consistency identity that forms the basis of the training objective.The method progresses from transition-rate prediction to posterior prediction.Source: §1 · p. 2
  • Projecting the identity onto per-coordinate marginals yields a tractable closed form using D × |S|-dimensional quantities.Sequence-level posterior prediction is otherwise intractable because it involves a generator matrix of dimension |S|^D × |S|^D.Source: §1 · p. 2
  • On OpenWebText, the method has the lowest generative perplexity among compared methods for 8–64 steps, with lower perplexity than Di4C+ReDi1 using 16× fewer steps.On ImageNet, it achieves comparable FID and Inception Score to existing diffusion-language-model acceleration methods.Source: §1 · p. 2

2. Preliminary

The preliminaries describe discrete flow matching as learning transition rates for a target probability path, often through coordinate-wise conditional paths and posterior prediction. They then review MeanFlow’s interval-average velocity, its self-consistency identity, and a target approximation that can destabilize training.

2.1. Discrete Flow Matching

  • Discrete flow matching learns a marginal transition rate from the conditional rate that generates a conditional path, yielding the target probability path.The conditional rate can be chosen to suit the desired path.Source: §2.1 · pp. 2–3
  • A common construction uses coordinate-wise conditional paths and rates, including mixture paths with uniform or masked source distributions.The mixture-path rate is specified using the posterior over the clean coordinate.Source: §2.1 · p. 3
  • A common training objective minimizes Bregman divergence between conditional and estimated generators, equivalently learning the posterior under generalized KL divergence.With a masked source and mixture path, the objective recovers the training objective for masked diffusion models.Source: §2.1 · p. 3

2.2. MeanFlow

  • MeanFlow defines an average velocity over a time interval and derives a self-consistency identity used to formulate a fixed-point training objective.The interval runs from t to r, and the velocity field uses noise at t=0 and data at t=1.Source: §2.2 · p. 4
  • MeanFlow’s target depends on an unknown oracle average velocity, so approximating it with the network prediction introduces bias and unstable training.A velocity-based reparameterization addresses this issue and trains through a standard velocity regression objective.Source: §2.2 · p. 4

3. Proposed Method

The Discrete Average Generator extends MeanFlow to CTMCs by defining an interval-average generator and deriving a self-consistency identity. The method turns that identity into a tractable training objective through prediction parameterizations and per-coordinate projection.

3.1. Average Generator

  • The average generator is defined over an interval, and its self-consistency identity expresses it as a correction of the instantaneous generator.The identity provides the basis for a fixed-point training objective.Source: §3.1 · p. 5

3.2. Training Objective of Discrete Average Generator

  • The U-prediction target can be self-referential and unstable, so the method reparameterizes the instantaneous generator for Q-prediction to eliminate that dependence.The method also parameterizes the average generator through per-coordinate posterior distributions.Source: §3.2 · p. 6

3.3. Tractable Training Objective of Discrete Average Generator

  • Per-coordinate projection yields a closed equation for the exact marginal of the finite-step kernel, avoiding simulation of the kernel over the interval.The projected transport term can be estimated without bias using a single forward pass shared across coordinates.Source: §3.3 · pp. 6–8
  • At the boundary r=t, the objective reduces to the instantaneous rate and is a strict generalization of Equation 4.The boundary equality holds when the predicted posterior matches p1|t.Source: §3.2 · p. 6, §3.3 · p. 8

4. Experimental Results

The method improves few-step generation across Potts simulations, OpenWebText, and ImageNet, with the strongest reported text result matching a 16× acceleration over Di4C+ReDi1. The experiments evaluate discrete-state simulations and tokenized text and images using task-specific metrics.

4.1. Experimental Setup

  • Evaluation spans exactly enumerable Potts simulations, 1024-token GPT-2 OpenWebText sequences, and 256-code VQGAN ImageNet sequences.For simulations, total variation is measured against the exact ground truth; text uses generative perplexity and entropy, while images use FID and Inception Score. OpenWebText training starts from Di4C+ReDi1 for 1k iterations; ImageNet starts from ReDi2 for 3 epochs.Source: §4.1 · p. 9
  • The baselines include distilled SDTT and Di4C, rectification-based ReDi, and their Di4C+ReDi1 combination; the proposed method does not distill a teacher trajectory.Source: §4.1 · p. 9

4.2. Experimental Results

  • The correction reduces total variation by 21% at K=4, 43% at K=8, and 67% at K=16, with the same ordering across all three seeds.The comparison is between networks trained with the baseline objective and Equation 21.Source: §4.2 · p. 10
  • At K=64, generative perplexity is 16.97 for our method versus 20.32 for Di4C+ReDi1, and it is lower than the baseline’s K=1024 value of 18.44.The paper reports this as a 16× acceleration; the comparison covers OpenWebText.Source: §4.2 · p. 10, Figure 3 · p. 10
  • At K=64, qualitative OpenWebText samples suggest fluent, topically coherent text without repetitive degeneration.The samples are unconditional; this inspection is reported in Appendix E.Source: §4.2 · p. 10
  • In ImageNet 256 × 256 generation with four network evaluations, our method reaches Inception Score 250 versus ReDi2’s 240, at comparable FID of 7.81 versus 7.86.The figure reports experimental results on ImageNet; values are in Table 4.Source: Figure 2 · p. 9, §4.2 · p. 10

5. Related Work

Prior acceleration work spans continuous generative models and discrete diffusion, using solvers, caching, distillation, consistency, and coupling rectification. This work instead learns a CTMC average generator directly on discrete states.

5. Related Work

  • Continuous-space acceleration includes higher-order solvers, caching, distillation, consistency models, rectified flow, shortcut models, and MeanFlow.Source: §5 · p. 10
  • Discrete approaches include distillation methods, coupling rectification, one-step distribution matching, and inference-time caching or parallel decoding.The cited examples include SDTT, Di4C, ReDi, Di[M]O, Fast-dLLM, and Jacobi Forcing.Source: §5 · p. 10
  • Unlike flow maps on a probability simplex, this method remains on discrete states and uses left multiplication by a CTMC generator as its transport term.The construction applies to CTMCs with coordinate-wise rates.Source: §5 · p. 11

6. Conclusion

The paper extends MeanFlow to CTMCs through a transition-kernel average generator and derives a tractable training objective. It situates this approach among continuous and discrete acceleration methods and includes a proof that the discrete flow-matching objective reduces to posterior cross-entropy under a coordinate-wise mixture path.

6. Conclusion

  • The conclusion reports up to 67% lower K-step sampler total variation in simulations, lowest OpenWebText generative perplexity from 8 to 64 steps with 16× acceleration, and the highest ImageNet Inception Score.The approach defines the average generator as the normalized transition-kernel increment; its per-coordinate marginal identity yields a closed equation whose transport term can be estimated without bias using one additional forward pass.Source: §6 · p. 11

A. Detailed Discussion of Related Work

  • Continuous acceleration methods include solvers, caching, distillation, consistency, and interval operators; discrete methods include trajectory distillation, coupling rectification, and accelerated inference, while this method learns a CTMC average generator without multi-step distillation.Source: A. Detailed Discussion of Related Work · p. 15

B.1. Proof of Training Objective for Discrete Flow Matching

  • Under the coordinate-wise mixture path, the training objective reduces to cross-entropy with respect to the posterior; restricting rates to off-diagonal entries makes the Bregman divergence well-defined.Source: §B.1 · pp. 15–16

B.2. Proof of Proposition 3.1 of Campbell et al. (2024)

The appendix derives marginal transition rates by comparing forward equations, reproduces the MeanFlow identity using differentiation, and defines a CTMC material derivative through Dynkin’s formula.

B.2. Proof of Proposition 3.1 of Campbell et al. (2024)

  • The proof identifies the marginal transition rate by comparing the conditional-rate equation with the Kolmogorov forward equation, then reindexes it to express rates from x to z.Source: §B.2 · p. 17

B.3. Proof of the MeanFlow Identity

  • The MeanFlow identity follows by differentiating the interval-scaled average velocity, expanding its total derivative, and rearranging terms.The total derivative combines the time partial derivative with the velocity-weighted spatial gradient.Source: §B.3 · pp. 17–18

C.1. Notation and the material derivative of a CTMC

  • For a time-dependent function on CTMC states, Dynkin’s formula gives its expected derivative as the expectation of its time derivative plus generator action, defining the discrete material derivative.The generator acts by left multiplication and is the discrete analogue of the continuous directional derivative.Source: §C.1 · p. 18

C.2. Proof of Proposition 1

The appendix establishes the average generator’s transition properties and develops its projected training and sampling procedures. The projection targets exact per-coordinate marginals, while the coordinate-factorized sampler cannot represent cross-coordinate dependence.

C.2. Proof of Proposition 1

  • The average generator is a valid generator whose single Euler step reproduces the exact transition law, and it converges to the instantaneous generator as r approaches t.These are the basic properties used to connect finite-interval transitions to instantaneous rates.Source: §C.2 · p. 19

C.3. Marginal projection and the tractable projected transport term

  • The projected identity targets the exact per-coordinate marginals of the finite-step kernel, and replacing the exact generator with the single-site model in its projected transport term is not an approximation.The objective can use instantaneous rates and single-coordinate perturbations rather than teacher trajectories over the interval.Source: §C.3 · pp. 20–21

C.4. The masked path

  • On the masked path, the training objective reduces to a weighted cross-entropy over masked coordinates.The paper identifies this as the objective used on ImageNet.Source: §C.4 · p. 21

C.5. Sampling with the Discrete Average Generator

  • The sampler takes an independent Euler step for each coordinate, so it can match exact per-coordinate marginals but cannot represent the true kernel’s cross-coordinate dependence.The paper does not claim a guarantee that exact marginals reduce K-step sampler error; setting r=t recovers the standard ancestral sampler as K grows.Source: §C.5 · pp. 21–22

C.6. The uniform-state (DUO) process

  • For the uniform-state process, the rate is not proportional to posterior mass, so the mixture-path posterior parameterization does not carry over.The appendix gives an unbiased projected-transport estimate using one draw and one forward pass.Source: §C.6 · p. 23

C.7. Proof of Training Objective for Discrete Average Generator

  • The proof reduces the posterior-parameterized objective to the stated training objective, with the masked-path case reducing to Equation 31.The derivation uses the per-coordinate target rate and the normalized posterior parameterization.Source: §C.7 · pp. 24–25

D.1. Experimental Setup of Simulations

The simulation setup uses an enumerable Potts target to evaluate sampler total variation exactly. The appendix also specifies training and evaluation settings and points to tables containing detailed benchmark values.

D.1. Experimental Setup of Simulations

  • The Potts simulation uses a four-coordinate, four-state target, allowing exact enumeration of the ground-truth distribution and sampler total variation.The setup uses coupling strength β=1.5 and fixed perturbations to break permutation symmetry.Source: §D.1 · p. 25

D.2. Implementation Details of OpenWebText and ImageNet

  • Training conditions the network on time and step size, with separate time-pair sampling designs for OpenWebText and ImageNet.A quarter of pairs use r=t, where the correction vanishes; ImageNet sampling hyperparameters are selected by sweep, while OpenWebText baselines come from Yoo et al. (2025).Source: §D.2 · p. 26

D.3. Detailed values for Section 4

  • Detailed simulation, OpenWebText, and ImageNet values are provided in Tables 2, 3, and 4, respectively.Table 2 reports Potts-model total variation means over three seeds.Source: §D.3 · p. 26, Table 2 · p. 26

E. Unconditional Text Generation Samples (𝐾= 64)

This appendix presents unconditional OpenWebText samples generated with the exact-marginalization sampler at 64 steps. The reported evaluation for that configuration includes generative perplexity and entropy, alongside illustrative text outputs.

E. Unconditional Text Generation Samples (𝐾= 64)

  • The displayed OpenWebText outputs use 1,024 tokens, 64 sampling steps, temperature 1.0, and the exact-marginalization sampler.The evaluated batch reports LLaMA-3.1-8B generative perplexity 16.97 and entropy 5.02.Source: E. Unconditional Text Generation… · p. 26
  • Another displayed sample discusses economic conditions in Egypt and U.S. government responses.The sample includes claims about economic hardship, oil production, and aid programs.Source: E. Unconditional Text Generation… · p. 29
  • One displayed sample is labeled “Internet Openness, Privacy Law & Climate Study” and is presented at K=64 and T=1.0.The associated text moves among internet openness, privacy, and climate-related topics.Source: E. Unconditional Text Generation… · p. 30

E.1. CFG Ablation on ImageNet

The ImageNet ablation finds that the best guidance weight depends on the sampling-step budget, with precision and recall moving in opposite directions as guidance increases.

E.1. CFG Ablation on ImageNet

  • FID is 7.81 at w=3.0 for four-step sampling and 6.90 at w=1.5 for eight-step sampling; precision rises and recall falls monotonically with w at both budgets.The authors attribute the shift toward lower guidance at eight steps to iterative refinement improving fidelity, allowing lower guidance to preserve diversity.Source: §E.1 · p. 32

E.2. Additional Qualitative Samples on ImageNet

The ImageNet appendix shows qualitative samples from a controlled CFG-weight sweep and a four-class comparison of the method with baselines. The highlighted 4-step setting is w=3.0, with FID 7.81.

E.2. Additional Qualitative Samples on ImageNet

  • The four-class qualitative comparison follows Yoo et al. (2025).Source: §E.2 · p. 32
  • With the random seed fixed per class, the CFG-weight sweep highlights w=3.0 as the optimal 4-step setting, with FID 7.81.The qualitative sweep varies w across 1.0, 1.5, 2.0, 3.0, 4.0, and 5.0 for nine ImageNet classes.Source: Figure 5 · p. 33
  • The 4-step ImageNet comparison appends the method below the baselines, reporting FID 7.81.The figure shows two samples per model for goldfinch, sea anemone, cheeseburger, and spotted salamander.Source: Figure 6 · p. 34
Loading 2609.38364v1…