Large Language Diffusion Models

That the core abilities of LLMs (scaling, in-context learning, instruction following) depend on autoregressive next-token prediction. Nie et al. argue they come from the generative modeling principle, maximum likelihood with a large Transformer and data, rather than from the left-to-right factorization.

The reversal curse: a model trained left to right struggles to produce what precedes a given text (e.g. the line before a line of a poem), even when it knows the forward direction.

Forward: each token is independently replaced by a mask with probability \(t\), from \(t = 0\) (clean text) to \(t = 1\) (all masked). Reverse: a Transformer mask predictor, with no causal mask, sees the partly masked sequence and predicts all masked tokens at once, step by step from \(t = 1\) back to \(t = 0\).

Cross-entropy on the masked positions only, weighted by \(1/t\): \(-\mathbb{E}\big[\frac{1}{t}\sum_i \mathbb{1}[x_t^i = \mathrm{M}] \log p_\theta(x_0^i \mid x_t)\big]\). This is an upper bound on the negative log-likelihood, so LLaDA is a principled generative model trained by (approximate) maximum likelihood.

BERT masks a fixed fraction (15%) of tokens; LLaDA samples the masking ratio \(t\) uniformly from 0 to 1 for each sequence, so it learns to generate from fully masked text. MaskGIT’s heuristic objective lacks the \(1/t\) weight and therefore the link to likelihood.

2.3T tokens with 4096-token sequences, about 0.13M H800 GPU hours, similar to an autoregressive model of that size. LLaDA cannot use a KV cache, so grouped-query attention (which mainly saves KV-cache memory) was replaced by plain multi-head attention with a smaller FFN to keep the parameter count comparable.

The prompt is left intact and only response tokens are masked (each with probability \(t\)); the loss is computed on masked response tokens. Short pairs are padded with end-of-sequence tokens, which the model learns to predict, so it can control the length of its answers. 4.5M pairs were used.

It starts from a fully masked response of chosen length and runs a set number of steps from \(t = 1\) to \(0\). At each step it predicts all masked tokens, then remasks a share of them, those predicted with the lowest confidence, so the confident tokens stay. More steps cost more but give better text.

From \(10^{20}\) to \(10^{23}\) pretraining FLOPs it matched them overall on six tasks, scaling even better on MMLU and GSM8K and closing the gap on weaker tasks like PIQA. An earlier study had estimated masked diffusion needs 16× more compute for the same likelihood, but likelihood is only an indirect measure.

The base model beat LLaMA2 7B on nearly all 15 zero/few-shot tasks and was overall competitive with LLaMA3 8B, stronger in math and Chinese. After SFT alone (no RL alignment) it trailed LLaMA3 8B Instruct slightly but followed instructions and held multi-turn conversations.

On 496 Chinese poem line pairs, it scored about the same forward (51.8%) and backward (45.6%). GPT-4o (82.7% vs 34.3%) and Qwen2.5-7B (75.9% vs 38.0%) were much stronger forward but fell far behind it in reverse, because LLaDA treats all positions alike.

Generation length must be set by hand; the comparison with autoregressive models on identical data stopped below \(10^{23}\) FLOPs; there is no KV cache or other system optimization, and sampling algorithms are still basic; and it has not had RL alignment. It is also smaller and trained on less data than leading LLMs.