Scaling Beyond Masked Diffusion Language Models

Paper: 2602.15014 Authors: Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang, Justin Deschenaux, Jingyu Liu, John Thickstun, Ante Jukic Categories: cs.LG, cs.CL

Abstract

This paper presents the first comprehensive scaling law study of discrete diffusion language models, comparing uniform-state, interpolating, and masked diffusion approaches. The authors demonstrate that masked diffusion models can achieve approximately 12% better FLOPs efficiency through a simple cross-entropy training objective. Critically, they find that perplexity, while informative within a diffusion family, can be misleading when comparing across different diffusion methods. Models with worse likelihood scaling may actually be preferable due to faster and more practical sampling characteristics. At 1.7B parameters, uniform-state diffusion remains competitive on likelihood benchmarks and outperforms both autoregressive and masked diffusion models on GSM8K mathematical reasoning tasks, despite having worse validation perplexity.

Key Contributions

• First scaling law analysis for uniform-state and interpolating discrete diffusion methods, providing empirical evidence across multiple model sizes
• Discovery that masked diffusion can be made ~12% more FLOPs-efficient using cross-entropy objectives instead of traditional diffusion losses
• Demonstration that perplexity is not a reliable cross-family metric, with speed-quality Pareto frontiers revealing different trade-offs
• Evidence that uniform-state diffusion at 1.7B parameters outperforms masked diffusion and autoregressive models on GSM8K reasoning tasks
• Open-source release of code, model checkpoints, and video tutorials for reproducibility

Methodology and Experimental Design

The research employs a systematic scaling approach, training models across multiple parameter counts to establish scaling laws for different diffusion families. The authors compare three primary discrete diffusion approaches: masked diffusion (which masks tokens during the forward process), uniform-state diffusion (which transitions tokens to a uniform distribution), and interpolating diffusion (which gradually interpolates between data and noise distributions).

For masked diffusion, the authors introduce a modified training objective using cross-entropy loss rather than the standard diffusion objective. This modification simplifies the training process while maintaining model quality, leading to the observed 12% FLOPs efficiency gain. The training procedure involves:

$\mathcal{L}_{\text{CE}} = -\mathbb{E}_{x_0, t, x_t} [\log p_\theta(x_0 | x_t, t)]$

where $x_0$ represents the original data, $x_t$ is the noised version at timestep $t$, and $p_\theta$ is the model’s predicted distribution.

The evaluation framework encompasses both traditional language modeling metrics (perplexity on validation sets) and downstream task performance, including mathematical reasoning (GSM8K), common sense reasoning, and other benchmarks. This dual evaluation strategy proves crucial for revealing the limitations of perplexity as a sole metric.

Results and Analysis

The scaling law experiments reveal distinct patterns for each diffusion family. Masked diffusion shows the best perplexity scaling, which explains its current dominance in the field. However, when examining the speed-quality Pareto frontier—plotting generation speed against output quality—uniform-state diffusion demonstrates competitive performance with significantly faster sampling.

At the 1.7B parameter scale, the results challenge conventional wisdom:

• Uniform-state diffusion achieves competitive perplexity despite simpler dynamics
• On GSM8K mathematical reasoning, uniform-state diffusion outperforms both masked diffusion and autoregressive baselines
• The cross-entropy trained masked diffusion models match or exceed standard masked diffusion performance while requiring fewer FLOPs

The GSM8K results are particularly striking, suggesting that perplexity and reasoning capability may not be as tightly coupled as previously assumed. The uniform-state model’s superior performance on mathematical reasoning, despite worse validation perplexity, indicates that different diffusion processes may learn different inductive biases that favor certain types of tasks.

The speed-quality analysis reveals that uniform-state diffusion can generate samples 2-3× faster than masked diffusion for comparable quality levels. This advantage stems from the simpler transition dynamics and more efficient sampling procedures inherent to uniform-state approaches.

Implications for Diffusion Language Modeling

This work has significant implications for the future direction of diffusion language models. The field has largely converged on masked diffusion due to its strong perplexity numbers, but this research suggests that such convergence may be premature. The findings indicate that:

Metric Selection Matters: Perplexity serves as a useful metric within a model family but fails as a universal comparison tool across different diffusion approaches. Researchers should adopt multi-faceted evaluation frameworks that include sampling efficiency, downstream task performance, and computational requirements.

Efficiency Considerations: The 12% FLOPs improvement for masked diffusion through cross-entropy training demonstrates that even established methods have room for optimization. This finding suggests that other diffusion families might similarly benefit from objective function redesign.

Task-Specific Advantages: The superior GSM8K performance of uniform-state diffusion hints at task-specific advantages for different diffusion processes. Future work should investigate which diffusion families excel at which types of reasoning or generation tasks.

Practical Deployment: For real-world applications where generation speed matters, uniform-state diffusion’s faster sampling may outweigh its perplexity disadvantage. The speed-quality Pareto frontier provides a more practical lens for model selection than perplexity alone.

Takeaways

1. Perplexity is informative within diffusion families but misleading across them—speed-quality trade-offs matter more for practical deployment
2. Masked diffusion can be made ~12% more FLOPs-efficient using simple cross-entropy objectives without sacrificing quality
3. Uniform-state diffusion remains competitive at scale (1.7B parameters) and outperforms masked diffusion on mathematical reasoning tasks
4. The speed-quality Pareto frontier reveals that models with worse likelihood scaling may be preferable due to faster sampling
5. The field should move beyond perplexity-only evaluation and adopt multi-dimensional assessment frameworks for diffusion language models
6. Different diffusion processes may encode different inductive biases, making them suitable for different downstream applications
7. Open-source release of models and code enables further research into scaling laws and diffusion method comparison