PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast Sampling via Discrete Non-Markov Diffusion Models with Predetermined Transition Time

Zixiang Chen, Huizhuo Yuan, Yongqian Li, Yiwen Kou, Junkai Zhang, Quanquan Gu

arXiv:2312.09193v3cs.LGcs.AIstat.ML

TL;DR

Training-free acceleration for discrete diffusion remains under-explored. DNDM uses a predetermined transition-time set to reduce neural-network evaluations, and experiments report faster sampling while preserving sample quality. Its infinite-step study offers insight into bridging discrete and continuous-time processes.

  • Problem

    Training-free sampling acceleration for discrete diffusion remains under-studied.

    Source: §1 · p. 2

  • Method

    DNDM uses a predetermined transition-time set to accelerate discrete diffusion sampling without retraining the network.

    Source: §1 · p. 2

  • Results

    DNDM is about 3× faster than baselines at T=50 and about 30× faster at T=1000 while preserving sample quality.

    Source: §1 · p. 2

  • Takeaways & Limitations

    DNDM-C’s infinite-step formulation offers insight into bridging discrete and continuous-time processes for discrete diffusion.

    Source: §5 · p. 10

    Empirical claims are limited to text generation with non-autoregressive models; other domains and autoregressive architectures were not explored.

    Source: Limitations · p. 14

Abstract

from arXiv · show

Discrete diffusion models have emerged as powerful tools for high-quality data generation. Despite their success in discrete spaces, such as text generation tasks, the acceleration of discrete diffusion models remains under-explored. In this paper, we propose discrete non-Markov diffusion models (DNDM), which naturally induce the predetermined transition time set. This enables a training-free sampling algorithm that significantly reduces the number of function evaluations (i.e., calls to the neural network), making the sampling process much faster. Furthermore, we study the transition from finite to infinite step sampling, offering new insights into bridging the gap between discrete and continuous-time processes for discrete diffusion models. Extensive experiments on natural language generation and machine translation tasks demonstrate the superior performance of our method in terms of both generation speed and sample quality compared to existing methods for discrete diffusion models.

1 Introduction

Discrete diffusion has gained interest, but training-free sampling acceleration remains under-studied. The paper introduces DNDM, using predetermined transition times to reduce neural-network evaluations and extend sampling to an infinite-step setting.

1 Introduction

  • The paper situates its contribution in a field where discrete diffusion acceleration is under-studied, especially for training-free methods.A recent reparameterized diffusion model improves text-generation speed and quality but uses a training-based algorithm.Source: §1 · p. 2
  • DNDM preserves the original discrete process’s marginal and conditional distributions while enabling acceleration for multinomial and absorbing diffusion.Its transition-time set supports a training-free sampling algorithm.Source: §1 · p. 2
  • At T=50, DNDM is about 3× faster than baselines, and at T=1000 it is about 30× faster while preserving sample quality.The transition set’s cardinality is less than T and approaches O(1) as T grows.Source: §1 · p. 2
  • DNDM-C studies infinite-step sampling, offering insight into bridging discrete and continuous-time processes for discrete diffusion models.With a pretrained network, it generates the real data distribution using O(1) network evaluations.Source: §1 · p. 2

2 Background

The background describes categorical discrete diffusion as a forward corruption process and a learned reverse process for generation. It focuses on multinomial and absorbing diffusion and notes that standard reverse sampling evaluates the network at every step.

2 Background

  • The paper considers one-hot categorical variables and focuses on multinomial and absorbing discrete diffusion models.Discrete diffusion uses categorical noise rather than the Gaussian noise common in continuous-state models.Source: §2 · p. 3
  • The Markov forward process gradually corrupts data with noise over T steps, with updates governed by Bernoulli retention and noise draws.The noise is either uniform over the vocabulary or concentrated on an absorbing state.Source: §2 · p. 3
  • Reverse generation starts from noise and repeatedly applies a learned predictor, while the stochastic reverse process requires a network evaluation at every step.The predictor is trained by maximizing an evidence lower bound; the paper leaves detailed loss discussion to the appendix.Source: §2 · p. 3

