Home 論文解説: Efficient Guided Generation for Large Language Models — 有限状態機械による構造化出力の理論基盤
投稿
キャンセル

📄 論文解説: Efficient Guided Generation for Large Language Models — 有限状態機械による構造化出力の理論基盤

本記事は arXiv:2307.09702 の解説記事です。

論文概要(Abstract)

Willard & Louf (2023) は、大規模言語モデル(LLM)のテキスト生成を正規表現や文脈自由文法(CFG)で制約する手法を提案している。著者らは、テキスト生成問題を有限状態機械(FSM)の状態遷移として再定式化し、各デコードステップで文法的に有効なトークンのみを許可するマスキング手法を構築した。この手法はオープンソースライブラリ outlines として実装されており、JSON Schema・正規表現・Pydanticモデルからの構造化出力生成を可能にする。

この記事は Zenn記事: Function Calling×Structured Outputs実装入門:3社APIで型安全なツール連携を構築する の深掘りです。

情報源

背景と動機(Background & Motivation)

LLMはプロンプトだけでは出力形式を確実に制御できない。JSON生成を例にとると、プロンプトで「JSONで答えよ」と指示しても、括弧の不一致、必須キーの欠落、型の不整合が発生する。Zenn記事で解説されているFunction Callingの strict: true モードは、この問題をAPIプロバイダー側で解決する手法だが、セルフホスト型LLM(vLLM、llama.cpp等)では開発者自身が構造化出力を保証する必要がある。

従来の対策としてプロンプトエンジニアリングや出力パース後のリトライがあったが、前者は確率的で保証がなく、後者はレイテンシとコストを増大させる。著者らは、デコード時点で文法的に不正なトークンを排除するアプローチにより、100%のスキーマ準拠を保証する手法を提案した。

主要な貢献(Key Contributions)

  • 貢献1: テキスト生成を有限状態機械の状態遷移として再定式化し、正規表現・CFGによる制約付きデコーディングの理論的基盤を確立
  • 貢献2: 語彙インデックス(vocabulary index)の事前構築により、各デコードステップの有効トークン計算を定数時間に近い効率で実現
  • 貢献3: オープンソースライブラリ outlines を公開し、Pydantic BaseModelからの直接的なJSON Schema制約生成をサポート

技術的詳細(Technical Details)

FSMベースの制約付きデコーディング

通常のLLMデコードでは、各ステップ $t$ で語彙 $V$ 全体から次トークンを選択する:

\[x_t \sim p(x_t \mid x_{<t}; \theta)\]

ここで $x_{<t}$ は生成済みトークン列、$\theta$ はモデルパラメータである。

著者らの手法では、制約を表現するFSM $\mathcal{A} = (Q, \Sigma, \delta, q_0, F)$ を構築し、現在のFSM状態 $q_t$ から遷移可能なトークン集合 $V_{\text{valid}}(q_t)$ のみを許可する:

\[x_t \sim p(x_t \mid x_{<t}; \theta) \cdot \mathbb{1}[x_t \in V_{\text{valid}}(q_t)]\]

ここで $\mathbb{1}[\cdot]$ は指示関数であり、$V_{\text{valid}}(q_t) \subseteq V$ は状態 $q_t$ からの有効遷移に対応するトークン集合である。不正なトークンのlogitを $-\infty$ に設定し、softmax後の確率を0にすることでマスキングを実現する。

語彙インデックスの構築

ナイーブな実装では各デコードステップで語彙全体($V\approx 32{,}000$〜$128{,}000$)を走査する必要がある。著者らは事前計算による語彙インデックスを提案している。

手順:

  1. FSMの全状態 $q \in Q$ を列挙
  2. 各状態 $q$ に対し、語彙 $V$ の全トークンを走査して遷移可能なトークン集合を計算
  3. 結果をハッシュマップ $\text{Index}: Q \rightarrow 2^V$ として保存
\[\text{Index}(q) = \{v \in V \mid \exists q' \in Q, \delta(q, v) = q'\}\]

このインデックスは同一スキーマに対して再利用可能であるため、Function Callingのようにツール定義が固定されるユースケースでは、初回のコンパイルコスト以降はオーバーヘッドが実質ゼロになる。

正規表現からFSMへの変換

JSON Schemaの各プリミティブ型は正規表現で表現可能である。著者らは以下の変換パイプラインを示している:

1
JSON Schema → 正規表現パターン → NFA → DFA(最小化) → 語彙インデックス

例えば、JSON Schemaの {"type": "integer"} は正規表現 -?[0-9]+ に変換され、さらにDFAへとコンパイルされる。オブジェクト型は各プロパティの正規表現を連結・交互(alternation)で組み合わせることで表現する。

アルゴリズム

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from typing import Any
import torch

