用中间分布桥接扩散与自回归模型,实现更快更准的文本生成。
Diffusion Language Model Parallel Decoding via Product-of-Experts Bridge
- 构建扩散与自回归模型的产物专家中间分布,降低采样难度。
- 相比标准扩散模型,速度提升5倍,逼近95%自回归模型性能。
- 适合需要高速高质生成的数学推理与代码生成任务。
扩散语言模型(DLM)通过并行解码提供显著速度优势,但缺乏词元依赖导致生成质量低于自回归(AR)模型。近期工作尝试通过重要性采样弥合差距,以DLM为提议分布,AR为目标分布。然而两者分布差异巨大,采样需大量粒子,计算成本高昂。本文提出PoE-Bridge框架,引入一个中间分布来缩小差距。该分布由DLM提议分布与AR目标分布的乘积专家(PoE)构造。先用DLM并行生成多个候选续写,再通过拒绝采样验证并引导候选向PoE分布靠拢;随后使用重要性采样进一步校正至AR目标。我们还提出混合温度采样以增强多样性,弹性拒绝窗口减少无效验证。实验表明,PoE-Bridge相较标准DLM解码速度提升5倍,且在数学推理和编程等挑战性任务上恢复至少95%的AR模型性能,有效缩小质量差距。代码已开源:https://github.com/juntongshi48/poe-bridge。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。