言語モデルからの高速制御生成：適応的加重棄却サンプリングを用いた手法

Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling

April 7, 2025

著者: Benjamin Lipkin, Benjamin LeBrun, Jacob Hoover Vigly, João Loula, David R. MacIver, Li Du, Jason Eisner, Ryan Cotterell, Vikash Mansinghka, Timothy J. O'Donnell, Alexander K. Lew, Tim Vieira

cs.AI

要旨

制約付き言語モデルからの生成において主流のアプローチは、局所的に制約されたデコーディング（LCD）であり、各タイムステップでトークンを逐次的にサンプリングし、制約が常に満たされるようにするものです。通常、これはトークンマスキングによって実現されます。つまり、語彙全体をループし、制約に適合しないトークンを除外します。このアプローチには2つの重要な問題があります。(i) すべてのトークンに対して制約を評価することは非常にコストがかかる可能性があります。言語モデルの語彙はしばしば10万トークンを超えるためです。(ii) LCDは文字列全体の分布を歪める可能性があり、局所的な情報のみに基づいてトークンをサンプリングするため、行き止まりのパスに導くことがあります。本研究では、これらの両方の問題に対処する新しいアルゴリズムを提案します。まず、生成の各ステップで語彙全体に対して制約を評価することを避けるために、通常は桁違いに少ない制約評価で済む適応的リジェクトサンプリングアルゴリズムを提案します。次に、このアルゴリズムを拡張して、非常に少ない追加コストで低分散かつ不偏な重要度重みの推定値を生成する方法を示します。これらの推定値は、以前に提案された逐次モンテカルロアルゴリズム内で使用することで、局所的な制約適用の近視眼的な挙動を補正することができます。テキストからSQL、分子合成、目標推論、パターンマッチング、JSONといった領域での広範な実証評価を通じて、我々のアプローチが最先端のベースラインを上回り、より広範なクラスの制約をサポートし、実行時間と性能の両方を改善することを示します。追加の理論的および実証的分析により、我々の手法の実行効率は、計算の動的な使用によって駆動され、制約なしと言語モデル間の乖離に応じてスケーリングし、その結果、より優れたモデルに対して実行時間の改善が大きくなることを示しています。

English

The dominant approach to generating from language models subject to some constraint is locally constrained decoding (LCD), incrementally sampling tokens at each time step such that the constraint is never violated. Typically, this is achieved through token masking: looping over the vocabulary and excluding non-conforming tokens. There are two important problems with this approach. (i) Evaluating the constraint on every token can be prohibitively expensive -- LM vocabularies often exceed 100,000 tokens. (ii) LCD can distort the global distribution over strings, sampling tokens based only on local information, even if they lead down dead-end paths. This work introduces a new algorithm that addresses both these problems. First, to avoid evaluating a constraint on the full vocabulary at each step of generation, we propose an adaptive rejection sampling algorithm that typically requires orders of magnitude fewer constraint evaluations. Second, we show how this algorithm can be extended to produce low-variance, unbiased estimates of importance weights at a very small additional cost -- estimates that can be soundly used within previously proposed sequential Monte Carlo algorithms to correct for the myopic behavior of local constraint enforcement. Through extensive empirical evaluation in text-to-SQL, molecular synthesis, goal inference, pattern matching, and JSON domains, we show that our approach is superior to state-of-the-art baselines, supporting a broader class of constraints and improving both runtime and performance. Additional theoretical and empirical analyses show that our method's runtime efficiency is driven by its dynamic use of computation, scaling with the divergence between the unconstrained and constrained LM, and as a consequence, runtime improvements are greater for better models.