Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling

生成符合某些约束的语言模型的主要方法是局部约束解码（LCD），它通过在每个时间步增量采样标记，确保约束从未被违反。通常，这是通过标记屏蔽实现的：遍历词汇表并排除不符合约束的标记。这种方法存在两个重要问题。（i）对每个标记评估约束可能代价过高——语言模型的词汇量通常超过100,000个标记。（ii）LCD可能会扭曲字符串的全局分布，因为它仅基于局部信息采样标记，即使这些标记可能导致无解路径。本文提出了一种新算法，解决了这两个问题。首先，为了避免在生成的每一步对整个词汇表进行约束评估，我们提出了一种自适应拒绝采样算法，该算法通常需要少几个数量级的约束评估。其次，我们展示了如何扩展该算法以极小的额外成本产生低方差、无偏的重要性权重估计值——这些估计值可以在之前提出的序贯蒙特卡洛算法中可靠地使用，以修正局部约束强制执行的短视行为。通过在文本到SQL、分子合成、目标推断、模式匹配和JSON领域的广泛实证评估，我们证明了我们的方法优于现有的最先进基线，支持更广泛的约束类型，并同时提升了运行时间和性能。进一步的理论和实证分析表明，我们的方法的运行效率得益于其动态计算使用方式，随着无约束和有约束语言模型之间的差异扩大而扩展，因此对于更好的模型，运行时间改进更为显著。

论文试图解决语言模型生成过程中受约束解码的两大问题：(i) 当前方法（如LCD）需要对整个词汇表进行逐词评估，计算成本高；(ii) 局部约束可能导致生成路径陷入死胡同，从而扭曲全局分布。这并非全新问题，但现有方法在效率和准确性之间难以平衡。

论文提出了一种新的算法，包含两个核心创新：(i) 一种自适应拒绝采样方法，通过动态选择候选集显著减少约束评估次数；(ii) 结合重要性权重估计，用于校正局部约束导致的偏差，同时保持低方差和无偏性。这种方法不仅提高了效率，还扩展了可支持的约束类型。

论文通过多个领域（如文本到SQL、分子合成等）的广泛实验验证了方法的有效性，并展示了其在运行时效率上的优势，特别是对于高质量模型。此外，作者提供了理论分析，证明算法的计算复杂度与未约束和约束模型之间的差异相关。代码尚未提及开源，但实验设计详尽，涵盖了多种数据集（如Spider、ZINC等）。未来研究可以进一步探索更复杂的约束条件和更大规模模型的应用。

近期相关工作包括：(1)「Constrained Decoding for Closed-Vocabulary Text Generation」，探讨了基于屏蔽技术的约束解码方法；(2)「Sequential Monte Carlo for Constrained Language Models」，提出了使用SMC框架改进约束生成；(3)「Guided Decoding with Pre-trained Language Models」，利用引导机制优化生成质量。这些研究主要集中在提高约束生成的灵活性或性能，而本论文的独特之处在于同时解决了效率和偏差问题。