技術要約：扱いやすい提案による局所的制約デコーディングにおけるバイアスの軽減

問題提起

大規模言語モデル（LLM）は、JSONスキーマ、有効なSQLクエリ、または特定の関数呼び出し形式といった、論理的または構造的な制約を厳密に遵守した出力を生成することに失敗することがよくあります。指示チューニングや選好最適化は一般的なアライメントを改善しましたが、「ハード」な制約を強制するには依然として信頼性に欠けます。

既存のアプローチは、主に2つの制限に直面しています。

1. 拒絶サンプリング（Rejection Sampling）: バイアスはありませんが、複雑な制約に対しては受理率が極端に低くなるため、計算量的に非現実的です。
2. 局所的制約デコーディング（Locally Constrained Decoding: LCD）: この標準的な手法は、各ステップで制約に即座に違反するトークンをマスクします。しかし、これは以下の2つの決定的な欠陥を抱えています。
   • バイアス: 未来の妥当性を考慮せずに現在のステップでマイオピック（近視眼的）にトークンをマスクするため、モデルの潜在的な分布を歪め、サンプリングバイアスを引き起こします。
   • 終了保証の欠如: LCDは局所的な妥当性のみを保証します。あるシーケンスは、固定されたトークン予算 $n$ 内で制約を満たすことができなくても、局所的には無限に妥当であり続ける可能性があります（例：閉じられていない中括弧を持つJSON）。

これらの問題に対処するため、LCDを提案分布として使用し、粒子（パーティクル）を再重み付けすることでターゲット分布に収束させる逐次モンテカルロ（SMC）フレームワークが提案されてきました。しかし、SMCの収束速度は提案の質によってボトルネックとなります。LCDは、未来の有効な補完に関する確率的情報を持たず、有限の予算内で制約充足を保証しないため、弱い提案となります。

手法

著者らは、制約付き分布 $p_{lm}(\cdot \mid C)$ からのSMCサンプリングのための、扱いやすい提案（tractable proposals）とポテンシャル関数を構築する汎用的なアプローチを提案しています。核心となる革新は、制約を有限オートマトンとして表現し、効率的なGPU実行のためにテンソル演算を活用することにあります。

1. 全域的制約デコーディング (Globally Constrained Decoding: GCD)

GCDは、生成されるすべてのトークンが残りのトークン予算 $n$ 以内で有効な補完につながることを保証することで、LCDを改善します。

• テンソル化された有限オートマトン: 有限オートマトン（FA）、具体的には非決定性有限オートマトン（NFA）および決定性有限オートマトン（DFA）としてエンコードされた制約を、テンソル表現（遷移行列 $T_{out}, T_{in}$ および放出行列 $B_f$）にコンパイルします。これにより、GPU上での効率的な実行が可能になります。
• フォワード・バックワード推論:
  • フォワードプロセス: プレフィックス $x_{1:t-1}$ を観測した後に到達可能なFA状態の集合を計算します。
  • バックワードプロセス: ある状態が、ちょうど $n-t$ ステップ後に受理状態に到達できるかどうかを計算します。
• グローバルマスキング: トークンは、残りの予算内で有効な継続を許容する状態へと導く場合にのみ許可されます。これにより、長さ $n$ の生成された任意のシーケンスが制約 $C$ を満たすことが保証されます。

2. 確率的全域的制約デコーディング (Probabilistic GCD: P-GCD)

GCDは制約充足を保証しますが、そのマスクはバイナリ（有効/無効）であり、LLMの下での有効な補完の尤度を無視しています。P-GCDは、確率的情報を統合することでこれに対処します。

• HMM蒸留: LLMの分布は、先行研究（Zhang et al., 2023）に従い、隠れマルコフモデル（HMM）によって近似されます。
• 回路乗算（Circuit Multiplication）: テンソル化されたFA（エッジ上の未正規化HMMと見なされる）は、回路乗算アルゴリズムを用いて蒸留されたHMMと掛け合わされます。これは、FAとHMMの共有された代数的構造を活用しています。
• 積分布: 得られる積分布 $p_{prod}$ は、FAのハードな論理制約と、LLMの確率的な好みの両方を捉えます。
• 提案とポテンシャル:
  • 提案 ($q_{pgcd}$): 制約付きの次トークン確率を用いて、各ステップでLMの分布を再重み付けします。
  • ポテンシャル ($\phi_{pgcd}$): SMCにおける粒子の重みを、LMのプレフィックス確率と、制約充足に関するHMMの先読み推定値の積によって決定します。

