論文要約：Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

1. 背景と課題 (Problem)

大規模言語モデル（LLM）をドメイン固有のタスクに適用する際、一般的に教師あり微調整（SFT）が用いられます。しかし、SFT に用いる下流データ（ドメイン固有データ）が、モデルの事前学習分布や汎用的な QA（質問応答）スタイルの生成分布と大きく乖離している場合、以下の問題が発生します。

• オフポリシー不安定性: 下流データは、現在のモデル方策（Policy）にとって低確率（低尤度）のサンプルであることが多く、SFT 時の勾配推定に大きな分散をもたらします。これにより、最適化が不安定になり、学習が収束しにくくなります。
• 破滅的忘却（Catastrophic Forgetting）: 下流分布への過度な適合（Over-shifting）により、モデルが汎用的な能力（一般知識や推論能力）を失う現象が引き起こされます。
• 既存のデータ書き換え手法の限界: 分布の乖離を埋めるための「データ書き換え（Data Rewriting）」手法は存在しますが、既存手法には以下の欠点があります。
  • 分布の不一致: 書き換えは「プロンプト誘導型の条件付き分布（入力＋解答＋書き換え指示）」からサンプリングされ、モデル本来の「QA スタイルの生成分布（入力のみ）」と一致しません。
  • 多様性の崩壊（Diversity Collapse）: 固定されたテンプレートやプロンプトに依存するため、生成される書き換えデータが画一的になり、モデルの多様な生成能力を損なう恐れがあります。

2. 提案手法 (Methodology)

著者らは、データ書き換えを「方策学習（Policy Learning）」の問題として再定義し、強化学習（RL）を用いて、モデルの自然な QA 生成分布に合致しつつ多様性を保つ書き換え方策を学習するエージェントを提案しました。

2.1 全体フレームワーク

1. Stage I: 書き換えエージェントの RL 学習
   • ベースモデル: 凍結されたインストラクションチューニング済みモデル（$\pi_0$）を使用。
   • 方策: LoRA アダプタを介してパラメータ効率よく学習される書き換え方策 $R_\phi$。
   • 最適化アルゴリズム: グループ相対方策最適化（GRPO）を使用。価値ネットワークを不要とし、グループ内の相対的な報酬で方策を更新します。
2. Stage II: データセット構築と下流 SFT
   • 学習済みのエージェントで書き換えデータを生成し、タスク整合性を検証。失敗した場合は元の専門家の解答にフォールバックする「Generate-Verify-Fallback」パイプラインで最終学習データセットを作成し、SFT を実施します。

2.2 報酬設計（Unified Objective）

強化学習の報酬関数は、以下の 3 つの要素を統合的に最適化するように設計されています。

1. タスク整合性（Task Consistency）: ハードゲート
   • 書き換えられた解答が最終答えの正解性と推論の妥当性を満たすかを確認します（ルールベース + LLM ジャッジ）。
   • 整合性が取れない場合（$r_{task}=0$）は、他の報酬を適用せず、学習を停止させます。これにより、誤ったデータによる学習ノイズを防ぎます。
2. 分布アライメント（Distributional Alignment）: QA スタイルへの適合
   • 書き換えられた解答 $\tilde{y}$ が、ベースモデル $\pi_0$ にとって「入力 $x$ のみ」を条件とした QA スタイルの生成分布において、より高確率（高尤度）であるかを評価します。
   • 負の対数尤度（NLL）を正規化し、モデルが自然に生成しやすい形式を報酬として最大化します。
3. 多様性（Diversity）: 崩壊防止
   • 同じ入力に対する複数の整合性のある書き換え候補間の意味的距離（埋め込みベクトルのコサイン類似度）を計算し、候補セット全体の多様性を最大化するよう報酬を与えます。
   • これにより、固定テンプレートによる画一化（Mode Collapse）を防ぎます。

2.3 技術的工夫

• LoRA による「パッチ」: 全体を再学習させるのではなく、凍結されたベースモデルの上に軽量な LoRA 層を乗せることで、局所的な修正（パッチ）を行い、方策のドリフト（Policy Drift）を抑制します。
• ハードゲートによる安定化: 不正解なサンプルに補助的な報酬（アライメントや多様性）を与えないことで、学習の不安定性を排除しています。

3. 主要な貢献 (Key Contributions)

1. 分布ミスマッチの解明と RL による解決: 既存のデータ書き換えパイプラインにおける「プロンプト誘導分布」と「QA 生成分布」の不一致を特定し、RL ベースのアプローチでこれを解消するフレームワークを提案しました。
2. 統合的な目的関数の設計: タスク整合性（ハード制約）、分布アライメント、多様性を同時に最適化するユニファイドな報酬設計により、高品質で分布整合性の高い書き換えデータセットを構築可能にしました。
3. 実験による有効性の実証: 複数のモデルサイズ（1B, 3B, 7B）およびバックボーン（Llama, Mistral）を用いた実験で、標準的な SFT と同等のドメイン内性能を維持しつつ、汎用ベンチマークでの忘却を大幅に軽減することを示しました。

4. 実験結果 (Results)

数学推論タスク（下流）と一般ドメインベンチマーク（忘却評価）を用いた評価結果は以下の通りです。

• 性能と忘却のトレードオフの改善:
  • 提案手法（Rewriting-agent）は、標準的な SFT と同等かそれ以上の数学推論性能向上（例：Mistral-7B で +55.23%）を達成しました。
  • 同時に、一般ドメインベンチマーク（MMLU など）での性能低下（忘却）を大幅に抑制しました（例：SFT は 19.17% 低下に対し、提案手法は 7.35% 低下に留まる）。
  • 既存の書き換え手法（SDFT, Mind the Gap）よりも、分布アライメントと多様性のバランスが取れており、全体的なスコア（OverallAvg）が最も高くなりました。
• 学習の安定性:
  • 提案手法で書き換えられたデータを用いた SFT は、学習損失が低く、より滑らかに収束することが確認されました。これは、書き換えデータがモデルの自然な生成分布に近いため、最適化が安定していることを示唆しています。
• アブレーション研究:
  • 分布アライメント報酬（$r_{dist}$）を除去すると汎用性の保持が最も大きく低下し、多様性報酬（$r_{div}$）の除去もパフォーマンスを低下させました。
  • 「ハードゲート」方式は、ソフトな報酬付与と比較して、タスク整合性のパス率と最終性能の両面で優位性を示しました。

5. 意義と結論 (Significance)

本論文は、オフポリシーな SFT における「分布の乖離」が破滅的忘却の主要因であることを再認識し、それを「データソース側」から RL を用いて修正する新しいパラダイムを提示しました。

• 実用性: 大規模モデルの再学習を伴わず、LoRA による軽量な修正で実現可能であり、計算コストと効果のバランスに優れています。
• 理論的貢献: 単なるデータ拡張ではなく、モデルの生成分布に適合した「教師信号の再設計」が、学習の安定性と汎用性の維持に不可欠であることを示しました。
• 将来展望: 数学推論以外のドメイン（コーディング、対話など）や、より大規模なモデルへの適用可能性が期待されます。

結論として、この RL ベースの書き換えエージェントは、下流タスクへの適応と既存能力の保持という、従来トレードオフであった課題を同時に解決する有効な手段を提供しています。