PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-Block Diffusion Language Models

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

arXiv:2606.29215v2cs.LGcs.CL

TL;DR

MultiBD seeks to overcome SingleBD’s sequential block bottleneck, but existing training states do not match bounded running-set inference with heterogeneous noise patterns. The paper introduces MultiTF and Block Buffer decoding; MBD-LLaDA2-Mini raises average TPF from 3.47 to 6.19 while improving average accuracy from 79.95% to 81.03%.

  • Problem

    Existing BD-LMs lack training states aligned with bounded MultiBD running-sets containing multiple noisy blocks with heterogeneous noise patterns, limiting practical inter-block parallelism.

  • Method

    The paper post-trains BD-LMs with MultiTF and uses Block Buffer decoding to align training with bounded noise-groups while preserving KV-cache reuse and static input shapes.

  • Results

    Average TPF rises from 3.47 to 6.19 and average accuracy from 79.95% to 81.03%; with DMax, average TPF reaches 9.34 with a 1.02 percentage-point accuracy drop.

  • Takeaways & Limitations

    Reliable MultiBD requires both training-time alignment with bounded running-set states and inference-time system support to improve decoding parallelism while maintaining generation quality.

  • Takeaways & Limitations

    Direct running-set execution produces variable token counts across steps and requests, making CUDA Graph capture and replay difficult.

Abstract

from arXiv · show

Block Diffusion Language Models (BD-LMs) improve diffusion-based text generation with KV caching and flexible-length generation. A natural next step is to extend them from Single-Block Diffusion (SingleBD) to Multi-Block Diffusion (MultiBD), where a running-set of consecutive blocks is decoded concurrently for inter-block parallelism. However, existing BD-LMs are mostly trained under teacher forcing, where the model observes only one noisy block conditioned on a clean prefix. While the recent diffusion forcing strategy introduces visibility among multiple noisy blocks, its training states still differ from MultiBD inference, where decoding operates on a bounded running-set with heterogeneous slot-wise noise patterns. To bridge this gap, we propose Multi-Block Diffusion Language Models (MBD-LMs), obtained by post-training BD-LMs with Multi-block Teacher Forcing (MultiTF). MultiTF integrates teacher forcing and diffusion forcing by training on bounded noise-groups conditioned on clean prefixes, with randomized noise-schedulers that better match MultiBD inference states. To make MultiBD practically executable, we further introduce an optimized decoding algorithm based on the Block Buffer mechanism that preserves prefix-cache reuse, keeps input shapes static, and translates increased decoding parallelism into wall-clock acceleration. Empirically, MBD-LLaDA2-Mini increases average Tokens Per Forward pass (TPF) from 3.47 to 6.19 and improves average accuracy from 79.95% to 81.03%; when combined with DMax, MBD-LLaDA2-Mini-DMax reaches an average TPF of 9.34 with only a 1.02% accuracy drop on math and code benchmarks.

1 Introduction

MBD-LMs extend block diffusion from sequential SingleBD to bounded MultiBD, aligning training and inference with MultiTF and enabling practical inter-block parallelism. A Block Buffer inference pipeline converts this parallelism into speedups while preserving caching and static input shapes.

  • Framework: MBD-LMs unify TF-trained and D2F-trained BD-LMs while identifying bounded MultiBD as an intermediate regime for reliable and efficient inference.SingleBD decodes one noisy block after clean cached prefixes, whereas MultiBD overlaps future-block refinement with KV-cache storage to enable inter-block parallelism.
  • Method: MultiTF post-trains BD-LMs on bounded groups of consecutive noisy blocks, using randomized heterogeneous noise schedules and group-aware attention to match practical MultiBD states.Groups vary in size and relative position to cover possible running-set configurations during inference.
  • Inference: The Block Buffer maintains fixed block slots, preserves KV and prefix caching, and keeps input shapes static for CUDA Graph capture while overlapping decoding with cache storage.Future blocks activate idle slots, and completed front blocks leave after committing their KV cache, translating higher TPF into wall-clock speedup.
  • Experiments: TPF rises from 3.47 to 6.19 (+78.4%) and average accuracy from 79.95% to 81.03% for MBD-LLaDA2-Mini versus LLaDA2-Mini.With DMax, MBD-LLaDA2-Mini-DMax reaches 9.34 average TPF (+47.1% over SingleBD) with only a 1.02 percentage-point accuracy drop.

