Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models
Marianne Arriola, Aaron Gokaslan, Justin T. Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, Volodymyr Kuleshov
ICLR 2025 Oral
cs.LG, cs.AI
2025-03-13
BD3-LMs do AR across blocks and masked diffusion inside, adding KV cache and variable-length decode. A 110M model hits ≤28.23 PPL on LM1B vs MDLM 31.78, sampling up to ~10k tokens.
Discrete diffusion LMs can denoise many tokens in parallel and are easier to steer than left-to-right models. In practice they still fail three operational tests.
Most recent discrete diffusion architectures emit a fixed-length vector, so they cannot stop when a reply is done. Bidirectional context also blocks KV caching, which makes decode more expensive than autoregression. Likelihood lags as well: on LM1B, MDLM sits at ≤31.78 test perplexity, while a matched 110M autoregressive Transformer reaches 22.83.
BD3-LMs, from Marianne Arriola, Volodymyr Kuleshov and colleagues at Cornell Tech (ICLR 2025 Oral), interpolate the two regimes. The model is autoregressive across blocks of L' tokens and runs masked discrete diffusion inside each block. L'=1 recovers AR; L' equal to the full length recovers ordinary diffusion.
A length-L sequence is split into B blocks. Log-likelihood factorizes over blocks; each block-conditional is a discrete denoising diffusion model that sees the previously generated clean blocks.
The backbone is a Transformer with a block-causal mask: tokens inside the current block attend bidirectionally, history is visible, future blocks are not. At decode time, past blocks write into a KV cache and the current block is sampled in parallel. That is the architectural fix for variable length and caching.
Training has a bookkeeping problem. Denoising the current block needs a noisy input, but the next block needs a clean encoding of the current one, so a naive loop runs every token through the net twice. The paper concatenates the clean sequence with a noisy copy and uses a custom mask: noisy tokens attend to other noisy tokens in their block and to clean tokens in earlier blocks. One forward pass covers the full block-diffusion loss, 20-25% faster than two passes, and overall training stays within 2x of vanilla diffusion. They also pretrain with L' equal to the full context for 850K steps, then finetune the target block size for 150K steps.
Gradient variance explains a surprising gap. For L'=1 the diffusion objective equals AR NLL in expectation, yet after 16B LM1B tokens the model still trails by about two perplexity points (≤25.56 vs 22.88). Masked diffusion computes cross-entropy only on masked positions, about half the tokens on average. Forcing a fully-masked forward process matches the AR objective, recovers 22.88 PPL, and drops NELBO variance from 1.52 to 0.11.
For L'>1 they clip the noise schedule, sampling mask rates from U[β,ω] instead of U[0,1], so the model rarely sees near-empty or fully-masked blocks that give weak, noisy gradients. Every 5K steps they grid-search β and ω to minimize NELBO variance. The best interval tracks block size: L'=4 likes heavier masking (around U[0.45,0.95] or U[0.5,1]); L'=16 sits closer to the middle. Clipped schedules beat linear, log, cosine, and square on both PPL and variance.
The net is 12 layers, width 768, 12 heads, 110M parameters, RoPE, no timestep conditioning. Context is 128 on LM1B and 1024 on OpenWebText, batch 512.
Diffusion numbers are NELBO upper bounds, not exact NLL.
On LM1B after 65B tokens:
| Method | Test PPL↓ |
| AR Transformer | 22.83 |
| Transformer-XL Base | 23.5 |
| D3PM (absorb) | ≤82.34 |
| SEDD | ≤32.68 |
| MDLM | ≤31.78 |
| BD3-LM L'=16 | ≤30.60 |
| BD3-LM L'=8 | ≤29.83 |
| BD3-LM L'=4 | ≤28.23 |
Smaller blocks move toward AR. The paper claims up to 13% over MDLM; the table's L'=4 number is about 3.5 points below 31.78.
On OpenWebText after 524B tokens: AR 17.54, SEDD ≤24.10, MDLM ≤22.98, BD3-LM L'=16/8/4 at ≤22.27 / ≤21.68 / ≤20.73. Zero-shot from the same OWT checkpoint, L'=4 reaches 42.52 on Pubmed, below AR's 48.59, and is the best diffusion model on Wikitext, LM1B, and AG News. It loses on PTB (96.81 vs MDLM 90.96 and AR 81.07), and does not beat MDLM on Lambada or Arxiv.
From 500 unconditional samples, SEDD cannot exceed the 1024 training context. BD3-LM L'=16 has median 798 and max 9982 tokens, about 10x. Sampling stops at [EOS] or when mean entropy of the last 256 tokens falls below 4, which cuts off degenerate run-ons. AR reaches median 4008 and max 131K, matching the training-set cap.
Generative perplexity under GPT-2 Large, 300 samples:
| Method | L=1024 Gen.PPL / NFE | L=2048 |
| AR | 14.1 / 1K | 13.2 / 2K |
| SEDD | 52.0 / 1K | — |
| MDLM | 46.8 / 1K | 41.3 / 2K |
| SSD-LM L'=25 | 37.2 / 40K | 35.3 / 80K |
| BD3-LM L'=4 | 25.7 / 1K | 23.6 / 2K |
SSD-LM is Gaussian block diffusion over embeddings, 400M parameters. When its NFE budget is cut to the same order as BD3, Gen.PPL collapses to about 281. The discrete route is better at an order of magnitude fewer network calls.
Variable length and KV caching were the two missing pieces if discrete diffusion is going to sit in a real decoder. BD3-LM adds both, and moves likelihood close enough that AR and diffusion can share a table. For people training diffusion LMs, the reusable parts are the block-causal mask, the clean/noisy concatenated trainer, and the per-block-size clipped schedule. The schedule trick also applies to plain MDLM and D3PM, which the authors note.
This is a 110M study, not a 7B chat model. Smaller blocks improve quality and give up in-block parallelism; larger blocks look more like full-sequence diffusion and cache more. The leftover AR gap is real: 28.23 vs 22.83 on LM1B.
Training is still more expensive than vanilla diffusion, and can remain close to 2x even after vectorization. Blocks are generated sequentially, so small L' gives back the speed and control advantages that made diffusion interesting. The best block size is task-dependent. The NELBO loosens as L' grows, so every diffusion PPL here is an upper bound.
There is no LLM-scale run and no instruction, code, or math eval. Zero-shot is mixed. The OWT variable-length setup drops [BOS]/[EOS] injection used by some public baselines, so those comparisons are not perfectly matched. GPT-2 Large Gen.PPL rewards surface fluency and should not be the only quality call. Code and weights are public; the claim currently holds at this scale.