技術要約: LARK (Learnability-Grounded Trajectory Selection)

問題提起

本論文は、より大きな教師モデルが生成した推論軌跡（reasoning trajectories）を、より小さな生徒用大規模言語モデル（LLM）に模倣させるよう微調整するパラダイムである「推論蒸留（reasoning distillation）」の課題に取り組んでいる。既存のデータ選択手法の多くは、質問レベルで動作するか、あるいは軌跡の品質、正確性、モデルの確信度といったヒューリスティックに依存しているが、決定的な要因である**学習可能性（learnability）**を見落としがちである。

具体的には、単一の質問に対して、複数の候補となる推論軌跡が存在する場合がある（異なる教師モデル、サンプリング実行、またはスタイルによって生成される）。たとえこれらの軌跡が同じ正解に到達したとしても、提供されるトレーニング信号は大きく異なる。ある軌跡が「高品質」または「自然」に見えるからといって、それが必ずしも生徒モデルにとって最も効率的に学習できるものであるとは限らない。ここで提示される中心的な問いは、「蒸留される推論軌跡が、生徒モデルにとって学習可能であることを保証するための、原理に基づいた軌跡選択基準を設計できるか？」というものである。

手法: LARK

著者らは、静的な品質メトリクスではなく、最適化理論に基づく定義に基づいて軌跡を選択する手法であるLARK（Learnability-grounded Anchor-time Ranking）を提案する。

1. 最適化目的関数としての学習可能性

LARKは、軌跡選択を方策最適化問題として定式化する。著者らは、生徒のトレーニング損失の減少率を特徴付ける学習可能性因子 $\rho(\theta; q)$ を定義する。この因子は、以下の通り定義される：
$$\rho(\theta; q) \triangleq \frac{\|\nabla_\theta \mathcal{L}(\theta; q)\|^2}{\mathcal{L}(\theta; q)}$$
ここで、$q$ は $K$ 個の候補軌跡に対する重みベクトルである。Polyak–Łojasiewicz (PL) 条件の下では、高い $\rho$ はより速い収束を意味する。

2. アンカータイム近似

全微調整プロセスにおける時間積分された学習可能性目的関数を直接最適化することは困難である。LARKは**アンカー相対条件（anchor-relative condition）**を導入する。これは、過剰パラメータ化されたモデルが「レイジー・トレーニング（lazy training）」領域にある場合、学習率が初期のアンカーパラメータ $\theta_{ref}$ における値に比例するという仮定に基づいている。これにより、目的関数を $\rho(\theta_{ref}; q)$ を最大化することによって近似することが可能になる。

3. ローカル線形化とフォワードパス・プロキシ

$\rho(\theta_{ref}; q)$ を直接最大化することは、単一の軌跡のみを選択する退化した「ワンホット」解に崩壊してしまうため、問題がある。これを避けるために、LARKは以下の手法を用いる：

• ローカル線形化（Local Linearization）: 標準的なSFT（選択なし）である一様分布 $p$ の周囲で、一次テイラー展開を用いて $\rho$ を展開する。
• フォワードパス・サロゲート ($\hat{g}_k$): 全ての候補に対してバックワードパス（逆伝播）を計算して正確な勾配を求める代わりに、LARKはフォワードパス（順伝播）のみを用いて計算可能なプロキシ $\hat{g}_k$ を導出する。このプロキシは、生徒モデルのトークンレベルの予測残差とクロスエントロピー損失に基づき、学習率の局所的な勾配を推定する。
  • 具体的には、$\hat{g}_k$ は、生徒が構造的に「自信を持って間違っている（confidently wrong）」状態にある軌跡を好む。つまり、損失 $\ell_k$ は高い（予測が不十分であることを示す）が、残差質量が少数の「惜しい（near-miss）」トークンに集中している（明確な修正信号を提供している）状態である。

4. $\chi^2$-正則化選択方策

