技術要約: PriFT (Prior-Support Guided Fine-Tuning)

問題提起

教師あり微調整（SFT）は、事前学習済み大規模言語モデル（LLM）をダウンストリームタスクに適応させるための標準的かつ計算効率の高い手法であり、多くの場合、強化学習（RL）の初期化として機能します。しかし、SFTはRLと比較して汎化性能のギャップが生じるという課題があります。オフポリシー手法であるSFTは、固定されたデモンストレーションデータをトークンごとに適合させます。デモンストレーション内のターゲット・トークンがモデルの事前学習分布と適合していない場合、このプロセスは過学習や事前学習知識の歪みを引き起こす可能性があります。

近年のアプローチでは、ターゲット・トークンの確率やエントロピーなどを用いて、モデルの予測分布との整合性に基づいてトークンに非一様な学習重みを割り当てることで、これを軽減しようとしています。しかし、これらの手法は通常、現在微調整中のオンラインモデルを用いてこれらの重みを計算しています。本論文は、これが決定的な欠陥を生んでいると主張しています。すなわち、オンラインモデルの予測分布は、学習中に事前学習分布から急速にドリフト（逸脱）していくという点です。その結果、重み付けの信号が最適化の軌跡と絡み合い、モデルが最初に好んだトークンに学習信号を集中させ、他のトークンを抑制するという自己強化的なダイナミクスを誘発します。これにより、サンプリングの多様性が低下し、さらなるオンポリシー最適化（RL）のためのモデルの能力が制限されてしまいます。

手法: PriFT

オンラインでの再重み付けによる不安定性を解決するために、著者らは**PriFT (Prior-support guided Fine-Tuning)**を提案します。核心となる洞察は、トークンの重みは進化し続けるオンラインモデルではなく、凍結された事前学習済み参照モデルから導出されるべきであるという点です。これにより、「事前サポート（prior support）」、すなわちタスク固有の適応を行う前に、ターゲット・トークンがモデルの知識によってどの程度支持されているかについての安定した推定が可能になります。

フレームワーク

PriFTは、凍結された事前学習モデル $\pi_{\theta_{pt}}$ から導出されたトークン重み $m_t$ を導入することで、標準的なSFT目的関数を修正します：
$$L_{PriFT}(\theta_{on}) = -\sum_{t=1}^{T} u(\pi_{\theta_{pt}}(\cdot | x, y_{<t}), y_t) \log \pi_{\theta_{on}}(y_t | x, y_{<t})$$
ここで、$u(\cdot)$ はトークン・ユーティリティ関数です。重みを凍結された参照モデルから計算することで、PriFTは重み推定を最適化の軌跡から切り離し、オンライン手法で見られる自己強化的なバイアスを防ぎます。

実装形態

本論文では、ユーティリティ関数 $u$ の2つの具体的な実装を紹介しています。

1. PriFT-prob: 事前学習済みのターゲット・トークン確率を直接使用します。
   $$m_t = \pi_{\theta_{pt}}(y_t | x, y_{<t})$$
   これは、ターゲット・トークンが事前分布の下でどの程度尤もらしかったかに基づいて勾配の大きさをスケーリングします。標準的なSFTの更新方向を維持しつつ、事前サポートが低いトークンの重みを下げます。

2. PriFT-mass: 鋭い分布を持つ「容易な」トークンへのバイアスを軽減するために、事前学習分布における累積確率質量を使用します。
   $$m_t = \mathbb{I}[u^{(MASS)}_t \geq 0.5]$$
   ここで、$u^{(MASS)}_t = \sum_{v: \pi_{\theta_{pt}}(v) \leq \pi_{\theta_{pt}}(y_t)} \pi_{\theta_{pt}}(v)$ です。
   この指標は、ターゲットよりも「可能性が低い」トークンに割り当てられた確率質量の割合を測定します。閾値0.5は、事前学習の確率質量の少なくとも半分によって支持されているトークンを選択します。PriFT-probとは異なり、これはすべての選択されたトークンが等しい重みを受け取るバイナリ選択マスクとなります。これにより、些細な高確信度トークン（例：機能語）への過度な強調を防ぎつつ、より困難なコンテキストからの情報量の多いトークンを保持します。

計算オーバーヘッド

