When data is scarce and compute is not, masked diffusion LMs beat autoregression

Diffusion Beats Autoregressive in Data-Constrained Settings

Mihir Prabhudesai, Mengning Wu, Amir Zadeh, Katerina Fragkiadaki, Deepak Pathak

cs.LG, cs.AI, cs.CV, cs.RO

2025-07-22

CMU trains ~200 models on repeated limited data: diffusion's data-reuse half-life is ~494 epochs vs ~31 for AR; at 100M unique tokens validation loss is 3.55 vs 3.71.

What problem this solves

Public high-quality text is flattening while compute is not. Robotics and healthcare never had internet-scale unique data to begin with. Autoregressive LMs train next-token left to right and look cheap in a single pass. Masked diffusion LMs generate by randomly masking and infilling with bidirectional context; prior scaling work said they need about 16× the compute to match AR validation NLL.

That 16× was measured in a one-epoch regime, each token seen once. Scaling compute there grows model size and unique data together, so it mixes compute efficiency with sample efficiency. CMU pins the unique-token budget and repeats the same data, asking which family wins when data, not FLOPs, is the constraint.

Method

Both families share a GPT-2-style Transformer with RoPE, English C4, GPT-2 BPE, and 2048-token sequences. Unique-token budgets are 25M, 50M, and 100M, up to 800 epochs, models from 7M to 2.5B, about 200 runs in total. Hyperparameters follow Muennighoff et al., originally tuned for AR, a mild AR-friendly bias.

AR uses a causal mask and next-token loss. Diffusion samples a mask rate r, independently replaces tokens with [MASK], reconstructs the masked positions with full attention, and weights the loss by 1/r as an ELBO. Because the mask is redrawn every example, the model trains on a large set of conditional prediction orders.

Scaling follows the data-constrained Chinchilla form: repeated data decays exponentially, and a fitted half-life RD marks when further epochs on the same tokens stop helping much.

Results

Near the single-epoch Chinchilla point, diffusion is worse: 10.65 versus 7.07 validation loss at 100M unique tokens. With repetition, AR starts overfitting around 50 epochs. Diffusion does not overfit inside the budget and reaches 3.55 at 500 epochs, below AR's best 3.71. Relative to each family's single-epoch start, that is a 67% drop versus 48%.

Fitted RD is about 494 for diffusion and 31 for AR. Repeated data stays nearly as useful as fresh data for about 4 epochs in AR and about 100 in diffusion. The compute at which diffusion crosses AR follows a power law in unique tokens, exponent about 2.174.

Downstream tracks the loss. Best-by-validation models at 100M unique tokens:

TaskARDiffusion
SciQ58.0568.67
LAMBADA10.9115.19
ARC-Easy35.6337.84
HellaSwag27.3730.24
PiQA60.9460.72

A 2.3B diffusion model trained for 130 epochs on 500M unique tokens under the predicted critical compute (still not converged) scores 79.13 versus 67.82 on SciQ and 22.30 versus 15.07 on LAMBADA. AR stays slightly ahead on PiQA.

Controls: attention dropout and zeroing token attention on AR do not stop overfitting. Training AR on N fixed sequence permutations does; at N=16 the 100-epoch loss approaches diffusion. Diverse conditional orders from random masking are the main sample-efficiency source.

Why it matters

The claim that diffusion LMs need 16× compute mixed sample efficiency with compute efficiency. When unique data runs out first, or never existed at internet scale, diffusion reuses the same tokens for an order of magnitude more epochs. The practical line in the paper is blunt: if you are compute-constrained, use AR; if you are data-constrained, use diffusion. Hybrids that interpolate order diversity against per-step supervision density are the next step the authors flag.

Limitations

Raw validation losses are not directly comparable across families because the entropy constants E0 differ; the mechanism section drops that term before comparing. At 100M unique tokens the models are small, several tasks sit near chance, and PiQA still favors AR. Hyperparameters were tuned for AR, so diffusion may have headroom, but that caveat does not rescue AR's overfitting. The 500M diffusion run stopped for compute and never showed an overfit boundary. Results are language modeling; transfer to robotics or clinical sequences is an argument, not a measurement. Multi-step diffusion sampling cost at inference is barely discussed.

Terms

Source

What people are saying

Related papers

All paper explainers