PaperlayerView exact revision on arXiv

Source-linked AI summary

Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data

Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, Chongxuan Li

arXiv:2406.03736v4cs.LGcs.CL

TL;DR

Absorbing discrete diffusion requires concrete-score estimation, whose theoretical structure and practical scaling were not fully understood. The paper reparameterizes that score as clean-data conditional probabilities, proposes time-independent RADD networks with caching, and unifies diffusion with AO-ARMs. RADD converges faster than SEDD while achieving state-of-the-art performance among diffusion models on five zero-shot language-modeling benchmarks at GPT-2 scale.

  • Problem

    Absorbing diffusion relies on concrete scores and scaling, while the theoretical relationship between these quantities and alternative autoregressive models is insufficiently understood.

  • Method

    The paper expresses the concrete score as clean-data conditional probabilities times an analytic time-dependent scalar, then removes time conditioning in RADD and relates diffusion objectives to AO-ARMs.

  • Results

    RADD achieves state-of-the-art performance among diffusion models on five zero-shot language-modeling benchmarks at GPT-2 scale, measured by perplexity.

  • Takeaways & Limitations

    Time-independent conditional modeling provides a simpler diffusion parameterization with caching-based sampling acceleration and alternative training and likelihood objectives.

  • Takeaways & Limitations

    The sequence model assumes independent dimensions, replacing the exponentially large joint transition process with separate one-dimensional diffusion processes.

Abstract

from arXiv · show

Discrete diffusion models with absorbing processes have shown promise in language modeling. The key quantities to be estimated are the ratios between the marginal probabilities of two transitive states at all timesteps, called the concrete score. In this paper, we reveal that the concrete score in absorbing diffusion can be expressed as conditional probabilities of clean data, multiplied by a time-dependent scalar in an analytic form. Motivated by this finding, we propose reparameterized absorbing discrete diffusion (RADD), a dedicated diffusion model without time-condition that characterizes the time-independent conditional probabilities. Besides its simplicity, RADD can reduce the number of function evaluations (NFEs) by caching the output of the time-independent network when the noisy sample remains unchanged in a sampling interval, which enables sampling acceleration. Built upon the new perspective of conditional distributions, we further unify absorbing discrete diffusion and any-order autoregressive models (AO-ARMs), showing that the upper bound on the negative log-likelihood for the diffusion model can be interpreted as an expected negative log-likelihood for AO-ARMs. Further, our RADD models achieve SOTA performance among diffusion models on 5 zero-shot language modeling benchmarks (measured by perplexity) at the GPT-2 scale. Our code is available at https://github.com/ML-GSAI/RADD.

1 INTRODUCTION

The paper develops RADD by reinterpreting absorbing diffusion's concrete score as clean-data conditional probabilities, enabling simpler parameterization, faster sampling, and a theoretical connection to AO-ARMs. RADD achieves state-of-the-art performance among discrete diffusion models on five zero-shot language-modeling benchmarks at GPT-2 scale.

  • Auto-regressive models dominate language modeling but sequential sampling is inefficient and limits reasoning in nonsequential orders.
  • Absorbing diffusion estimates concrete scores for masked-token denoising, while prior scaling tricks lacked a complete theoretical explanation.
  • The concrete score equals a clean-data conditional probability multiplied by an analytic time-dependent scalar, explaining the scaling trick as a reparameterization.
  • Absorbing diffusion and AO-ARMs have equivalent training objectives, allowing diffusion likelihood bounds to be interpreted as expected AO-ARM negative log-likelihoods.
  • RADD removes time conditions so its network models time-independent conditional probabilities, simplifying the model and training target.
  • RADD caches network outputs when noisy samples remain unchanged, reducing function evaluations and accelerating sampling.
  • RADD achieves state-of-the-art performance among discrete diffusion models on five zero-shot language-modeling benchmarks at GPT-2 scale, measured by perplexity.

2 BACKGROUND

