Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

Abstract: High-throughput long-context generation is one of the central challenges for LLMs. Generation is typically memory-bandwidth-bound rather than compute-bound: each decoding step must stream the accumulated key/value (KV) cache from memory, so bandwidth demand grows with context length while only one token is emitted. Two parallel approaches have therefore emerged: reducing memory access with efficient attention variants and linear-time mixers such as Mamba, or increasing parallel computation by generating blocks of tokens at once. However, technical challenges arise when combining these two ideas. Earlier hybrid diffusion models such as DiffuMamba use bidirectional Mamba mixing, including a reverse-direction scan relative to causal generation. This reverse scan needs to scan the entire sequence, so its states are not prefix-only and cannot be precisely reused as a cache even when diffusion is performed block by block. We propose a BDLM Mamba--attention hybrid that addresses this challenge by restricting the reverse Mamba scan to the active denoising block, which enables exact caching across blocks. In an 87M-parameter DCLM sweep, BDLM Mamba-H achieves the best C4-en validation perplexity compared to BDLM attention and full-sequence baselines. At 350M parameters, it remains competitive with BDLM attention. For long-context inference, BDLM Mamba-H reaches 19.7x the throughput of full-sequence DiffuMamba-H at 65K tokens and 3.7x the throughput of BDLM attention at 262K, showing that Mamba hybrids are a potential long-context diffusion architecture.

• The paper's main contribution is a hybrid BDLM (Mamba-H) that confines reverse Mamba mixing to the active block, enabling efficient cacheable inference for long-context generation.
• It utilizes an all-block objective and timestep-conditioning factorization to maintain clean prefix caching and support scalable autoregressive block factorization.
• Empirical results demonstrate that BDLM Mamba-H achieves up to 19.7x speedup and competitive perplexity at various scales, validating its performance for long-context tasks.

Hybrid Block Diffusion LLMs with Partial Bidirectionality: An Authoritative Summary

Introduction and Motivation

Long-context generation with LLMs remains memory-bandwidth-bound due to the requirement to stream the accumulated KV cache at each autoregressive decoding step, leading to an unfavorable scaling as context increases. Two principal directions have emerged for mitigating this bottleneck: (1) attention variants and state-space mixers such as Mamba that reduce memory access complexity, and (2) blockwise or masked diffusion language modeling that enables parallel generation of token blocks through iterative denoising.

Combining these avenues has proven nontrivial. Prior work, like DiffuMamba, interleaves bidirectional Mamba mixing into masked diffusion denoisers, but requires full-sequence reverse scans, precluding effective prefix caching for blockwise generation. This constraint inhibits native block-cached inference and scalable long-context throughput.

This paper introduces a hybrid Block Diffusion LLM (BDLM) architecture—BDLM Mamba-H—where reverse-direction Mamba mixing is strictly localized to the active denoising block, while prefix-only forward states can be cached and reused across blocks. This design realizes efficient cache-aligned hybrid blockwise diffusion, supporting performant long-context inference without retracing the full context at every step, and demonstrates strong empirical performance across multiple settings.

Technical Contributions

The principal contributions are as follows:

• Cache-aligned All-block Objective: The BDLM Mamba-H architecture implements a training objective that directly aligns the generative sampler’s cache with training: attention layers cache KV states, Mamba layers cache forward boundary states, and all-block losses train these objects at all block positions in parallel via an autoregressive block factorization.
• Partial Reverse Mamba Hybridization: The reverse-direction Mamba scan is restricted to within the current block, allowing block-local bidirectional mixing while ensuring that cached prefix states depend only on completed blocks and remain reusable—a strict requirement for high-throughput block-cached inference.
• Timestep-conditioning Factorization: Timestep embeddings are factored such that only the active block receives timestep modulation, preventing contamination of the clean prefix cache by step-dependent signals and supporting timestep-invariant prefix caching crucial for native blockwise sampling.
• Empirical Validation and Scaling: Experiments demonstrate that BDLM Mamba-H achieves the best validation perplexity among 87M-parameter DCLM models and remains competitive with BDLM attention at 350M scale. Inference throughput benchmarks confirm substantial superiority in the long-context regime.

Architectural Details and Blockwise Diffusion

BDLMs factor sequence likelihood autoregressively over blocks, allowing each block’s denoising process to be performed conditioned only on the clean, cached prefix. The hybrid Mamba-attention layers are scheduled, with attention layers acting at defined intervals and Mamba layers providing linear-time state transitions via selective state-space mechanisms.

The bidirectionality of denoising is addressed by splitting the Mamba scan: the forward scan is applied to the entire clean prefix plus the current block, cached up to block boundaries, while the reverse scan is strictly localized to the current block. This ensures that only within-block representations depend on bidirectional context, with the prefix cache remaining time-agnostic and reusable across the block sequence.

Prefix-caching is formalized via a cacheable layer factorization, where:

$$\mathcal{C}^{\ell}(c) = P^{\ell}(c), \qquad h_{b,t}^{\ell+1} = B^{\ell}(h_{b,t}^{\ell}; \mathcal{C}^{\ell}(c), e_t)$$

Here, $P^{\ell}(c)$ caches the representation for the prefix independently of timestep $t$. Only $B^{\ell}$ applies timestep modulation, strictly to the active block.

Experimental Evaluation

Validation perplexities on C4-en and Paloma-C4 are reported at both 87M and 350M scales. BDLM Mamba-H exhibits the lowest NELBO-derived perplexity at 87M—61.6 versus 76.5 for BDLM attention and 83.7/87.6 for full-sequence baselines. At 350M, BDLM Mamba-H (C4-en PPL 38.2) closely matches BDLM attention (37.0), supporting the claim that cache-compatible hybridization does not compromise denoising quality.

Throughput benchmarks show important distinctions in the ability to scale to extreme sequence lengths:

• At 65K tokens, BDLM Mamba-H achieves a 19.7x speedup compared to full-sequence DiffuMamba-H.
• At 262K tokens, BDLM Mamba-H provides a 3.7x throughput gain over BDLM attention.

Throughput performance and the crossover against BDLM attention as sequence length grows are highlighted in the following comparison:

Figure 1: 700M random-initialized single-A100 inference throughput comparison for BDLM Mamba-H, BDLM attention, and baseline architectures across sequence lengths up to 262K tokens.

The methodology aligns with real-world inference workloads, utilizing block size 256 with BF16 and cacheable prefix states, benchmarking both synthetic and random-initialization settings on contemporary A100 hardware with up to 262K context.

Practical Implications and Future Directions

The BDLM Mamba-H architecture demonstrates that native cache-aligned diffusion LMs can support performant long-context generation, substantially reducing memory bandwidth constraints by eliminating the need to stream or recompute full-sequence KV caches for each inference step. This enables practical scaling of diffusion-based language modeling to lengths previously impossible or prohibitively expensive for full-sequence denoisers.

The presented all-block objective and architectural factorizations set a precedent for building cache- and SR-compatible hybrid state-space diffusion models, suggesting a general approach for integrating efficient mixers in diffusion LMs more broadly. The partial reverse Mamba hybrid also provides an explicit mechanism to balance bidirectional context for local denoising while supporting scalable caching—potentially generalizable to other SSMs or linear-time architectures.

Limitations include slower training throughput due to all-block objective evaluation but are offset by the long-context inference gains. The learning rate transfer—using Complete(d)P—was not exhaustively optimized for dLLMs, leaving open questions for hyperparameter transfer and further throughput-versus-quality scaling law investigations.

Conclusion

This work demonstrates that hybrid blockwise diffusion LLMs combining partially bidirectional Mamba layers and attention with a cache-aligned all-block objective yield strong language modeling performance and substantial long-context inference throughput gains. The design maintains competitive validation perplexity against attention-only BDLMs and is scalable to sequence lengths unfeasible for full-sequence diffusion denoisers. The partial reverse Mamba construction provides a principled architecture for realizing efficient, cacheable, and high-capacity diffusion LMs, with broad practical relevance for the deployment of next-generation long-context models.