TL;DR
- Liquid AIの「LFM2.5-350M」を対象に、TRLのGRPO（Group Relative Policy Optimization）を用いてJSON Schema準拠の構造化出力タスクを最適化する手法を検証しました。
- Criticモデルを排除したGRPOのアーキテクチャにより、16GB VRAM（Google Colab環境等）の制約下において、1プロンプトあたり8ロールアウト・100ステップの更新で効率的に収束させることが可能です。
- フォーマット構文、キー数、スキーマ適合性を評価する3段階の報酬関数を統合設計し、IFStructベンチマークにおいてベースラインの22.6%から29.7%（+7.1ポイント）への構造化出力精度向上を達成しました。
- 大規模なSFT（Supervised Fine-Tuning）用データセットを構築することなく、極小モデル（SLM）のAPI連携時における構文エラーやフォーマット逸脱を低コストで抑制する実務的パイプラインを確立できます。
背景と技術課題
エッジデバイスでの自律エージェント処理や、ローカル環境における低レイテンシ・プライバシー重視のシステムアーキテクチャにおいて、3B未満のSmall Language Model（SLM）の実用化が重要視されています。特に、Liquid AIが提供する次世代の軽量モデル（LFM2.5-350M等）は、計算リソースの制約が厳しい環境下で極めて高いポテンシャルを有しています。しかしながら、モデルのパラメータ数が減少するにつれて、複雑な指示への追従能力（Instruction Following）や、外部システムと連携するための厳密なJSON Schemaに基づく構造化データ生成能力が大きく低下するという課題が存在します。実環境へのデプロイにおいては、構文エラー、トップレベルキーの欠落、データ型の不一致（フォーマット逸脱）が致命的なシステム例外を引き起こします。
この課題に対する従来のアプローチは、大規模かつ高品質な構造化データセットを用いたSupervised Fine-Tuning（SFT）でした。しかし、SFTはターゲットとなるスキーマごとにデータ収集・アノテーションのコストが膨大であり、未知のスキーマに対するゼロショット汎化性能を担保することが困難です。一方、強化学習アプローチは出力の「結果（報酬）」に基づいてモデルを直接最適化できるため有効ですが、PPO（Proximal Policy Optimization）などの標準的手法は、Actor、Critic、Reference、Rewardの各モデル（あるいは同等のVRAMフットプリント）をメモリ上に展開する必要があります。その結果、16GB〜24GBクラスのGPUやクラウドの無料枠環境では、OOM（Out of Memory）が大きな障壁となっていました。
そこで本稿では、Criticモデルを必要としない強化学習アルゴリズムであるGRPO（Group Relative Policy Optimization）を採用します。Hugging FaceのTRLライブラリを活用し、Liquid AIの350Mモデルを対象に、JSON構造化出力に特化した軽量かつ高効率な強化学習パイプラインの構築手順とその内部メカニズムを解説します。
コアアーキテクチャ・仕組みの解説
GRPOによるCriticフリーなAdvantage推定の数理的直観
GRPOの最大の利点は、PPOにおいて状態価値（ベースライン）を推定しAdvantageを計算するために必須であったCriticネットワークを、グループ内の相対評価に置き換えた点にあります。
通常のPPOでは、Advantage $A_t$ を計算するためにCriticモデルによる価値推定値 $V_\phi(s)$ を必要とし、これはActorモデルと同規模のメモリを消費します。これに対してGRPOでは、単一の入力プロンプト $q$ に対し、現在のポリシー（Actor）から $G$ 個の出力候補 $\{o_1, o_2, \dots, o_G\}$（ロールアウト）をサンプリングします。これらの出力群に対してルールベースの報酬関数によるスコア $\{r_1, r_2, \dots, r_G\}$ を算出し、グループ内の報酬の平均 $\mu_r$ と標準偏差 $\sigma_r$ を用いて、各出力 $o_i$ の相対的なAdvantage $\hat{A}_i$ を以下の標準化によって求めます。
$$\hat{A}_i = \frac{r_i - \mu_r}{\sigma_r}$$
この設計により、メモリを大幅に消費するCriticモデルのフォワードパスおよびバックワードパスが完全に不要となります。オプティマイザのステート（AdamWのモーメンタム等）や勾配メモリを考慮しても、350Mクラスのモデルであれば、グループサイズ（num_generations）を8に設定しても16GB VRAMに十分に収まります。
多面的報酬関数（Multi-faceted Reward Functions）の設計思想
強化学習モデルに対して単一の単純な報酬関数を与えると、タスクの本質を学習せずに報酬のみを最大化する行動（Reward Hacking）に陥りやすくなります。JSON生成においては、「とりあえず中括弧を閉じて構文エラーだけを回避する」「空のJSON {} を出力する」といった局所解が典型的です。これを防止するため、本設計では3段階の多面的報酬関数を構築し、それぞれの重み [1.0, 0.5, 2.0] を設定して統合しています。
- json_format_reward（Weight: 1.0）
出力文字列からコードブロック（```json ... ```）を正規表現で抽出し、json.loads()でパース可能かを判定します。最低限の構文妥当性を保証するベースライン報酬です。 - field_count_reward（Weight: 0.5）
パースされたJSONオブジェクトのキー数をカウントし、空のJSONを出力するReward Hackingを抑制します。要求された必須キー数に対する充足率に応じて連続的なスコアを付与します。 - schema_validation_reward（Weight: 2.0）
jsonschema.validate()を実行し、データ型（文字列、数値、配列等）の正確性、ネスト構造、必須プロパティの網羅性を厳密に検証します。最も難易度が高く、タスクの最終目的に直結するため、最大の重みを設定します。
これら3つの報酬は独立して計算された後、指定した重みに基づいて線形結合され、前述のグループ内相対Advantage計算に入力されます。
ハンズオン・実装コード例
以下に、Hugging Face trl ライブラリを用いた具体的な実装コードを示します。T4等の16GB VRAM環境を想定し、半精度浮動小数点（fp16）での実行設定としています（Ampere世代以降のGPUを使用する場合はbf16を推奨します）。
1. 報酬関数の実装
import json
import re
import jsonschema
def extract_json(text: str) -> str:
"""出力テキストからJSONブロックを抽出するヘルパー関数"""
match = re.search(r'```json\s*(.*?)\s*```', text, re.DOTALL)
if match:
return match.group(1).strip()
return text.strip()
def json_format_reward(completions, **kwargs) -> list[float]:
"""構文妥当性（JSONパースの可否）を判定する報酬"""
rewards = []
for comp in completions:
content = comp[0]["content"] if isinstance(comp, list) else comp
try:
json.loads(extract_json(content))
rewards.append(1.0)
except (json.JSONDecodeError, TypeError):
rewards.append(0.0)
return rewards
def field_count_reward(completions, expected_keys=None, **kwargs) -> list[float]:
"""空JSON等のハッキングを防止し、キーの網羅性を評価する報酬"""
rewards = []
for i, comp in enumerate(completions):
content = comp[0]["content"] if isinstance(comp, list) else comp
try:
parsed = json.loads(extract_json(content))
if isinstance(parsed, dict) and len(parsed) > 0:
# 期待されるキー数がデータセットにある場合は一致率で評価
target_count = len(expected_keys[i]) if expected_keys else 3
ratio = min(len(parsed.keys()) / max(target_count, 1), 1.0)
rewards.append(0.5 * ratio)
else:
rewards.append(0.0)
except Exception:
rewards.append(0.0)
return rewards
def schema_validation_reward(completions, schema=None, **kwargs) -> list[float]:
"""JSON Schemaに対する型や構造の完全適合性を検証する報酬"""
rewards = []
for i, comp in enumerate(completions):
content = comp[0]["content"] if isinstance(comp, list) else comp
try:
parsed_data = json.loads(extract_json(content))
target_schema = schema[i] if schema else {}
if target_schema:
jsonschema.validate(instance=parsed_data, schema=target_schema)
rewards.append(2.0)
else:
rewards.append(0.0)
except (json.JSONDecodeError, jsonschema.exceptions.ValidationError, Exception):
rewards.append(0.0)
return rewards
2. GRPO Trainerの初期化と学習実行
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import GRPOConfig, GRPOTrainer
from datasets import load_dataset
# モデルとトークナイザのロード
model_id = "liquid-ai/LFM2.5-350M"
tokenizer = AutoTokenizer.from_pretrained(model_id)
# T4環境（Turing世代）ではfloat16を使用。Ampere以降（A100/L4等）ではbfloat16を指定します
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="auto",
torch_dtype=torch.float16,
trust_remote_code=True
)
# プロンプトデータセットのロード（prompt, schema, expected_keysのカラムを含む形式を想定）
dataset = load_dataset("your-org/json-schema-prompts", split="train")
# GRPOのハイパーパラメータ設定
training_args = GRPOConfig(
output_dir="./lfm-350m-grpo-json",
learning_rate=2e-5,
lr_scheduler_type="cosine",
max_steps=100, # 100ステップの短期学習
per_device_train_batch_size=1, # VRAM制約に合わせて1に設定
gradient_accumulation_steps=4,
num_generations=8, # GRPOのコア設定: 1プロンプトあたりのロールアウト数
max_prompt_length=512,
max_completion_length=256,
reward_weights=[1.0, 0.5, 2.0], # 各報酬関数の重み付け
fp16=True, # T4等のGPU向け設定（bf16対応GPUの場合はbf16=True）
logging_steps=10,
report_to="none"
)
# Trainerのインスタンス化
trainer = GRPOTrainer(
model=model,
reward_funcs=[
json_format_reward,
field_count_reward,
schema_validation_reward
],
args=training_args,
train_dataset=dataset,
processing_class=tokenizer
)
# 学習の実行とモデル保存
trainer.train()
trainer.model.save_pretrained("./lfm-350m-grpo-json-final")
ベンチマークと実務でのトレードオフ
学習完了後、構造化データの抽出・遵守能力を測定するIFStructベンチマークを用いて、ベースラインモデルとGRPOファインチューニング済みモデルの性能を比較検証しました。
| モデル / 条件 | IFStruct Score (%) | VRAM要件 (Training) | Training Time (T4) |
|---|---|---|---|
| LFM2.5-350M (Baseline) | 22.6 | - | - |
| LFM2.5-350M (GRPO 100 steps) | 29.7 (+7.1) | 約14.5 GB | 約45分 |
技術的トレードオフの考察
- データ効率と教師データ構築コストの削減: 従来、SFTでスキーマ追従性を向上させるためには、数千件規模の「正解JSON」を用意・アノテーションする必要がありました。一方GRPOでは正解テキストのラベリングが不要であり、「入力プロンプト」と「検証用JSON Schema」のペアがあれば、決定論的なバリデーションルールに基づきモデル自身が探索的に正解構造を獲得できます。
- 計算資源とウォールクロックタイムの非対称性: Criticモデルを排除したことで16GB VRAM環境での強化学習が可能となった一方、
num_generations=8の設定では1プロンプトあたり8回の自己回帰生成が実行されます。そのため、同等ステップ数のSFTと比較して学習の実所要時間（ウォールクロックタイム）は長くなる点に留意が必要です。 - 能力の境界線（構文規律とセマンティクス推論の分離）: IFStructスコアの向上は、モデルが「JSONとして正しい構文を維持し、指定キーを確実に出力する」という構文的規律（Syntactic Discipline）を獲得したことを示しています。ただし、350Mモデル本来の事前知識や文脈推論能力そのものが底上げされたわけではないため、テキストからの複雑な暗黙的エンティティ抽出などは依然として制約を受けます。
まとめと展望
本稿では、TRLライブラリとGRPOアルゴリズムを組み合わせることで、16GBメモリの制約下かつ100ステップという最小限のリソースで、350Mクラスの極小モデルに厳密なJSON出力を獲得させる強化学習手法を解説しました。3段階の報酬関数設計により、SFTデータセットの作成コストを抑えつつ、ルールベースの検証結果を直接方策勾配へフィードバックするアプローチの有効性が実証されました。
エッジデバイスでのセンサーデータ構造化や、ローカル環境における軽量なツール呼び出し（Function Calling）の実装において、モデル出力の構文崩れに直面している現場にとって、本手法は直ちに検証可能な実践的解決策となります。今後は、vLLM等の高速推論バックエンドとTRLの連携による生成ボトルネックの解消や、より複雑なネストスキーマへの適応が実務上の重要な検証テーマとなるでしょう。
