研究人员提出了一种混合块扩散语言模型（BDLM），结合了 Mamba 和注意力架构，解决了早期混合扩散模型中固有的缓存挑战。通过将反向 Mamba 扫描限制为仅激活的去噪块，该方法实现了跨块的精确 KV 缓存复用，克服了全序列扫描的前缀-only 限制。

  • BDLM Mamba-H 架构将反向扫描限制在激活的去噪块上，从而在逐块扩散过程中实现精确缓存。
  • 在 DCLM 上的 87M 参数遍历中，BDLM Mamba-H 相比 BDLM attention 和全序列基线取得了最佳的 C4-en 验证困惑度。
  • 在 350M 参数下，该模型在性能指标方面与 BDLM attention 保持竞争力。
  • 对于长上下文推理，BDLM Mamba-H 在 65K tokens 时的吞吐量达到全序列 DiffuMamba-H 的 19.7 倍，在 262K tokens 时达到 BDLM attention 的 3.7 倍。

这种方法证明了 Mamba 混合架构是长上下文扩散的可行方案，在保持竞争力困惑度的同时提供了显著的吞吐量提升。