A Visual Introduction to Masked Diffusion Language Models

Autoregressive — One Token at a Time

Once

there

was

small

cat

named

Milo

lived

tall

red

house

big

green

hill

Every

day

sat

the

door

and

saw

the

birds

fly

One

day

bird

came

and

sang

song

for

him

Block Diffusion — Multiple Tokens Simultaneously

Once[MASK]
there[MASK]
was[MASK]
a[MASK]

small[MASK]

cat[MASK]

named[MASK]

Milo[MASK]

.[MASK]

He[MASK]

lived[MASK]

in[MASK]

a[MASK]

tall[MASK]

red[MASK]

house[MASK]

on[MASK]

a[MASK]

big[MASK]

green[MASK]

hill[MASK]

.[MASK]

Every[MASK]

day[MASK]

he[MASK]

sat[MASK]

by[MASK]

the[MASK]

door[MASK]

and[MASK]

saw[MASK]

the[MASK]

birds[MASK]

fly[MASK]

by[MASK]

.[MASK]

One[MASK]

day[MASK]

a[MASK]

bird[MASK]

came[MASK]

and[MASK]

sang[MASK]

a[MASK]

song[MASK]

for[MASK]

him[MASK]

.[MASK]

Figure 1. Autoregressive vs. block diffusion, driven by one shared reveal clock. Top: tokens generated strictly left-to-right (AR); pending slots are a subtle dashed underline since AR has no [MASK] sentinel. Bottom: tokens grouped into fixed-size blocks; within a block, tokens un-mask in random order, and blocks complete strictly left-to-right (block diffusion). Pending slots show [MASK]. The bottom row finishes about 4× sooner than the top row on the same shared clock, and both then rest before the animation auto-cycles.

Introduction

Autoregressive language models are the workhorse of modern AI. Practically every major LLM in production today, from GPT to Claude to Gemini, is autoregressive at its core: they generate text one token at a time, left to right, each token conditioned on everything that came before. A remarkable amount of engineering effort has gone into making this loop fast. Techniques like kernel optimization, KV caching, and speculative decoding [?] shave cost off every generated token, and modern GPU architectures and custom accelerators like TPUs are co-designed with this particular workload in mind, all to squeeze more tokens per second out of the same fundamentally sequential recipe. But at some point it's worth asking: what comes next? What are the alternatives to autoregression?

Masked diffusion language models [?] offer one such alternative. Where an autoregressive model builds a sequence left to right one token at a time, a masked diffusion model starts from a fully-masked sequence and progressively unmasks tokens in any order, potentially several at once. The most immediately appealing consequence is parallelism: because multiple tokens can be unmasked in a single reverse step, generation is no longer bottlenecked on producing one token at a time. There are other reasons to be interested. The model can revisit and revise tokens it already committed to, a technique called remasking [?] . And there is evidence that a non-autoregressive unmasking order is a better inductive bias for problems whose dependency structure isn't left-to-right [?] . But the parallelism argument is the one that promises a clean systems win, and it is what first drew attention to this paradigm.

But applied naively, they aren't actually faster. Unmasking many tokens per step sounds like a straightforward speedup, and it would be, if inference cost were just "number of forward passes." It isn't. Masked diffusion models leverage bidirectional attention, which is strictly more powerful than the causal attention that autoregressive models use, but at the cost of violating core assumptions required for KV caching. KV caching is the ability to reuse computed keys and values from previous passes, made possible by the one-way flow of information from left to right. The rest of this post is about that tradeoff: what assumptions enable efficient AR inference, why masked diffusion violates these assumptions, and how the field is putting the pieces back together.

Masked Language Modeling

Before we get to masked diffusion, it is worth revisiting the idea it builds on: masked language modeling. The task is simple to state. Take a sentence, replace a few of its tokens with a special [MASK] symbol, and ask a neural network to predict the original tokens at the masked positions. For each masked position, the network emits a full categorical distribution over the vocabulary, and we take the argmax (or sample from it) to fill the slot in.

This is exactly the pretraining objective of BERT [?] . For each training sentence $\mathbf{x}$ , BERT samples a random subset of positions $M \subset \{1, \ldots, L\}$ — canonically 15% of them — and replaces each of those tokens with [MASK] to produce a corrupted sequence $\tilde{\mathbf{x}}$ . A bidirectional transformer $\mathbf{x}_\theta$ then predicts a categorical distribution over the vocabulary at every masked position, and the loss is cross-entropy on those positions only:

\mathcal{L}_{\mathrm{MLM}} = -\, \mathbb{E}_{\mathbf{x},\, M} \left[ \sum_{\ell \in M} \log p_\theta(\mathbf{x}^\ell \mid \tilde{\mathbf{x}}) \right]

