技術的サマリー：ReAD - 大規模言語モデルのための強化学習に基づく能力蒸留

1. 問題定義

本論文は、大規模言語モデル（LLM）における能力蒸留の課題に取り組む。これは知識蒸留（KD）の特殊な形態であり、特定の能力（推論、コーディング、数学など）における大規模な教師モデルの性能を小規模な学生モデルが模倣するように訓練することで、サービスコストを削減するものである。

既存の手法は通常、能力を独立した訓練目標として扱い、ある能力の最適化が他の能力に悪影響を及ぼさないと仮定している。しかし、著者は重要なギャップを特定している：固定されたトークン予算の下では、単一の能力を改善することが、しばしば体系的かつ予算依存の能力間転移を引き起こし、対象能力の向上が他の有用な能力の犠牲の上に成り立つという点である。これにより、以下の 2 つの非効率性が生じる：

1. 逓減収益性：単一の目標に割り当てられた追加トークンは、次第に小さくなる限界収益をもたらす。
2. 有害な波及効果：ある能力向けに行われた更新が、下流タスクの成功に必要な非対象能力の性能を低下させる。

核心的な問題は問題 2.3として定式化される：相互依存する能力の集合 $C$ に対して固定された蒸留予算 $B$ をどのように配分すれば、他の能力への悪影響（波及効果）を最小化しつつ、下流タスクの有用性 $U_\tau$ を最大化できるか。

2. 手法：ReAD フレームワーク

著者は、能力間の相互依存性を明示的にモデル化し、トークン予算を適応的に配分するフレームワークであるReAD（Reinforcement-guided cApability Distillation）を提案する。ReAD は以下の 3 つの主要な段階で動作する：

2.1 タスクに不可欠な能力の特定

ReAD は、与えられた下流タスク $\tau$ に対してタスク要求ベクトル $r_\tau \in \Delta^{|C|}$ を構築する。

• メカニズム：軽量な識別器 $g_\phi$（Transformer エンコーダに MLP を接続したもの）が、タスク記述カードを $r_\tau$ にマッピングする。
• 訓練：識別器は「局所的な監督信号」を用いてオフラインで事前訓練される。学生モデルに対して小規模かつ疎な介入を行い、その結果生じる能力プロファイルの変化（$\Delta s$）とタスク有用性の変化（$\Delta U$）を測定することで、特定のタスクに有益な能力を予測することを学習する。
• 役割：$r_\tau$ は配分の事前分布として機能し、タスク有用性を向上させる能力を優先し、これらの特定の次元における低下を「有害な波及効果」として扱う。

2.2 即時の能力指向データ生成

静的なデータセットを使用する代わりに、ReAD は蒸留プロセス中に動的に監督データを生成する。

• テンプレート：各能力 $c$ に対して、テンプレートライブラリが指示の足場、スロット、出力制約を定義する。
• 難易度制御：推論ステップ数やコードの複雑さなどの要因に基づき、決定論的な難易度スコア $d_c(x)$ を計算する。システムは、訓練が進むにつれて易しい例から難しい例へと移行するカリキュラムを採用する。
• サンプリング：各決定ステップにおいて、現在の配分ポリシーに基づいて能力ラベルをサンプリングし、プロンプトをインスタンス化し、教師モデルが蒸留のための完了を提供する。

2.3 不確実性認識コンテキストバンディトによる配分

ReAD は、予算配分をコンテキストバンディトによって解かれる逐次決定問題として定式化する。

• コンテキスト（$x_t$）：タスク要求ベクトル $r_\tau$、現在の学生モデルの能力プロファイル $s_{probe}(S_t)$、残存予算 $b_t$、および最近の配分履歴を含む。
• 行動空間：以前の配分の摂動と、上位 $k$ 個の能力に焦点を当てた疎な行動を含む、有限の配分ベクトル集合 $A(\tau)$。
• 代理報酬（$\hat{R}_t$）：蒸留ステップ後、システムは以下を計算する：
  $$\hat{R}_t = r_\tau^\top \Delta s_{probe} - \beta \cdot \text{Spill}_t - \lambda \cdot \text{cost}_t$$
  ここで、$\text{Spill}_t$ はタスクに不可欠な能力における後退に対してペナルティを課す。