The background introduces continuous-time discrete diffusion, its reverse process and concrete score, absorbing transitions, and any-order autoregressive modeling. It also describes independent-dimension approximations used to make sequence diffusion tractable.

  • 2.1 CONTINUOUS TIME DISCRETE DIFFUSION MODEL: Continuous-time discrete diffusion evolves categorical states through a time-dependent transition-rate matrix and reverses the process using a learned reverse rate.
  • 2.1 CONTINUOUS TIME DISCRETE DIFFUSION MODEL: With Qt parameterized as σ(t)Q, the transition solution has a matrix-exponential form, permitting one-step sampling between arbitrary times.
  • 2.1 CONTINUOUS TIME DISCRETE DIFFUSION MODEL: Absorbing diffusion replaces tokens with a mask state [M], and its matrix structure is used for empirical performance and sampling acceleration.
  • 2.1 CONTINUOUS TIME DISCRETE DIFFUSION MODEL: Because the forward transition rates are known, reverse-process learning can focus on estimating the concrete score.
  • 2.1 CONTINUOUS TIME DISCRETE DIFFUSION MODEL: Denoising score entropy trains the score network, but its loss is only an upper bound on negative log-likelihood with an unknown gap.
  • 2.1 CONTINUOUS TIME DISCRETE DIFFUSION MODEL: For sequences, independent dimensions share the transition matrix, avoiding the exponential state space of joint sequence transitions.
  • 2.2 ANY-ORDER AUTOREGRESSIVE MODELS: AO-ARMs model joint distributions autoregressively across all variable orders using weight sharing, masks, and an expected negative log-likelihood objective.

3 REPARAMETERIZED ABSORBING DISCRETE DIFFUSION

This section reparameterizes the absorbing-diffusion concrete score as time-independent clean-data conditionals, enabling simpler modeling and cached sampling. It also establishes equivalence between absorbing diffusion and any-order autoregressive objectives.

  • 3.1 Reparameterizing the Concrete Score: Theorem 1 expresses the absorbing-diffusion concrete score as a clean-data conditional distribution multiplied by an analytic time-dependent term.This reparameterization explains the scaling trick and separates time dependence from the learned conditional distribution.
  • 3.1 Reparameterizing the Concrete Score: RADD removes time conditioning and uses a time-independent network to model conditional probabilities of clean data directly.The implementation minimally modifies SEDD by removing the time-conditioning input and using a softmax output.
  • 3.2 Efficient Samplers to Reduce NFEs: At most d sampling intervals can change a d-token sequence, so RADD caches the network output whenever the noisy sample remains unchanged.Because SEDD is time-conditioned, it cannot use the same caching strategy to reduce NFEs.
  • 3.2 Efficient Samplers to Reduce NFEs: RADD’s caching strategy theoretically reduces expected NFEs, especially at larger sampling-step counts, and experiments show faster convergence trends.The expected-NFE curve agrees with experiments, while the model achieves slightly better perplexity with less sampling time.
  • 3.3 Unifying Absorbing Discrete Diffusion and Any-Order Autoregressive Models: The absorbing diffusion objective is equivalent to the any-order autoregressive objective when the final total noise level σ̄(T) → +∞.The diffusion denoising loss can therefore be interpreted through any-order autoregressive training objectives.

4 EXPERIMENTS

