近日,英伟达、香港大学与麻省理工学院的研究团队联合发布了一种名为 Fast-dLLM 的创新技术,旨在提升扩散语言模型的推理效率。与传统的自回归模型不同,扩散语言模型采用逐步去除文本噪声的方式生成文本,因此能够在一次迭代中生成多个单词,整体效率更高。然而,在实际应用中,许多开源扩散语言模型的推理速度仍然不如自回归模型,主要受限于缺乏键值(KV)缓存支持以及并行解码时生成质量下降。
KV 缓存是自回归模型中常用的加速推理技术,通过存储和重用之前计算的注意力状态,显著减少重复计算,从而提高生成速度。但由于扩散语言模型采用双向注意力机制,直接应用 KV 缓存并不容易。Fast-dLLM 架构的创新之处在于将文本生成过程划分为多个块,每个块包含一定数量的 token。通过这种块状生成方式,模型可以在生成一个块之前预先计算并存储其他块的 KV 缓存,进而避免重复计算。
尽管 KV 缓存机制有效提升了推理速度,但在并行解码时,生成质量往往会下降。这是因为扩散型模型在解码时假设条件独立性,然而标记之间可能存在复杂的依赖关系。为了解决这一问题,Fast-dLLM 提出了一种基于置信度的平行解码策略。在每个解码步骤中,模型会计算每个标记的置信度,并选择置信度超过阈值的标记进行解码。这种策略确保了高置信度的情况下能够安全进行并行解码,从而维持生成文本的连贯性和准确性。
为验证 Fast-dLLM 的性能,研究人员在 NVIDIA A10080GB GPU 上对 LLaDA 和 Dream 两种扩散语言模型进行了全面评估,涵盖了数学推理和代码生成等任务。在 KV 缓存机制的测试中,块大小为32时,模型的吞吐量达到了54.4tokens/s,准确率为78.5%。并行解码测试中,采用动态阈值策略优于固定 token 数基线。整体来看,LLaDA 模型在 GSM8K 任务中仅用 KV Cache 加速3.2倍,并行解码加速2.5倍,两者结合的速度提升达到8.1倍,生成长度达到1024时的端到端加速更是高达27.6倍。所有测试结果显示,Fast-dLLM 在加速的同时,保持了生成质量的稳定。
划重点: