🤖 machine learning

BlockGen: Flexible Blockwise Sequence Modeling with Hybrid Samplers

This paper introduces BlockGen, a flexible blockwise sequence modeling framework that combines masked and uniform-state diffusion with AR-informed predictor-corrector sampling to demonstrate that while uniform diffusion outperforms masked diffusion in few-step block-by-block generation, the performance gap narrows or reverses with increased sampling steps and specific block sizes.

Original authors: Justin Deschenaux, Caglar Gulcehre

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Justin Deschenaux, Caglar Gulcehre

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to write a story, but you have two different ways of doing it.

The Old Way (Autoregressive): You write one word at a time, from left to right. Once you write a word, you lock it in and move to the next. This is fast and reliable, but if you make a mistake in the first sentence, you can't easily go back and fix it without rewriting the whole thing. Also, because you can only write one word at a time, it takes a long time to finish a long story.

The New Way (Diffusion): Imagine you start with a page full of random gibberish (like "xkq#z@!"). You have a smart editor who looks at the whole page at once and tries to fix a few words to make them make sense. Then they look again and fix a few more. They repeat this process, slowly turning the gibberish into a story. This is great because the editor can look at the whole picture and fix mistakes anywhere on the page. However, doing this "all at once" is usually slower and the final story might not be as perfect as the old way.

Recently, researchers found a third way: Block-by-Block. Instead of writing one word at a time or fixing the whole page at once, you write (or fix) the story in small chunks, like paragraphs. You finish paragraph 1, lock it in, and then move to paragraph 2. This gives you the speed of the old way but the flexibility of the new way.

The Problem the Paper Solves

The authors, Justin Deschenaux and Caglar Gulcehre, noticed two big issues with how people were testing these "Block-by-Block" models:

The "Random Guess" Problem: When the model makes a mistake in a paragraph, the old correction methods just picked random words to fix. It's like a teacher telling a student, "Fix a random word in your essay," instead of pointing out the specific sentence that makes no sense.
The "One-Size-Fits-All" Problem: Previous studies tested these models using only one specific block size (e.g., always 16 words at a time). But maybe a block of 4 words works better for some tasks, and a block of 32 works better for others. No one had tried mixing them.

The Solution: BlockGen

They built a new system called BlockGen. Think of it as a super-flexible writer that can handle paragraphs of any length.

Mixing the Sizes: Instead of training the model to only write 16-word chunks, they trained it on a mix of sizes: 1 word, 2 words, 4 words, up to 16 or 32 words.
The Magic Trick: Because the model learned to write in all these different sizes, it learned a secret: it can write a single word perfectly (like the old "Autoregressive" way) just as well as it can fix a whole paragraph. This allows the model to act as its own "verifier."

The New Strategy: ARPC (The "Smart Editor")

The paper introduces a new way to generate text called ARPC (Autoregressive-Informed Predictor-Corrector).

Here is how it works with a creative analogy:
Imagine you are writing a paragraph using the "Block-by-Block" method. You generate a draft of the whole paragraph.

The First Pass: The model writes the whole paragraph quickly.
The "Smart Check": Before you lock the paragraph in, the model takes a quick look at its own work. It asks, "If I were writing this word-by-word (the old reliable way), what would I have written here?"
The Correction: If the model's quick "word-by-word" guess is very different from what it just wrote, it knows that specific word is likely wrong. It only rewrites those specific bad words.

This is much smarter than the old method, which would just pick random words to rewrite. It's the difference between a teacher saying "Fix a random word" vs. "Fix the sentence where you used the wrong verb."

What They Found

The authors tested this on math problems (GSM8K) and general writing (OpenWebText).

The "Uniform" vs. "Masked" Debate: In the world of diffusion, there are two main types of "editors":
- Masked: The editor can only replace words with a special "blank" token. Once a blank is filled, it's locked.
- Uniform: The editor can change any word to any other word at any time.
- Previous finding: When fixing the whole page at once, the "Uniform" editor was better.
- BlockGen finding: When working block-by-block, the "Uniform" editor is still better if you just do a simple pass. However, when you use the new "Smart Check" (ARPC), the "Masked" editor actually becomes slightly better at high-quality tasks (like math) because it's more precise.
Speed vs. Quality: The "Smart Check" (ARPC) allows the model to get much closer to the quality of the old, slow "word-by-word" writers, but it does it much faster because it fixes whole chunks at once.
The Trade-off: The paper admits that training this flexible model is more expensive (it takes more computer power) than training a standard model. Also, their models are still smaller than the massive AI models used by big tech companies today, so they can't claim this works for every possible task yet.

In a Nutshell

The paper introduces BlockGen, a flexible AI that learns to write in chunks of varying sizes. By combining this flexibility with a "Smart Check" system (ARPC) that uses the model's own knowledge to find and fix only the specific mistakes, they created a method that is faster than writing word-by-word but just as accurate, and smarter than previous "fix-it-all" methods. They showed that while the "Uniform" approach is generally strong, using this smart correction method changes the rules, making the "Masked" approach surprisingly competitive for complex tasks.

Technical Summary: BlockGen: Flexible Blockwise Sequence Modeling with Hybrid Samplers

Problem Statement