Experiments evaluate RADD’s training variants, sampling efficiency, and zero-shot language-modeling perplexity against absorbing-diffusion baselines. RADD improves efficiency through caching and achieves strong performance across five benchmarks, while the study notes scope limitations.

  • Experimental setup: RADD uses a time-independent network with fewer parameters than SEDD and is trained under four alternative objectives.The objectives are DSE, t-DCE, λ-DCE, and any-order autoregressive loss.
  • Accelerated generation: Caching makes RADD more efficient than SEDD during sampling while preserving comparable sample quality.The reported efficiency gain is attributed to reducing evaluations when the noisy sample remains unchanged.
  • Zero-shot perplexity: Scaled SEDD outperforms unscaled SEDD, while RADD-DSE without time-conditioning outperforms time-conditioned SEDD-Scale under similar parameter counts.These ablations support the proposed explanation of the scaling trick and time-independent parameterization.
  • Zero-shot perplexity: All RADD losses outperform SEDD on average across the five datasets, despite theoretical objective equivalence holding only in expectation.Finite-data gradient-estimation differences can lead the models toward distinct local optima.
  • Zero-shot perplexity: RADD achieves state-of-the-art performance among diffusion models on five zero-shot language-modeling benchmarks at the GPT-2 scale.The benchmarks include LAMBADA, WikiText2, PTB, WikiText103, and 1 Billion Words.
  • Limitations: The evaluation primarily uses GPT-2-scale models, and the model generates only full-length outputs rather than variable-length sequences.The authors identify scaling and output-length flexibility as future-work boundaries.

B PROOF OF THEOREM 1

The proof derives analytic conditional and joint distributions for absorbing diffusion, then divides them to factorize the concrete score. This establishes a time-dependent analytic factor multiplied by a clean-data conditional distribution.

  • Proof strategy: The proof exploits a continuous-time Markov chain with an absorbing state and independence across diffusion dimensions.This independence simplifies the conditional and joint-distribution analysis.
  • Conditional distribution: The one-dimensional forward conditional distribution permits only retaining the original token or transitioning it to the absorbing mask state.The transition probabilities are analytic functions of the integrated noise schedule.
  • Joint distribution: The multidimensional noisy-data distribution factors into an analytic time-only term and the clean-data joint distribution over unmasked tokens.The clean-data component is independent of time.
  • Concrete score: Theorem 1 expresses the concrete score as a time-independent clean-data conditional distribution multiplied by an analytic time-dependent term.The factorization follows by dividing the derived joint distributions.

C PROOF OF THEOREM 2

The proof transforms absorbing-diffusion losses through masking-probability reparameterization and permutation-based sampling. It concludes that the diffusion objective is equivalent to an any-order autoregressive objective under complete masking at the terminal time.

  • Theorem 2: Theorem 2 states that absorbing discrete diffusion and any-order autoregressive objectives are equivalent when the final total noise level approaches infinity.This condition ensures every token is masked with probability one.
  • Loss transformation: The diffusion objective is first simplified from DSE to t-DCE, then reparameterized using λ(t), the forward-process probability that a token is masked.The change of variables uses λ(t) = 1 −e^-σ̄(t).
  • Loss transformation: The t-DCE loss becomes λ-DCE after analytically simplifying the transformed integral.This produces a loss expressed over masking probability rather than diffusion time.
  • Connection to AO-ARMs: The λ-DCE expectation is reorganized by the number of masked tokens and uniformly sampled masking configurations.Permutation notation then connects the resulting expression to any-order autoregressive training.
  • Connection to AO-ARMs: The resulting expression is mathematically equivalent to the any-order autoregressive objective and supports efficient parallel optimization.The paper also states that diffusion models can sample like AO-ARMs and vice versa.
  • Reverse-process boundary: For reverse sampling, exact multidimensional conditionals are inaccessible when multiple tokens are masked, so existing samplers assume within-interval dimension independence.The assumption enables parallel updates but does not represent the full joint conditional directly.

D.4 EQUIVALENCE OF TWEEDIE τ-LEAPING AND EULER METHOD UNDER LOG-LINEAR NOISE SCHEDULE

The Euler and Tweedie τ-leaping samplers make the same token-update decisions under a log-linear noise schedule. Their equivalence follows because both use the same unmasking probability in that schedule.

  • Sampling behavior: Both samplers keep unmasked tokens unchanged and independently decide whether masked tokens are unmasked within each interval.When unmasking occurs, the token value is sampled from a conditional distribution.
  • Sampling behavior: The samplers differ generally only in the analytic form of their unmasking probability ψ(t, s).The two expressions are not equivalent for arbitrary noise schedules.
  • Equivalence condition: Under a log-linear noise schedule, the Euler and Tweedie τ-leaping unmasking probabilities simplify to the same expression.This equality establishes equivalence of the two sampling methods under that schedule.
  • NFE analysis: The number of function evaluations is treated as a random variable determined by the sampling method and noise schedule.Token changes across intervals determine the evaluation count.
  • NFE analysis: For unconditional generation, independent token changes make the number of changed dimensions in each interval binomially distributed.For conditional generation, the corresponding binomial parameter uses the generating sequence length.

E DISCUSSION FOR MEAN PARAMETERIZATION AND RADD

The paper argues that mean prediction in absorbing diffusion is equivalent to learning time-independent conditional distributions of clean data, and that related sampling parameterizations are equivalent.

  • Mean parameterization: The reverse distribution’s analytic form motivates setting s = 0 to derive the mean-parameterization relationship.
  • Mean parameterization: Mean prediction is equivalent to learning conditional distributions on clean data.The resulting mean parameterization should be time-independent and corresponds to the reparameterized cθ(xt).
  • RADD: The time-independent mean parameterization is denoted µθ(xt) and is equivalent to the reparameterized cθ(xt).Empirical results from prior work are cited as supporting this theory.
  • Sampling equivalence: The Tweedie τ-leaping sampler for score parameterization is equivalent to sampling with mean prediction.This equivalence is established by comparing the corresponding sampling formulations.
  • Sampling equivalence: αt is the probability that a token remains unmasked at time t, equal to e−¯σ(t) for score parameterization.This identification makes Eq. (E.2) and Eq. (D.8) equivalent.

F COMPARISON WITH CHEN ET AL. (2024)

RADD differs from Chen et al. (2024) in when it updates tokens and in its applicable diffusion settings, while its caching strategy can reduce function evaluations for absorbing models.

  • Discrete-time models: Chen et al. (2024) pre-samples token-specific change times and updates only at those predetermined discrete steps.For absorbing diffusion, each token changes once, giving NFEs = |Tchange| ≤ min(d, T).
  • Continuous-time models: In continuous-time models, Chen et al. (2024) applies updates at d sorted change points, resulting in NFEs = d.Reducing NFEs below d was not investigated for that continuous-time method.
  • Comparison: RADD updates tokens only when they change by using a time-independent parameterization, unlike Chen et al.’s predetermined time-point updates.The difference leads to different applicable scenarios.
  • Applicable scenarios: Chen et al. (2024) applies to absorbing and multinomial diffusion, whereas RADD is specifically designed for absorbing diffusion.RADD’s cache strategy can be applied straightforwardly in continuous time to reduce NFEs below d because its input is time-independent.
  • AO-ARMs: AO-ARMs model all univariate conditionals with a weight-sharing network and predict d − k conditionals in parallel for a condition set of size k.They are trained over uniformly sampled variable orders and sampled autoregressively from masked initialization.

H COMPARISON TO PRIOR WORKS CONCERNING EQUIVALENCE DISCUSSION

The paper contrasts earlier qualitative or direct equivalence arguments with its rigorous alternative proof based on time-independent reparameterization, extending the analysis across four losses.

  • Prior work: Earlier work described the absorbing-diffusion and AO-ARM loss relationship qualitatively without a rigorous proof.Austin et al. characterized the losses as similar but not exactly identical.
  • Equivalence proof: The paper rigorously proves the equivalence between absorbing discrete diffusion and AO-ARMs using a continuous-time framework and time-independent parameterization.
  • Derivation: Its derivation first reduces the ELBO to LT DSE(x0), then transforms LT DSE(x0) into LAO through step-by-step substitutions.
  • Equivalence proof: The paper provides an alternative proof that leverages the time-independent properties of its reparameterization formulation.
  • Loss analysis: The analysis extends the equivalence result to four distinct losses and is accompanied by experiments in Tables 1 and 2.The paper states that this comprehensive loss-function exploration was not conducted in prior work.

