训练具有部分双向性的混合块扩散语言模型

高通量长上下文生成是 大语言模型 的核心挑战之一。生成通常受内存带宽限制而不是受计算限制：每个解码步骤都必须从内存中传输累积的键/值 (KV) 缓存，因此带宽需求随着上下文长度的增加而增长，同时仅发出一个词元。因此出现了两种并行方法：通过有效的注意力变体和线性时间混合器（例如 Mamba）来减少内存访问，或者通过立即生成词元块来增加并行计算。然而，当将这两种想法结合起来时，就会出现技术挑战。早期的混合扩散模型（例如 DiffuMamba）使用双向 Mamba 混合，包括相对于因果生成的反向扫描。这种反向扫描需要扫描整个序列，因此它的状态不是仅前缀的，并且即使逐块进行扩散也不能精确地重用为缓存。我们提出了一种 BDLM Mamba-注意力混合体，通过将反向 Mamba 扫描限制在主动降噪块上来解决这一挑战，从而实现跨块的精确缓存。在 87M 参数 DCLM 扫描中，与 BDLM 注意力和全序列基线相比，BDLM Mamba-H 实现了最佳的 C4-en 验证困惑度。在 350M 参数下，它与 BDLM 注意力仍然具有竞争力。对于长上下文推理，BDLM Mamba-H 在 65K 词元时达到全序列 DiffuMamba-H 吞吐量的 19.7 倍，在 262K 时达到 BDLM 注意力吞吐量的 3.7 倍，表明 Mamba 混合体是一种潜在的长上下文扩散架构。