2 Preliminaries

Diffusion language models generate text by iteratively denoising masked sequences, while block diffusion models apply this process autoregressively within blocks and enable KV caching. Existing Teacher Forcing and D2F training regimes do not fully match MultiBD inference states, and native D2F is incompatible with BD-LM prefix caching.

  • Diffusion Language Models: DLMs formulate generation as iterative denoising: tokens are independently masked, the model predicts clean values at masked positions, and inference progressively fills high-confidence masks.Training uses weighted masked-token cross-entropy evaluated only on masked tokens.
  • Block Diffusion Language Models: BD-LMs partition sequences into blocks and model them autoregressively, using block-causal attention so SingleBD decoding can reuse KV caches.Each block is decoded through a DLM process while attending to itself and preceding blocks.
  • Teacher forcing: Teacher Forcing corrupts only the current block and predicts it from clean prefix blocks, so it learns a state that is conceptually incompatible with MultiBD inference.The mismatch arises because MultiBD concurrently decodes multiple noisy blocks rather than one noisy block conditioned on a clean prefix.
  • Discrete diffusion forcing: D2F exposes the model to multiple noisy suffix blocks with ordered noise ratios, but its training states still differ from MultiBD inference and its native arbitrary-length clean prefix prevents direct prefix-cache reuse.Native D2F uses full attention on the clean prefix; enforcing cache-compatible block-causal attention causes larger quality degradation, motivating MultiTF.

3 Methodology

MBD-LMs align training with MultiBD inference by post-training BD-LMs on bounded, heterogeneous noise-groups, while Block Buffer execution preserves cache reuse and enables static-shape parallel decoding.

  • Multi-Block Diffusion: MultiBD concurrently decodes a bounded running-set of consecutive blocks while retaining earlier blocks as a clean prefix in the KV cache.Each running-set contains active noisy blocks and completed blocks awaiting caching, with slot-wise mask ratios describing their states.
  • Multi-Block Diffusion: The running-set formulation unifies existing regimes: SingleBD is the |Rs| = 1 case, whereas expanded running-sets with monotonic schedulers resemble block-causal D2F training states.This relationship concerns training-state construction only; D2F remains a training paradigm, while MultiBD is an inference regime.
  • Multi-Block Teacher Forcing: MultiTF post-trains BD-LMs using inference-like noise-groups that match MultiBD’s bounded running-set structure and heterogeneous slot-wise noise patterns.Later noise-groups are conditioned on clean earlier groups, extending teacher forcing from one noisy block to bounded consecutive groups.
  • Multi-Block Teacher Forcing: MultiTF constructs systematic and random group-layouts to cover bounded group sizes and different group-relative positions corresponding to possible MultiBD running-sets.The maximum training group size is Gmax, while systematic shifts provide structured coverage and random layouts increase layout diversity.
  • Static-Shape MultiBD Execution: Block Buffer decouples the logical running-set from the physical input, supplying static shapes for CUDA Graph replay despite changing processed-token counts.The mechanism preserves inter-block parallelism, overlaps decoding with KV-cache storage, and maintains prefix-cache reuse.

4 Experiments

