カスタムエージェント訓練：モデル蒸留ライブチュートリアル

モデル蒸留で小規模AIを大規模モデル並みに鍛える：オンポリシーとオフポリシーの決定的違い

Hugging Face のライブチュートリアルでは、高コストな大規模言語モデル（LLM）に依存せず、小規模・低予算モデルでも高性能なエージェントを構築するための「モデル蒸留」技術が詳述されました。教師モデルから学生モデルへ知識を転送するこの手法は、特定のタスクにおける性能向上の鍵であり、オンポリシーとオフポリシーという2つの主要アプローチの違いを理解することが実装成功の第一歩となります。

蒸留とは何か：小規模モデルのパフォーマンスを最大化する技術

モデル蒸留（Distillation）は、大規模で高コストな「教師モデル」の知識を、小さく高速な「学生モデル」へ転送するプロセスです。これは単なるデータ圧縮ではなく、特定のタスクに特化して小規模モデルの能力を引き出すための標準的な機械学習手法です。

蒸留には主に3つの目的があります。

1. 圧縮: 大規模モデルを小型化する。
2. 能力転送: 専門分野の知識を小規模モデルに移す。
3. 合成データ生成: 教師モデルを使って学習データを自動生成する。

重要なのは、蒸留が必ずしも悪意ある「攻撃（Distillation Attack）」や大規模モデル特有の現象ではないという点です。2015年のヒントン氏による論文に遡る古典的な手法であり、API の背後にあるモデルから出力を抽出して別のモデルを訓練する一般的なプロセスです。

学習の4つの軸：ハードラベルとソフトラベルの違い

蒸留を理解するには、以下の4つの軸で分類することが有効です。特に「教師モデルからの信号の種類」が最初の重要な分岐点となります。

ハードラベル（Hard Labels）vs ソフトラベル（Soft Labels）

• ハードラベル: 先週紹介された教師データによる監督微調整（SFT）に相当します。教師モデルが生成したテキストそのもの（トークン列）を学生モデルが模倣して学習します。これは「正解」だけを学ぶため、情報が限定的です。
• ソフトラベル: 教師モデルの出力ロジット（確率分布）を使用します。単に最初の単語だけでなく、「次に何が来る可能性があったか」「なぜその単語が選ばれなかったか」という詳細な情報まで含みます。

ソフトラベルを用いることで、学生モデルは教師モデルの「思考の過程」や「他の選択肢との比較」を学習でき、より豊かな言語表現を習得できます。

この違いは、ブラックボックス（ロジットにアクセスできない）かホワイトボックス（ロジットにアクセスできる）かの観点とも重なります。ソフトラベルを使うには教師モデルへの内部アクセスが必要ですが、その分学習効率と精度が向上します。

時間軸の選択：オフポリシーとオンポリシーの決定的違い

蒸留手法を分類する最も重要な軸は、「教師モデルがトレーニングループ内にあるかどうか」です。これにより「オフポリシー」と「オンポリシー」に大別されます。

オフポリシー蒸留（Off-policy Distillation）

事前に準備されたデータセットを用いて学習する方法です。教師モデルが生成した過去の軌跡（トレス）を学生モデルが模倣します。

• 特徴: 教師モデルはトレーニングループ外にあり、オフラインでデータを生成してから学習を開始します。
• メリット: 計算リソースの管理が容易で、合成データ生成パイプラインと相性が良いです。
• デメリット: 閉鎖型API（内部ロジットへのアクセスがない）からの学習には適していません。また、教師モデルが完璧な戦略しか示さない場合、学生モデルは「失敗する状況」や「困難な局面での判断」を学ぶ機会を失う可能性があります。

オンポリシー蒸留（On-policy Distillation）

学生モデル自身が生成したテキストを教師モデルで評価し、そのフィードバックを即座に学習ループに組み込む方法です。

• 特徴: 教師モデルがトレーニングループ内に常駐します。学生モデルの出力に対してリアルタイムでスコアリングや修正が行われます。
• メリット: 強化学習（RL）における「報酬信号が希薄（スパース）」という問題を解決できます。将棋で言えば、ゲーム終了時の勝敗だけでなく、各手ごとの評価を即座に得られるため、学習効率が格段に向上します。
• 実装の鍵: 逆KLダイバージェンス（Reverse KL Divergence）を用いた損失計算が核心となります。これにより、学生モデルのミスを教師モデルが直接修正するループが形成され、より安定した学習が可能になります。

TRL ライブラリによる実践的デモ：小規模モデルの実装

Hugging Face のTRL（Transformer Reinforcement Learning）ライブラリを活用すれば、これらの理論を実際のコードで再現できます。ライブチュートリアルでは、以下の手順で実験が行われました。

1. 環境構築: TRL ライブラリをインポートし、教師モデルと学生モデルを設定する。
2. ループの実行: オンポリシー設定において、学生モデルがテキストを生成→教師モデルが評価→逆KLダイバージェンスで損失計算→パラメータ更新というサイクルを回す。
3. 能力の統合: 異なるチェックポイント（学習段階）を組み合わせて、小規模モデルに複数の能力を付与する手法を実演しました。

このデモは、個人開発者や小規模チームが、ローカル環境や低予算でも高性能なAIエージェントを構築・運用できる道筋を示しています。特にオンポリシー手法の解説は、LLM の推論効率と学習精度を両立させるための標準的なプラクティスとして、今後は業界全体で普及していく可能性が高いでしょう。

まとめ

モデル蒸留は、大規模モデルへの依存を減らしつつ高性能なエージェントを実現するための強力な手段です。オフポリシーが「過去の成功例の模倣」であるのに対し、オンポリシーは「リアルタイムでの試行錯誤と修正」を通じて学習効率を最大化します。TRL などのツールを活用し、自社のリソースに合った蒸留手法を選択することで、より効率的でコストパフォーマンスの高いAI開発が可能になります。