セレクターが元のデータ分布から離れすぎて（これが汎化性能を損なう）、学習可能性基準を「ハック」することを防ぐため、LARKは$\chi^2$-正則化項を課す。これにより、閉形式の予算パラメータ化された選択ルールが導かれる：

• 軌跡はスコア $\hat{g}_k$ によってランク付けされる。
• 上位 $B$ 個の軌跡が選択される。
• ハードな上位 $B$ 個の切り捨てとは異なり、LARKは $(B+1)$ 番目の軌跡に対するスコアの差分に比例したソフトな重みを割り当てる。この重み付けルールは、一様事前分布からの $\chi^2$ ダイバージェンスを制御しながら、学習可能性の改善の下限を最大化するように導出されている。

主な貢献

1. 定式化: 著者らは、推論軌跡の選択を、アンカータイム学習率 $\rho$ を原理的な目的関数として導入した、学習可能性に基づいた方策最適化問題として定式化した。
2. アルゴリズム: フォワードパス・サロゲート $\hat{g}_k$ と $\chi^2$-正則化を用いた、閉形式の予算パラメータ化された選択ルールを導出した。これにより、高価な候補ごとのバックワードパスを回避しつつ、推定誤差と汎化に関する境界を理論的に保証する。
3. 実証的検証: LARKは、複数のベースモデル（Qwen-2.5-7B/1.5B, Llama-3.2-3B）および推論ベンチマーク（AIME, AMC, GPQA, MATH-500）において、既存のベースライン（ヒューリスティック、品質ベース、および生徒側スコアリング手法を含む）を一貫して上回ることを示した。

実験結果

• パフォーマンス: LARKは、テストされたすべてのモデル予算設定において、最高の平均 Acc@5 を達成した。例えば、Qwen-2.5-7B で予算 $B=3$ の場合、LARK は最強のベースラインである RSR に対して 7.07 パーセントポイントの向上を実現した。
• 効率性: LARK はフォワードパスのみの手法であり、最も低コストな生徒側ベースラインと同等の計算コストである。一方で、正確な勾配ベースの手法は、大幅に多くの計算を必要とする。
• 診断分析:
  • 損失の減衰: LARK によって選択された軌跡は、トレーニング中の SFT 損失の最も速い減少を誘発した。
  • 予測の妥当性: LARK スコア $\hat{g}_k$ は、教師モデルごとに集計した際、ダウンストリームの精度と強い正の相関（ピアソン $r=0.69$）を示した。
  • トークンレベルの挙動: 選択された軌跡は、正解トークンの確率が低い一方で、「間違いトークンの集中度」が高いという特徴を示した。これは、LARK が、生徒が間違ってはいるものの、そのエラーが構造化されており修正可能であるような例を選択していることを裏付けている。

意義と主張

本論文は、学習可能性が、軌跡の品質や難易度とは独立した、効率的な推論蒸留のための明確かつ効果的な原理であることを主張している。選択を生徒モデルの最適化ダイナミクスに根ざさせることで、LARK は最も効率的な学習信号を提供する軌跡を特定する。

著者らは、本手法が以下の特性を持つことを強調している：

• 汎化性能の維持: $\chi^2$-正則化を通じて、選択方策が狭い範囲の「容易すぎる」または「ハックされた」軌跡に過学習することを防ぐ。
• 理論的裏付け: 選択ルールは、明示的な誤差境界を伴う、最適化理論（PL 条件、勾配流）の第一原理から導出されている。
• 実用性: 選択ルール自体のハイパーパラメータ調整を必要とせず（正則化強度は予算 $B$ によって決定される）、バックワードパスによる計算オーバーヘッドを回避している。

本研究は、将来的な拡張として、よりノイズの多い候補プールやより広い推論ドメインにこの学習可能性に基づいた視点を適用できる可能性を示唆しているが、現在の範囲は数学的問題解決における推論蒸留に厳密に限定されている。