LLMが構造的な出力を行う方法を調べてみた — デコーディング手法とその最適化

こんにちは、iimonでエンジニアをしている須藤です。

最近、アプリケーションでLLMを利用していて、JSON出力がどの程度安定しているのか気になっていました。
また、ClaudeやOpenAIが「100%のスキーマ準拠」を謳っていて、その裏側の仕組みも面白そうだったので、今回調べてみることにしました。

特にデコーディングの段階でどのように制約をかけるか、そしてその仕組みがどう最適化されているかを中心に見ていきます。

※ 本記事で主に扱うのは「文法的に正しい JSON をどう保証するか」です。型や minimum, pattern, format などの意味制約は本記事では触れません。

LLMが文法に従うことで得られるメリット

• LLMの出力を後続の処理に渡しやすくなる
• LLMが苦手な計算処理などをツールを使って処理しやすくなる

たとえば、文章から属性を抽出してコード上で分岐処理したり、計算処理を Python 側に逃がしたりしやすくなる
このように、構造化出力は周辺ツールやアプリケーションロジックとの接続を滑らかにします

従来の方法で構造（文法）に従わせるのが難しい理由

• プロンプトを工夫したり、fine-tuning を行ったりしても
  • LLMが確率的にトークンを選択している以上、100%文法に合致させるのが難しい
• 複数回に分けて応答させて、生成プロセス外で組み立てる
  • 単純な構造以外では、複雑さやコストが増大する

そこで、構造上許容されるもののみを生成候補として残すようにしたのが、Constrained Decoding という手法です

トークン生成プロセスを見る

LLMのトークン生成は単純化すると

1. プロンプトを受け取る
2. プロンプトをトークンにする
3. LLMに手順2を入れて、スコア(logitsベクトル)を出す
4. スコアにSoftmaxをかけて、確率分布(distribution)を出す
5. 特定のデコーディングアルゴリズムで、トークンを選択する
6. 手順5で選択したトークンを手順2のトークン列に追加する
7. トークンEOSが出るまたは、トークンのmax長まで、手順3~6を繰り返す
8. トークンをデトークナイズして応答を文字列で返す

概略図

token
↓ LLM
logits 4|-2|3|2.1|-1
↓ Softmax
distribution 0.66|2e-3|0.24|0.10|4e-3
↓ sampler
sampled token 4

Constrained Decoding では
手順3で出したスコアにマスクをかけてスコアを-infinityにすることで、確率を0にしている

このようにトークン生成プロセスでマスクをかけて、確率分布を操作し、出力するトークンに制約をかける
また decoding とは別のステップで制約 (mask) を適用している

Python で概念的に実装してみると

def llm_generation(prompt):
    current_tokens = tokenize(tokenizer, prompt)
    for i in range(max_token):
        scores = LLM(current_tokens)   # 次のtokenのlogits（softmax前のスコア）
        # マスク処理
        mask = generate_mask(current_tokens, tokenizer)
        masked_scores = apply_mask(mask, scores)
        token_i = decode(masked_scores)
        if token_i == EOS:
            break
        current_tokens.append(token_i)
    return detokenize(tokenizer, current_tokens)

マスクを作る方法

では、そのマスクを実際にどう作るのか、また計算量をどう抑えるかを見ていきます。

有限状態機械を使った、事前計算

（このセクションは主に Outlines 論文に基づく）

マスクを作るためには、文法とすでに生成済みのトークンに基づき、生成するトークンが有効かを判断する必要がある
LLMの語彙全てに対して評価する場合、語彙サイズをNとして、O(N)のコストが発生する

文法を有限状態機械(FSM, Finite State Machine)として表現すると
生成済みのトークンは状態を決定して、状態ごとに有効なトークンを紐づけることができる
各状態で許容されるトークン集合は事前計算してインデックスにすることで、毎ステップ語彙全走査を避けて効率的にマスクを生成できる

ではどうやって有限状態機械で文法を表現するのかを見ていきます

構文仕様（制約）の記述

簡単な規則であれば正規表現で表すことができる
ただし、参照や再帰などの柔軟な構造は表現できない
そこで制約を表現し、検査する仕組みとして CFG が有効となる

