将线性注意力改造为扩散语言模型

扩散语言模型 (dLLM) 通过并行解码加速推理，为自回归模型提供了一种有前景的替代方案。最近的 dLLM 通常使用分块半自回归解码，自回归生成块，同时并行对每个活动块内的标记进行去噪。然而，尽管有 KV 缓存，每个去噪步骤仍然会关注所有先前的块，从而反复产生前缀注意成本。受这一瓶颈的推动，我们询问是否可以通过线性化先前块的注意力来进一步加速 dLLM 推理。我们引入了块混合注意力，它在主动降噪块内保留精确的 softmax 注意力，同时对之前的块应用线性注意力。我们证明，这种混合注意力可以改造为具有最小 后训练 的预训练 dLLM：LLaDA-Hybrid 替换了 LLaDA~2.1（一个 16B 开源 dLLM）中 20 个注意力层中的 6 个，很大程度上遵循 LoLCAT（Zhang 等人，2024）。转换只需要大约 60 小时，同时保持基准性能：HumanEval 上为 72.0% vs. 75.6%，MBPP+ 上为 63.0% vs. 57.7%，CMATH 上为 86.7% vs. 88.3%。通过 Triton 实现，LLaDA-Hybrid 的解码吞吐量提高了 1.7 倍，并在耗尽内存之前支持更多并发请求，这表明预训练的 dLLM 可以有效地线性化，以实现更快的推理。我们的代码位于：https://github.com/Diuven/LLaDA-Hybrid。