论文

DMax:dLLM 的主动并行解码

DMax: Aggressive Parallel Decoding for dLLMs

模型推理批处理与并行

摘要

我们提出了 DMax,一种高效扩散语言模型 (dLLM) 的新范例。它可以减少并行解码中的错误累积,从而在保持生成质量的同时实现积极的解码并行性。与通过二进制掩码到词元转换进行解码的传统掩码 dLLM 不同,DMax 将解码重新表述为从掩码嵌入到词元嵌入的渐进式自我细化。我们方法的核心是 同策略 统一训练,这是一种新颖的训练策略,可以有效地统一屏蔽和统一 dLLM,使模型能够从屏蔽输入及其自身的错误预测中恢复干净的标记。在此基础上,我们进一步提出软并行解码。我们将每个中间解码状态表示为预测词元嵌入和掩码嵌入之间的插值,从而实现嵌入空间中的迭代自我修正。跨各种基准的大量实验证明了 DMax 的有效性。与原始的 LLaDA-2.0-mini 相比,我们的方法将 GSM8K 上的 TPF 从 2.04 提高到 5.47,同时保持精度。在 MBPP 上,它将 TPF 从 2.71 提高到 5.86,同时保持相当的性能。在两个 H200 GPU 上,我们的模型在批量大小为 1 时平均达到 1,338 TPS。代码位于:https://github.com/czg1225/DMax