CFG(Context Free Grammar)

• 規則の集まりで構造を表現し、規則の中に他の規則への参照を含めることができる
  • 再帰的構造を表現可能
• 文法上更に展開できる記号である nonterminal（例: <expr>）と、それ以上展開できない terminal（例: "+", "(", NUMBER）で構成される
• start symbol（規則の出発点となる nonterminal）に対しての rule を定義していく
  • 例: <expr> ::= <expr> "+" <term> | <term>
• EBNF (Extended Backus Naur Form, CFG をコンパクトに書くための記法) 等で表現される

これで、文法的に許容される terminal 集合を求められるようになりました

文法の制約とトークンの制約の接続

（このセクションは主に SynCode 論文に基づく）

LLMのトークンは文法上の字句(例 "+","(")とは対応しておらず
複数トークンで字句を表現したり、逆に1つのトークンが複数の字句にまたがる可能性がある

つまり、以下のような可能性がある

• 現在のトークンが示す字句の種類が後から変化する
• 1つのトークンが次の字句まで指定してしまう（例: 1} というトークンは、数値 1 と構造記号 } の2つの字句を同時に指定してしまう）

対応方法

1. 生成済みのトークンをデコードして文字列にし（lexerで字句に分解）、文法上確定済みの部分 (lexical) と、変化しうる部分 (remainder) に分ける
2. lexical から、文法上の次に続きうる terminal 系列を計算する（これを accept sequence と呼ぶ）
   • LR parser（構文解析器）で、現在の parser state から shift（次の入力を受け入れる操作）可能な terminal を並べることを繰り返す
3. remainder を、accept sequence ごとの DFA（決定性有限オートマトン）に入れて状態として表現する
   • ここで「live state」は、その状態から先に進めば受理状態に到達できる状態を指す
   • 具体的には、terminal に一部マッチして将来完成できる／ちょうど完成（受理状態そのもの）／完成して残りの文字列が次の terminal に現れる、いずれかの状態
   • live state でないなら、その accept sequence のルートは破棄
4. DFA の state と accept sequence の組み合わせから、受理状態に到達できるトークンを判定してマスクを取得
5. 取得した accept sequence ごとの mask を union して最終的な mask を得る

手順4の DFA の状態と accept sequence の組み合わせが許容するトークンは事前に計算し、インデックスを持つことで効率化できる

具体例

以下は説明のために単純化した字句規則を用いた例であり、JSON number の厳密な仕様をそのまま表しているわけではない。

