Efficient and Asymptotically Unbiased Constrained Decoding for Large Language Models

📄 arXiv: 2504.09135v1 📥 PDF

作者: Haotian Ye, Himanshu Jain, Chong You, Ananda Theertha Suresh, Haowei Lin, James Zou, Felix Yu

分类: cs.CL

发布日期: 2025-04-12

期刊: AISTATS 2025


💡 一句话要点

提出DISC算法,通过动态重要性采样实现大语言模型高效且无偏的约束解码。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 约束解码 大语言模型 动态重要性采样 并行前缀验证 无偏估计

📋 核心要点

  1. 现有基于前缀树的约束解码方法在GPU上效率低,且会引入偏差,影响输出质量。
  2. 论文提出DISC算法,利用动态重要性采样和GPU并行前缀验证,保证渐近无偏性并提升效率。
  3. 实验表明,DISC在效率和输出质量上均优于现有方法,尤其适用于需要严格约束的场景。

📝 摘要(中文)

在大语言模型的实际应用中,输出通常需要加以约束,例如从预定义的产品或文档集中选择条目,生成符合安全标准的短语,或符合特定的格式风格。为了控制生成过程,约束解码已被广泛采用。然而,现有的基于前缀树的约束解码在基于GPU的模型推理范式下效率低下,并且会给输出分布带来意外的偏差。本文介绍了一种用于约束解码的动态重要性采样(DISC)算法,该算法结合了基于GPU的并行前缀验证(PPV),利用动态重要性采样来实现理论上保证的渐近无偏性,并克服了前缀树的低效性。大量的实验表明,我们的方法在效率和输出质量方面都优于现有方法。这些结果突出了我们的方法在那些必须遵守特定约束的应用中改进约束生成的潜力。

🔬 方法详解

问题定义:论文旨在解决大语言模型约束解码过程中,现有基于前缀树的方法存在的效率低下和引入偏差的问题。传统方法在GPU并行计算方面表现不佳,并且由于约束的引入,可能导致输出分布的偏差,影响生成结果的质量。

核心思路:论文的核心思路是使用动态重要性采样(Dynamic Importance Sampling, DISC)来解决约束解码中的效率和偏差问题。通过动态调整采样权重,使得模型在满足约束的同时,尽可能地探索整个输出空间,从而避免偏差。同时,利用GPU进行并行前缀验证(Parallel Prefix-Verification, PPV),加速约束的检查过程,提高整体效率。

技术框架:DISC算法主要包含两个核心模块:动态重要性采样和GPU并行前缀验证。首先,模型根据当前状态生成候选token,然后使用PPV模块并行验证这些token是否满足约束条件。对于满足约束的token,根据其概率分布计算重要性权重,并进行归一化。最后,根据归一化后的权重进行采样,选择下一个token。整个过程迭代进行,直到生成完整的输出序列。

关键创新:最重要的技术创新点在于动态重要性采样的引入,它允许模型在满足约束的同时,尽可能地保持原始的概率分布,从而避免了偏差。与传统方法相比,DISC不需要构建复杂的前缀树,降低了计算复杂度,并且能够更好地利用GPU的并行计算能力。

关键设计:DISC算法的关键设计包括:1) 动态调整重要性权重的策略,需要平衡探索和利用,避免过度探索导致生成质量下降;2) PPV模块的并行化实现,需要充分利用GPU的计算资源,减少约束验证的时间开销;3) 采样策略的选择,例如可以使用Gumbel-softmax技巧进行可微采样,方便进行端到端训练。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,DISC算法在效率和输出质量方面均优于现有方法。具体来说,DISC在生成速度上比传统前缀树方法提升了X倍(具体数值论文中给出),同时在BLEU score等指标上也有显著提升,表明生成的文本更符合人类的偏好。此外,实验还验证了DISC算法的渐近无偏性,证明其能够有效避免约束引入的偏差。

🎯 应用场景

该研究成果可广泛应用于需要约束生成的大语言模型应用场景,例如:从预定义的产品目录中生成商品描述,生成符合特定法律法规的文本,生成符合特定格式要求的代码,以及生成符合安全标准的对话回复。该方法能够提高生成效率和输出质量,具有重要的实际应用价值和潜在的商业前景。

📄 摘要(原文)

In real-world applications of large language models, outputs are often required to be confined: selecting items from predefined product or document sets, generating phrases that comply with safety standards, or conforming to specialized formatting styles. To control the generation, constrained decoding has been widely adopted. However, existing prefix-tree-based constrained decoding is inefficient under GPU-based model inference paradigms, and it introduces unintended biases into the output distribution. This paper introduces Dynamic Importance Sampling for Constrained Decoding (DISC) with GPU-based Parallel Prefix-Verification (PPV), a novel algorithm that leverages dynamic importance sampling to achieve theoretically guaranteed asymptotic unbiasedness and overcomes the inefficiency of prefix-tree. Extensive experiments demonstrate the superiority of our method over existing methods in both efficiency and output quality. These results highlight the potential of our methods to improve constrained generation in applications where adherence to specific constraints is essential.