Source-linked AI summary
Scaling Beyond Masked Diffusion Language Models
Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang, Justin Deschenaux, Jingyu Liu, John Thickstun, Ante Jukic
TL;DR
The paper asks whether Masked diffusion is truly dominant and whether perplexity suffices to compare diffusion families. It studies three discrete diffusion families with compute-matched scaling and sampling-centric evaluations, finding that sampling efficiency and downstream performance can favor models with worse likelihood scaling. At 1.7B parameters, Uniform-state diffusion remains competitive on likelihood-based evaluations and outperforms the other compared models on GSM8K after fine-tuning.
Problem
Perplexity is incomplete for comparing diffusion families because it does not reflect inference-time behavior and different diffusion processes induce different likelihood bounds.
Method
The paper conducts compute-matched scaling studies of Masked, Uniform-state, and interpolating diffusion models and combines likelihood-based evaluation with sampling-centric speed-quality analysis.
Results
At 1.7B parameters, Uniform-state diffusion remains competitive on likelihood-based evaluations and outperforms AR, MDLM, and Eso-LM on GSM8K after supervised fine-tuning despite worse validation perplexity.
Takeaways & Limitations
Perplexity is informative within a diffusion family but can mislead across families, so practical comparison should jointly consider likelihood, sampling efficiency, and downstream performance.
Takeaways & Limitations
The study leaves the behavior and real-world advantages of diffusion language models at much larger scales for future investigation.
Abstract
from arXiv · showhide
Diffusion language models are a promising alternative to autoregressive models due to their potential for faster generation. Among discrete diffusion approaches, Masked diffusion currently dominates, largely driven by strong perplexity on language modeling benchmarks. In this work, we present the first scaling law study of uniform-state and interpolating discrete diffusion methods. We also show that Masked diffusion models can be made approximately 12% more FLOPs-efficient when trained with a simple cross-entropy objective. We find that perplexity is informative within a diffusion family but can be misleading across families, where models with worse likelihood scaling may be preferable due to faster and more practical sampling, as reflected by the speed-quality Pareto frontier. These results challenge the view that Masked diffusion is categorically the future of diffusion language modeling and that perplexity alone suffices for cross-algorithm comparison. Scaling all methods to 1.7B parameters, we show that uniform-state diffusion remains competitive on likelihood-based benchmarks and outperforms autoregressive and Masked diffusion models on GSM8K, despite worse validation perplexity. We provide the code, model checkpoints, and video tutorials on the project page: http://s-sahoo.github.io/scaling-dllms
1. Introduction
Diffusion language models may offer faster generation than autoregressive models, but perplexity does not fully capture their inference-time trade-offs. This study compares three diffusion families through scaling and speed-quality analyses, finding that sampling efficiency and downstream performance can favor models with worse likelihood scaling.
- Motivation: Diffusion models generate by iteratively refining sequences in parallel, potentially enabling faster decoding than token-by-token autoregressive generation.Discrete diffusion methods have narrowed much of the perplexity gap to autoregressive models.
- Motivation: Perplexity is incomplete for cross-family comparison because different diffusion processes induce different likelihood bounds and inference-time behaviors.Uniform-state diffusion can excel in few-step generation, while Masked diffusion can benefit from additional sampling compute.
- Approach: The study systematically compares Masked, Uniform-state, and interpolating diffusion through compute-matched scaling and sampling-centric speed-quality evaluations.The evaluation combines likelihood-based metrics with throughput and sample quality across sampling steps.
- Contributions: 12%: a low-variance objective brings MDLM’s compute multiplier to within about 12% of autoregressive models while shifting compute-optimal checkpoints toward smaller models.The shift toward smaller checkpoints reduces inference cost.
- Contributions: The study presents the first systematic IsoFLOP scaling analysis of Uniform-state diffusion and interpolating diffusion models.The representative models are Duo and Eso-LM, compared alongside MDLM.
- Results: Models with worse perplexity scaling can dominate the speed-quality frontier, and at 1.7B parameters Duo outperforms AR, MDLM, and Eso-LM on GSM8K after supervised fine-tuning.Duo remains competitive on likelihood-based downstream evaluations despite worse validation perplexity.
2. Background
Discrete diffusion models corrupt clean sequences toward a simple prior and learn reverse processes for generation. This section contrasts Masked, Uniform-state, and interpolating diffusion through their objectives, sampling behavior, and inference trade-offs.
- Discrete diffusion: Discrete diffusion transforms clean sequences into a simple prior through forward corruption, then learns a reverse process that reconstructs data.The forward process is typically position-wise independent, with a noise schedule controlling the corruption level.
- Interpolating diffusion: Interpolating diffusion combines clean-token and categorical-prior components through a monotone noise schedule, with α0 controlling the interpolation between MDLM-like and AR-like behavior.Eso-LM uses a decoder-only denoiser with causal attention on shuffled sequences and supports KV caching without sacrificing parallel generation.
- Masked diffusion: Masked diffusion uses a mask prior, progressively replacing tokens with [MASK], while its denoiser uses bidirectional attention.Its reverse posterior is convenient because the mask is absorbing, but bidirectional attention prevents KV caching.
- Uniform-state diffusion: Uniform-state diffusion moves tokens toward a uniform categorical prior rather than [MASK], allowing reverse sampling to revise token values repeatedly.This self-correction supports few-step sampling and guided generation, while time-conditioned denoising improves validation perplexity and sample quality.
3. Scaling Laws
Compute-matched scaling shows that diffusion families have comparable scaling exponents to autoregressive models but substantially different constant-factor compute gaps. Likelihood alone misses sampling advantages: Duo and Eso-LM can occupy different regions of the speed-quality frontier through few-step generation and KV caching.
- Scaling-law setup: Matched-FLOP IsoFLOP studies compare AR, MDLM, Duo, and Eso-LM using shared compute budgets, model grids, and scaling-law fits.The study uses exact training FLOPs and applies the same compute-optimality procedure across model families.
- Masked diffusion: MDLM requires ≈16× more compute than AR to match validation loss under true ELBO training, while its compute-optimal checkpoints use ≈2× fewer parameters.MDLM’s validation-loss scaling slope is comparable to AR, but its constant-factor compute gap is much larger.
- Masked diffusion: 12%: low-variance MDLM training reduces the compute multiplier from ≈16× to ≈14× relative to AR and shifts compute-optimal models toward smaller sizes.The smaller models also reduce sampling cost at inference time.
- Uniform-state diffusion: Duo requires ≈23× more compute than AR to match perplexity, yet supports faster few-step inference through self-correction and uses ≈2× fewer compute-optimal parameters.These results illustrate a likelihood-efficiency trade-off within the speed-quality comparison.
- Interpolating diffusion: Eso-LM requires ≈32× more compute than AR to match perplexity but supports KV caching and uses ≈2× fewer parameters at matched compute.Its full-diffusion configuration is evaluated with α0 = 1.
- Speed-quality trade-off: Duo dominates throughput ranges [200,400] ∪ [600,∞], while Eso-LM dominates [400,600]; thus perplexity alone can mis-rank practical methods.The comparison varies sampling steps and evaluates sample quality with Generative Perplexity.
4. Scaling to the Billion-Parameter Regime
The paper scales AR and three diffusion families to 1.7B parameters under a shared large-scale training and evaluation setup. AR leads overall likelihood-based performance, while Duo leads the diffusion models on several tasks and surpasses AR, MDLM, and Eso-LM on GSM8K after supervised fine-tuning.
- Experimental setup: All models use 1.7B parameters, a 2048-token context, and 2.1T training tokens under a common large-scale pretraining protocol.The models are trained on Nemotron-Pre-Training-Dataset mixtures using 64 H100 GPUs.
- Likelihood-based benchmarks: AR achieves the best overall performance on the likelihood-based downstream benchmark suite.Among diffusion models, MDLM leads ARC-e, BoolQ, and SIQA, while Duo leads OBQA, PIQA, and RACE.
- Maths and reasoning: Duo outperforms AR, MDLM, and Eso-LM on GSM8K after supervised fine-tuning.The evaluation uses the augmented GSM8K training set and reports throughput on the full test set with batch size 1.
- Maths and reasoning: Confidence-based sampling effectively collapses to left-to-right generation, so the GSM8K comparison generates one token at a time from every model.In this memory-bound setting, AR latency is expected to be comparable to diffusion models despite KV caching.
5. Discussion and Conclusion
The study argues that perplexity-based scaling is insufficient for comparing diffusion families because sampling mechanisms create practical trade-offs. Across 1.7B models, Duo can outperform stronger-likelihood methods on math and reasoning after fine-tuning, while larger-scale studies remain needed to test broader advantages.
- Discussion: Compute-matched scaling across Masked, Uniform-state, and interpolating diffusion shows comparable scaling exponents to AR but family-dependent constant-factor gaps.MDLM has the strongest likelihood scaling, whereas Duo and Eso-LM require substantially more compute to match AR perplexity.
- Discussion: Perplexity is informative within a diffusion family but can mislead across families with different diffusion processes and sampling capabilities.Few-step generation in Duo and KV caching in Eso-LM are practical advantages not captured by validation perplexity alone.
- Discussion: At 1.7B parameters, AR remains strongest on likelihood-based metrics, while Duo can outperform AR and other diffusion models on math and reasoning after supervised fine-tuning.This result persists despite Duo’s weaker perplexity scaling.
- Future directions: Larger-scale diffusion studies are needed to clarify emergent behaviors, long-range reasoning, and whether distinct generative mechanisms yield consistent real-world advantages.The paper frames these questions as future work rather than established outcomes.
Impact Statement
The paper aims to advance machine learning through diffusion language-model research. It notes potential applications to synthetic text and biological-sequence design, along with associated benefits and risks.
- Impact: The work concerns machine learning advances involving synthetic text generation and possible biological-sequence design applications.The paper identifies both potential benefits and risks for these applications.
A. Additional Experiments
The additional experiments compare diffusion-model behavior across sampling, scaling, and downstream evaluation settings. They highlight the low-variance objective’s effect on MDLM scaling and report GSM8K accuracy and model configurations for the broader study.
- Gen. PPL and entropy are evaluated against time discretization for MDLM, Eso-LM, and Duo sampling configurations.
- MDLM trained with the low-variance training loss yields compute-optimal models with fewer parameters than MDLM trained with the ELBO.
- GSM8K accuracy is reported after supervised fine-tuning with a cosine schedule and maximal and minimal learning rates.
- The scaling-law study uses specified Transformer configurations for AR, MDLM, Duo, and Eso-LM.