3 Discrete Non-Markov Diffusion Models (DNDM)

DNDM replaces stepwise noise draws with shared noise and token-specific transition times, preserving key distributions while making the reverse process deterministic given those times. Its transition-time set supports finite-step acceleration and an infinite-step variant.

3.1 Forward and Reverse Process

  • DNDM’s non-Markov forward process uses time-invariant noise and preserves q(xt|x0), which in turn preserves q(x0|xt) and allows the Markov-trained predictor to remain applicable.The transition time is when a token switches from data to noise; the definition also applies to multinomial noise, unlike the cited absorbing-state-only definition [Hoogeboom et al. (2021a)].Source: §3.1 · p. 4
  • Given a transition time, DNDM’s reverse step is deterministic, and its ELBO is equivalent to the original process’s up to a constant.Sampling transition times in advance identifies the steps where tokens change.Source: §3.1 · p. 4

3.2 Accelerated Reverse Sampling

  • Each token has an independent transition time distributed according to the schedule; for a linear schedule, transition times are uniform over 1 to T.The method evaluates the network only at times in the transition set, reducing evaluations; experiments report fewer than half as many as Markov sampling.Source: §3.2 · pp. 5–6
  • DNDM-Multi and DNDM-Absorb accelerate multinomial and absorbing sampling, while DNDM-k adds top-k transition-time information to improve sample quality.For sequences, each token receives its own transition time.Source: §3.2 · pp. 5–6

3.3 Continous-time (Infinite Step) Reverse Sampling

  • DNDM-C takes a scale-invariant schedule to the continuous limit, orders the token transition times, and updates at those times rather than at infinitely many intermediate steps.The resulting method uses N network evaluations; unlike ARDM, it covers multinomial as well as absorbing diffusion and finite as well as continuous time, but does not include ARDM’s log-likelihood-based parallelization.Source: §3.3 · pp. 6–7

4 Experiments

The experiments evaluate DNDM on conditional machine translation and unconditional text generation. They report faster sampling alongside quality results, including comparisons across sampling methods and step counts.

4.1 Conditional Text Generation

  • Conditional evaluation uses BLEU and wall-clock time across IWSLT14 DE-EN, WMT14 EN-DE, and WMT16 EN-RO, comparing against RDM and RDM-k variants.Experiments use shared BPE vocabularies and an encoder-decoder architecture; reported time and BLEU are averaged over five runs except the RDM baseline.Source: §4.1 · p. 7
  • DNDM-based methods significantly accelerate machine-translation sampling, and their runtime grows only marginally with more steps, unlike RDM-based methods’ linear growth.The reported speedup increases with the number of sampling steps.Source: §4.1 · pp. 7–8
  • For continuous-time sampling, DNDM-C improves BLEU by up to 0.3 on IWSLT14 and WMT16, while WMT14 results show both gains and drops.The WMT16 improvement has one exception: absorbing diffusion without top-k selection.Source: §4.1 · p. 8, Table 2 · p. 8, Table 3 · p. 9
  • On IWSLT14, DNDM’s BLEU score grows nearly linearly with log time while the baseline curve is flat; multinomial diffusion gains 1 BLEU in under 60 seconds.For absorbing diffusion, DNDM outperforms RDM before RDM reaches 50 steps.Source: §4.1 · p. 8, Figure 1 · p. 9

4.2 Unconditional Text Generation

  • The 100-step DNDM-k-Multi example shows noise being removed as time runs from 100 to 0, with most transitions near the start and English text emerging.Figure 2 tracks BLEU on the left and text at different time steps on the right.Source: Figure 2 · p. 10, §4.2 · p. 10
  • On unconditional text generation, DNDM-based methods outperform vanilla multinomial sampling in both sampling time and perplexity, running 5 times faster on text8 and 14 times faster on enwik8.Table 4 compares sampling methods on both datasets.Source: §4.2 · p. 10, Table 4 · p. 10