3. 効率的な実装

著者らは、巨大なクロネッカー積行列を明示的に実体化することを避けるために、テンソルネットワーク演算（具体的にはクロネッカー・ベクター恒等式）を利用しています。これにより、遷移ステップの計算コストを $O((HE)^2)$ から $O(H^2E + HE^2)$ （ここで $H$ はHMMの隠れ状態サイズ、$E$ はFAのエッジ数）に削減し、効率的な自己回帰生成を可能にしています。

主な貢献

1. GCDおよびP-GCDフレームワーク: 予算内で制約充足を保証する、扱いやすい提案とポテンシャルを提供する、全域的制約デコーディングおよび確率的GCDの導入。
2. GPU上のテンソル化オートマトン: OutlinesやXGrammarといった既存のフレームワークのCPUバウンドな制限を克服し、効率的なGPUネイティブ推論のための有限オートマトン（NFAを含む）のテンソル化を初めて提案。
3. NFAのサポート: 複雑な制約（例：可変パラメータを持つ関数呼び出し）において、NFAがDFAよりも指数関数的に簡潔になり得ることを実証。これにより、DFAの構築がメモリオーバーフローを引き起こす可能性がある実世界のエージェントアプリケーションへのスケーラビリティを実現。
4. SMCとの統合: (P-)GCDがLCDベースの提案と比較して、SMCの収束を大幅に加速させ、ターゲット分布に到達するために必要な粒子数を大幅に減少させることを示す。

実験結果

著者らは、xLAM（関数呼び出し）、CommonGen（キーワードベースの生成）、Spider（Text-to-SQL）の3つのベンチマークで(P-)GCDを評価しました。

• 関数呼び出し (xLAM):
  • (P-)GCDは、同じ粒子数においてLCDよりも大幅に高い精度を達成しました。
  • P-GCDは、GCDよりも速く真の分布（拒絶サンプリング）に収束し、確率的情報の利点を浮き彫りにしました。
  • 制約充足: GCDおよびP-GCDは、すべての粒子予算において100%の制約充足率を達成しましたが、LCDは中程度の予算（例：JSONで~95%）でも制約充足に失敗しました。
  • 効率性: LCDおよびGCDのオーバーヘッドはわずか（1.08倍および1.10倍の減速）でした。P-GCDはHMM計算のため低速（3.11倍から10.2倍）でしたが、実行可能な範囲内に留まりました。
• キーワード生成 (CommonGen): P-GCDを提案とポテンシャルの両方に使用することで、LCDと比較してBLEUスコアと収束速度が大幅に向上しました。
• Text-to-SQL (Spider): GCDおよびP-GCDは、さまざまな粒子予算において、特に中程度の値において、LCDよりも実行精度を向上させました。
• NFA vs. DFA: 複雑なスキーマ（例：可変長関数パラメータ）において、NFAはサイズが線形に増加するのに対し、DFAは指数関数的に増加することを実験で確認しました。Spider SQL文法の構成では、DFA構築はトークンレベルでメモリ不足（OOM）エラーを引き起こしましたが、NFAは実行可能な状態（約32k状態）を維持しました。

意義と主張

本論文は、局所的制約デコーディングが相当なサンプリングバイアスを導入し、有限のトークン予算内で制約充足を保証できないことを主張しています。提案された(P-)GCDは、以下の解決策を提供します。

1. 妥当性の保証: 生成される任意のシーケンスの長さ $n$ が制約を満たすことを保証し、LCDの失敗ケースを排除します。
2. バイアスの低減: 確率的情報（P-GCD経由）とグローバルな先読みを取り入れることで、提案はターゲットとなる条件付き分布をより良く近似します。
3. 効率性の向上: SMCフレームワークにおいて、(P-)GCDはLCDベースのアプローチよりも大幅に少ない粒子数でターゲット分布に収束し、高品質な制約付き生成の計算コストを削減します。

本研究は、論理的制約（オートマトン経由）と確率的モデリング（HMM経由）をGPU効率の高い方法で統合することにより、厳格な構造的遵守を必要とする関数呼び出しやコード生成などのLLMアプリケーションにおける、信頼性の高い制約充足型生成への基礎的な一歩となるものです。