~/satyajit

architectures / diffusion

Masked-diffusion language models, from first principles

mdjsonmcp

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!

non-autoregressive · bidirectional · denoise a masked sequence over T steps
[M] maskedfilledjust unmaskedstep 0[M][M][M][M][M][M]all [M]step 1the[M][M]on[M][M]4× [M]step 2thecat[M]on[M]mat2× [M]step Tthecatsatonthemattext ✓unmask a fewre-predictunmask a fewre-predictunmask a fewre-predictbidirectional attentionthecatsatontheeach token attends both left & right (not causal)

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, p(x)=∏ip(xi∣x<i)p(x) = \prod_i p(x^i \mid x^{<i}), 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 t∈[0,1]t \in [0, 1] each token independently survives with probability αt\alpha_t and is masked otherwise:

q(xti∣x0i)=αt 1 ⁣[xti=x0i]+(1−αt) 1 ⁣[xti=M]q(x_t^i \mid x_0^i) = \alpha_t\,\mathbf{1}\!\left[x_t^i = x_0^i\right] + (1 - \alpha_t)\,\mathbf{1}\!\left[x_t^i = \mathrm{M}\right]

with α0=1\alpha_0 = 1, clean, and α1=0\alpha_1 = 0, all masks. A masked token stays masked as tt grows: the state absorbs. LLaDA uses αt=1−t\alpha_t = 1 - t, so tt is the expected fraction masked; Dream-Cubed, which runs the same process over Minecraft block ids, masks with probability sin⁡(πt/2)\sin(\pi t / 2).

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 tt to an earlier ss, a masked position is revealed with probability (αs−αt)/(1−αt)(\alpha_s - \alpha_t)/(1 - \alpha_t), which is (t−s)/t(t - s)/t on the linear schedule, and its value is the clean token. A model of that one distribution, pθ(x0i∣xt)p_\theta(x_0^i \mid x_t), 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 −αt′/(1−αt)-\alpha_t'/(1 - \alpha_t). On the linear schedule the weight is 1/t1/t, and the loss is LLaDA's:

L(θ)=− Et, x0, xt ⁣[1t∑i=1L1 ⁣[xti=M]log⁡pθ ⁣(x0i∣xt)]\mathcal{L}(\theta) = -\,\mathbb{E}_{t,\,x_0,\,x_t}\!\left[\frac{1}{t}\sum_{i=1}^{L}\mathbf{1}\!\left[x_t^{i} = \mathrm{M}\right]\log p_\theta\!\left(x_0^{i} \mid x_t\right)\right]

Draw tt uniformly, mask each token with probability tt, predict the originals at every masked position at once, and weight by 1/t1/t. 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 1/t1/t 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 tLtL 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:

PartShapeParameters
Embedding126,464 × 4,096517,996,544
Attention, per block4 × 4,096 × 4,09667,108,864
MLP, per block3 × 4,096 × 12,288150,994,944
Two RMSNorms, per block2 × 4,0968,192
32 blocks6,979,584,000
Final RMSNorm4,0964,096
Output head4,096 × 126,464517,996,544
Total8,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 ii still predicts token i+1i + 1, 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 LL masks to the prompt and choose a number of steps NN. At each step, from tt to s=t−1/Ns = t - 1/N:

  1. Run the whole sequence, prompt and masks, through the network once, and read a prediction and its probability at every masked position.
  2. Commit enough predictions that about sLsL masks remain. A committed token is never revisited.
  3. 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 N=LN = L the model commits one token per pass, in an order it chooses; with N=L/kN = L/k it commits kk. 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 jj, the input at jj 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.

Block diffusion: left to right again, in blocks

Block Diffusion (BD3-LM) splits the sequence into BB blocks of L′L' tokens and factorises across them autoregressively,

log⁡pθ(x)=∑b=1Blog⁡pθ ⁣(xb∣x<b),\log p_\theta(x) = \sum_{b=1}^{B} \log p_\theta\!\left(x^{b} \mid x^{<b}\right),

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 L′=1L' = 1 it is autoregressive; at L′=LL' = L 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 [0,1][0, 1]. On OpenWebText at 524B tokens the bound tightens as blocks shrink: ≤ 22.27 at L′=16L' = 16, ≤ 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.

share