PriFTは、トークンの重みまたはマスクを事前計算するために、凍結された事前学習モデルの単一のフォワードパスを必要とします。これは相対的に約 $1/(3E)$ （$E$ はエポック数）のオーバーヘッドを加えますが、これはオフラインで行うことが可能です。毎ステップの最適化中に参照分布へのアクセスを必要とするKL正則化手法とは異なり、PriFTはスカラーの重みまたはバイナリマスクを保存するだけで済むため、大幅なメモリオーバーヘッドを回避できます。

主な貢献

1. 信号ドリフトの特定: 著者らは、オンラインのトークン統計が事前学習分布から急速にドリフトし、最初に好まれたトークンに学習を集中させ、汎化性能を低下させる自己強化ループを生み出すことを示しました。
2. 凍結された参照による安定した再重み付け: 凍結された事前学習モデルからトークン重みを導出することが、様々な再重み付けルール（確率ベースおよびエントロピーベース）において、オンラインモデルを使用する場合よりも一貫して優れていることを示しました。
3. PriFTフレームワーク: 重み付けの信号を最適化の軌跡から切り離し、凍結された参照から重みを事前計算する、シンプルで汎用的なフレームワークを提案しました。
4. 実証的検証: PriFTが様々なSFTベースラインの中で最先端の結果を達成し、その後のRLトレーニングのための優れた初期化を提供することを示す広範な実験を行いました。

実験結果

著者らは、Qwen2.5-Math、Qwen3、LLaMA-2などのモデルを用い、数学的推論、コード生成、および医学的質問回答においてPriFTを評価しました。

• 数学的推論: Qwen2.5-Math-7Bにおいて、PriFT-massは平均精度（Avg@16）40.67、Pass@16 69.60を達成し、最強の先行ベースラインであるTALRをそれぞれ3.97ポイントおよび8.75ポイント上回りました。注目すべきは、オンライン再重み付け手法（DFTなど）はしばしばPass@16を犠牲にしてAvg@16を向上させる傾向があるのに対し、PriFTは両方の指標を同時に向上させた点です。
• コード生成: Qwen3-4B-Instructにおいて、PriFTのバリアントは標準的なSFTや他のトークン再重み付けベースラインを上回り、特に困難なLiveCodeBenchのスプリットで顕著でした。PriFT-massは、元のチェックポイントと比較してHumanEval+での性能低下が最も小さく、既存の能力の保持に優れていることが示唆されました。
• 医学的QA: PriFT-massは、MMLU-MedicalおよびMedMCQAベンチマークにおいて、他のトークン再重み付け手法を上回り、最高の全体性能を達成しました。
• RLの初期化: PriFTで初期化されたモデルは、標準的なSFTやオンライン再重み付け手法（DFT）で初期化されたモデルと比較して、その後のRLトレーニング（DAPO）後に強いパフォーマンスを示しました。例えば、Qwen3-8B-Baseにおいて、PriFT-massはベースラインのSFT初期化と比較して、RL後のAvg@16で9.57ポイントの向上をもたらしました。

意義と主張

本論文は、PriFTがトークン再重み付けSFTの根本的な限界、すなわち重み付け信号と最適化の軌跡の絡み合いに対処していると主張しています。事前学習モデルを「事前サポート」の安定したソースとして利用することで、PriFTは以下を実現します：

• サンプリングの多様性の保持: 狭い範囲の推論パスへの確率質量の自己強化的な集中を回避し、より広いエントロピー分布を維持することで、より高いPass@k性能を維持します。
• 分布シフトの軽減: オンライン再重み付け手法と比較して、微調整中の事前学習モデルからのKLダイバージェンスを最小限に抑えます。
• RLへの準備態勢の強化: モデルのさらなるオンポリシー最適化のための能力を維持することで、RLのためのより強力な初期化を提供します。これは、「優れたSFT」とは単に即時的な性能だけでなく、将来の学習に必要な多様性を維持することでもあることを示唆しています。

著者らは、中規模のオープンモデルと特定のタスクドメインに焦点を当てているという限界についても言及しており、より大きなモデルやより広範なユーティリティ関数の探索は今後の課題としています。彼らはPriFTがすべての汎化問題を解決すると主張しているのではなく、SFTパラダイムに対する堅牢で汎用的な改善策として位置付けています。