Training Hybrid Block Diffusion Language Models with Partial Bidirectionality
Proposed a partial bidirectional BDLM Mamba-attention hybrid, boosting long-context inference throughput with parameters from 87M to 350M.
Key Findings
Methodology
This work combines diffusion language models with the Mamba linear-time mixer, restricting the reverse Mamba scan to active denoising blocks to enable exact cache reuse. Training employs a full-block objective, aligning training caches with inference. The model at 87M parameters achieves a C4 validation perplexity of 61.6, outperforming baselines, and remains competitive at 350M. For long-context inference, BDLM Mamba-H attains 19.7x throughput over full-sequence DiffuMamba-H at 65K tokens and 3.7x over BDLM attention at 262K, demonstrating scalability.
Key Results
- At 87M parameters, BDLM Mamba-H achieves a validation perplexity of 61.6 on C4-en, better than BDLM attention (76.5) and full-sequence baselines (83.7). At 350M, perplexity remains comparable (~38.2), with MCQA accuracy reaching 43.3%. In long inference, throughput at 65K tokens exceeds DiffuMamba-H by 19.7x, and at 262K tokens, it surpasses BDLM attention by 3.7x.
- Training uses cache-aligned full-block objectives, ensuring the caches built during training match those during inference. The approach effectively mitigates the bottleneck caused by full-sequence models, enabling efficient long-context generation.
- Limiting reverse Mamba to active blocks prevents non-prefix reverse states, thus maintaining cache reusability. The method demonstrates significant improvements in inference speed while preserving model quality, especially for very long sequences.
Significance
This research addresses the critical challenge of long-context generation in large language models by integrating partial bidirectionality with linear-time mixing. It offers a scalable solution that balances inference efficiency and predictive accuracy, crucial for applications like document summarization, long dialogue, and multi-modal understanding. The approach paves the way for deploying large models in resource-constrained environments without sacrificing performance, thus impacting both academia and industry. The demonstrated throughput gains at multi-hundred-thousand token lengths mark a significant step toward practical long-text AI systems.
Technical Contribution
The key innovation is restricting the reverse Mamba scan to the active denoising block, enabling exact cache reuse during training and inference. The full-block training objective aligns the model’s cache construction with inference, while the timestep-free prefix cache design enhances flexibility. The transfer of learning rates across model scales using Complete(d)P-style rules facilitates effective scaling from 87M to 350M parameters. Additionally, the model incorporates timestep conditioning factors that are only applied within active blocks, maintaining cache invariance. These contributions collectively enable high throughput for long sequences while maintaining competitive perplexity.
Novelty
This work is the first to introduce a partial bidirectional Mamba hybrid that confines reverse scanning to active blocks, ensuring cache reusability in diffusion models. Unlike previous full-sequence bidirectional models, it maintains prefix-only states for efficient caching. The combination of full-block training and local bidirectionality represents a novel approach to scaling long-context diffusion models, bridging the gap between efficiency and quality in long-text generation.
Limitations
- Despite efficiency gains, training with the hybrid architecture still incurs higher computational costs (~25-30% slower than full-sequence models at 350M).
- The restriction of reverse scan to active blocks may impact global consistency in extremely long sequences, especially in tasks requiring full context understanding.
- Hardware dependency on high-memory GPUs like A100 may limit broader applicability without further optimization.
Future Work
Future research will focus on optimizing the reverse scan scope, exploring dynamic block sizes, and integrating multi-modal data. Enhancing the timestep conditioning mechanism for better controllability and extending the architecture to multi-task and multi-modal settings are promising directions. Additionally, efforts to reduce training costs and adapt the approach to diverse hardware environments will be pursued.
AI Executive Summary
The rapid growth of large language models has brought long-context generation to the forefront, yet existing architectures face significant bottlenecks due to memory bandwidth constraints. Autoregressive models, while powerful, require streaming key/value caches for each token, making long sequences computationally expensive. To address this, researchers have explored efficient attention variants and parallel block generation, but combining these approaches introduces technical challenges, especially in maintaining cache efficiency.
This paper introduces a novel hybrid architecture—partial bidirectional diffusion language models with Mamba attention—that restricts the reverse Mamba scan to active denoising blocks. This design ensures that caches built during training can be exactly reused during inference, significantly boosting throughput for long sequences. The training employs a full-block objective aligned with inference cache structures, ensuring consistency and efficiency. Experiments on models ranging from 87M to 350M parameters demonstrate that the BDLM Mamba-H achieves superior perplexity scores and remarkable inference speedups, reaching 19.7x faster than full-sequence diffusion models at 65K tokens and 3.7x faster than attention-based models at 262K tokens.
These advancements mark a significant step toward practical deployment of large-scale language models capable of handling extremely long contexts. The approach balances the need for high-quality predictions with computational efficiency, opening new avenues for applications like long-form summarization, extended dialogues, and multi-modal understanding. While the method shows promising results, challenges remain in further reducing training costs and extending the architecture's flexibility. Future work will focus on refining the reverse scan scope, optimizing hardware utilization, and exploring multi-modal extensions, aiming to make long-context AI systems more accessible and scalable.
Deep Dive
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.