Experiments show that MultiTF converts BD-LMs into MBD-LMs with substantially higher decoding parallelism while preserving or improving accuracy. Gains persist with DMax, across backbones, and in wall-clock throughput, provided training matches bounded, heterogeneous MultiBD inference states.

  • Main analysis: MBD-LLaDA2-Mini raises average TPF from 3.47 to 6.19 (+78.4%) and average accuracy from 79.95% to 81.03% versus native SingleBD.Without DMax, its TPF is comparable to LLaDA2-Mini-DMax under SingleBD (6.19 vs. 6.35), while accuracy is higher (81.03% vs. 79.59%).
  • Compatibility with T2T-enhanced decoding: Combining MultiTF with DMax increases average TPF from 6.35 to 9.34 (+47.1%) with only a 1.02 percentage-point average accuracy drop.The result demonstrates that MBD-LMs can stack with Token-to-Token-enhanced decoding.
  • Effect of train–inference alignment: MultiTF alignment is necessary because training-free MultiBD increases TPF but can reduce accuracy, while reliable inference requires bounded active sets and heterogeneous slot-wise noise patterns.MultiBD typically maintains about two active blocks, occasionally expanding to three or four; MultiTF also generalizes to SDAR, where TPF rises from 2.54 to 4.46 and accuracy from 69.00% to 69.74%.
  • Ablation study: The full MultiTF configuration raises TPF from 6.57 to 9.87 and AUP from 536.89 to 805.34, while average accuracy changes only from 84.67% to 84.59%.Combining systematic and random noise-group layouts gives the highest TPF and AUP, whereas D2F-style scheduling raises TPF to 8.76 but drops accuracy to 79.34%.
  • Throughput analysis: MBD improves wall-clock throughput: LLaDA2-Mini reaches 745.92 versus 517.16 Avg. TPS, matching the predicted 1.44× scaling from TPF and step latency.TPF rises from 3.47 to 6.19 while step latency increases from 7.07 ms to 8.78 ms; static shapes add resident and dummy tokens that reduce token efficiency.

5 Related Work

Diffusion language models enable parallel denoising but face serving limitations, while block-causal variants add KV caching and flexible-length generation at the cost of sequential blocks. This work targets MultiBD as a distinct inter-block parallelism regime and introduces MultiTF to align training with its bounded, heterogeneous multi-block states.

  • Diffusion Language Models: DLMs generate text through iterative denoising and parallel token refinement, but fully bidirectional models lack natural KV caching and flexible-length generation.Representative models include LLaDA (Nie et al., 2025), Dream (Ye et al., 2025), and LLaDA2.x (Bie et al., 2025, 2026).
  • Block Diffusion Language Models: BD-LMs (Arriola et al., 2025; Bie et al., 2025; Cheng et al., 2025) use block-causal generation to enable KV caching and intra-block parallel decoding.SingleBD decodes one noisy block conditioned on a clean cached prefix, but processes blocks sequentially and underuses inter-block parallelism.
  • Multi-Block Diffusion: MultiBD refines a bounded running-set of consecutive blocks concurrently, expanding BD-LMs beyond SingleBD’s sequential block processing.The running-set defines the broader inference regime studied in this work.
  • Efficient Diffusion Language Models: Efficient DLM research explores distillation, scheduling, caching, and parallel decoding, including D2F’s noisy-block visibility and methods targeting the accuracy–parallelism trade-off.Related methods include D2F (Wang et al., 2025), DMax (Chen et al., 2026), d3LLM (Qian et al., 2026), LightningRL (Hu et al., 2026), dParallel (Chen et al., 2025), Fast-dLLM (Wu et al., 2025), and LoPA (Xu et al.).
  • Positioning of This Work: Unlike inference-time heuristics or token-level parallelism alone, this work treats MultiBD as a target regime and proposes MultiTF to train inference-like multi-block states.The approach emphasizes bounded running-set structure and heterogeneous slot-wise noise patterns as train–inference alignment factors, post-training BD-LMs into MBD-LMs.

6 Conclusion

MBD-LMs enable reliable MultiBD inference by aligning training states with bounded running-set decoding and heterogeneous slot-wise noise. MultiTF and Block Buffer jointly improve parallelism and throughput while preserving generation quality through KV caching and static-shape execution.

  • 6 Conclusion: MBD-LMs unify BD-LMs for reliable MultiBD inference, addressing SingleBD’s sequential bottleneck through inter-block parallelism.MultiBD requires training states aligned with bounded running sets and heterogeneous slot-wise noise patterns.
  • 6 Conclusion: MultiTF post-trains BD-LMs with bounded noise-groups, a Group-Aware Dual-Stream Mask, and randomized block-level noise-schedulers.These components bridge the mismatch between existing training states and MultiBD inference states.
  • 6 Conclusion: The Block Buffer inference engine enables static-shape execution while preserving KV caching and prefix-cache reuse.Experiments on math and code benchmarks show improved decoding parallelism and realized throughput while maintaining generation quality.

