PaperlayerView exact revision on arXiv

Source-linked AI summary

Generalized Interpolating Discrete Diffusion

Dimitri von Rütte, Janis Fluri, Yuhui Ding, Antonio Orvieto, Bernhard Schölkopf, Thomas Hofmann

arXiv:2503.04482v2cs.CLcs.AIcs.LG

TL;DR

Autoregressive and masked diffusion models struggle to revise already generated tokens, motivating a more flexible discrete-diffusion framework. The paper introduces GIDD and a corresponding ELBO, then applies masking plus uniform noise to obtain compute-matched state-of-the-art diffusion language modeling, improved sample quality, and self-correction. The authors note that training language models at overall state-of-the-art scale remains out of scope.

  • Problem

    Autoregressive generation and masked diffusion cannot revise filled tokens, while designing diffusion trajectories with flexible combinations of noise is technically difficult.

  • Method

    GIDD generalizes interpolating discrete diffusion through flexible mixing schedules and derives a diffusion ELBO, with experiments combining masking and uniform noise.

  • Results

    The method achieves compute-matched state-of-the-art diffusion language-modeling performance, while uniform noise improves sample quality and enables self-correction.

  • Takeaways & Limitations

    GIDD provides a flexible discrete-diffusion framework in which suitable noise can give diffusion models the ability to correct their own mistakes.

  • Takeaways & Limitations

    Training language models at the overall state-of-the-art scale is out of scope, and random cropping may overemphasize shorter documents.

Abstract

from arXiv · show

While state-of-the-art language models achieve impressive results through next-token prediction, they have inherent limitations such as the inability to revise already generated tokens. This has prompted exploration of alternative approaches such as discrete diffusion. However, masked diffusion, which has emerged as a popular choice due to its simplicity and effectiveness, reintroduces this inability to revise words. To overcome this, we generalize masked diffusion, deriving a new family of general interpolating discrete diffusion (GIDD) which offers greater flexibility in the design of the noising processes. Leveraging a novel diffusion ELBO, we achieve compute-matched state-of-the-art performance in diffusion language modeling. Exploiting GIDD's flexibility, we explore a hybrid approach combining masking and uniform noise, leading to improved sample quality and unlocking the ability for the model to correct its own mistakes, an area where autoregressive models notoriously have struggled. Code: https://github.com/dvruette/gidd/

1. Introduction

Autoregressive generation is effective but cannot revise earlier errors, while discrete diffusion decomposes generation into revisable denoising steps. GIDD generalizes masked diffusion and combines masking with uniform noise to support self-correction and improved sample quality.

  • Autoregressive models require one model invocation per token and cannot correct earlier errors or incompatibilities.
  • Discrete diffusion progressively removes information from the entire sample and reverses this degradation through a freely chosen number of denoising steps.
  • GIDD generalizes masked diffusion to interpolating discrete processes with flexible time-varying noising distributions, including closed-form transitions and a diffusion ELBO.
  • Mask-only GIDD achieves compute-matched state-of-the-art diffusion language-modeling performance through a reweighted training objective.
  • Combining masking and uniform noise improves sample quality and enables models to iteratively correct their own mistakes.

2. Discrete Diffusion Models

Discrete diffusion models generate samples by degrading discrete data through a Markov chain and learning to reverse that process. Masked diffusion simplifies this process by replacing tokens with [MASK], but filled tokens cannot later be revised and contribute no training loss.

  • Discrete diffusion gradually degrades a data sample through a Markov chain toward an easy-to-sample prior, then learns its reverse transitions.
  • For sequence data, tokens are commonly noised independently, while the model infers the original sequence from the noisy input during inference.
  • Masked diffusion progressively replaces tokens with [MASK] and trains the model to fill in missing tokens from context.
  • Masked diffusion cannot change a token after it is filled, allowing errors or incompatibilities to persist through the final sample.
  • Only masked tokens carry a loss signal because unmasked tokens are noise-free, reducing the effective batch size and potentially slowing convergence.

3. Generalized Interpolating Diffusion

