Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

2026-07-02

AI Summary

This work addresses the memory bandwidth bottleneck in large language models during long-context generation, caused by frequent reads from the key-value (KV) cache. To overcome this, the authors propose a hybrid architecture that integrates Mamba with attention mechanisms. By introducing a local bidirectional scanning scheme, the backward Mamba state is constrained to depend only on the current denoising block, enabling precise cross-block cache reuse—a capability absent in conventional hybrid diffusion models. Built upon the BDLM framework and trained with block-level diffusion, the method achieves state-of-the-art perplexity on the C4-en validation set with an 87M-parameter model, while a 350M-parameter variant remains competitive. At context lengths of 65K and 262K, inference throughput reaches 19.7× that of DiffuMamba-H and 3.7× that of BDLM attention, respectively.

Abstract

High-throughput long-context generation is one of the central challenges for large language models. 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.

Problem

• long-context generation
• hybrid diffusion models
• memory bandwidth
• KV cache reuse
• bidirectional scanning

Innovation

• block diffusion
• Mamba
• partial bidirectionality
• KV cache reuse
• long-context generation