A Theoretical View of MultiTF

MultiTF is a coverage-based surrogate for the ideal MultiBD training objective rather than a direct guarantee of improved downstream accuracy. Its approximation gap is controlled by running-set coverage mismatch and noise-ratio distribution mismatch.

  • Objective Interpretation: MultiTF approximates ideal MultiBD training by matching inference-state coverage and slot-wise noise-ratio distributions, not by directly guaranteeing downstream accuracy.The theoretical view explicitly frames MultiTF as a surrogate whose gap depends on these distributional mismatches.
  • Running-Set Coverage: Systematic shifts cover every consecutive running-set of size 2 through Gmax exactly once, while random layouts add non-regular noise-group-size combinations.For each fixed group size, shifted layouts also place every block in each group-relative logical slot.
  • Objective Mismatch Bound: The objective gap is bounded by MδR + Ltδt, where δR measures running-set mismatch and δt measures conditional noise-ratio mismatch.The bound follows from bounded loss for the running-set term and Lipschitz dependence on noise ratios for the conditional term.
  • Implications: Reducing δR and δt tightens the gap to ideal MultiBD inference; systematic shifts reduce support mismatch, random layouts add diversity, and chain-uniform scheduling reduces noise-ratio mismatch.The scheduler produces heterogeneous slot-wise noise gaps aligned with MultiBD states.

B MultiTF Training Implementation Details

This appendix details MultiTF post-training, defining its training and inference terminology and specifying the framework and data used for different model families.

  • B MultiTF Training Implementation Details: MultiTF post-trains BD-LMs into MBD-LMs using noise-groups, group-layouts, and noise-schedulers on the training side, versus Block Buffers and slots at inference.Gmax denotes the maximum noise-group size, Λ the set of group-layouts, λ one group-layout, and Hm one noise-group.
  • B MultiTF Training Implementation Details: VeOmni (Ma et al., 2025) is used as the training framework for MultiTF.
  • B MultiTF Training Implementation Details: SDAR models use reasoning/code data from prior studies, while LLaDA2.x and DMax-enhanced models use the corresponding mixtures from their base recipes.The cited prior studies are Boizard et al. (2025) and jtatman (2025).

B.1 Group-Layout Construction · B.2 Chain-uniform Noise-Scheduler · B.3 Group-Aware Dual-Stream Mask

MultiTF constructs diverse consecutive noise-group layouts, assigns each group a chain-uniform sequence of monotonic randomized mask ratios, and uses a group-aware dual-stream attention mask. Together, these components approximate MultiBD’s bounded running-set visibility while preventing information leakage.

  • B.1 Group-Layout Construction: MultiTF partitions consecutive blocks into noise-groups whose layout matches a possible MultiBD running-set.Each group-layout λ=(H1,…,H|λ|) partitions the block sequence into consecutive groups.
  • B.1 Group-Layout Construction: Systematic shifted layouts cover every consecutive running-set of a fixed length, up to boundary effects, across all group-relative positions.For group size g, selecting the shift h=(a−1) mod g places each length-g running-set {a,…,a+g−1} in one shifted layout.
  • B.1 Group-Layout Construction: Random layouts add non-regular group-size combinations and boundary patterns without replacing systematic layouts’ coverage guarantee.Random groups are sampled sequentially until the full block sequence is covered.
  • B.2 Chain-uniform Noise-Scheduler: For each noise-group, MultiTF samples a floor from the lower noise range, then recursively samples each block’s mask ratio up to an effective upper bound.Each sampled ratio becomes the floor for the next block, producing monotonic but randomized slot-wise noise levels.
  • B.2 Chain-uniform Noise-Scheduler: The chain-uniform scheduler creates larger and more variable block-level noise-ratio gaps than fixed-step D2F, matching heterogeneous active blocks during MultiBD inference.The noise-transition margin ratio ρ is independent of the random scheduler power-law bias γrand used in the ablation.
  • B.3 Group-Aware Dual-Stream Mask: MultiTF concatenates noisy and clean sequences and applies a group-aware dual-stream attention mask with noisy diagonal, noisy-to-clean offset-causal, and clean block-causal components.These components are denoted MGD, MGOC, and MBC, respectively.
  • B.3 Group-Aware Dual-Stream Mask: Noisy tokens attend only within their noise-group and same-or-preceding blocks, groups condition on earlier clean prefixes, and clean tokens never attend to noisy tokens.This implements the required visibility pattern without information leakage.

