Post-Training is About States, Not Tokens: A State Distribution View of SFT, RL, and On-Policy Distillation

技術的サマリー：ポストトレーニングを状態分布の形成として捉える

問題提起

大規模言語モデル（LLM）のポストトレーニング手法（特に教師あり微調整（SFT）、強化学習（RL）、蒸留）の現在の分析は、主に最適化目的（最尤法、方策勾配、KL 発散など）に焦点を当てています。しかし、この目的中心の視点では、いくつかの実証的現象を説明することが困難です。

• なぜ穏やかな SFT は最小限の忘却でパフォーマンスを向上させることができる一方で、攻撃的（「ストレス」）な SFT は壊滅的な忘却と目標パフォーマンスの低下を引き起こすのか。
• なぜ RL は、スパースまたはノイズの多い報酬を使用しているにもかかわらず、しばしば安定した改善をもたらすのか。
• なぜ蒸留における学生モデルが、教師が唯一の監督源である場合でも、劣化した教師を上回るパフォーマンスを発揮することがあるのか。

本論文は、これらの不一致が生じるのは、標準的な分析が、監督が適用される状態分布という重要な変動軸を見落としているためであると主張します。自己回帰モデルにおいて、「状態」とはプロンプトと生成されたプレフィックスの組み合わせとして定義されます（$s_t = (x, y_{<t})$）。異なるポストトレーニング手法は、根本的に異なる状態分布（固定データセットのプレフィックス対学習者によるロールアウト）に信号を適用するため、モデルの挙動、局所性、保持能力に大きな影響を与えます。

手法

著者は、ポストトレーニングを状態分布の形成として捉える統合フレームワークを提案します。彼らは、あらゆるポストトレーニング手法を 2 つの明確な軸に分解します。

1. 状態ソース（$q_k$）： 訓練コンテキスト（状態）がどこから来るか。
2. 信号ソース（$S$）： その状態においてどのような監督（トークン、報酬、ログオッズ）が提供されるか。

理論的定式化：

• SFT： オフポリシー状態適合。モデルは固定データセットの状態（$d_{data}$）におけるトークン損失を最小化します。データセットの軌道がモデル自身のロールアウトから乖離する場合、これはミスマッチを生じさせ、露出バイアスやストレス下での不安定性を招きます。
• RL： オンポリシー局所改善。モデルは現在の方策（$d_{\pi_\theta}$）から軌道をサンプリングし、報酬に基づいて更新を行います。更新は、モデルが実際に訪問する状態に自然に局在化されます。
• オンポリシー蒸留（OPD）： 学生が状態（$d_{\pi_S}$）をサンプリングするが、教師が信号（ログオッズまたは継続）を提供するハイブリッド手法。これにより、状態ソースと信号ソースが分離され、学生は教師のグローバルな軌道分布を継承することなく、局所的な修復を学習できます。

実験設定：

• モデル： Qwen3-0.6B-Base に LoRA アダプターを適用。
• ハードウェア： 単一の RTX 3090 GPU。
• タスク：
  • 目標： GSM8K（数学的推論）。
  • 保持： TruthfulQA および MMLU（壊滅的忘却の測定）。
• 指標： 目標精度、保持率、忘却の大きさ、および状態ドリフト（ロールアウト状態の語彙的特徴における最大平均不一致（MMD）を介して測定）。

実験条件：

1. 穏やかな SFT： GSM8K での標準的な訓練。
2. ストレス SFT： 劣化を誘発するための攻撃的な訓練（5 エポック、高い学習率、高い LoRA ランク）。
3. OPD： 劣化した「ストレス SFT」モデルを信号ソースとして使用して学生を訓練するが、状態は学生自身の方策からサンプリングする。
4. オンポリシー RL： GSM8K におけるグループ相対的な正確な答え報酬を用いた軽量 RL。

主要な結果

1. SFT は文脈依存である： 穏やかな SFT は、保持損失を無視できるレベルで GSM8K を改善しました（0.448 $\to$ 0.512）。逆に、ストレス SFT は目標タスク（0.420）と保持タスク（TruthfulQA/MMLU が大幅に低下）の両方を劣化させました。これは、SFT が本質的に破壊的なのではなく、モデルの自然な状態分布と整合しないオフポリシーの圧力の下でそうなることを示しています。
2. OPD は劣化した教師を上回る： 「ストレス SFT」モデルを教師として使用して OPD で訓練された学生は、すべての指標で教師を上回りました（GSM8K: 0.466 対 0.420; MMLU: 0.430 対 0.364）。これは、学生が自身の状態で訓練され、局所的な条件付きガイダンスのみを受ける場合、教師のグローバルな軌道の失敗を継承しないことを確認しています。
3. RL は局所性を保持する： 軽量なオンポリシー RL は、ほぼベースレベルの保持を維持しながら GSM8K を改善しました（0.472）。これは、RL が学習者が到達可能な状態に基づいた局所的な改善手順として機能するという見解を支持しています。
4. ドリフト指標は不十分である： ほぼ同一のスカラー状態ドリフト（MMD）を持つ手法が、保持結果において大きく異なっていました。例えば、ストレス SFT とストレス教師からの OPD は、類似した MMD 値（約 0.0109）を示しましたが、保持率（0.82 対 0.95）は大きく異なりました。これは、分布シフトの大きさだけでなく、訓練状態のソースと局所性の方が重要であることを示しています。

貢献

1. 統合フレームワーク： SFT、RL、OPD を共通の「状態分布視点」の下で定式化し、状態ソースと信号ソースを明示的に分離しました。
2. 制御された実証パイプライン： 目標精度、保持、忘却、ロールアウト状態のドリフトを測定する単一 GPU 実験環境を実装し、状態分布の効果を分離しました。
3. 状態中心の主張に対する証拠： 以下の実証的証拠を提供しました。
   • OPD は、状態サンプリングと信号生成を分離することで、パフォーマンスを回復し、劣化した教師を上回ることができる。
   • オンポリシー RL は、その局所性により、最小限の保持損失で目標パフォーマンスを向上させる。
   • スカラードリフト指標（MMD など）は、状態ソースを考慮しない限り、ポストトレーニングのダイナミクスを完全に記述できない。

意義と主張

本論文は、ポストトレーニングの挙動が、監督信号と訓練状態分布との相互作用によって支配されると主張します。著者は、「訓練状態のソース、局所性、学習者依存性」が、監督信号そのものの形式と同様に重要であると論じています。

• SFT の再定義： 壊滅的忘却は単に最尤関数の機能ではなく、密なオフポリシー更新がモデルを、自身の将来の状態分布と互換性のない領域に押しやる際に生じます。
• 蒸留の再定義： 学生が教師を上回ることは、教師の誤りがしばしば教師自身の軌道分布と結びついているためです。学生の状態に対して教師に問い合わせることで、学生はグローバルな軌道の失敗を継承することなく、有用な局所的な修復にアクセスできます。
• RL の再定義： RL の安定性は、そのオンポリシー性によるものであり、更新がモデルが実際に到達し、回復可能な状態にのみ適用されることを保証しています。

本論文は、目的レベルの分析だけでは説明できない安定性とパフォーマンスの逆転を説明するためには、状態中心の視点が必要であると結論付けています。著者は、これらの知見は小規模な研究に基づいており、決定的なベンチマーク主張というよりも、機械的な仮説に対する証拠として捉えるべきであると強調しています。