例えば {"count": 2 まで生成している場合（最終的に {"count": 2} の形を作ると仮定）
lexerは { "count" : 2 に分解して
2はintになるかもしれないし、.が続くとfloatにもなりうる
lexicalは LBRACE STRING COLON（2より前の確定部分）となり、
accept sequence は {INT, RBRACE}, {FLOAT, RBRACE} などの候補（各々が順序付き terminal の系列）として求められる
accept sequence の先頭の terminal が FLOAT で続いて RBRACE が期待される場合
FLOAT は正規表現で、[0-9]+.[0-9]+と表せる
つまり状態は、以下のように分けられる

1. 何も読んでいない
2. .の前の数字部分だけ読んだ
3. .を読んだ
4. .の後の数字まで読んだ (FLOAT 受理)
5. FLOAT 完了後、次の terminal (RBRACE) の DFA へ

これを図で表すと以下のようになる

また次に続くトークンに対し、accept sequence の DFA に入れたときに、受理状態になるかを見ることで許容されるトークンかを判別できる

remainder を DFA の状態として扱うこの設計は、近似処理のような側面があり、状態で直接表現されないトークン同士が隣接した場合に不整合が生じないのか気になりますが、実験上は SynCode はかなり精度高く生成できているようです。

soundnessとcompleteness

ここまでで見たマスクの作り方には、形式的な意味で「どこまで完璧か」という性質がある。
soundness（mask が妥当な候補を落とさない）とcompleteness（mask が不正な候補を残さない）の観点で整理しておく（SynCode 論文参考）。

例えば、accept sequence の長さが 1 の場合、次に来るトークンが次の terminal まで進んだときに、それが許容できるかを甘く見るしかない

• 物差しが短すぎて測れない

つまり、妥当な候補は落とされない（sound）が、不正な候補を完全には弾けない（incomplete）。
そして accept sequence が語彙中の最大トークン長を超えるなら、甘く見る必要がなくなり、completeness も成立する。

accept sequence は長くするほど completeness に近づくが、計算が重くなるというトレードオフがある。
このトレードオフを抑える工夫として、次章ではXGrammarの最適化を見ていきます

最適化事例 XGrammar

（この章は XGrammar 論文に基づく）

前章では 生成済み文字列のうち文法的に確定した部分を parse して、現在の文法的文脈を解釈していたが
次は、stack と状態機械を利用した方法を見てみます。
XGrammar 論文では、Llama-3 系モデル + H100 GPU の環境で、Outlines や llama.cpp と比較して、mask 生成は JSON Schema で約 3.5 倍・CFG で約 10 倍、エンドツーエンドの推論時間も JSON で最大 14 倍・CFG で最大 80 倍程度の高速化が報告されている。

XGrammarの最適化は大きく以下の3カテゴリに分けられる。

カテゴリ	やること	主な手法
実行時のオーバーヘッド削減	mask 生成のコストを下げる	context independent / dependent 分離、adaptive cache、context expansion、persistent execution stack
文法自体の最適化	文法表現を簡略化して曖昧性を減らす	rule inlining、node merging
実行基盤の最適化	LLM 推論との統合	mask 生成と推論の overlap、jump forward decoding

まずは基礎となる PDA（Pushdown automata）から見ていきます。

Pushdown automata (PDA) とは

• CFG を読むための実行機械で、スタックで親や子の state を積んでいく
• 2種類の遷移
  • 具体的な文字を読む遷移(character edge)
  • 別の規則に入る遷移(rule reference edge)

イメージ

arrayがあり、その中のstringをLLMが出力した場合
stack には array の state が積まれ、その上に string の state が push される
stack top（現在処理中の state）は string の state となり、" が出力されると string の state が pop され、array の状態機械に戻る

byte レベルの PDA アプローチ

文法上の語句と LLM の扱うトークンが一致しない問題に対して、ここではトークンを byte として、PDA に流すことで、受理可能な byte 列かどうかをそのまま判定することができるようにしている
字句レベル・トークン単位で対応を取るのではなく、文法を byte 表現で扱うことで両者を整合的に追跡可能にしている

実行時のオーバーヘッドを減らす工夫

context independent / dependent token

token には2つの種類がある

• context independent token
  • 局所的な状態だけで判断できるもの
  • 事前に計算可能
  • 大多数はこっち
• context dependent token
  • 実行中のスタック全体を見ないと判断できないもの
  • JSONの場合、1%未満程度

つまり大部分は事前に計算して、実行時には残ったトークンを判断することで効率化できる

非決定的文法

• 配列の中の最後以外の文字なのか、最後の文字なのかという2つの状態があり得る
• 状態が複数取りうる場合、状態は複数の stack 集合として扱う

遷移には3つの種類がある

1. 子ルールに入って stack を push する場合
2. ルール内で前進し、stack topの位置だけが変わる場合
3. 親ルールに戻って stack を pop する場合

1, 2 は context independent token、3 は context dependent token と言える

キャッシュの保存 (adaptive token mask cache)

• context independent token なら、stack top node をキーとするキャッシュが作れる
• token は「accepted な context independent」「rejected な context independent」「context dependent」の3つに分類できる
• 多くの場合、accepted か rejected のどちらかに大きく偏っている（accept-heavy / reject-heavy）
• よって accept/reject のうち少数派と、context dependent token の2つだけを保存すれば、storage を最適化できる(adaptive storage)
  • Llama-3.1 と JSON 文法の組み合わせでは、storage を約 0.2% まで削減

依存範囲が少ないことで事前計算できるトークンがあり、それが大きな効率化につながっていることがわかった

次は stack top 以外にも依存するトークンの最適化を見ていきます

Context expansion

• ここでも偏りを最適化のきっかけにする
• context dependent token が親側に戻ったときに、取りうる文字列が実際には限られていることを利用する

流れ

• 親側で受理されうる文字列集合(expanded suffix)を前計算し、
• 親側に戻るとき、expanded suffix で始まらないものを破棄し、通過したもののみ従来の full stack 処理を行う
• Llama-3.1 と JSON 文法では 90% の context dependent token を減らせた

名前の通り、子が見ている文脈を親の文脈まで拡張することで、independentの範囲を広げていますね

非決定文法の場合は複数の stack を持ちうるとありましたが、stack 同士で重複が大きく storage 的に非効率になる問題や、full stack なトークン処理を stack 数の分だけ毎回行う必要があるという問題があります。

これを解決するため、木構造で一度作成された複数の stack を永続的に共通化 (persistent execution stack) して、各 stack ではどの枝かをもつようにすれば
前計算でトークンを処理するときに、辞書順に処理すると隣り合うトークンの最大共通接頭辞を再利用することができ計算量を大幅に減らせる
stack で持つのが pointer だけになるので、pointer を付け替えるだけで rollback もしやすくなる

• read, reader, ready の順で処理するとき、reader は read 時点の stack に rollback して、残りの "er" を処理するだけで良くなる

文法自体の最適化

rule inlining

• CFG には要素数の少ないルールが有り、実行時に rule に入り込んで調べる必要があり、曖昧性が増してしまう
• そこで他の rule を参照しない rule を親の rule に inline で入れ込むことで、token checking 効率と context expansion の効率も上がる
• オートマトンが爆発しないようにサイズ上限は設ける

node merging

• 同じ文字を読んだとき、複数候補に分岐してしまうものがあり、実質同じ挙動になるなら最初からまとめる
• 分岐は token check や mask merge の数を増やしてしまう

実行基盤最適化 (LLM serving engine)

Overlapping mask generation and LLM inference

XGrammar では、mask generation と LLM inference をオーバーラップさせることで、structured generation の追加コストを隠すことを狙っている。
これは両者の計算資源と依存関係をうまく分離できることを利用した最適化といえる。

例えば以下のような性質が成立する場合、両者を並列実行し、同期後に sampling を行うような最適化が考えられる。

• mask generation は典型的には CPU 側で行える
• LLM inference は典型的には sampling 以外を GPU 側で行える
• そして両方とも過去に生成したトークンにしか依存しない

これによって mask generation による end to end レイテンシはかなり隠すことができる (near zero overhead structured generation)

jump forward decoding

• 文法上、次のトークンが一意に確定できる場合は、LLM 推論をスキップして直接トークンを出力に追加する
• 例: def func() まで生成された後、文法上必ず : が続くことが確定している場合、LLM を呼ばずに : を直接付加できる
• PDA は決定性なので「次に来うる状態が1つだけ」かを判定でき、これで該当箇所を発見する
• LLM 推論呼び出し自体を省略するので、レイテンシ削減に直接効く（複数トークン分まとめてスキップできることも）

文法制約下での確率分布のバイアス

（この章は Grammar-Aligned Decoding 論文に基づく）

ここまで見てきたような、文法に違反するトークンをマスクして除外する方式は GCD (Grammar-Constrained Decoding) と呼ばれる手法群にあたる。文法整合性を高い信頼性で担保できる点に強みがあるが、マスク後の分布がLLM本来の確率分布とどれくらい一致しているかは別の話です。

LLM は Constrained Decoding されることを前提にせずに確率分布を生成している。
マスクして一部のトークンを0にし残りを再正規化した分布は、「文法という条件下でLLMが本来選びたかった分布」とは別物になっているのではないか。

つまり、文法には従っているけど、それがLLM的にもっともらしい選択ができているのかは別問題。
そこにはバイアスが存在し、本来は文法制約を条件としたLLMの条件付き分布が望ましい。

こうしたバイアスを補正しようとする方向性は GAD (Grammar-Aligned Decoding) と呼ばれている。

Expected Future Grammaticality (EFG)

EFGとは

ある prefix から LLM の確率に従って続きを生成したとき、最終的に文法を満たして完了できる確率

なぜGCDではズレるのか

GCDは「その場で文法に合うか」だけを見て候補を残すため、今は合法でも将来的に文法違反が確定する経路を選んでしまうことがある。
言い換えると、文法が許容されるなら1、されないなら0という荒い重み付けを行っているとみなせる。
ここがbiasの源になる。

どう近似するのか

EFGを計算するには、生成後に末尾から確率を遡って伝播させる。
prefix 毎に EFG を計算し、分岐があるなら各遷移先について P_LLM × EFG を合計する。
初期値は、文法上有効なら EFG = 1、無効なら EFG = 0 とする。

EFGが考慮されない場合の例

• 0から始まる場合は文法的に00のみ許容され、1から始まる場合は任意の数字を選択できるとして、1で終わる数字を出力させたとき
• 文法的制約が確率に反映されていない場合は、00を選択してしまう
• つまりはじめに0を選択するLLMの局所確率で選択して、文法の制約付きの局所確率を扱えていない

EFG決定の具体例

以下は EFG の直感を示すための簡易的な例であり、実際のデコーディング過程を厳密に再現したものではない。

以下の図では、ノード内に「prefix」と「EFG値」、辺のラベルに「LLM確率 P_LLM」を書く。

例1: 分岐がない場合（"00" のみ許容）

文法が "00" のみを許容し、"00" で EOS（完了）する LLM 確率が 1e-8、各ステップ P_LLM = 0.45 と仮定する。

文法を満たす経路が1本しかないため、LLMの各段階の遷移確率がそのままEFGの減衰になり、最終的に非常に小さな値（≈ 2e-9）に潰れる。
GCD の近似（EFG = 1 扱い）では、この小ささを見落としてしまう。

例2: 分岐がある場合

"1" でEOSできるLLM確率が0.3、次も "1" を出すLLM確率も0.3とする。分岐先の深い探索は行わず、EFG≈1と近似する。

より多くの将来分岐を網羅するほど、近似値は真のEFGに近づく。

EFGを使ったサンプリング(Adaptive Sampling with Approximate Expected Futures, ASAp)の実験によると、バイアス削減には効果的だが、近似の収束が遅い場合があり、必ずしもタスクの正解率向上には繋がらないようです

最後に

確率的にトークンを選ぶ LLM でも、デコーディング時に無効な候補を除外することで、少なくとも文法レベルでは強い制約をかけられる。
その中核は、マスクで無効なトークンの確率を 0 にし、有効なものだけからサンプリングするというシンプルな仕組みだった。
具体的には以下のような流れだった。

• 文法の表現: 正規表現/FSM・CFG・PDA で構造を記述
• 効率化: 事前計算インデックス、キャッシュ、stack 共有、CPU/GPU 並列化
• 限界と課題: accept sequence の長さと completeness のトレードオフ、マスク後の確率分布の歪み (GAD/EFG)

昨今、LLM の能力を引き出すにはハーネスエンジニアリング、コンテキストエンジニアリング、権限管理など「外側の仕組み」が注目されていると思いますが、Constrained Decoding はそのうち、トークン生成プロセスに対する明確な例だなあと思いました。また、LLM 単体で完結させるのではなく、周辺の制約・誘導の設計が実用化において重要だと再認識できました。

さらに、Constrained Decoding が reasoning に与える影響、diffusion LLM への適用、schema への拡張なども、学んでみたら面白そうだなあと思いました。

間違っている内容などありましたら、ご指摘お願いします。

最後までお付き合いいただきありがとうございます！ 弊社ではエンジニアを募集しております。 ご興味がありましたらカジュアル面談も可能ですので、下記リンクより是非ご応募ください！

• iimon採用サイト
• Wantedly

参考資料

• Outlines: Efficient Guided Generation for Large Language Models — マスクの事前計算（FSM）
• SynCode: LLM Generation with Grammar Augmentation — 文法とトークンの接続
• XGrammar: Flexible and Efficient Structured Generation Engine for Large Language Models — 最適化事例
• Grammar-Aligned Decoding — 確率分布のバイアス