Source-linked AI summary
Scaling Behavior of Discrete Diffusion Language Models
Dimitri von Rütte, Janis Fluri, Omead Pooladzandi, Bernhard Schölkopf, Thomas Hofmann, Antonio Orvieto
TL;DR
The paper addresses the limited understanding of how discrete diffusion language models scale relative to autoregressive models. It studies masked, uniform, and hybrid diffusion with SNR-based interpolation and tuned batch-size and learning-rate choices, finding favorable compute-bound scaling and especially token-efficient compute-optimal scaling for uniform diffusion.
Problem
The scaling behavior of DLMs across noise types and relative to ALMs remains insufficiently explored, despite the importance of scaling laws for large-scale training.
Method
The paper compares masked, uniform, and hybrid diffusion using SNR-based interpolation while systematically tuning batch size, learning rate, model size, and training duration.
Results
DLM variants scale comparably in compute-bound settings, while uniform diffusion requires more parameters and less data for compute-optimal training and matches predicted trends at 3B and 10B scales.
Takeaways & Limitations
DLMs, particularly uniform diffusion, are promising alternatives to ALMs when training compute is abundant relative to available data.
Takeaways & Limitations
The phenomenological power-law fit for batch-size behavior is not expected to remain accurate near the irreducible loss.
Abstract
from arXiv · showhide
Modern LLM pre-training consumes vast amounts of compute and training data, making the scaling behavior, or scaling laws, of different models a key distinguishing factor. Discrete diffusion language models (DLMs) have been proposed as an alternative to autoregressive language models (ALMs). However, their scaling behavior has not yet been fully explored, with prior work suggesting that they require more data and compute to match the performance of ALMs. We study the scaling behavior of DLMs on different noise types by smoothly interpolating between masked and uniform diffusion while paying close attention to crucial hyperparameters such as batch size and learning rate. Our experiments reveal that the scaling behavior of DLMs strongly depends on the noise type and is considerably different from ALMs. While all noise types converge to similar loss values in compute-bound scaling, we find that uniform diffusion requires more parameters and less data for compute-efficient training compared to masked diffusion, making them a promising candidate in data-bound settings. We scale our uniform diffusion model up to 10B parameters trained for $10^{22}$ FLOPs, confirming the predicted scaling behavior and making it the largest publicly known uniform diffusion model to date.
1 INTRODUCTION
DLMs offer parallel generation and token revision, but their scaling behavior depends strongly on diffusion noise and differs from ALMs. The paper studies this behavior with interpolated noise types and careful hyperparameter tuning, finding favorable large-scale scaling for DLMs.
- DLMs can generate multiple tokens per denoising step and revise every token throughout generation, addressing two limitations of ALMs.The number of denoising steps T is independent of sequence length N.
- The paper refines prior scaling-law analyses by tuning batch size and learning rate while comparing compute- and token-bound scaling across masked, uniform, and hybrid diffusion.Its hybrid family smoothly interpolates between masked and uniform diffusion using an SNR-defined transition.
- Uniform diffusion scales more favorably with increased compute, requiring more parameters and less data for compute-optimal training than masked diffusion.The result is especially relevant when compute grows faster than available pre-training data.
- All examined noise types have largely similar compute-bound scaling, while large DLMs closely follow predicted trends and can be competitive with ALMs at scale.The study validates predictions with 3B masked and uniform models and a 10B uniform model trained at substantially larger compute budgets.
2 METHOD
The method formulates discrete diffusion as a unified categorical corruption-and-denoising process and expresses its likelihood bound through log-SNR. A hybrid mixing distribution then interpolates between masked and uniform noise.
- 2.1 DISCRETE DIFFUSION MODELS: Discrete diffusion trains a denoising model to reverse a categorical Markov corruption process that transforms clean discrete data into a pure-noise prior.The forward process progressively adds noise, while the backward model removes it.
- 2.2 GENERALIZED INTERPOLATING DISCRETE DIFFUSION: GIDD unifies masked and uniform diffusion by interpolating categorical transitions between the initial data state and a noise distribution.Specialized noise processes arise under appropriate mixing schedules.
- 2.2 GENERALIZED INTERPOLATING DISCRETE DIFFUSION: The authors reformulate GIDD in terms of signal-to-noise ratio, yielding a simpler and more flexible likelihood bound aligned with continuous-state diffusion theory.The ELBO becomes an importance-sampling procedure over log-SNR values and the associated forward noising process.
- 2.4 A UNIVERSAL HYBRID MIXING DISTRIBUTION: A log-SNR-controlled mixing distribution smoothly transitions from masked to uniform noise, with parameters governing the transition point and speed.For positive transition-shape parameter a, varying b approaches pure masking or pure uniform noise.
- Independent per-token noise levels provide an anisotropic alternative to the global noise level typically used in sequence diffusion models.The cited approach was proposed to stabilize autoregressive rollouts and later accelerate DLM inference.
3 ESTIMATING SCALING LAWS
The paper rederives DLM scaling laws without fixing batch size, learning rate, or irreducible loss assumptions. It uses CompleteP and systematic sweeps to characterize compute- and token-optimal training choices.
- 3 ESTIMATING SCALING LAWS: Prior DLM scaling analyses fixed learning rate and batch size and fit power laws without an irreducible-loss offset, so this paper rederives scaling laws while dropping those assumptions.The laws are estimated directly on negative ELBO, an upper bound on negative log-likelihood.
- 3 ESTIMATING SCALING LAWS: CompleteP is adopted to transfer optimal learning rates across model width and depth, with base initialization and learning-rate values calibrated through smaller-model sweeps.The transferred values work well, while the base learning rate still requires adjustment for batch size.
- 3 ESTIMATING SCALING LAWS: Learning-rate annealing is omitted from the main scaling procedure because it permits a single run to cover training horizons, while ablations preserve optimal hyperparameters.The reported performance difference from omitting annealing is a constant factor.
- 3 ESTIMATING SCALING LAWS: The compute- and token-optimal batch size scales almost linearly with training-token count and shows no dependence on target loss in these experiments.No saturation appears even at batch sizes of 10^6 tokens, suggesting a higher critical batch size than reported for ALMs.
4 EXPERIMENTS AND RESULTS
Experiments show that diffusion scaling depends on noise type and carefully tuned training hyperparameters. Uniform diffusion scales more favorably in token-constrained settings, while noise types converge in compute-constrained settings; batch-size and step-count relationships also impose practical limits.
- Training setup and hyperparameters: The optimal batch size depends primarily on training tokens and follows an almost linear power law, while the optimal learning rate follows a power law in batch size.The batch-size scaling exponent is 0.82 and the learning-rate exponent is 0.34.
- Scaling laws and compute-optimal frontier: All noise types approximately converge under compute-constrained scaling, while uniform diffusion out-scales the alternatives in token-constrained settings.The extrapolations remain accurate up to 50× larger compute budgets than those used for fitting.
- Scaling laws and compute-optimal frontier: Uniform noise has the best diffusion scaling behavior, requiring more parameters and less data for compute-optimal training while slightly improving compute-bound loss scaling.The reported scaling laws also call for faster model-size scaling than autoregressive laws from prior literature.
- Ablating the effect of learning rate annealing: Learning-rate annealing leaves optimal batch size and learning rate largely unchanged while improving final loss by 2.45% ± 0.138% across settings.This constant-factor improvement extrapolates to the largest runs.
- Relation between batch size and step count: For a fixed model size and target loss, effective minimum batch size and step count constrain attainable loss even when the learning rate is tuned.The phenomenological power-law description is not expected to remain accurate near the irreducible loss.
5 CONCLUSION
The study finds that DLMs scale competitively with autoregressive models, with uniform diffusion especially favorable when data is constrained. DLMs therefore remain promising alternatives with parallel generation, token revision, and test-time compute scaling.
- All examined DLM variants scale comparatively in compute-bound settings, while uniform diffusion scales most favorably in token-bound environments.
- Uniform diffusion performs below masked diffusion at small scales, but its shrinking likelihood gap and more parameter-heavy compute-optimal scaling support greater capacity at scale.
- DLMs can enable parallel generation, revise previously generated tokens, and provide straightforward ways to scale test-time compute.
- The authors conclude that DLMs, particularly uniform diffusion, are promising candidates for next-generation language models.
A.1 DOWNSTREAM EVALUATIONS
Downstream performance broadly correlates with training ELBO, with uniform diffusion tending to perform better on reasoning-heavy tasks and masked diffusion on knowledge-heavy tasks. Uniform diffusion also supports adaptive denoising steps beyond the completion length.
- Uniform diffusion performs comparatively better on reasoning-heavy benchmarks, whereas masked diffusion performs slightly better on knowledge-heavy benchmarks.
- Confidence-based sampling gives masked diffusion a noticeable boost over classic ancestral sampling on GSM8k.
- Uniform diffusion can use more adaptive denoising steps than there are completion tokens, unlike masked diffusion without remasking.
A.2 RELATION BETWEEN BATCH SIZE AND STEP COUNT
Batch size and optimizer step count exhibit a hyperbolic iso-loss relationship, so equivalent losses can be reached through different batch-size and step-count combinations.
- Equivalent losses follow a hyperbolic relationship between batch size and optimizer step count.
- The experiments use longer runs for smaller batches and shorter runs from batch size 256 onward to reduce computational cost.
A.4 OPTIMAL HYPERPARAMETERS
Optimal hyperparameters depend primarily on training tokens, batch size, training horizon, and noise type rather than model size. Uniform noise favors somewhat larger batches, while learning-rate behavior follows predictable batch- and horizon-dependent trends.
- Optimal batch size: More uniform noise consistently favors larger optimal batches, although the difference is small and unlikely to affect practical power-of-two batch-size choices.
- Optimal hyperparameters: Available observations suggest that optimal batch size and learning rate are more likely than not unaffected by model size.
- Optimal batch size: A critical batch size likely exists above which increasing batch size yields diminishing returns, but it is not reached in these experiments.
- Optimal batch size: Optimal batch size scales almost linearly with training tokens and shows no strong dependence on model size.
- Optimal learning rate: Optimal learning rate increases with larger batches and decreases with longer training horizons, while remaining largely independent of model size.
A.5 TOKENIZER
The tokenizer converts between bytes and tokens using measured bytes-per-token statistics, while related passages describe learning-rate scaling and the greater difficulty of uniform diffusion.
- 4.2278 B/tok and 0.23653 tok/B determine the tokenizer’s conversion between negative log-likelihood in nats and bits per byte.
- η∗ appears to follow a power law in training steps for each batch size, with model sizes distinguished by markers.
- Uniform diffusion is strictly more difficult than masked diffusion because it must account for which tokens are noisy as well as predict their replacements.
C PROOF OF PROPOSITION 1
The proof reformulates the diffusion objective using signal-to-noise ratio and shows that the ELBO is invariant to the induced SNR distribution and forward-process marginals. A separate figure passage contrasts scaling fits based on two FLOP-estimation methods.
- SNR is defined as α/(1−α), with log-SNR λ = log(α/(1−α)), using the mixing rate to measure preserved signal.
- Changing variables from time t to log-SNR λ rewrites the GIDD ELBO in terms of the log-SNR differential.
- The ELBO is invariant to both the SNR distribution and forward-process marginals, which serve to approximate the integral through importance sampling.
- Fitting on interpolated data produces tighter confidence bounds and better scaling exponents than the alternative fit shown for the FLOP-estimation methods.
D SCALING COEFFICIENTS
The scaling analysis compares FLOP estimates, smoothing choices, and power-law forms, finding that interpolation generally improves fit quality and that an intercept is usually unnecessary. The reported compute-optimal trends differ because the classic FLOP approximation underestimates executed computation.
- Scaling coefficients are reported across FLOP approximations, smoothing variants, and noise types, with figures showing fitted trends and confidence intervals.
- The classic M = 6P approximation systematically underestimates training FLOPs, producing scaling laws that differ significantly from the more accurate estimate.
- Raw optimal observations can deviate substantially from compute-optimal trends because the model-size sweep contains few unique sizes.
- Interpolated squared-fit values generally yield better R2 fits than raw observations because smoothing reduces the brittleness caused by sparse model sizes.
- Fits with an intercept usually give E ≈ 0, supporting the assumption f(C) = AC^α without an irreducible term in this setting.