🤖 machine learning

Multi-Block Diffusion Language Models

This paper introduces Multi-Block Diffusion Language Models (MBD-LMs), which bridge the training-inference gap in diffusion-based text generation through a novel Multi-block Teacher Forcing strategy and an optimized Block Buffer decoding algorithm, achieving significant gains in both generation speed (Tokens Per Forward pass) and accuracy.

Original authors: Yijie Jin, Jiajun Xu, Yuxuan Liu, Chenkai Xu, Yi Tu, Jiajun Li, Dandan Tu, Xiaohui Yan, Kai Yu, Pengfei Liu, Zhijie Deng

Published 2026-06-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yijie Jin, Jiajun Xu, Yuxuan Liu, Chenkai Xu, Yi Tu, Jiajun Li, Dandan Tu, Xiaohui Yan, Kai Yu, Pengfei Liu, Zhijie Deng

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

The Big Picture: Fixing the "Assembly Line" Bottleneck

Imagine a factory that builds cars (text).

Old Way (Autoregressive): The factory builds one car at a time. They finish the engine, then the wheels, then the paint, and only then do they start the next car. It's reliable, but slow.
Diffusion Way (The Newer Factory): Instead of building from scratch, this factory starts with a pile of scrap metal (random noise) and gradually refines it into a perfect car. This allows them to work on many parts of the car at the same time (parallel processing), which is much faster.

The Problem:
The current "Diffusion Factory" (called Block Diffusion) has a clever trick: it builds cars in batches (blocks). However, it still has a bottleneck. It finishes one batch, puts it in storage (caching), and then starts the next batch. It's like a relay race where the runner must stop completely to hand off the baton before the next runner can even start running. This creates "waiting bubbles" where the factory sits idle.

The Solution: The "Multi-Block" Relay

The authors propose a new method called Multi-Block Diffusion (MBD).

The Analogy: The Overlapping Relay
Imagine a relay race where the runners don't wait for the previous runner to cross the finish line before they start.

Runner A is finishing their lap.
Runner B is already sprinting on the next lap.
Runner C is getting ready in the next lane.

They are all running at the same time! This is Multi-Block Diffusion. Instead of finishing one block of text before starting the next, the model refines several blocks of text simultaneously. This creates a "pipeline" where work is constantly happening, drastically speeding up the process.

The Hurdle: Training vs. Reality

There was a major problem with trying to do this.

The Training: The models were trained like a strict teacher (Teacher Forcing). The teacher showed the student: "Here is a perfect paragraph, now fix this one messy sentence." The student never practiced fixing multiple messy sentences at once.
The Reality: When the model tried to run the "Multi-Block" race (fixing 2 or 3 blocks at once), it stumbled because it had never practiced that specific scenario. It was like asking a student who only practiced single math problems to suddenly solve a complex equation with three variables at once.

The Fix: "Multi-Block Teacher Forcing" (MultiTF)

To fix this, the authors created a new training method called Multi-block Teacher Forcing (MultiTF).

The Analogy: The Simulation Drill
Instead of just showing the student one messy sentence, the teacher now shows them a whole row of messy sentences (a "noise group") and says, "Fix all of these at once."

They don't just make them messy in a predictable order; they make them messy in a chaotic, realistic way that matches what happens during the actual race.
This "drill" teaches the model how to handle multiple blocks of text simultaneously without getting confused.

The Engine: The "Block Buffer"

Even with a trained model, running this on a computer is tricky. Computers like to work with fixed sizes (like a tray with exactly 4 slots). If you try to add a new block dynamically, the computer has to stop and rearrange everything, which slows it down.

The Analogy: The Fixed Tray
The authors built a special "Block Buffer" mechanism.

Imagine a conveyor belt with 4 fixed slots.
Instead of adding a new slot when you need it, you just activate an empty slot that was already there.
When a block is finished, it slides off the belt into a "storage locker" (the cache), and a new empty slot slides in at the back.
This keeps the conveyor belt size exactly the same, allowing the computer to run at maximum speed without stopping to rearrange the belt.