Geometrically, cross-entropy is a force pulling the model's predicted distribution $\mathbf{x}_\theta^\ell$ toward the one-hot target $\mathbf{x}^\ell$ . As training progresses, whatever mass the model is currently spreading over near-miss vocabulary words gets pulled onto the single correct word.

Model Prediction

p_\theta(\mathbf{x}^\ell \mid \tilde{\mathbf{x}})

Target

\mathbf{x}^\ell

This loss is fantastic for representation learning: fill-in-the-blank is a hard enough task that the transformer has to build genuinely useful features to solve it, which is why BERT is such a strong starting point for downstream classifiers and encoders. But BERT was never designed to generate text, and it turns out that a 15%-MLM model can't do it, for a subtle reason.

How would we go about generating a novel sequence with a model trained like this? The natural thing to try is to hand it an input that is entirely [MASK] and ask it to fill in every position at once:

But this input is drawn from a distribution the model has never encountered during training: every training example BERT ever saw had 85% of its tokens still visible as context. With no real tokens to condition on, the model has nothing to anchor its predictions to — the outputs are essentially untrained behavior. And the problem is not unique to the 100% case. To generate a sequence from scratch we need a model that behaves sensibly all the way from “fully masked” to “almost done”, which is a continuum of masking rates that BERT's fixed 15% schedule simply never visits.

This is the gap masked diffusion closes. The core observation is small but transformative: instead of training on a single fixed masking rate, train the same architecture on a whole family of masking rates — every rate from 0% to 100%. The resulting model can be handed an input at any level of corruption and produce a sensible prediction, and that turns it from a representation learner into a generative model.

Masked Diffusion Models

A different way to see the two paradigms: freeze the decoder after each step and stack the intermediate lines from top to bottom. Row 0 is the empty (fully-masked) sequence; row k is the sequence after k tokens have been generated; the last row is the fully-decoded line. The autoregressive block reveals slots strictly left-to-right; the masked diffusion block reveals slots in a random order — you can see this by tracking which column gains a word from one row to the next.

Autoregressive

The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 

Masked Diffusion

The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 
The  little  cat  sat  on  the  warm  mat 

Figure 9. Decoding trajectories for autoregressive (left) and masked-diffusion (right) generation on a short line. Each row is one decoding step; rows fade in top-to-bottom on a shared clock. Pending slots are colored rectangles sized to the eventual token so the columns stay aligned across rows.

Turning “train on every masking rate” into an actual algorithm takes a bit of scaffolding. We need to say how to corrupt a sequence to any target rate (the forward process), how to invert one step of that corruption at inference time (the reverse process), and how to weight the different rates against each other when we train (the training loss). Once those three pieces are in place, generation is just the iterative unmasking loop from the previous section — run in reverse from a fully-masked start. We follow the MDLM formulation throughout [?] .

Masking as the Forward Process

Let $x \in \mathcal{V}^L$ be a clean sequence of length $L$ over vocabulary $\mathcal{V}$ , and let $\mathbf{m}$ denote the special absorbing mask token. A monotone schedule $\alpha_t : [0, 1] \to [0, 1]$ with $\alpha_0 = 1$ (all clean) and $\alpha_1 = 0$ (fully masked) defines the noising marginal:

We write these per-position distributions as $\mathrm{Cat}(x;\, p_1, \ldots, p_K)$ , a categorical distribution over $K$ outcomes with probabilities $p_k$ summing to one.

q(\mathbf{z}_t \mid \mathbf{x}) = \prod_{\ell=1}^L \mathrm{Cat}\!\left(\mathbf{z}_t^\ell;\; \alpha_t\, \mathbf{x}^\ell + (1 - \alpha_t)\, \mathbf{m}\right)

Each position is independently either kept as its original token with probability $\alpha_t$ or absorbed into $\mathbf{m}$ with probability $1 - \alpha_t$ . Once a position is masked, it stays masked at all later times — that is what makes the process absorbing.

The absorbing forward process. A short sequence is unrolled vertically, one row per forward-process timestep. Row 0 is fully clean; each subsequent row corrupts one or two more positions to the [MASK] sentinel. Once a position is masked it stays masked in every row below — the mask token is absorbing, and each column has a single point of no return.

Reverse Process

The model $\mathbf{x}_\theta(\mathbf{z}_t, t)$ is trained to predict a categorical over the clean token at each masked position (the SUBS parameterization: mask probability is fixed at zero, and unmasked positions carry over unchanged). Plugging $\mathbf{x}_\theta$ into the true posterior of the forward process gives the reverse step $p_\theta(\mathbf{z}_s \mid \mathbf{z}_t)$ for $s < t$ :

