# Masked-diffusion language models, from first principles

> Satyajit Ghana — Head of Engineering @ Inkers Technology
> canonical: https://ai.thesatyajit.com/architectures/masked-diffusion-lm
> architecture: Masked-diffusion language model (diffusion, 2025)
> date: 2026-09-26
> tags: diffusion, language-models, llm, non-autoregressive, explainer

An autoregressive language model, the [Transformer](/architectures/transformer) as almost everyone deploys it, factorises text left to right, $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](https://arxiv.org/abs/2406.07524) (Sahoo et al., 2024) made the objective simple, [LLaDA](https://arxiv.org/abs/2502.09992) (Nie et al., 2025) trained it from scratch at 8 billion parameters, and [Block Diffusion](https://arxiv.org/abs/2503.09573) (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](https://arxiv.org/abs/2107.03006): add one token, `[MASK]`, to the vocabulary, and corrupt a sequence by replacing tokens with it. At noise level $t \in [0, 1]$ each token independently survives with probability $\alpha_t$ and is masked otherwise:

$$
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 $\alpha_0 = 1$, clean, and $\alpha_1 = 0$, all masks. A masked token stays masked as $t$ grows: the state absorbs. LLaDA uses $\alpha_t = 1 - t$, so $t$ is the expected fraction masked; [Dream-Cubed](/articles/dream-cubed), which runs the same process over Minecraft block ids, masks with probability $\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 $t$ to an earlier $s$, a masked position is revealed with probability $(\alpha_s - \alpha_t)/(1 - \alpha_t)$, which is $(t - s)/t$ on the linear schedule, and its value is the clean token. A model of that one distribution, $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 $-\alpha_t'/(1 - \alpha_t)$. On the linear schedule the weight is $1/t$, and the loss is LLaDA's:

$$
\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 $t$ uniformly, mask each token with probability $t$, predict the originals at every masked position at once, and weight by $1/t$. That is [BERT](/architectures/encoder-bert)'s masked-language-model loss ([paper](https://arxiv.org/abs/1810.04805)) 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/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 $tL$ 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](/articles/illada-diffusion-language-model) 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](https://arxiv.org/abs/2508.15487) (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 $i$ still predicts token $i + 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 $L$ masks to the prompt and choose a number of steps $N$. At each step, from $t$ to $s = 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 $sL$ 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 = L$ the model commits one token per pass, in an order it chooses; with $N = L/k$ it commits $k$. 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 $j$, the input at $j$ 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](/articles/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 $2 \times 32 \times 4{,}096 \times 2$ 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 $k$ tokens from one pass samples each from its own marginal, $\prod_i p_\theta(x^i \mid x_t)$, 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 $B$ blocks of $L'$ tokens and factorises across them autoregressively,

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