GIDD extends masked diffusion by parameterizing discrete noising with a mixing rate and a time-dependent mixing distribution, enabling broader diffusion trajectories. Its ELBO remains theoretically sound, recovers masked diffusion as a special case, and supports self-correction during sampling.

  • Motivation: Combining masking and uniform noise can corrupt every token, enabling revision of already-unmasked tokens and a more informative correction-oriented training task.
  • Generalized Interpolating Diffusion: GIDD generalizes interpolating diffusion to arbitrary suitable time-varying interpolants, avoiding the need to solve a separate inverse design problem for each noise combination.
  • Forward Process: A GIDD mixing schedule combines α_t, which controls signal-to-noise over time, with π_t, which specifies the distribution toward which data is noised.
  • Forward Process: The framework constructs Markov chains with the desired marginals and derives instantaneous and cumulative transition formulas for arbitrary valid schedules.
  • ELBO: Setting π_t to the masking token makes the GIDD ELBO equivalent to the masked-diffusion ELBO.
  • ELBO: The GIDD ELBO has a global minimum of zero when the model and true noised distributions match, so the schedule does not theoretically limit the best possible model.
  • Sampling: A fixed-point self-correction procedure resamples differing tokens and commits the candidate with the highest model likelihood until convergence.

4. Mixing Schedule

The mixing schedule combines masking with configurable uniform noise while retaining the masked prior, parameterized by p_u to control the midpoint uniform-token fraction. The section also derives and stabilizes ELBO weighting, with clamping and dynamic reweighting addressing extreme weights at low and high noise levels.

  • Mixing Schedule: The schedule retains the masked prior while introducing configurable uniform noise, with p_u specifying the expected maximum uniform-token fraction at t = 1/2.The mixing rate and distribution are designed to reach this target at the midpoint.
  • Validation: GIDD with p_u = 0.0 recovers masked diffusion, whose validation PPL matches MDM closely, while suitable weighting yields significant gains especially when p_u > 0.The final best setting uses dynamic loss weights and weight decay and is called GIDD+.
  • ELBO Weighting: GIDD's ELBO weights grow exponentially near very low and high noise levels, where denoising provides little training signal.The extreme behavior occurs for noise-free, masked, and uniformly corrupted tokens.
  • ELBO Weighting: Clamping the weights to w_max limits the extreme ELBO values, with preliminary experiments selecting w_max = 1.Clamping mostly affects mask and uniform-token weights.
  • ELBO Weighting: Dynamic reweighting preserves relative differences among masked, uniform, and noise-free tokens while controlling the maximum loss weight.The relative weights are determined empirically, and reweighting is equivalent to changing the training-time noise sampling or schedule.

5. Experiments

Experiments evaluate GIDD on language modeling, ablations, generation, self-correction, and zero-shot understanding. Hybrid uniform noise improves generative quality and self-correction, while mask-only GIDD+ remains strongest among diffusion models for validation and benchmark performance.

  • Ablation Study: GIDD with pu = 0.0 matches the MDM training trajectory and final validation PPL of 24.37, while closely matching the compute-matched MDLM baseline.The comparison uses SMALL models trained on OpenWebText.
  • Ablation Study: Uniform noise slightly worsens validation perplexity but improves expressivity, reflecting the greater difficulty of jointly modeling masking and uniform corruption.The highest-noise setting also shows signs of improved scaling behavior.
  • Ablation Study: Custom weighting schemes improve both mask-only and hybrid settings, with dynamic weighting significantly closing the performance gap.The best configuration combines dynamic loss weights with 0.02 weight decay and is called GIDD+.
  • Unconditional Generation: At 32 denoising steps, GIDD+ SMALL with pu = 0.1 reaches generative PPL 387, versus 904 for pu = 0.0 and 1302 for MDM.Uniform-noise models show better generative PPL despite slightly worse validation PPL, especially at low inference compute.
  • Self-Correction: For BASE models, pu = 0.2 lowers generative PPL from 214 to 93.3 and raises self-accuracy from 62.0% to 73.5% as more tokens are replaced.Mask-only models do not improve generative PPL or self-accuracy through the same self-correction procedure.
  • Self-Correction: Self-correction improves clarity, grammaticality, factuality, writing style, and creativity for hybrid-noise models, but deteriorates all categories for mask-only models.The evaluation uses LLM-based grading before and after self-correction.
  • Benchmark Performance: The best diffusion model, GIDD+ with pu = 0.0, surpasses MDM and GPT2-small on zero-shot benchmark accuracy but remains below the retrained Llama baseline.It is trained on only 131B tokens, compared with diffusion models trained for twice as long.

6. Related Work

The paper situates GIDD within discrete diffusion and related approaches that adapt diffusion, flow matching, or denoising objectives to discrete data.

  • Discrete diffusion research has introduced ELBOs for discrete Markov chains, continuous-time extensions, alternative score-matching objectives, and simplified mask-only objectives.
  • Random cropping rather than chunking may over-emphasize short documents and possibly cause overfitting to shorter sequences.The authors describe the observed difference as small and potentially attributable to run-to-run variance.
  • Adaptive denoising orders and adaptive loss weights have been explored to address limitations of mask-only diffusion models.
  • Other approaches adapt Gaussian diffusion in embedding space, discrete flow matching, or Bayesian flow networks to discrete data.
  • Combining masking and uniform noise was previously popularized by BERT for representation learning.