5 Conclusion and Future Work

The conclusion presents DNDM as a discrete-state diffusion model with accelerated sampling and an infinite-step variant. The paper’s experiments focus on non-autoregressive text generation, with audio and image generation left for future exploration.

5 Conclusion and Future Work

  • DNDM combines a discrete non-Markov diffusion model with an accelerated sampling algorithm for discrete-state generation.The acceleration reduces neural-network function evaluations without sacrificing sample quality.Source: §5 · p. 10
  • DNDM-C provides an infinite-step sampling algorithm and new insights into bridging discrete- and continuous-time processes for discrete diffusion models.Source: §5 · p. 10
  • The study focuses on text generation using non-autoregressive models; applying the method to audio and image generation is proposed as future work.Source: §5 · p. 10

Broader Impact

The broader-impact statement identifies synthetic-data generation as a possible future application and notes that sensitive-data tasks may require proper controls.

Broader Impact

  • The authors believe the work may enable future applications of synthetic data generation, potentially leading to positive impacts.Source: Broader Impact · p. 14
  • The experiments are reported to demonstrate state-of-the-art performance in generative-model acceleration.Source: Broader Impact · p. 14
  • The authors say proper controls may be needed when applying the method to tasks involving sensitive data.Source: Broader Impact · p. 14

Limitations

The paper bounds its empirical claims to non-autoregressive text generation and discrete-Markov comparisons, leaving other settings and poor-quality hyperparameter regions unexplored. It also cautions that infinite-step DNDM-C is not guaranteed to outperform the 1000-step algorithm.

Limitations

  • The empirical claims cover non-autoregressive text, while audio, image generation, and autoregressive GPT architectures remain unexplored.The paper leaves these settings as future work.Source: Limitations · p. 14
  • Infinite-step DNDM-C is not guaranteed to produce better sample quality than the 1000-step accelerated algorithm, and training estimation errors may accumulate across steps.The paper identifies the assumption of optimal neural-network training as an ideal case that may not hold in practice.Source: Limitations · p. 14
  • The paper focuses on comparisons with discrete Markov diffusion and does not explore hyperparameter regions with poor sample quality.Its stated focus is acceleration while maintaining good sample quality.Source: Limitations · p. 14

A Related Work

  • The review places DNDM alongside continuous-time solvers and discrete innovations, including Zheng et al. (2023)'s reparameterized loss and sampling technique.It describes the paper's contributions as parallel to these discrete-model studies.Source: A Related Work · pp. 14–15

B Additional details of Discrete Diffusion

B.1 Comparison between D3PM and DNDM

This appendix compares DNDM with Markov discrete diffusion and DDIM, then describes its finite- and continuous-time training objectives. It highlights selective network evaluations, transition-time-based de-randomization, and a limitation of the continuous-time loss decomposition.

B.1 Comparison between D3PM and DNDM

  • The Markov and non-Markov processes yield the same q(x_t|x_0), although DNDM's fixed w makes its states history-dependent and prevents deriving this conditional through Markov transition kernels.Source: §B.1 · p. 15
  • For absorbing diffusion, DNDM and D3PM share a forward process, but DNDM precomputes transition times and evaluates the network only when a token needs to change, skipping sampling steps.The transition probabilities are otherwise identical when a masked token transitions.Source: §B.1 · p. 16
  • Unlike DDIM, DNDM fully de-randomizes after sampling transition time τ; its transition at τ selects x_0, and DNDM is not a special case of DDIM.Source: §B.1 · p. 16

B.2 Training Objective

  • The infinite-step average loss can provide a better checkpoint than the finite-step averaged loss on some tasks.Source: §B.2 · p. 17

B.3.1 Finite Time DNDM