Autoregressive (AR) Transformers dominate sequence modeling but suffer from sequential generation limits (throughput/latency) and causal attention constraints that hinder reasoning tasks requiring bidirectional context. Discrete diffusion models offer an alternative by iteratively refining noisy sequences, allowing parallel updates. However, recent comparisons between Masked Diffusion Models (MDMs) and Uniform-State Diffusion Models (USDMs) have yielded conflicting insights depending on the sampling regime:

Ancestral Sampling: USDMs often outperform MDMs in downstream tasks despite worse perplexity, attributed to the ability of USDMs to self-correct mistakes (tokens can transition between any states), whereas MDMs fix tokens once unmasked.
Blockwise Generation: Prior work largely compares full-sequence diffusion. It remains unclear if the USDM advantage holds when tokens are generated block-by-block (semi-autoregressive), a setting that enables KV caching and faster inference.
Informed Correctors: Existing comparisons often use "un-informed" predictor-corrector (PC) samplers that re-inject noise at random positions. It is unknown how "informed" correctors (targeting likely errors) affect the MDM vs. USDM ranking, particularly without auxiliary verifier models.

Methodology: BlockGen

The authors introduce BlockGen, a blockwise sequence model trained over a mixture of block sizes, instantiated with either masked or uniform diffusion within blocks.

1. Mixture Formulation and Training

BlockGen defines the sequence density as a mixture over $M$ block sizes $S = \{s_1, \dots, s_M\}$ :
$p_{\theta}^{\text{BlockGen}}(x) = \sum_{i=1}^M \gamma_i p_{\theta}^{(s_i)}(x)$
where $p_{\theta}^{(s_i)}$ is a valid density factorized into blocks of size $s_i$ .

Shared Denoiser: Instead of training separate models, a single shared denoiser $\hat{x}_\theta$ takes the block size $L'$ as an input, adjusting attention patterns accordingly.
Likelihood Bounds: The paper derives two bounds for the mixture density. A tighter log-sum-exp bound is used for evaluation, while a looser mixture likelihood bound (an expectation over block sizes) is used for training to maintain computational efficiency.
Stratified Sampling: To reduce variance during multi-GPU training, block sizes are selected via stratified sampling rather than i.i.d. sampling.

2. AR-Informed Predictor-Corrector (ARPC)

A key contribution is the ARPC sampler, which leverages the BlockGen mixture to perform informed corrections without an auxiliary verifier.

Mechanism: During generation, after standard ancestral diffusion steps, the sampler executes a "clean proposal" step. It then performs a second forward pass in AR mode ( $L'=1$ ) using the same denoiser to score the proposed tokens.
Scoring: Tokens with low AR log-likelihood are identified as "unlikely" and re-noised (corrected) to a previous noise level.
Advantage: This allows the model to use its own AR predictions to guide diffusion corrections, effectively interleaving AR and diffusion paradigms within a single model.

3. Block-Level Attention

BlockGen utilizes block-causal attention: tokens within a block attend bidirectionally, while the block attends causally to all preceding blocks. This enables KV caching across decoded blocks, preserving the efficiency benefits of semi-autoregressive generation.

Key Results

1. Mathematical Reasoning (GSM8K)

Evaluated on GSM8K using models trained on TinyGSM:

Ancestral Sampling: In the block-by-block setting, USDMs outperform MDMs, especially in the few-step regime (low NFE). This confirms that the full-sequence USDM advantage carries over to blockwise generation.
ARPC Sampling: The performance gap between USDMs and MDMs narrows and reverses at high NFE. Under ARPC with 16-token blocks, MDMs slightly outperform USDMs on GSM8K accuracy.
ARPC vs. Baselines: ARPC consistently outperforms both ancestral sampling and Entropy-Informed PC (EIPC) across most NFE budgets. ARPC with MDMs approaches the performance of AR sampling with greedy decoding.

2. Language Modeling (OpenWebText & LM1B)

Likelihood: Training over a mixture of block sizes significantly improves perplexity compared to fixed-block models. BlockGen with masked diffusion achieves 17.5 Val. PPL on OpenWebText (vs. 16.7 for AR and 21.6 for the best fixed-block BDM), narrowing the gap between diffusion and autoregressive models.
Sample Quality:
- Under ancestral sampling, USDMs achieve lower Generative Perplexity (Gen. PPL) than MDMs.
- Under ARPC, the trend reverses at higher NFE (e.g., 64 NFE/block), where masked-ARPC achieves a better Gen. PPL frontier than uniform-ARPC.
- ARPC narrows the gap to AR but does not surpass AR in Gen. PPL within the practical low-temperature regime.

Significance and Claims

The paper claims to resolve two open questions regarding discrete diffusion:

Blockwise Validity: The advantage of uniform-state diffusion over masked diffusion observed in full-sequence generation does carry over to block-by-block generation under ancestral sampling.
Informed Correction: The ranking between MDMs and USDMs is regime-dependent. While USDMs excel with simple ancestral or un-informed samplers, informed correctors (specifically ARPC) can close this gap and even reverse the ranking, favoring masked diffusion at higher compute budgets (NFE).

The authors emphasize that BlockGen provides a flexible framework that interpolates between AR and pure diffusion, allowing a single model to act as both a blockwise sampler and an AR verifier. They note that while their empirical claims are robust within the tested budgets and samplers, they do not establish if these trade-offs persist at the scale of large language models (LLMs), as their experiments are limited to 170M-parameter backbones. The work suggests that the choice of diffusion paradigm (masked vs. uniform) is not universally superior but depends heavily on the sampling strategy and compute budget.