7. Conclusion

The paper introduces GIDD and applies it to diffusion language modeling, emphasizing self-correction, compute-matched perplexity, and adjustable inference-time computation.

  • GIDD extends discrete diffusion with flexible noising choices, while reweighted ELBO training achieves state-of-the-art perplexity for compute-matched diffusion models.
  • With suitable noise, diffusion models can naturally self-correct, addressing a capability with which next-token prediction has struggled.
  • Denoising steps and self-correction iterations provide ways to trade off inference speed and accuracy.
  • The extreme scale required to train overall state-of-the-art language models is outside the scope of this work.

Impact Statement

The paper reports potential consequences for language modeling and diffusion-model use, while noting performance trade-offs associated with uniform noise and model capacity.

  • Increasing uniform noise decreases trained-model likelihood because the model must consider and potentially replace every token in its context.
  • The authors hypothesize that larger models may be less affected by uniform noise because the more difficult training task may expose capacity limits.
  • The capacity-scaling experiment varies model sizes TINY, SMALL, and BASE while holding the training horizon constant and testing uniform-noise levels 0.0, 0.1, and 0.2.
  • Uniform-noise models consistently lag behind mask-only models on downstream performance, although performance improves with model size.

C. Self-Correction Step

Self-correction iteratively identifies and replaces suspected erroneous tokens in denoised samples. Uniform-noise training improves sample quality especially at low denoising-step counts, while mask-only baselines lack comparable ability.

  • Self-Correction Step: Self-correction is a fixed-point iteration applicable to fully or partially denoised samples.
  • Self-Correction Step: The algorithm queries the model for tokens it considers wrong, then replaces one token at a time to avoid reintroducing conflicts.
  • Downstream Performance: Downstream performance improves with model size, but hybrid noise models lag behind mask-only counterparts across scales.
  • Self-Correction Results: GIDD+ with p_u = 0.0 has weak but present self-correction, whereas MDM has none and worsens when self-correction is applied.
  • Sample Quality: Uniform-noise training produces significantly better generative PPL than no-uniform-noise training, especially with few denoising steps.

E. Training Details

The models use a DiT architecture with GPT2 tokenization across three sizes, trained on 512-token contexts with specified optimization and noise-sampling procedures. Additional figures examine self-correction and denoising-step effects.

  • Model configurations: Models use DiT architectures with GPT2 tokenization in TINY, SMALL, and BASE configurations.The configurations contain 28.4M, 92.1M, and 321.2M non-embedding parameters, respectively.
  • Additional analyses: Small-model self-correction performs best at p_u = 0.1 rather than p_u = 0.2, suggesting the preferred uniform-noise ratio depends on model size.The small-model figure reports the same overall trend as the base models.
  • Additional analyses: Generative perplexity decreases monotonically as denoising steps increase, with uniform noise especially benefiting the low-step regime.The figure describes this trend as an overall sample-quality benefit from uniform noise during training.
  • Training setup: Training uses 512-token contexts and batches of 512 for 500k steps, totaling 131B training tokens.Some comparison models are trained for twice as long, totaling 262B training tokens.
  • Optimization: Optimization uses Adam with β = (0.9, 0.99), learning rate 5 · 10^-4, warmup, cosine decay, and gradient clipping at norm 1.0.Weight decay is 0.0 for ablations and 0.02 for the final GIDD+ configuration.
  • Noise and preprocessing: Noise times are sampled uniformly from (10^-4, 1 − 10^-4) using low-discrepancy sampling, with loss weights clipped to 10^4.Long sequences are randomly windowed to 512 tokens, while short sequences are padded to that length.

F. Evaluation Details