B.3.2 Continous Time DNDM

  • The continuous-time derivation handles multiple discrete noise distributions, including multinomial diffusion, beyond an absorbing-only setting.Source: §B.3.2 · p. 19
  • The continuous-time ELBO decomposes over transition times, but IID transition times prevent further decomposition into a loss related to token-position information, unlike Hoogeboom et al. (2021a).Source: §B.3.2 · pp. 18–19

C Choice of the Transition Time

Transition order affects generation quality, and the transition-time schedule determines where transitions are concentrated. A validation-tuned Beta distribution can approximate schedule shapes and generally performs near the reported choices.

C Choice of the Transition Time

  • For a schedule αt = g(t/T), transition probability is approximately (1/T)|g′(t/T)|, concentrating transitions where |g′| is large.The linear schedule yields uniform transition times; the cosine schedules concentrate them differently, including toward the middle of the range for the cosine-square schedule.Source: C Choice of the Transition Time · pp. 20–21
  • A validation-tuned Beta distribution can approximate transition-time shapes, and its translation-task performance is better than most alternative schedules while their NFE efficiencies are similar.Across the tested Beta parameters, most choices remain close to the optimum, though validation-selected schedules are not always test-set maxima.Source: C Choice of the Transition Time · p. 21, Table 5 · p. 21
  • Updating tokens with transition time τ ≥ t, rather than only τ = t, lets tokens be corrected over time and is described as more robust than Algorithm 1 at nearly the same sampling time.This variation uses the fact that xt = x0 when τ > t.Source: C Choice of the Transition Time · p. 21
  • Left-to-right transitions outperform right-to-left transitions across all datasets and step counts.The authors report that earlier transitions of left-side tokens lead to better generation quality.Source: C Choice of the Transition Time · p. 22, Table 6 · p. 22

D Discussion on the Number of Function Evaluations (NFE).

DNDM requires neural-network evaluations only at distinct transition times, so its NFE equals the size of the transition set. Repeated transition times reduce that count, with the expected count depending on the transition-time distribution.

D Discussion on the Number of Function Evaluations (NFE).

  • NFE equals |T|, the number of distinct transition times, because the algorithm skips neural-network evaluations at steps outside T.The transition set satisfies 1 ≤ |T| ≤ min{N, T}.Source: D Discussion on the Number of Function… · p. 22
  • The expected transition-set size is [1 − CT,N,Dτ] · T, with the equality condition in the bound holding only for a uniform transition-time distribution.Here pi = P(τ = i); the reported result states that the bound is an equality if and only if Dτ is uniform.Source: D Discussion on the Number of Function… · pp. 22–23
  • For T = 25, NFE is about half the original count; for T = 50, about one-third; and for T = 1000, less than one-twentieth.These are reported average reductions compared with the original counts.Source: D Discussion on the Number of Function… · p. 23

E Discrete Non-Markov Diffusion Model with Top-k Transition Time (DNDM-K).

DNDM-K combines sampled transition times with top-K token selection. It evaluates the network only when the number of tokens remaining to decode changes, then updates selected, not-yet-updated tokens.

E Discrete Non-Markov Diffusion Model with Top-k Transition Time (DNDM-K).

  • DNDM-K selects tokens through a two-step process rather than choosing each updated token directly from its sampled transition time.The method is inspired by using neural-network token scores to prioritize updates.Source: E Discrete Non-Markov Diffusion Model… · pp. 24–25
  • At step t, Kt counts tokens with τn ≥ t; when Kt changes, the algorithm selects top-scoring tokens and updates those not previously updated.When Kt does not change, it carries the sentence forward unchanged.Source: E Discrete Non-Markov Diffusion Model… · p. 25
  • Because DNDM-K evaluates the network only when Kt changes, it provides an accelerated sampling algorithm.The procedure updates tokens with the highest Kt scores that have not yet changed.Source: E Discrete Non-Markov Diffusion Model… · p. 25

F.1 Conditional Text Generation