I ALGORITHMS FOR TRAINING AND SAMPLING

The paper specifies training and sampling procedures for AO-ARMs and discrete diffusion, including a simplified time-independent RADD architecture and implementation details for caching and numerical sampling.

  • AO-ARM training: AO-ARM training samples data, a variable order π, and a uniformly distributed cardinality l, then masks variables at or beyond l.
  • AO-ARM sampling: AO-ARM sampling initializes all positions with [M], samples an order π, and iteratively generates variables according to that order.
  • Discrete diffusion: Discrete diffusion training and sampling require a network cθ, noise schedule σ, a time range, and a step size Δt.The unconditional sampler initializes xT with all [M] tokens and caches cθ(xt).
  • RADD architecture: RADD uses a decoder-only transformer with rotary positional encoding, removes time-conditioning components, and applies a final softmax to produce a conditional distribution.
  • Training details: The experiments use EMA 0.9999, gradient-norm clipping at 1, and warmup during the first 2500 iterations.
  • Training details: The authors did not perform hyperparameter search because of limited computational resources, so further tuning may improve performance.
  • Sampling details: For unconditional generation, a log-linear noise schedule makes Euler and Tweedie τ-leaping equivalent for RADD, while SEDD perplexity is measured with Tweedie τ-leaping.

J.4 FURTHER EVALUATION OF GENERATIVE PERPLEXITY

The further evaluation compares RADD and SEDD across sampling precision, steps, batch sizes, and autoregressive decomposition orders. RADD maintains comparable perplexity while its caching strategy consistently reduces sampling time, and forward decomposition performs best among tested orders.

  • SEDD vs RADD: RADD and SEDD exhibit comparable perplexity under both fp64 and fp32 for the same number of sampling steps.Under fp64, perplexity converges at large steps, whereas fp32 continues decreasing because precision errors act like annealing.
  • SEDD vs RADD: Up to 2.5 to 3 times faster sampling is achieved by RADD with large sampling steps while maintaining similar perplexity to SEDD.Even with 1024 sampling steps, the cache strategy enables about 1.5 times acceleration.
  • Batch size and hardware utilization: RADD sampling time decreases substantially from batch size 1 to 4, then shows minimal additional reduction beyond batch size 4.This pattern indicates that hardware utilization is nearly maximized at batch size 4.
  • SEDD vs RADD: RADD consistently outperforms SEDD in sampling speed across identical batch sizes and timesteps.The caching mechanism reduces redundant computations, with especially strong gains at larger batch sizes and higher timesteps.
  • Sampling as any-order autoregressive models: Standard AO-ARM sampling closely matches converged fp64 τ-leaping perplexity in large steps, while forward decomposition performs best among tested orders.The random decomposition averages perplexity across randomly generated orders.
  • Extended zero-shot evaluation: Training RADD-small for 1000k iterations can slightly improve performance compared with the 400k-iteration main-text setup.The extended evaluation attributes this improvement to increased exposure to training data.

K.1 ADDITIONAL SAMPLES

The appendix presents additional unconditional and conditional samples from RADD-λ-DCE small and medium models. Conditional examples identify prompt tokens in blue, while the samples use 1024 steps under a log-linear noise schedule.

  • Additional samples: All additional samples are generated with 1024 steps under a log-linear noise schedule.
  • Additional samples: Additional unconditional samples are shown for RADD-λ-DCE small and medium models.These examples correspond to Figures 2 and 4.
  • Additional samples: Additional conditional samples are shown for RADD-λ-DCE small and medium models, with prompt tokens marked in blue.These examples correspond to Figures 3 and 5.
Loading 2406.03736v4…