q(\mathbf{z}_s^\ell \mid \mathbf{z}_t^\ell, \mathbf{x}^\ell) = \begin{cases} \delta(\mathbf{z}_s^\ell = \mathbf{z}_t^\ell) & \text{if } \mathbf{z}_t^\ell \ne \mathbf{m}, \\[4pt] \mathrm{Cat}\!\left(\mathbf{z}_s^\ell;\; \dfrac{(1 - \alpha_s)\, \mathbf{m} + (\alpha_s - \alpha_t)\, \mathbf{x}^\ell}{1 - \alpha_t}\right) & \text{if } \mathbf{z}_t^\ell = \mathbf{m}. \end{cases}

Unmasked positions carry over deterministically. Masked positions either stay masked with weight $(1 - \alpha_s) / (1 - \alpha_t)$ or unmask to the predicted clean token with weight $(\alpha_s - \alpha_t) / (1 - \alpha_t)$ .

Training Loss

With the forward and reverse processes fixed, the training objective writes itself. We want the model to be accurate at every masking rate visited by the schedule, so we take an expectation of the per-rate MLM loss from the previous section over $t$ as well as over the corruption $q(\mathbf{z}_t \mid \mathbf{x})$ at that rate. The continuous-time NELBO [?] collapses (after the SUBS parameterization) to exactly this — a weighted cross-entropy on the masked positions:

