BDLM Mamba-H restricts reverse scan to enable exact caching in hybrid diffusion models

Researchers propose a hybrid Block Diffusion Language Model (BDLM) combining Mamba and attention architectures that addresses the caching challenges inherent in earlier hybrid diffusion models. By restricting the reverse Mamba scan to only the active denoising block, the method enables exact KV cache reuse across blocks, overcoming the prefix-only limitation of full-sequence scans.

• The BDLM Mamba-H architecture restricts reverse scanning to the active denoising block, enabling precise caching during block-by-block diffusion.
• In an 87M-parameter sweep on DCLM, BDLM Mamba-H achieves the best C4-en validation perplexity compared to BDLM attention and full-sequence baselines.
• At 350M parameters, the model remains competitive with BDLM attention in terms of performance metrics.
• 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.

This approach demonstrates that Mamba hybrids are a viable architecture for long-context diffusion, offering significant throughput improvements while maintaining competitive perplexity.