Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

Technical Summary: Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

Problem Statement
High-throughput long-context generation in large language models (LLMs) is primarily constrained by memory bandwidth rather than compute. In standard autoregressive decoding, every step requires streaming the accumulated key/value (KV) cache from memory, causing bandwidth demand to scale linearly with context length while only emitting one token. While two parallel approaches exist to mitigate this—reducing memory access via linear-time mixers (e.g., Mamba) and increasing parallelism via block generation—combining them presents a specific technical hurdle.

Earlier hybrid diffusion models, such as DiffuMamba, utilize fully bidirectional Mamba scans, including a reverse-direction pass over the entire denoising window. This reverse scan depends on active tokens, meaning its internal states are not prefix-only. Consequently, these states cannot be precisely cached and reused across blocks, negating the efficiency gains of block diffusion. The core challenge is to design a hybrid architecture that supports block-wise generation with exact caching while retaining the benefits of bidirectional mixing.

Methodology
The authors propose BDLM Mamba-H, a hybrid architecture combining Block Diffusion Language Models (BDLMs) with Mamba and attention layers. The core innovation is a partial bidirectionality mechanism that restricts the reverse-direction Mamba scan to the active denoising block only, rather than the full sequence.

Key methodological components include:

1. Cache-Aligned Training Objective: The model employs an "all-block" objective where completed prefix blocks are stored as:
   • KV states for attention layers.
   • Forward-only Mamba boundary states (recurrent states) for Mamba layers.
     The reverse scan is confined to the current active block ($b_t$), ensuring that the prefix cache ($C_\ell(c)$) remains a function of the clean prefix alone, independent of the reverse diffusion timestep ($t$).
2. Timestep Factorization: To maintain cache validity, the architecture factors the layer computation into a timestep-free prefix operator and a timestep-conditioned active-block continuation. This prevents the clean-prefix cache from becoming dependent on the specific reverse diffusion step, a requirement for valid reuse during inference.
3. Training Protocol: The model is trained using a discrete diffusion objective over blocks. The prefix cache is constructed by scanning the clean prefix, and downstream block losses directly train the representations stored in these prefix states, aligning the training computation with the inference-time sampler's cache interface.

Key Contributions

• Cache-Aligned Training: The authors formulate a BDLM all-block objective where cached objects (KV states and Mamba prefix states) match inference-time block caches exactly. This allows downstream block losses to train the prefix scan used for recurrent state reuse.
• Controlled Validation and Scale-Up: The authors conduct a comprehensive sweep at 87M parameters on the DCLM dataset and scale selected models to 350M parameters using a Complete(d)P-style learning-rate transfer rule.
• Long-Context Throughput Gains: The study evaluates inference throughput at extended context lengths, demonstrating that the hybrid approach significantly outperforms both full-sequence baselines and pure attention-based block diffusion models as context length increases.

Results

• Validation Quality (87M Parameters): BDLM Mamba-H achieved the best C4-en validation perplexity (61.6) among the 87M models, outperforming BDLM attention (76.5), full-sequence DiffuMamba-H (83.7), and full-sequence attention (87.6).
• Validation Quality (350M Parameters): At the 350M scale, BDLM Mamba-H remained competitive with BDLM attention, achieving a C4-en perplexity of 38.2 compared to 36.97 for BDLM attention.
• Throughput Performance:
  • At 65K tokens, BDLM Mamba-H achieved 19.7x the throughput of the full-sequence DiffuMamba-H baseline.
  • At 262K tokens, BDLM Mamba-H achieved 3.7x the throughput of BDLM attention.
  • While full-sequence models are faster at short contexts (8192 tokens) due to the overhead of the all-block objective, the native BDLM models with caching demonstrate superior scalability in long-context regimes.

Significance and Claims
The paper claims that BDLM Mamba-H represents a viable architecture for long-context diffusion. By restricting the reverse Mamba scan to the active block, the model enables exact caching across blocks, solving the state-reuse problem inherent in previous bidirectional diffusion hybrids. The authors assert that this approach allows hybrid Mamba-attention models to retain competitive validation quality while offering substantial throughput improvements for long-context inference, positioning block partial-reverse hybrids as a practical direction for scaling diffusion language models.

The paper remains modest regarding optimality, noting that the learning rate transfer rules used for scaling are strong candidates based on existing parameterization laws but not proven optimal for discrete diffusion language models (dLLMs). Additionally, the authors acknowledge that while the block objective improves long-context generation, it currently incurs a training throughput penalty (approx. 25–34% slower) compared to full-sequence baselines on current hardware configurations.