The Results: Faster and Smarter

The paper tested this on math and coding tasks. Here is what they found:

Speed: The new method (MBD) processes significantly more "tokens" (words/ideas) per second than the old method.
- Analogy: If the old factory made 3 cars an hour, the new factory makes 6 cars an hour.
Quality: Usually, when you try to go faster, you make more mistakes. However, because they trained the model specifically for this "Multi-Block" style (using MultiTF), the new factory didn't just go faster; it actually made fewer mistakes than the old factory.
Compatibility: This new method works well with other speed-boosting tricks (like DMax) that were already being used, stacking up to make the system even faster.

Summary

The paper introduces a way to make AI text generators faster by having them work on multiple chunks of text at the same time. They solved the problem by:

Training the AI to handle multiple chunks at once (MultiTF).
Building a computer system that keeps the workload steady and efficient (Block Buffer).

The result is a text generator that is both faster and more accurate than previous versions.

Technical Summary: Multi-Block Diffusion Language Models (MBD-LMs)

1. Problem Statement

Diffusion Language Models (DLMs) offer native parallel decoding as an alternative to autoregressive generation but historically struggle with efficient serving due to a lack of KV caching and dynamic-length generation support. Recent Block Diffusion Language Models (BD-LMs) addressed this by introducing block-causal generation, enabling Single-Block Diffusion (SingleBD). In SingleBD, the model decodes one noisy block conditioned on a clean, cached prefix, allowing for KV caching and intra-block parallelism. However, blocks are still processed sequentially, creating "storing bubbles" where the system waits to cache a completed block before decoding the next, locking inter-block parallelism.

While Discrete Diffusion Forcing (D2F) introduced visibility among multiple noisy blocks to enable Multi-Block Diffusion (MultiBD) and inter-block parallelism, a significant train–inference mismatch remains. D2F trains on long, monotonic noisy suffixes, whereas practical MultiBD inference operates on a bounded running-set of consecutive blocks with heterogeneous slot-wise noise patterns. Existing BD-LMs trained under Teacher Forcing (TF) or D2F do not align with these specific inference states, leading to suboptimal performance when attempting to decode multiple blocks concurrently. Furthermore, naive MultiBD implementations often rely on dynamic input shapes, which are incompatible with static-shape execution mechanisms like CUDA Graph capture, hindering practical throughput gains.

2. Methodology

The paper proposes Multi-Block Diffusion Language Models (MBD-LMs), a unified framework and training strategy to bridge the gap between existing BD-LMs and efficient MultiBD inference.

2.1 Unified Formulation: MBD-LMs

The authors define MBD-LMs as a generalization of BD-LMs where a running-set ( $R_s$ ) of consecutive blocks is decoded concurrently. Unlike SingleBD (where $|R_s|=1$ ) or full-sequence D2F, practical MBD-LMs operate on a bounded intermediate regime where $|R_s|$ is small (e.g., 2–4 blocks) to balance parallelism with computational efficiency.

2.2 Multi-block Teacher Forcing (MultiTF)

To align training states with the bounded running-set inference regime, the authors introduce Multi-block Teacher Forcing (MultiTF), a post-training method for existing BD-LMs.

Noise-Group Construction: Instead of corrupting a single block (TF) or a long monotonic suffix (D2F), MultiTF constructs noise-groups: bounded, consecutive groups of noisy blocks conditioned on a clean prefix.
Group Layouts: The training data is partitioned using both systematic layouts (enumerating all shifts and group sizes up to a maximum $G_{max}$ ) and random layouts to ensure diverse coverage of running-set sizes and block-relative positions.
Chain-Uniform Noise Scheduler: Within each noise-group, MultiTF employs a chain-uniform noise scheduler. Unlike D2F's monotonic schedule, this method samples a group-level floor and then samples each subsequent block's noise ratio conditioned on the previous block's ratio. This creates heterogeneous slot-wise noise gaps that better match the empirical noise patterns observed during MultiBD inference.
Group-Aware Dual-Stream Mask: A specialized attention mask allows noisy blocks within a group to attend to each other while maintaining block-causal visibility to the clean prefix, preventing information leakage from future noisy tokens to the clean prefix.

