Tail-Aware Top-$k$ On-Policy Distillation

技術要約：Tail-Aware Top-k On-Policy Distillation (TA-OPD)

問題提起
On-policy distillation (OPD) は、学生言語モデルを、学生自身の軌跡に沿って教師の次トークン分布に適合するように訓練する手法であり、off-policy distillation に固有の訓練と推論の間の乖離に対処するものである。一般的な手法である Normalized Top-k OPD は、教師の top-$k$ トークンに限定された学生と教師の分布間の逆カルバック・ライブラー（KL）ダイバージェンスを最小化する。

本論文は、この正規化された目的関数における決定的な欠陥を特定している。それは、テイル確率（top-$k$ トークンの外側にある総確率質量）に関する情報を破棄してしまう点である。理論的解析および経験的証拠は、正規化された目的関数を最小化することが、学生のテイル確率とエントロピーを着実に増大させる系統的なバイアスを生み出すことを示している。学生のテイル確率が増大するにつれ、教師の top-$k$ セット外のトークンをサンプリングする頻度が高まり、教師による監督が信頼できなくなる接頭辞へと漂流していく。この現象は、特に学生と教師の間に大きな能力差がある場合に、ダウンストリームの精度低下を招く。

手法：Tail-Aware Top-k OPD (TA-OPD)
テイル確率の信号消失を解決するために、著者らは明示的な テイル・トークン ($v_{tail}$) を導入することで、この欠落した情報を復元する手法である TA-OPD を提案する。

1. 拡張されたトークン集合: 単に top-$k$ トークン上で確率を正規化するのではなく、TA-OPD は拡張集合 $S^+_t = S^k_t \cup \{v_{tail}\}$ を定義する。ここで、$S^k_t$ は教師の top-$k$ トークンを表す。
2. 確率の割り当て:
   • $S^k_t$ 内のトークンは、元の確率 $p_t(v)$ を保持する。
   • テイル・トークン $v_{tail}$ には、総テイル確率が割り当てられる：$p_t(v_{tail}) = p^{tail}_t = \sum_{v \notin S^k_t} p_t(v)$。
   • 同様に、教師の分布も $q^{tail}_t$ で拡張される。
3. 損失関数: この手法は、拡張された集合上での逆 KL ダイバージェンスを最小化する：
   $$\ell^{TA}_t = \sum_{v \in S^k_t} p_t(v) \log \frac{p_t(v)}{q_t(v)} + p^{tail}_t \log \frac{p^{tail}_t}{q^{tail}_t}$$
   この定式化は、学生と教師のテイル確率を明示的に比較し、学生のテイル確率を教師のそれに近づける。

理論的特性

• 勾配ダイナミクス: top-$k$ セットの外側にあるロジットに対してゼロの勾配しか提供しない正規化された目的関数とは異なり、TA-OPD はテイルの対数比を通じてテイル・ロジットを更新する。さらに、top-$k$ ロジットの勾配はもはや合計ゼロにはならず、これにより、学生がテイルを過剰に重み付けている場合に、テイルに対して top-$k$ ロジットを明示的に引き上げることを可能にする。
• 下界: TA-POD 損失は、フル語彙の逆 KL ダイバージェンスのタイトな下界であることが証明されている。その差は、テイル確率に比例する非負の残差項である。実際には、TA-OPD がテイル確率をゼロに近づけるにつれて、この残差は消失し、目的関数はタイトな近似となる。
• 不偏バリアント (SC-TA-OPD): サンプリングされたトークンの確率が利用可能な場合、著者らは Sample-Corrected TA-OPD を導出している。このバリアントは、サンプリングされたトークンに基づくバイアス補正項を追加し、値と勾配の両方において、フル語彙の逆 KL ダイバージェンスの 不偏推定値 を提供する。

主な貢献

1. テイル確率の増大の特定: 本論文は、正規化された top-k OPD が学生のテイル確率とエントロピーの着実な増大を引き起こし、性能低下を招くことを理論的および経験的に示した。
2. TA-OPD の提案: テイル・トークンを通じてテイル確率の信号を復元し、top-k 正規化に伴うエントロピーの増大を防ぐ新しい蒸留手法を提案した。
3. 理論的保証: TA-OPD がフル語彙の目的関数のタイトな下界であることの証明、および不偏なサンプル補正版の導出を行った。
4. 経験的な優位性: TA-OPD が、特に学生と教師の能力差が大きいシナリオにおいて、ベースラインを凌駕することを広範な実験で示した。

実験結果
著者らは、3 つの学生・教師ペア（例：OpenThinker3-7B から蒸留された Qwen2.5-7B-Instruct）を用い、数学的推論ベンチマーク（MATH500, Minerva, OlympiadBench, AMC, AIME24/25）および分布外（OOD）タスク（ARC-c, MMLU-Pro）において TA-OPD を評価した。

• 性能向上: TA-OPD は一貫して最高の平均精度を達成した。顕著な改善は以下の通りである：
  • Llama-3.1-8B において、数学ベンチマークの平均で +8.05 ポイント（正規化された top-k OPD と比較）。
  • Llama-3.1-8B において、MATH500 で +19.53 ポイント（31.30% から 50.83% へ向上）。
  • OOD パフォーマンスにおける大幅な向上（例：Qwen2.5-7B-Instruct の ARC-c 精度が 23.07% から 76.30% へ向上）。
• 訓練ダイナミクス: TA-OPD は訓練を安定させ、学生のエントロピーを 1.5 未満に、テイル確率を 0 付近に維持したが、正規化された top-k OPD はエントロピーを ~6 まで、テイル確率を ~0.6 まで上昇させた。
• 能力差への感度: この手法は、学生と教師の能力差が大きい場合に最も効果的である。能力差が小さい場合、性能向上は減少する。
• パラメータ感度: TA-OPD は小さな $k$（例：$k=2$）でも有効であり、$k \ge 2$ であれば $k$ に対して鈍感であることを示した。

意義と主張
本論文は、TA-OPD が、教師とのアライメントを維持するために不可欠なテイル確率の信号を保持することにより、既存の top-k 蒸留手法の根本的な欠陥に対処していると主張している。この手法は、教師への追加のクエリを top-$k$ の確率以外に必要とせず、実装が容易な実用的な改善として提示されている。著者らは、TA-OPD は大きな差があるシナリオで実質的な利点を提供する一方で、$k$ が増加するか、あるいは能力差が縮小すると、その優位性が減少することを指摘している。なお、計算上の制約により、本研究は 8B パラメータまでのモデルを対象としている。