Masked-diffusion language model · 2025 · Diffusion · 9 min
- diffusion
- non-autoregressive
- bidirectional
- language-models
- llm
- explainer
A 1:49 narrated explainer, drawn in code. Every number and picture in it is this page's own; the sources are below.
› transcript
Hi, I'm Dewdrop! This language model starts from blanks and fills them in, not left to right. Train one Transformer to fill in any fraction of blanks, then write by unmasking the confident ones. Training is fill in the blanks. Mask each token with probability t, and predict every original at once. t covers every ratio, and dividing by t makes the loss a bound on likelihood. LLaDA's eight billion parameters have exactly the shapes of a left-to-right model. That's why Dream could start from an existing one. Append masks to the prompt, and run the whole sequence through once. Every masked slot gets a prediction and a confidence, reading both sides. Commit the most confident. Those are never revisited. Remask the rest, and run again. A left-to-right model reads only the past, which never changes, so it caches keys and values. Here every token reads every other, so each commit moves every key and value. Each step is a full pass. Diffusing the whole answer means full passes and a fixed length. Block diffusion runs left to right over blocks. Finished blocks never change, so they cache exactly. Speed comes from committing several tokens per pass: about three with a confidence threshold, nearly six with self-checking. A Transformer without its causal mask, trained to fill any blanks. Its speed is how many it can commit per pass. A weighted fill-in loss, confident unmasking, and blocks that bring the cache back. Every source is in the full article. I'm Dewdrop. Bye!
Unlike a left-to-right autoregressive model that commits one token at a time, a masked-diffusion LM starts from an all-[M] sequence and, under bidirectional attention, unmasks a few tokens each step and re-predicts the rest — refining the whole sentence in parallel over T steps.
An autoregressive language model, the Transformer as almost everyone deploys it, factorises text left to right, , and pays one forward pass per generated token. A masked-diffusion language model learns something broader: to fill in any subset of blanks in a sequence from everything around them, in both directions. It generates by starting from a row of blanks and filling several per pass. MDLM (Sahoo et al., 2024) made the objective simple, LLaDA (Nie et al., 2025) trained it from scratch at 8 billion parameters, and Block Diffusion (Arriola et al., 2025) put left-to-right order back at the level of blocks. This page builds the model from the noise process up.
Noise for tokens is a mask
Gaussian noise means nothing for a token id. The discrete version that works is the absorbing state of D3PM: add one token, [MASK], to the vocabulary, and corrupt a sequence by replacing tokens with it. At noise level each token independently survives with probability and is masked otherwise:
with , clean, and , all masks. A masked token stays masked as grows: the state absorbs. LLaDA uses , so is the expected fraction masked; Dream-Cubed, which runs the same process over Minecraft block ids, masks with probability .
Two properties drive everything below. The only corruption is erasure, so a visible token is always correct. And the reverse process is the forward one read backward: going from to an earlier , a masked position is revealed with probability , which is on the linear schedule, and its value is the clean token. A model of that one distribution, , the clean token at a masked position given the partly masked sequence, is enough to run the chain.
The objective is a weighted masked-LM loss
MDLM derives the continuous-time evidence bound for this process. By copying visible tokens through and never predicting [MASK], it reduces the bound to a cross-entropy on the masked positions weighted by . On the linear schedule the weight is , and the loss is LLaDA's:
Draw uniformly, mask each token with probability , predict the originals at every masked position at once, and weight by . That is BERT's masked-language-model loss (paper) with two changes. BERT masks a fixed 15% of tokens; this masks every ratio from almost none to all, so the model learns to write from a blank page as well as to fill one gap. And the weight makes the average an upper bound on the negative log-likelihood, so the model is a generative model whose perplexity can be bounded and compared. With about masked positions per sequence, the weight also keeps lightly masked sequences from being outvoted. MDLM shows that the continuous-time bound does not depend on the schedule's functional form, which is why the linear and sine schedules train the same objective.
The bound still trails autoregression. On OpenWebText at 524B training tokens, Block Diffusion's table puts an autoregressive model at a perplexity of 17.54 and MDLM at ≤ 22.98.
The network is a Transformer without the causal mask
There is almost nothing new in the architecture. The denoiser is the Transformer stack with the causal mask removed, so every position attends to every other, and with [MASK] as one more input token. The head predicts a token at every position; the loss reads only the masked ones. LLaDA 8B is 32 blocks of width 4,096 with 32 heads, a SwiGLU MLP 12,288 wide, RMSNorm, RoPE, no biases, and a vocabulary of 126,464 with separate input and output tables, all from its released config:
| Part | Shape | Parameters |
|---|---|---|
| Embedding | 126,464 × 4,096 | 517,996,544 |
| Attention, per block | 4 × 4,096 × 4,096 | 67,108,864 |
| MLP, per block | 3 × 4,096 × 12,288 | 150,994,944 |
| Two RMSNorms, per block | 2 × 4,096 | 8,192 |
| 32 blocks | 6,979,584,000 | |
| Final RMSNorm | 4,096 | 4,096 |
| Output head | 4,096 × 126,464 | 517,996,544 |
| Total | 8,015,581,184 |
That is the 8.02B the iLLaDA write-up lists. LLaDA takes no timestep input: the fraction of masks in the input already says how noisy it is. MDLM's backbone kept a DiT-style time embedding and reports that training without it changes little.
Because the weights have an autoregressive model's shapes, a diffusion LM need not start from scratch. Dream 7B (Ye et al., 2025) starts from Qwen2.5-7B and trains on 580 billion tokens. It keeps the autoregressive shift, so the hidden state at position still predicts token , and it gives each masked token its own noise level according to how much clean context is near it.
Sampling: unmask, keep the confident, repeat
Generation runs the chain backward. Append masks to the prompt and choose a number of steps . At each step, from to :
- Run the whole sequence, prompt and masks, through the network once, and read a prediction and its probability at every masked position.
- Commit enough predictions that about masks remain. A committed token is never revisited.
- Put
[MASK]back everywhere else and go again.
The exact reverse process picks the positions to reveal at random. LLaDA's low-confidence remasking keeps the most confident predictions and re-masks the rest, which leaves the exact process and works better. With the model commits one token per pass, in an order it chooses; with it commits . For its instruct model LLaDA also fills the answer in blocks from left to right, semi-autoregressive remasking, because its fine-tuning data was padded with end-of-sequence tokens that a confidence-first sampler commits early, cutting answers short.
Bidirectional context is the point. A position can be predicted from both sides, which suits infilling and constraints fixed in advance. On reversed poem completion, LLaDA's paper reports 45.6 for its instruct model against 34.3 for GPT-4o, which leads forward, 82.7 to 51.8. The same property gives Dream-Cubed exact inpainting: a block the user fixes is never masked.
Why the KV cache breaks, and what parallel decoding costs
An autoregressive model can cache keys and values because the causal mask makes each position depend only on earlier tokens, which never change. Here attention runs both ways. When a step commits a token at position , the input at changes, and so does the hidden state of every position that attends to it, which is every position, at every layer above the first. LLaDA's paper puts it plainly: "LLaDA is incompatible with KV caching", and it uses plain multi-head attention for that reason. Each step is a full forward pass over prompt and answer.
LLaDA's multiplying weights, the 32 blocks' matrices plus the output head, are 7,497,318,400: about 15.0 GFLOPs per position per pass. With a 512-token prompt and a 512-token answer, one step over 1,024 positions is about 15.4 TFLOPs, and 512 steps of one token each are about 7.9 PFLOPs. An autoregressive model of the same size with a KV cache passes each of the 1,024 positions through once, about 15.4 TFLOPs in all. The compute ratio is the number of steps. The bet is on the other side of the ledger: one pass does 1,024 positions' work for a single read of the weights, so if each step commits many tokens, a decoder limited by memory bandwidth can still finish first.
Two approximations close the gap, each with a cost.
- Approximate caches. Keys and values drift slowly, so they can be reused for a while and refreshed selectively. Fast-dLLM refreshes block by block. Flash-dLLM recomputes only a window of masks plus the most-attended decoded tokens each step, with a kernel that writes keys and values straight into the cache, and reports 148.0 to 210.6 tokens/s for LLaDA-1.5 on one A100. A cached position costs bytes, 512 KiB in bf16, the same as Llama 2 7B's; iLLaDA's 8 key-value heads for 32 query heads make its cache a quarter of full multi-head attention's.
- Parallel decoding. Committing tokens from one pass samples each from its own marginal, , while the true joint couples them: two masked slots that could read "hot dog" or "ice cream" can come out "hot cream". Confidence thresholds limit the damage; Fast-dLLM commits a position only above a threshold, usually 0.9. On GSM8K with 512 generated tokens, the Flash-dLLM paper reports 1.0 token per step for greedy decoding, 2.8 with the threshold and 5.7 with its self-verification.
Block diffusion: left to right again, in blocks
Block Diffusion (BD3-LM) splits the sequence into blocks of tokens and factorises across them autoregressively,
with masked diffusion inside each block, given the finished ones. Attention is block-causal: a token sees its own block in both directions and every earlier block, never a later one. Finished blocks never change, so their keys and values cache exactly, as in an autoregressive model, and generation runs to any length. At it is autoregressive; at it is MDLM. Training runs noisy and clean copies of the sequence in one pass, with a mask that lets each noisy block see the clean blocks before it; the paper also finds the diffusion loss's gradient noisy and samples mask ratios from a clipped range rather than all of . On OpenWebText at 524B tokens the bound tightens as blocks shrink: ≤ 22.27 at , ≤ 21.68 at 8 and ≤ 20.73 at 4, against MDLM's ≤ 22.98 and autoregression's 17.54.
LLaDA's semi-autoregressive sampler and Fast-dLLM's block-wise cache are the inference-time half of the same idea; Block Diffusion trains for it.
What it is good and bad at
Good. Context from both sides, for infilling, editing and fixed constraints. Several tokens per pass. An order of generation the model chooses. One objective from pre-training through fine-tuning.
Bad. No exact KV cache without blocks, so every step costs a pass over the whole sequence. Parallel commits sample marginals, not the joint. The answer length is fixed unless blocks or extension add to it. And the likelihood bound trails autoregression at equal tokens.
What changed since the papers
LLaDA trained on 2.3 trillion tokens and fine-tuned on 4.5 million pairs. iLLaDA scaled the recipe to 12 trillion tokens with 8 key-value heads and reached base-model parity with Qwen2.5-7B, averaging 63.9 to its 63.3 (Dream 7B 61.4, LLaDA 51.1), while trailing it by 10 points after instruction tuning, 67.1 to 77.1. Inference work has moved from exact samplers to approximate caches and verified parallel decoding, and Block Diffusion shows how to train for decoding in blocks. The network is still a Transformer with its causal mask removed; what changed is the objective, the sampler and the cache.