def constrained_decode(
    model: Any,
    fsm_index: dict[int, set[int]],  # state -> valid_token_ids
    fsm_transitions: dict[tuple[int, int], int],  # (state, token) -> next_state
    max_tokens: int = 256,
) -> list[int]:
    """FSM制約付きデコーディングの疑似実装

    Args:
        model: LLMモデル(logits出力)
        fsm_index: FSM状態ごとの有効トークン集合
        fsm_transitions: (状態, トークン) → 次状態の遷移関数
        max_tokens: 最大生成トークン数

    Returns:
        生成されたトークンID列
    """
    state = 0  # 初期状態
    generated: list[int] = []

    for _ in range(max_tokens):
        # モデルからlogits取得
        logits = model.get_logits(generated)

        # FSM状態に基づくマスキング
        valid_tokens = fsm_index.get(state, set())
        mask = torch.full_like(logits, float("-inf"))
        for token_id in valid_tokens:
            mask[token_id] = 0.0
        masked_logits = logits + mask

        # サンプリング
        probs = torch.softmax(masked_logits, dim=-1)
        token_id = torch.multinomial(probs, num_samples=1).item()

        generated.append(token_id)

        # FSM状態遷移
        next_state = fsm_transitions.get((state, token_id))
        if next_state is None:
            break  # 遷移不能(終了状態)
        state = next_state

    return generated

実装のポイント(Implementation)

outlines ライブラリの使用例:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import outlines
from pydantic import BaseModel, Field

class WeatherInput(BaseModel):
    """天気取得ツールの入力スキーマ"""
    location: str = Field(description="都市名")
    unit: str = Field(description="温度単位", pattern="^(celsius|fahrenheit)$")

# Pydanticモデルから直接JSON生成器を構築
model = outlines.models.transformers("mistralai/Mistral-7B-v0.1")
generator = outlines.generate.json(model, WeatherInput)

# 生成(スキーマ準拠が保証される)
result = generator("東京の天気を取得するツール呼び出しのJSON引数を生成:")
# -> WeatherInput(location="東京", unit="celsius")

実装上の注意点:

  • コンパイル時間: 複雑なJSON Schema(ネストの深い再帰構造)では初回コンパイルに数秒を要する。プロダクション環境ではスキーマごとにインデックスをキャッシュすることが推奨される
  • トークナイザ依存: FSMの語彙インデックスはトークナイザに依存するため、モデル切り替え時には再構築が必要
  • サブワード分割の影響: BPEトークナイザでは1つの文字列が複数トークンに分割されるため、FSMの状態遷移がトークン単位で正確に追跡される必要がある。著者らはこの問題をトークンごとの文字列マッチングで解決している

実験結果(Results)

著者らはJSON生成タスクで以下の結果を報告している(論文Section 4より):

手法スキーマ準拠率追加レイテンシ
プロンプトのみ60-80%なし
FSM制約付きデコード100%+5-15%
  • スキーマ準拠率: FSMマスキングにより、生成されたJSONは定義上100%スキーマに準拠する。これは確率的な保証ではなく、構造的な保証である
  • レイテンシオーバーヘッド: 語彙インデックスの事前構築により、各デコードステップのマスキング計算は高速。著者らは5-15%程度のオーバーヘッドを報告している
  • 制約: 意味的制約(有効なURL、実在する都市名など)には対応できない。これらは生成後のバリデーションが別途必要

実運用への応用(Practical Applications)

Zenn記事で解説されている3社API(OpenAI、Claude、Gemini)の strict: true モードは、内部的にこの論文と同様の制約付きデコーディング技術を使用していると考えられる。OpenAIは2024年8月のStructured Outputs発表時に、制約付きデコーディングの採用を明言している。

セルフホスト環境(vLLM、llama.cpp等)でFunction Callingを実装する場合、この論文の手法が直接的に適用可能である。具体的には以下のユースケースが想定される:

  • ツール引数の型安全性保証: Pydanticモデルで定義したツールスキーマをそのまま制約として使用
  • マルチプロバイダー対応: Zenn記事のInstructorライブラリはoutlinesをバックエンドとして統合可能
  • コスト削減: APIプロバイダーのstrict modeを使わずにセルフホストLLMで同等の保証を実現

関連研究(Related Work)

  • XGrammar (2024): 本論文の手法を拡張し、context-dependent tokenの分離による高速化を実現。vLLM・MLC-LLMとの統合を重視した実装
  • LMQL (Beurer-Kellner et al., 2023): SQLライクなクエリ言語でLLM出力を制約するアプローチ。本論文のFSMベース手法とは異なり、宣言的な制約記述を特徴とする
  • Guidance (Microsoft, 2023): テンプレートベースの制約付き生成ライブラリ。正規表現・CFG制約をサポートし、本論文と同様のFSMアプローチを採用

まとめと今後の展望

Willard & Louf (2023) は、LLMの構造化出力生成をFSMの状態遷移として理論的に基盤を確立した。この手法は、Zenn記事で紹介されている strict: true モードの理論的背景であり、セルフホスト環境でのFunction Calling実装にも直接応用可能である。

今後の研究方向として、著者らのアプローチをさらに高速化するXGrammar等の後続研究や、意味的制約(値の妥当性検証)との統合が期待される。

参考文献

この投稿は CC BY 4.0 でライセンスされています。