B.4 MultiTF Objective and Model-specific Training Recipes … B.5 Sorted-uniform Scheduler Baseline

MultiTF standardizes bounded multi-block training inputs while allowing each base BD-LM to retain its model-specific objective or recipe. The framework supports default masked-token CE, DMax self-denoising, and SDAR block-wise weighting, while sorted-uniform provides a monotonic scheduler baseline.

  • B.4 MultiTF Objective and Model-specific Training Recipes: MultiTF constructs training states from a layout λ, noisy sequence xλ_t, clean sequence x0, and Group-Aware Dual-Stream Mask Aλ.Different base BD-LMs reuse these MultiTF input sequences while retaining their own model-specific training recipes.
  • B.4.1 Default MultiTF CE Objective: The default MultiTF objective applies masked-token cross-entropy to masked positions in the noisy portion of Xλ for BD-LMs trained with standard masked-token CE.This is the default objective when the original model recipe uses standard masked-token cross-entropy.
  • B.4.2 DMax-enhanced Models: OPUT Self-denoising: The OPUT self-denoising procedure predicts masked tokens by argmax from a no-gradient pass, replaces them in xλ_t, and returns the partially self-denoised sequence.This sequence is then used as the input for the self-denoising loss branch.
  • B.4.2 DMax-enhanced Models: OPUT Self-denoising: DMax-enhanced models reuse MultiTF inputs and add an OPUT branch that trains on partially self-denoised noisy sequences alongside the original noisy branch.The self-denoising branch performs a no-gradient forward pass, replaces masked positions with argmax predictions, and computes its loss afterward; gradients flow only through the subsequent loss computation.
  • B.4.3 SDAR Models: Block-wise Noise-weighted CE: SDAR models reuse MultiTF inputs and Group-Aware Dual-Stream Masks but replace global masked-token CE with block-wise noise-weighted CE.Each block’s loss is normalized by that block’s mask ratio, extending diffusion loss across the full block sequence while preserving SDAR’s per-block weighting.
  • B.4.3 SDAR Models: Block-wise Noise-weighted CE: SDAR’s block-wise normalization differs from Equation B.11, which normalizes globally over all masked positions in the noisy portion.A small constant ϵ provides numerical stability in the block-wise objective.
  • B.5 Sorted-uniform Scheduler Baseline: The sorted-uniform scheduler independently samples mask ratios uniformly within each noise-group, sorts them ascending, and assigns them to blocks to create monotonic block-level noise.Unlike the chain-uniform scheduler, its adjacent-slot gaps arise only from uniform order statistics and are not explicitly encouraged to be large.

C MultiBD Inference Implementation Details … C.4 Thresholds control activation and token updates.

The appendix implements MultiBD with a fixed-size Block Buffer that preserves prefix KV-cache reuse and static tensor shapes while representing dynamic running-set states. Separate thresholds coordinate block activation, context sharing, and token updates to stabilize concurrent refinement.

  • C MultiBD Inference Implementation Details: The optimized MultiBD algorithm executes the running-set with a static physical input shape while preserving prefix KV-cache reuse.This design targets practical execution of concurrent block decoding without changing the cached prefix representation.
  • C.1 A dynamic running-set prevents static-shape execution.: A direct dynamic running-set appends future blocks and removes completed blocks, but changing active-token counts prevents consistent execution shapes across steps and requests.The procedure still exposes inter-block parallelism, but its variable input size complicates efficient execution.
  • C.2 A fixed Block Buffer implements MultiBD states.: The fixed Block Buffer uses Nbuf physical slots, activating dummy slots for future blocks and retaining static tensor shapes for CUDA Graph capture and replay.Completed front blocks are committed to the prefix KV cache and removed logically without extending the physical input sequence.
  • C.2 A fixed Block Buffer implements MultiBD states.: Figure 6 shows that fully block-causal D2F improves cache compatibility but substantially hurts accuracy, motivating the optimized static-shape design.The figure contrasts cache compatibility with accuracy under fully block-causal D2F.
  • C.3 Block states advance the fixed Block Buffer.: Block Buffer slots transition among dummy, active, to-cache, and in-cache states, reproducing logical running-set evolution without changing the physical input shape.Dummy slots reserve capacity, active slots participate in forward passes, to-cache blocks await commitment, and in-cache blocks leave the active running-set.
  • C.4 Thresholds control activation and token updates.: Separate thresholds govern block activation, stability, semi-completion context sharing, and M2T/T2T token updates within the fixed Block Buffer.τadd and τstable gate future-block activation, τsemi enables top-1 preceding-block context before full caching, and τM2T with optional τT2T controls updates.
  • C.4 Thresholds control activation and token updates.: Separate M2T and T2T thresholds stabilize concurrent refinement because M2T adds new content, whereas T2T overwrites tentative content before commitment.The distinction reduces error propagation across the running-set by matching thresholds to different update reliability profiles.

C.5 Prefix Caching and Fully Block-Causal D2F

Fully block-causal D2F improves compatibility with prefix caching but sharply reduces accuracy, showing that cache compatibility cannot be achieved by simply restricting attention. Block Buffer MultiBD instead preserves immutable cached prefixes while refining active blocks in parallel.

  • Prefix caching: Native D2F breaks prefix-cache reuse because noisy prefix blocks are not organized as immutable causal prefix pages, unlike SingleBD.SingleBD recomputes only the current noisy block while reusing cached KV states for completed clean blocks.
  • Fully block-causal D2F: The fully block-causal variant completes arbitrary token-level prefixes to a block boundary, treats the boundary block as the first noisy suffix block, and computes loss only on masked suffix positions.The boundary block keeps its prefix positions clean while noising only its suffix positions; later blocks are fully suffix-noised.
  • Fully block-causal D2F: 77.60% to 69.60%: fully block-causal D2F substantially reduces accuracy versus prefix-full D2F, motivating MultiTF rather than a simple attention-mask restriction.The result indicates that D2F benefits from stronger prefix-full visibility.
  • Block Buffer MultiBD: Block Buffer MultiBD preserves prefix caching by making committed blocks immutable cached context while recomputing active blocks and keeping future dummy slots invisible until activation.This separates cached prefix blocks from active Block Buffer slots, enabling prefix KV reuse during parallel refinement.

D Experimental Details

This appendix specifies the inference conventions and MultiTF post-training hyperparameters used in the experiments. It defines the model variants and scheduler parameters used in the reported settings and ablations.

  • D Experimental Details: SingleBD (Native) denotes each BD-LM’s original single-block inference, while MultiBD (training-free) omits post-training and MBD-* denotes the corresponding MultiTF-post-trained model.The notation distinguishes native inference, training-free MultiBD inference, and MultiTF-post-trained variants.
  • D Experimental Details: Table 5 reports MultiTF post-training hyperparameters, with “—” marking parameters that are not applicable.The table provides the configuration used for MultiTF post-training.
  • D Experimental Details: t_low and t_high define the mask-ratio range, ρ determines the effective upper bound t_eff, and N_rand counts random group-layouts per sample.These parameters characterize the MultiTF scheduler and grouping configuration.
  • D Experimental Details: The random-scheduler ablation uses an independent power-law bias γ_rand and does not use ρ or the chain-uniform scheduler.Thus, γ_rand is separate from the margin-ratio parameter and chain-uniform scheduling.
Loading 2606.29215v2…