arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2606.08048cs.CLcs.LGcs.PF
收藏

通过专家乘积桥接的扩散语言模型并行解码

Diffusion Language Model Parallel Decoding via Product-of-Experts Bridge

Stanford University(斯坦福大学)

Juntong Shi, Brian L. Trippe, Jure Leskovec, Stefano Ermon, Minkai Xu

AI总结:

提出PoE-Bridge框架,通过专家乘积构建中间分布,结合扩散语言模型并行解码和自回归模型质量,实现5倍加速并恢复至少95%的AR性能。

摘要:

扩散语言模型(DLM)通过并行解码提供了显著的速度优势,但与自回归(AR)模型相比,缺乏令牌依赖性限制了生成质量。最近的进展试图通过重要性采样来弥合差距,其中DLM作为提议分布,AR作为目标分布。然而,由于它们分布之间的巨大差距,采样需要大量粒子,因此计算成本高昂。在本文中,我们引入了PoE-Bridge,一种新颖的解码框架,通过引入中间分布来弥合差距,从而大幅提高生成速度和准确性。该分布被构建为DLM提议和AR目标的专家乘积(PoE)。借助中间分布,我们首先使用DLM并行起草多个续写,然后应用拒绝采样验证起草的令牌,并将结果候选向PoE移动。接着,我们使用重要性采样进一步将PoE对齐的候选向AR目标校正。我们还提出了若干改进技术,包括用于增强多样性的混合温度采样和用于减少浪费验证的弹性拒绝窗口。实验上,PoE-Bridge在标准DLM解码方法上实现了显著提高的准确性,速度提升5倍,并恢复了目标AR模型至少95%的性能,在具有挑战性的数学推理和编码任务上高效地推进了大部分质量差距。我们的代码可在https://github.com/juntongshi48/poe-bridge获取。

英文摘要:

Diffusion language models (DLMs) offer substantial speed advantages through parallel decoding, but the lack of token dependencies limits generation quality compared to autoregressive (AR) models. Recent progress attempts to bridge the gap via importance sampling, with DLM being the proposal and AR being the target. However, due to the huge gap between their distributions, the sampling requires a large number of particles and is thus expensive to compute. In this paper, we introduce PoE-Bridge, a novel decoding framework that drastically improves generation speed and accuracy by introducing an intermediate distribution to bridge the gap. The distribution is constructed as a Product-of-Experts (PoE) of the DLM proposal and the AR target. With the intermediate distribution, we first use the DLM to draft multiple continuations in parallel, then apply rejection sampling to verify the drafted tokens and move the resulting candidates toward the PoE. We then use importance sampling to further correct the PoE-aligned candidates toward the AR target. We further propose several improved techniques, including mixed-temperature sampling for enhanced diversity and elastic rejection windows for reducing wasted verification. Empirically, PoE-Bridge achieves significantly improved accuracy with $5\times$ speedup over the standard DLM decoding approach, and recovers at least 95% of the target AR model's performance, efficiently advancing most of the quality gap on challenging mathematical reasoning and coding tasks. Our code is available at https://github.com/juntongshi48/poe-bridge.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