Evaluation reserves validation data, estimates likelihoods with the ELBO, and assesses samples through examples, LLM grading, and generative perplexity. The evaluation uses Gemma 2 9B as the primary grading model.

  • Validation and likelihood: Validation uses the final 100k OpenWebText samples, with longer validation examples cropped to random 512-token windows.The reserved portion is approximately 1.25% of the training set.
  • Validation and likelihood: Downstream evaluation uses likelihood-based multiple-choice tasks with 128 evenly spaced time samples to estimate the ELBO.Per-token likelihood includes context and completion but excludes padding.
  • Sample evaluation: Qualitative examples report corrections to grammar, coherence, and context-appropriate word choice for GIDD+ BASE with p_u = 0.2 and τ = 0.1.These examples are presented before and after self-correction.
  • Generative perplexity: The evaluation prompt asks graders to assess clarity and coherence, grammaticality, factuality, writing style, and related sample qualities.Figure 7 provides the prompt used for LLM-based sample-quality evaluation.
  • Generative perplexity: Generative perplexity uses Gemma 2 9B as a more capable grading model intended to better approximate the natural-language distribution than the evaluated models.The paper describes this as a tradeoff between language-modeling accuracy and evaluation efficiency.
  • Sample evaluation: LLM grading evaluates clarity, grammaticality, factuality, writing style, and creativity using scores from 1–10.GPT-4o is instructed to provide a justification and grade for each category in JSON format.

G. Evaluating Generative Perplexity of Diffusion Models

Generative perplexity is used as a proxy for sample quality, but its interpretation depends strongly on the grading model and sampling procedure. GPT2-large can assign better perplexity to samples whose diversity and quality have deteriorated.

  • Metric motivation: Generative perplexity treats high likelihood under a grading model as evidence that generated samples are higher quality or more likely under the data distribution.This interpretation assumes the grading model is a good proxy for the true distribution.
  • Metric limitations: Generative perplexity measured by GPT2-large decreases consistently as the min-p cutoff increases, but this trend does not correlate with subjective quality.The paper identifies this as a major limitation of GPT2-large generative perplexity.
  • Metric limitations: Prior reliance on GPT2-large is problematic because the model exhibits failure modes typical of small language models.These limitations motivate using more capable grading models for meaningful comparisons.
  • Sampling effects: Gumbel-max sampling can implicitly mask probabilities below approximately 10^-8, motivating explicit min-p regularization experiments.The paper reports generative perplexity near 90 around p = 10^-7 for SMALL models under GPT2-large evaluation.
  • Metric limitations: Samples generated with min-p cutoffs can receive lower GPT2-large perplexity while exhibiting drastically worse quality, low diversity, and repetition.The paper reports that larger cutoffs produce suspiciously low perplexity despite these failures.
  • Sampling effects: Unigram entropy can reveal catastrophic diversity collapse, but it remains limited for detecting subtler diversity losses when min-p ≤ 10^-6.The paper notes that Gemma 2 9B is a better, though imperfect, proxy than GPT2-large.

H. Proofs

The appendix proves properties of GIDD transition rates and ELBOs, including their continuous-time forms, weighting identities, and equivalence to masked diffusion in a special case. It also establishes that the ELBO is minimized when model and true marginals match.

  • Transition identities: The appendix derives cumulative transition expressions by recursively unrolling the α and β mixing coefficients.The derivation establishes α_t|s + β_t|s = 1 and that π_t|s is a probability vector.
  • Continuous-time rates: The continuous-time GIDD forward rate matrix is derived by first-order expansion of the transition probabilities as the time increment approaches zero.The rate depends on time derivatives of the mixing functions.
  • ELBO derivation: The continuous-time ELBO is obtained by expanding forward and backward transition terms, substituting them into the discrete-time ELBO, and taking the time-step limit.The derivation starts from a modified continuous-time ELBO that retains constant terms useful for cancellation.
  • ELBO derivation: The GIDD ELBO uses a pointwise IS-divergence weighting and is derived by substituting the GIDD forward and backward rates into the continuous-time ELBO.The appendix defines the weighting function and samples t uniformly over (0, 1).
  • Alternative ELBO: A full IS-divergence yields a looser alternative ELBO that may be easier to implement or lower variance, although the paper does not test those benefits.The full divergence is the sum of nonnegative pointwise divergences.
  • Optimality: The GIDD ELBO reaches its global minimum exactly when the model and true forward marginals agree for every x, t, and z.The proof uses nonnegativity of the KL and IS divergences and handles the zero-weight case through continuity arguments.

I. Unconditional Generation Samples

The samples contrast a mask-only model with a masking-plus-uniform-noise model under self-correction. The hybrid model improves sample quality, whereas self-correction does not help—and can worsen—the mask-only model.

  • The mask-only model achieves the best results on language understanding benchmarks.
  • For the mask-only model, self-correction does not improve sample quality and instead makes it worse in the LLM-evaluation experiment.The examples use a self-correction temperature of τ = 0.1.
  • The masking-plus-uniform-noise model is the best model in terms of sample quality.
  • The hybrid model identifies and corrects mistakes, improving sample quality during self-correction both qualitatively and quantitatively.
Loading 2503.04482v2…