2.3 Optimized Inference Engine: Block Buffer

To enable practical, high-throughput MultiBD inference, the authors design an optimized pipeline based on the Block Buffer mechanism:

Static-Shape Execution: Instead of dynamically appending blocks to the input (which breaks CUDA Graph compatibility), the engine maintains a fixed-size physical buffer ( $N_{buf}$ ).
State Transitions: Blocks within the buffer transition through four states: dummy (inactive), active (being decoded), to-cache (completed but not yet cached), and in-cache (committed to KV cache).
Mechanism: Future blocks enter the buffer by activating existing dummy slots rather than extending the input sequence. Completed blocks are committed to the KV cache and replaced by new dummy slots at the tail. This preserves the static input shape required for efficient GPU execution while maintaining the logical running-set for inter-block parallelism.
Prefix Cache Preservation: This design strictly separates the immutable cached prefix from the active buffer, ensuring compatibility with standard BD-LM prefix caching semantics, which native D2F often violates.

3. Key Contributions

Unified MBD-LM Formulation: The paper formulates MBD-LMs as a unified framework parameterized by a bounded running-set, identifying practical MultiBD as an intermediate regime between SingleBD and full-sequence D2F.
MultiTF Post-Training: A novel training strategy that aligns BD-LMs with MultiBD inference by training on bounded noise-groups with heterogeneous noise patterns, significantly reducing the train–inference mismatch.
Optimized Inference Engine: A Block Buffer-based pipeline that enables static-shape execution for CUDA Graphs, preserves prefix KV caching, and overlaps decoding with cache storing to achieve wall-clock speedups.

4. Experimental Results

The authors evaluated MBD-LMs on mathematical reasoning (GSM8K, MATH500) and code generation (MBPP+, HumanEval+) benchmarks using LLaDA2 and SDAR model families.

Parallelism and Accuracy: MBD-LLaDA2-Mini increased the average Tokens Per Forward pass (TPF) from 3.47 (SingleBD) to 6.19 (+78.4%) while improving average accuracy from 79.95% to 81.03%.
Combination with DMax: When combined with the DMax decoding strategy, MBD-LLaDA2-Mini-DMax achieved an average TPF of 9.34 (a +47.1% improvement over SingleBD DMax) with only a 1.02% accuracy drop.
Throughput (TPS): The optimized inference engine translated TPF gains into wall-clock throughput. MBD-LLaDA2-Mini-DMax achieved 926.67 TPS on average, compared to 779.49 TPS for the SingleBD DMax baseline.
Ablation Studies:
- Train-Inference Alignment: Training-free MultiBD (applying MultiBD inference to unmodified models) improved TPF but degraded accuracy. MultiTF post-training recovered and improved accuracy, confirming the necessity of state alignment.
- Scheduler Design: The proposed chain-uniform scheduler outperformed D2F-style monotonic schedulers and other uniform variants, highlighting the importance of heterogeneous noise gaps.
- Layout Diversity: Combining systematic and random group layouts yielded the best trade-off between TPF and accuracy.

5. Significance

The paper argues that reliable and efficient Multi-Block Diffusion requires a holistic approach combining training-time state alignment and inference-time system support.

It demonstrates that simply applying MultiBD inference to existing models is insufficient due to distributional mismatches in noise patterns and running-set sizes.
It proves that MBD-LMs can achieve inter-block parallelism without sacrificing the KV caching and prefix reuse benefits that make BD-LMs efficient in the first place.
The work establishes that the "bounded running-set" is the practical sweet spot for diffusion-based generation, offering a viable path to accelerating DLMs beyond the sequential bottlenecks of SingleBD while avoiding the complexity and cache incompatibility of full-sequence diffusion.