\mathcal{L}_{\mathrm{MDLM}} = -\, \mathbb{E}_t\, \mathbb{E}_{q(\mathbf{z}_t \mid \mathbf{x})} \left[ \dfrac{\alpha'_t}{1 - \alpha_t} \smash{\underbrace{\sum_{\ell \in M_t} \log p_\theta(\mathbf{x}^\ell \mid \mathbf{z}_t)}_{\mathcal{L}_{\mathrm{MLM}}(\mathbf{z}_t)}} \right]

Here $M_t = \{\ell : \mathbf{z}_t^\ell = \mathbf{m}\}$ is the set of positions masked at time $t$ , so the underbraced term is literally the MLM loss from the previous section, evaluated on the corrupted sequence $\mathbf{z}_t$ . The MDLM objective is that same MLM loss, averaged over all rates $t$ in the schedule and weighted by $\alpha'_t / (1 - \alpha_t)$ . Unmasked positions contribute zero, since carry-over guarantees a perfect prediction there.

Relation to Continuous Diffusion

In continuous diffusion, the forward process transforms data into a Gaussian, destroying information, and the reverse process aims to generate new data by reversing this information-destroying process.

t=0

t=1

In masked diffusion, the forward process destroys information through discrete masking, and likewise generation is framed as reversing this forward masking process.

Once there was a small cat named Milo . He lived in a tall red house on a big green hill . Every day he sat by the door and saw the birds fly by . One day a bird came and sang a song for him .

t=0

t=1

Accelerated Decoding

Applied naively, masked diffusion is not actually faster than autoregression. The naïve promise — "unmask many tokens at once, skip the token-by-token loop" — only holds if we ignore how autoregressive inference is actually made cheap in practice. AR is fast not because it does less work per token in the abstract, but because the systems layer underneath it exploits a very specific property of the model: past hidden states never depend on future tokens. A masked diffusion model doesn't have that property, and once we look at what that costs on real hardware, the raw parallelism advantage evaporates. This section walks through the pieces: the causal-attention mask underlying AR transformers, how KV caching exploits it to make inference cheap, why masked diffusion breaks the caching assumption, and finally how block diffusion recovers a version of caching without giving up parallel within-block decoding.

Causal Attention

Autoregressive transformers restrict attention to a specific pattern. The query at position t can only attend to keys at positions $\le t$ — every row of the attention matrix is masked to a lower-triangular shape. That means information flows strictly one way, from past to future: a token at position t knows about tokens 1…t, and no token to its right ever influences it. This is causal attention, and it is the property that makes efficient AR inference possible.

Figure 6. Causal self-attention (left) versus bidirectional attention (right). Each row is a query; each column is a key. Under the causal mask, every query attends only to past-and-current keys and the matrix is lower-triangular. Under bidirectional attention every query can see every key, so the full square fills in — there is no past-only structure to cache.

KV Caching

The one-way flow of information is what lets autoregressive inference reuse work across generation steps. Because past keys and values never depend on tokens that come later, they never need to be recomputed when new tokens are generated — they are safe to cache. The keys and values computed for each past token are stored in memory once and read back on every subsequent step, instead of being recomputed from scratch. The figure below stages the pieces before we walk through the mechanism: on the left, a High Bandwidth Memory block containing one KV cache per transformer layer; on the right, the transformer stack itself, with input tokens up top and the residual stream flowing down through three self-attention layers. The dashed cell in each row marks the token about to be generated.

Figure 5. Autoregressive transformer forward pass with key-value caching. During the prefill phase, the model runs all prompt tokens through the network in parallel. Each attention layer computes key/value tensors for every prompt position and writes them into the layer's KV cache in High Bandwidth Memory. During the decoding phase, tokens are generated one at a time: at each new position, the model reads the cached keys and values for every previous position (instead of recomputing them from scratch) and appends its own newly computed K/V to the cache. Caching turns per-step work from

O(L)

forward passes into

O(1)

, which is what makes autoregressive inference cheap enough to be practical.

Why Masked Diffusion Breaks KV Caching

Masked diffusion breaks the causal-attention assumption. Because the encoder is bidirectional, every position's hidden representation is a function of the whole current context. When a single [MASK] commits to a word at one reverse step, the hidden representation at every other position — including positions that stay masked — changes as well. Any keys or values a cache stored while position j was masked are out of date the moment j commits. The figure below shows two forward passes of the same transformer on an eight-token sentence: one before the commit, one after. Between the two passes, the colored square that stands in for the hidden state at each position shifts — every position, not just the one that was unmasked.

Figure 7. Two forward passes of the same masked transformer. In pass 1 the model sees an eight-token sentence with two [MASK]s and produces a hidden state at each position (the colored square beneath each token). In pass 2, one mask has been committed to little and the sequence is fed through again — and now the colored square at every position is different, including the one that stayed masked. Because attention is bidirectional, every hidden state depends on the whole context, so committing a single token changes the representation of every other token.

It's worth distinguishing masked diffusion from any-order autoregression here, since they can look similar at a glance and the difference matters for caching. Any-order autoregression fixes an unmasking order ahead of time and then decodes strictly along that order — each position conditions only on positions earlier in the chosen order. That model class is strictly weaker than full masked diffusion: it can't adapt the unmasking order to the sample, can't unmask multiple positions in parallel per step, and can't revise a committed token. What it can do is cache KVs along the fixed order, because along that order the model is still causal. Full masked diffusion gives up that fixed order to buy adaptive parallel decoding, and pays for it with a broken cache.

Block Diffusion

Block diffusion is the compromise between the two extremes. Tokens are grouped into fixed-size blocks; within a block, tokens are generated in a random order — masked-diffusion style, with each pending slot shown as a gray rectangle. Across blocks, generation is strictly left-to-right: block k+1 only starts once block k is fully revealed. The result is diffusion-style parallel filling inside each block, autoregressive ordering between blocks — and, because once a block is fully unmasked its KVs are frozen with respect to any later block, KV caching works along the block-level order. Within-block inference cost is $O(B^2)$ per reverse step, and $B \ll L$ , so total inference cost drops close to AR asymptotically while keeping parallel decoding and adaptive order inside each block.

Block Diffusion — One Block per Step

Once[MASK]
there[MASK]
was[MASK]
a[MASK]

small[MASK]

cat[MASK]

named[MASK]

Milo[MASK]

.[MASK]

He[MASK]

lived[MASK]

in[MASK]

a[MASK]

tall[MASK]

red[MASK]

house[MASK]

on[MASK]

a[MASK]

big[MASK]

green[MASK]

hill[MASK]

.[MASK]

Every[MASK]

day[MASK]

he[MASK]

sat[MASK]

by[MASK]

the[MASK]

door[MASK]

and[MASK]

saw[MASK]

the[MASK]

birds[MASK]

fly[MASK]

by[MASK]

.[MASK]

One[MASK]

day[MASK]

a[MASK]

bird[MASK]

came[MASK]

and[MASK]

sang[MASK]

a[MASK]

song[MASK]

for[MASK]

him[MASK]

.[MASK]

Figure 8. Block-diffusion generation on the same paragraph. Tokens are grouped into blocks of 4, laid out in the same 12-column grid as Figure 1. Every token in a block un-masks simultaneously in one decoding step; blocks themselves complete strictly left-to-right (autoregressive at the block level). A light-orange highlight follows the block currently being decoded.

Are Masked Diffusion Language Models Actually Faster?

Introduction

Masked Language Modeling

Masked Diffusion Models

Masking as the Forward Process

Reverse Process

Training Loss

Relation to Continuous Diffusion

Accelerated Decoding

Causal Attention

KV Caching

Why Masked Diffusion Breaks KV Caching

Block Diffusion

Landscape and Frontier

Acknowledgements

References

Comments