For conditional text generation, the paper selects Beta schedules using validation data and reports a WMT14 limitation tied to sparse timestep exposure in the pretrained checkpoint. Additional experiments include unconditional-generation settings and computational-time comparisons.

F.1 Conditional Text Generation

  • Finite-step transition schedules approximate the schedule with a Beta distribution whose parameters are selected by validation-set grid search.For continuous-step sampling, the authors choose between Beta(100, 4) and Beta(17, 4) using validation performance.Source: §F.1 · p. 25
  • On WMT14, a checkpoint trained on 50 steps may not handle unseen timesteps reliably when sampling with 1000 or infinitely many steps.The authors note that the longer training time and poorer WMT14 training performance make the checkpoint less dependable on those unseen steps.Source: §F.1 · p. 26

F.2 Unconditional Text Generation

  • The unconditional text-generation setup uses recovered multinomial-diffusion checkpoints and trains 12-layer Transformers on text8 and enwik8 for 500 epochs.The reported batch sizes are 256 for text8 and 128 for enwik8.Source: §F.2 · p. 27

G Additional Experiments

  • In the additional experiments, absorbing diffusion and RDM-absorbing diffusion have computational time that grows linearly with sampling steps.Figure 4 plots computational time against the number of sampling steps for absorbing diffusion, RDM-absorbing diffusion, DNDM-Absorb, and DNDM-T-Absorb.Source: G Additional Experiments · p. 27, Figure 4 · p. 27

G.1 Continuous Training

The appendix reports continuous-training comparisons, a WMT16 comparison with Mask-Predict, and a multinomial text-generation example. The displayed text-generation process converges at t = 39.

G.1 Continuous Training

  • Continuous training samples timestamps uniformly from [0, 1], and the authors report room to improve sampling scores relative to discrete training with continuous sampling.The comparison uses IWSLT14 and WMT16; Table 12 marks items that improve, but the supplied passages give no values.Source: Table 12 · p. 27, §G.1 · p. 28

G.2 Comparison with more generative models

  • On WMT16, DNDM runs faster and achieves a better BLEU score than Mask-Predict at a similar number of function evaluations.Experiments use the same GPU and machine setup.Source: Table 13 · p. 28, §G.2 · p. 28

G.3 Samples from the multinomial text models

  • For DNDM-Multi trained on IWSLT14, the full generation process shown in Figure 5 has 100 steps, and the generated sentence converges after t = 39.Source: §G.3 · p. 28

NeurIPS Paper Checklist

  • The checklist affirms that the abstract and introduction accurately reflect the paper’s stated scope of fast sampling via DNDM.The checklist points to the three contribution points at the end of the introduction.Source: NeurIPS Paper Checklist · p. 29

2. Limitations

These excerpts include a checklist confirmation that the paper has a limitations section, alongside statements about proofs, reproducibility, and experimental reporting. They do not specify the content of the limitations section.

2. Limitations

  • The checklist says a limitations section was added before the appendix, but these passages do not state its specific limitations.The accompanying checklist guidance recommends discussing assumptions, scope, computational efficiency, and other potential limitations; it is guidance rather than a paper-specific limitation.Source: §2 · p. 29
  • The authors state that theorems 3.1, 3.5, and D.1 are supported by rigorous proofs.The checklist also indicates that theorem statements and proofs are provided.Source: §2 · pp. 29–30
  • The authors report that public datasets are used and that their codebase will be made public upon acceptance.They also say training and evaluation code was submitted with the paper.Source: §2 · pp. 30–31
  • The authors say appendix sections D, E, and F provide training and test details, including experimental setup and procedures.The checklist response asserts that the paper supplies details needed to understand the results.Source: §2 · pp. 30–31
  • The checklist response states that confidence intervals are provided for experiments.Source: §2 · p. 31
  • The reported experiments use a single NVIDIA RTX A6000 GPU with 48 GB of memory.Source: §2 · p. 32
Loading 2312.09193v3…