• ポリシー更新：報酬回帰器（MLP）のアンサンブルが、各候補行動に対する期待報酬と不確実性（$\sigma$）を推定する。次の配分は、Upper Confidence Bound (UCB) 規則を用いて選択される：
  $$w_{t+1} = \arg \max_{w} \left( \mu(x_{t+1}, w) + \kappa \sigma(x_{t+1}, w) \right)$$
  これにより、ReAD は高報酬の配分を選択する利用（exploitation）と不確実性を減らす探索（exploration）のバランスを取り、訓練が進むにつれて飽和しているか、または高い波及効果を持つ能力から予算をシフトさせることができる。

3. 主要な貢献

1. 能力間相互作用の実証的研究：著者は、単一の能力のための蒸留が、他の能力全体に体系的かつ予算依存のパフォーマンス再分配を引き起こすことを示す制御された研究を提供する。これはしばしば、予算が増加するにつれて悪化する負の転移を引き起こす。
2. 予算制約付き能力蒸留の定式化：彼らは、能力が独立しているという仮定を超えて、有用性を最大化しつつ副作用を制御するために、相互依存する能力間で固定予算を配分する問題として定式化する。
3. 理論的解析：本論文は、共有表現における非直交な更新方向に起因して能力間転移が生じ、逓減収益性が単一の能力への継続的な配分を無駄にするという、局所的な理論的説明を提供する。
4. ReAD フレームワーク：タスク要求を推論し、即時にターゲットデータを生成し、不確実性認識バンディトを用いてリソースを適応的に配分する、強化学習に基づくフレームワークの提案。

4. 実験結果

実験は、トークン予算 20M および 150M の条件下で、Llama-3.3-70B を Llama-3.1-8B に、Qwen2.5-72B を Qwen2.5-14B に蒸留して行われた。

• 有用性の向上（RQ1）：ReAD は、評価された 8 つの能力（一般、制御性、推論、数学、コード、ツール使用、LCU、多言語）すべてにおいて、SFT/KD 目的と Resp/CoT/Logit シグナルを組み合わせた 6 つの強力なベースラインを一貫して上回った。それは、制御性や数学などのボトルネック能力において顕著な改善を伴う、最も強力な平均能力プロファイルを達成した。
• 予算効率（RQ2）：ReAD は、優れた利益 - 波及効果のトレードオフを示した。単一能力 KD は高い目標達成率を達成したが、対象外での著しい劣化を引き起こした。ReAD は、強い目標達成率を維持しつつ、負の転移を大幅に削減した。
• 適応的スケジューリング（RQ3）：静的な混合、貪欲ヒューリスティック、グリッドサーチされた配分との比較により、ReAD の適応的バンディトスケジューリングが、特に飽和と波及効果がより顕著になる大規模予算（150M）において、すべての代替案を上回ることが示された。
• 一般化：ReAD は、新しい能力ヘッドを必要とせず、保持された下流タスク（XSTest）への正の転移を示し、ベースラインと比較して安全拒否および不安全拒否の指標の両方を改善した。

5. 意義と主張

本論文は、ReAD が、能力蒸留を独立したタスクの集合として扱うことから、LLM 能力の内在的な相互依存性を考慮する結合最適化問題へと転換することを示すものであると主張する。

• 効率性：逓減収益性と能力間干渉を明示的にモデル化することで、ReAD は「予算の浪費」を削減し、限られた訓練トークンを有害な副作用を引き起こすのではなく、真のタスク関連有用性を提供する更新に費やすことを保証する。
• 堅牢性：このフレームワークは、異なる教師 - 学生ペア（Llama および Qwen ファミリー）全体で堅牢であり、特定のモデルアーキテクチャに依存しないことが示されている。
• 実用性：この手法は軽量なプローブと再利用可能な要求識別器を使用するため、リソース制約のある MLaaS 環境での展開に対して計算的に実行可能である。

著者は、固定予算下での効果的な能力蒸留には、タスクに関連する能力、能力間相互作用、および動的な予算配分についての推論を同時に必要とするとし、そのギャップを ReAD が埋めると結論づけている。