Trust-Region Behavior Blending for On-Policy Distillation

技術要約：オンポリシー蒸留における信頼領域行動ブレンディング（Trust-Region Behavior Blending）

問題提起

オンポリシー蒸留（On-Policy Distillation; OPD）は、オフライン知識蒸留におけるプレフィックスの不一致（prefix mismatch）に対処するため、現在の自身のポリシーからサンプリングされたプレフィックスを用いて学生モデルを訓練し、より強力な教師モデルに適合させる手法である。しかし、このアプローチは学習の初期段階において決定的な脆弱性を導入する。すなわち、弱い学生モデルは低品質なプレフィックスを生成してしまう。学生の軌跡分布が劣悪な場合、教師による指導が低品質または「使用不可能な」プレフィックスに対して適用されてしまい、学習シグナルの有効性が制限される。

このプロセスを安定化させるための既存の試みは、多くの場合、収集をオフポリシーへと移行させる（例：教師トークンの注入やターゲット分布の変更）。しかし、これらはターゲットとなる訓練分布を乱す可能性がある。一方で、純粋な学生のロールアウトは分布を維持するものの、初期品質の低さに苦しむ。核心となる課題は、収集プロセスを恒久的にオフポリシーへシフトさせたり、プレフィックスごとの基本的な蒸留目的を変更したりすることなく、いかにして初期ロールアウトの品質を向上させるかである。

手法：信頼領域行動ブレンディング（TRB）

著者らは、ロールアウト収集の初期フェーズにおける行動ポリシーを制御するために設計されたウォームアップ手法として、**信頼領域行動ブレンディング（Trust-Region Behavior Blending; TRB）**を提案している。TRBは、現在の学生ポリシーを中心とした「信頼領域」内に厳格に制約された上で、教師のポリシーに可能な限り近い行動ポリシーを見出すという原理に基づいて動作する。

1. プレフィックスごとの最適化
各生成プレフィックス $h$ において、TRBは制約付き最適化問題を解くことで、行動ポリシー $\mu^*(\cdot | h)$ を定義する。目的関数は、学生ポリシー $\pi_S$ に対するカルバック・ライブラー（KL）ダイバージェンスを制約条件としつつ、教師ポリシー $\pi_T$ に対するKLダイバージェンスを最小化することである。

$$\mu^*(\cdot | h) = \arg \min_{\mu} D_{KL}(\mu \parallel \pi_T) \quad \text{s.t.} \quad D_{KL}(\mu \parallel \pi_S) \le \epsilon$$

ここで、$\epsilon$ は学生からの許容される局所的な偏差を表す。この定式化により、サンプリング分布は教師によって支持されるトークンへと引き寄せられるが、学生を中心とした信頼領域内に留まることが保証される。

2. 閉形式解（Closed-Form Solution）
この最適化は、行動ポリシーが学生と教師のポリシー間の幾何学的な補間となる閉形式解を与える。
$$\mu_\beta(a | h) = \frac{\pi_S(a | h)^{1-\beta} \pi_T(a | h)^\beta}{Z_\beta(h)}$$
係数 $\beta^*(h)$ は、KL制約を満たす $[0, 1]$ 内の最大の実現可能な値である。教師がすでに信頼領域内にある場合、$\mu^*$ は教師となる。$\epsilon=0$ の場合、それは学生へと回帰する。

3. アニーリング・ウォームアップ
TRBは、初期のウォームアップフェーズにおいてのみ適用される。KL予算 $\epsilon$ は、期間 $K$ にわたって線形にゼロへとアニーリングされる。
$$\epsilon_k = \epsilon_0 \left(1 - \frac{k}{K}\right)$$
これにより、学習の開始時には強力な教師のガイダンスを提供してプレフィックスの品質を向上させ、学生が改善するにつれて、ロールアウト収集を純粋な学生サンプリングへと徐々に回帰させることが可能になる。極めて重要な点は、プレフィックスごとの逆KL蒸留損失自体は変更されず、データ収集ポリシーのみが変更される点である。

主な貢献

• 新しいウォームアップ戦略： TRBは、単にトークンを注入したり損失関数を変更したりするのではなく、学生を中心とした明示的なKL制約の下で教師に導かれたポリシーを最適化することにより、初期のロールアウト挙動を制御する手法を導入した。
• オンポリシー特性の保持： 教師と学生のポリシーを永続的に混合したり、ターゲット分布を変更したりする手法とは異なり、TRBは最終的に訓練が純粋なオンポリシー・ロールアウトに戻ることを保証し、メインの訓練フェーズにおけるターゲット分布を保持する。
• 理論的定式化： 本論文は、補間係数に関するKLダイバージェンスの単調性を根拠として、教師のガイダンスと学生の安定性のバランスをとる行動ポリシーの閉形式解を提供している。

実験結果

著者らは、2つの数学的推論蒸留設定においてTRBを評価した：

1. Qwen3-1.7B-Base を Qwen3-8B から蒸留。
2. Qwen3-0.6B-Base を Qwen3-4B から蒸留。

パフォーマンス：

• TRBは、両方の設定において比較されたすべての手法の中で最も高い平均パフォーマンスを達成し、バニラなOPD、固定-$\epsilon$ ブレンディング、Veto、SKD（教師のインターリーブ注入）、温度アニーリング、およびSFTウォームアップを上回った。
• Qwen3-1.7Bの設定において、TRBは 33.2% の平均pass@1を達成した（バニラなOPDは32.3%）。
• Qwen3-0.6Bの設定において、TRBは 44.4% を達成した（バニラなOPDは44.0%）。

分析：

• 学習初期のダイナミクス： 軌跡分析（図2）によれば、一部の手法（SKDなど）は初期に急速に上昇するものの、必ずしも最高の最終パフォーマンスを維持するわけではない。TRBは、最終的に優れた結果をもたらす「速いスタート」を提供する。
• エントロピーとプレフィックス品質： ウォームアップフェーズ中、TRBは訪問されたプレフィックスにおける教師トークンの平均エントロピーを減少させており、これはより高品質なガイダンスを示している。ステップ0プローブ（図4）は、TRBの下で生成されたプレフィックスが、教師または学生のいずれかによって継続された場合に、バニラなOPDのプレフィックスよりも高い成功率を示すことを実証しており、TRBが初期の状態分布をより学習可能な軌跡へと正常にシフトさせていることを示唆している。

意義と主張

本論文は、TRBが、オンポリシーの目的から永久に逸脱することなく、初期のデータ分布の品質を向上させることで、オンポリシー蒸留の脆弱性に効果的に対処することを主張している。

• 限定的な範囲： 著者らは、本研究がQwen3-Baseのペアを用いた2つの数学的推論設定と、正解性に基づく評価プロトコルに限定されていることを明示している。特定のウォームアップスケジュールが、変更なしに他のドメインや教師・学生間のギャップに転用できるとは主張していない。
• コストのトレードオフ： 本論文は、オンラインでの教師デコーディングと学生・教師の共存が必要となるため、ウォームアップフェーズ中に訓練コストが増加することを認めている。しかし、これらのコストは一時的なものであり、ウォームアップ期間終了後は標準的でより効率的なOPDの実行プロファイルに戻ることを指摘している。
• 核心的な洞察： 結果は、教師に導かれたオフポリシーの挙動が最も有益となるのは、学生が訪問したプレフィックスが教師と大きく乖離している場合であることを示唆している。この不一致が解消されると、継続的なオフポリシーのガイダンスは収穫逓減となる可能性があり、行動ポリシーを学生へとアニーリングする戦略を支持している。