オフライン強化学習の安定化！Data Rewritingとは？

紹介論文

今回紹介する論文はMind the Gap: Data Rewriting for Stable Off-Policy Supervised
Fine-Tuningという論文です。

この論文を一言でまとめると

本記事では、オフライン強化学習におけるデータ書き換え（Data Rewriting）に着目し、そのメカニズムと効果をわかりやすく解説します。ポリシーギャップを解消し、学習の安定化と性能向上を実現するData Rewritingの可能性を探りましょう。

オフライン強化学習（SFT）における課題：ポリシーギャップとは？

大規模言語モデル（LLM）の性能を飛躍的に向上させる手法として、教師ありファインチューニング（Supervised Fine-Tuning; SFT）が広く用いられています。しかし、SFTは一見シンプルな手法に見えて、実はオフライン強化学習という側面があり、特有の課題を抱えています。それが、ポリシーギャップです。

SFTをオフポリシー学習として捉える

SFTは、大規模言語モデルを特定のタスクに適応させるための効果的な手法です。しかし、SFTをオフポリシー学習として捉えると、その本質が見えてきます。オフポリシー学習とは、ある行動ポリシー（ここではエキスパートのデモンストレーション）によって収集されたデータを用いて、別のターゲットポリシー（学習によって最適化したいモデルのポリシー）を学習する枠組みです。

SFTでは、エキスパートのデモンストレーションが、固定された行動ポリシーから得られたデータに相当します。一方、学習の目的は、このデータを用いて、モデル自身のポリシー（ターゲットポリシー）を最適化することです。ここに、オフポリシー学習としてのSFTの構造があります。

ポリシーギャップの定義

SFTにおけるポリシーギャップとは、エキスパートのデモンストレーションという固定された行動ポリシーと、学習が進むにつれて変化していくモデル自身のポリシー（ターゲットポリシー）との間に生じるずれのことです。エキスパートのデータは、あくまで特定時点での最適な行動を反映したものであり、モデルの学習が進むにつれて、モデルがより良い行動を獲得したり、異なる戦略を採用したりする可能性があります。その結果、エキスパートのデータが必ずしもモデルの現在のポリシーにとって最適なデータとは言えなくなってしまうのです。

ポリシーギャップが問題となる理由

ポリシーギャップは、SFTの学習において様々な問題を引き起こします。主な問題点は以下の通りです。

• 高分散と不安定性：重要度サンプリング（Importance Sampling; IS）は、オフポリシー学習において、行動ポリシーとターゲットポリシーのずれを補正するための標準的な手法です。しかし、ポリシーギャップが大きいと、ISの重みが極端に偏り、分散が著しく増大します。その結果、学習が不安定になり、収束が遅れたり、最悪の場合、発散したりする可能性があります。
• 過学習：モデルがエキスパートのデモンストレーションに過度に適合してしまう現象も、ポリシーギャップが原因で発生します。モデルは、自身のポリシーとは異なるデータに無理やり合わせようとするため、汎化性能が低下し、未知のデータに対する対応力が弱まってしまいます。

既存のアプローチとその限界

ポリシーギャップに対処するための既存のアプローチとしては、以下のようなものがあります。

• KL正則化：モデルのアップデートを制限し、ターゲットポリシーが行動ポリシーから大きく逸脱しないようにする。
• 信頼領域法（Trust Region）：モデルのアップデート幅を制限し、学習の安定化を図る。
• クリッピング：重要度サンプリングの重みを一定範囲に制限し、分散の増大を抑制する。

これらの手法は、いずれも受動的にアップデートを抑制することで安定化を図るものですが、ポリシーギャップそのものを積極的に縮小するものではありません。そのため、根本的な解決には至らず、依然として学習の不安定性や汎化性能の低下といった問題が残ります。

そこで、本記事で紹介する「Data Rewriting」という手法は、このポリシーギャップに正面から取り組み、データセット自体を書き換えることで、より安定的なSFTを実現することを目指します。

Data Rewriting：ポリシーギャップを解消する革新的なアプローチ

オフライン強化学習（SFT）における最大の課題、それはポリシーギャップです。前のセクションでは、このポリシーギャップがなぜ問題なのかを解説しました。ここでは、その解決策として提案されているData Rewritingという革新的なアプローチについて、その基本的なアイデアと仕組みを詳しく見ていきましょう。

Data Rewritingの基本的なアイデア：データセットを書き換える？

Data Rewritingの核心は、その名の通り、学習に使用するデータセットを書き換えるという点にあります。従来のSFTでは、エキスパートのデモンストレーションデータ（固定された行動ポリシー）をそのまま学習に使用していました。しかし、Data Rewritingでは、モデルが生成したデータを取り入れ、よりターゲットポリシーに沿ったデータセットを構築することで、ポリシーギャップを積極的に縮小しようと試みます。

具体的には、以下の手順でデータセットを書き換えます。

1. モデルからの応答をサンプリング: 学習中のモデルに問題を与え、複数の回答を生成させます。
2. 正解を保持、不正解を再解決: 生成された回答のうち、正解しているものはオンポリシーデータとしてそのまま保持します。不正解だった場合は、正解例を参考に、再度問題を解くように促します。この際、モデルは単に正解をコピーするのではなく、自身の言葉で説明するように求められます（digest-and-retell）。
3. 再解決が失敗したら: 再度問題を解かせても正解が得られない場合は、元のエキスパートのデモンストレーションデータを採用します（フォールバック）。

重要度サンプリング（IS）との組み合わせ：残存するミスマッチを軽減

Data Rewritingによって、トレーニングデータはターゲットポリシーに近づきますが、完全に一致するわけではありません。そこで、重要度サンプリング（IS）を組み合わせることで、残存するミスマッチを軽減します。ISは、各データの重要度を調整することで、トレーニングデータの偏りを補正する役割を果たします。

Data Rewritingによってデータの質を高め、ISによって学習の安定性を高める。この2つの組み合わせが、Data Rewritingの強みと言えるでしょう。

Data Rewritingのメリット：学習の安定化と汎化性能の向上

Data Rewritingは、従来のSFTに比べて、以下のようなメリットがあります。

* 学習の安定化: ターゲットポリシーとのずれが小さくなるため、重要度サンプリングの分散が低減し、学習が安定します。
* 汎化性能の向上: モデルが自身の言葉で問題を解決する能力が促進されるため、未知のデータに対する汎化性能が向上します。
* 柔軟な組み合わせ: KL正則化などの既存の最適化手法と組み合わせることで、さらなる性能向上が期待できます。

Data Rewritingの注意点：計算コストとデータセットの質

Data Rewritingは強力なアプローチですが、注意点もあります。

* 計算コスト: モデルに複数回の応答を生成させるため、計算コストが増加する可能性があります。
* データセットの質: 書き換えられたデータの質が低い場合、かえって性能が低下する可能性があります。

これらの点に注意しながら、Data Rewritingを効果的に活用していくことが重要です。

次のセクションでは、Data Rewritingの具体的なプロセスを、3つのステージに分けて詳しく解説します。

Data Rewritingの具体的なプロセス：3つのステージ

Data Rewritingの核心は、ポリシーギャップを埋めるために、データセットを賢く書き換えることです。このプロセスは、まるで熟練の職人が素材を吟味し、最適な形に作り変えるように、3つの段階を経て進められます。それぞれのステージで何が行われ、データセットがどのように変化していくのか、詳しく見ていきましょう。

3段階のアラインメント階層

Data Rewritingは、SFTデータをよりターゲットポリシーに沿った混合分布へと変換するデータ書き換え演算子Tを適用します。この演算子Tは、以下の3つのステージからなるアラインメント階層を適用します。

* 自己整合（Self-alignment）
* ガイド付き整合（Guided-alignment）
* フォールバック（Fallback）

自己整合（Self-alignment）：まずは自力で挑戦！

最初のステップは自己整合です。これは、モデルが自らの力で問題解決を試みる段階です。具体的には、以下の手順で行われます。

1. 各入力x（問題文など）に対して、学習済みのモデル（πθ）から複数の応答をサンプリングします。
2. サンプリングされた応答の中に、問題を正しく解決するものがあるかどうかを確認します。
3. もし正解が見つかった場合、そのうちの1つをオンポリシーデータとしてランダムに保持します。

自己整合の段階では、モデルは過去の学習経験を活かし、自力で正解を導き出すことを目指します。このプロセスを通じて、モデルは自身の強みと弱みを認識し、より効率的な学習戦略を立てることができるようになります。

ガイド付き整合（Guided-alignment）：ヒントを参考に再挑戦！

自己整合がうまくいかなかった場合、次のステップはガイド付き整合です。ここでは、モデルは外部からのヒントを参考に、問題解決に再挑戦します。具体的な手順は以下の通りです。

1. 自己整合が失敗した入力xに対して、参照解法（reference solutions）と呼ばれる、正解への導き方を示す情報を提供します。
2. 参照解法を参考に、モデル（πθ）に再度応答を生成させます。この際、モデルは単に参照解法をコピーするのではなく、その内容を理解し、自身の言葉で説明する（digest-and-retell）ように促します。
3. 自己整合と同様に、複数の応答をサンプリングし、正しいものがあれば、書き換えられたデータとして1つの正しい応答をランダムに保持します。

ガイド付き整合は、モデルが困難な問題に立ち向かうための足がかりを提供します。参照解法を参考にすることで、モデルは新たな視点やアプローチを獲得し、より高度な問題解決能力を身につけることができます。

フォールバック（Fallback）：最終手段はエキスパートの知識！

自己整合とガイド付き整合の両方が失敗した場合、最後のステップはフォールバックです。ここでは、モデルは元のエキスパートのデモンストレーション、つまり教師データに立ち返ります。

フォールバックは、モデルが完全に手がかりを失った場合に、学習をストップさせないための安全策として機能します。エキスパートの知識を参考にすることで、モデルは再び学習の軌道に戻り、さらなる成長を目指すことができます。

最終的なデータセット：多様な知識の宝庫！

これらの3つのステージを経て、最終的に得られるデータセットD’は、以下の要素から構成される、多様な知識の宝庫となります。

* オンポリシーの例（Dself）：モデルが自力で生成した、高品質な正解データ。
* 書き換えられた例（Dretell）：参照解法を参考に、モデルが再構築したデータ。よりターゲットポリシーに沿った内容になっています。
* エキスパートデータ（Dexpert）：フォールバックとしてのみ含まれる、教師データ。

このデータセットは、モデルがより効果的に学習を進めるための、強力な学習資源となります。

階層的プロセスの効果：着実にポリシーギャップを縮小！

この階層的なプロセスにより、トレーニングデータは徐々にターゲットポリシーへと近づき、最適化を開始する前にポリシーギャップを効果的に縮小することができます。

Data Rewritingは、単にデータを書き換えるだけでなく、モデルの学習プロセスそのものを改善する、革新的なアプローチなのです。

実験結果から見るData Rewritingの効果：数学的推論ベンチマークでの性能向上

ここまで、Data Rewritingのアイデア、仕組み、そして具体的なプロセスについて解説してきました。このセクションでは、いよいよ実験結果を見ていきましょう。Data Rewritingが実際にどれほどの効果を発揮するのか、数学的推論ベンチマークでの性能向上を詳しく見ていきます。

実験設定：モデルとベンチマーク

論文では、Data Rewritingの効果を検証するために、以下の設定で実験が行われました。

• データセット: NuminaMath CoTデータセットを使用
• モデル:
  • Qwen2.5-Math-7B: 数学に特化したモデル
  • Llama-3.1-8B-Instruct: 指示に従う汎用モデル
• 評価ベンチマーク:
  • Math500
  • Minerva Math
  • OlympiadBench
  • AIME 2024
  • AMC 2023

これらのモデルとベンチマークを用いることで、Data Rewritingの効果を様々な角度から検証しています。

Data RewritingはSFTとDFTを一貫して改善

実験結果の結論から述べると、Data Rewriting（DR）は、Qwen2.5-Math-7BとLlama-3.1-8B-Instructの両方において、標準的なSFT（Supervised Fine-Tuning）とDFT（Dynamic Fine-Tuning）の性能を一貫して改善しました。これは、Data Rewritingがオフライン強化学習の安定化に貢献することを示す強力な証拠です。

Qwen2.5-Math-7Bでの驚異的な性能向上

数学特化モデルであるQwen2.5-Math-7Bでは、Data Rewritingの効果が特に顕著に現れました。

• DR+SFT: Data RewritingとSFTを組み合わせることで、平均精度が30.33%まで向上しました。
• DR+DFT: Data RewritingとDFTを組み合わせると、平均精度はさらに42.03%に達しました。

注目すべきは、DR+DFTが5つ全てのベンチマークで最高のスコアを達成したことです。特に、AIME 2024とAMC 2023での改善幅が大きく、Data Rewritingが難易度の高い問題に対する性能向上に貢献していることが示唆されます。

Llama-3.1-8B-Instructでも効果を発揮

指示チューニングされた汎用モデルであるLlama-3.1-8B-Instructでも、Data Rewritingは性能向上に貢献しました。

• DR+SFT: SFTの精度を11.39%から19.14%に改善
• DR+DFT: DFTの精度を22.54%から24.05%に改善

Qwen2.5-Math-7Bほどの大きな改善は見られなかったものの、Data Rewritingが様々な種類のモデルに対して有効であることが確認できました。

Data Rewritingはポリシーギャップを効果的に低減

実験結果の詳細な分析から、Data Rewritingがポリシーギャップを効果的に低減していることが明らかになりました。書き換えられたデータは、元のデータよりもターゲットポリシーに近い分布を持つことが確認されています。これにより、重要度サンプリングの分散が減少し、学習が安定化したと考えられます。

アブレーション分析：Data Rewritingの各要素の効果

Data Rewritingの効果をさらに詳しく分析するために、アブレーション分析が行われました。具体的には、以下の3つの設定で実験が行われました。

• 自己整合（Self-Alignment）のみ
• ガイド付き整合（Guided-Alignment）のみ
• 自己整合とガイド付き整合の両方

その結果、自己整合とガイド付き整合の両方を組み合わせた完全なData Rewritingが、最も良い結果を達成することがわかりました。これは、自己整合とガイド付き整合が相補的な役割を果たし、より効果的にポリシーギャップを低減することを示唆しています。

まとめ

このセクションでは、論文で報告されている実験結果を詳しく見てきました。Data Rewritingは、標準的なSFTやDFTと比較して、数学的推論ベンチマークで優れた性能を示すことが確認できました。特に、数学特化モデルであるQwen2.5-Math-7Bでの効果が大きく、難易度の高い問題に対する性能向上に貢献していることが示唆されました。また、アブレーション分析の結果から、Data Rewritingの各要素が相補的な役割を果たし、より効果的にポリシーギャップを低減することが示されました。

次のセクションでは、Data Rewritingの利点と限界をまとめ、今後の研究の方向性について考察します。

Data Rewritingの可能性と今後の展望：より安定的なオフライン強化学習へ

Data Rewritingは、オフライン強化学習（SFT）におけるポリシーギャップという課題に対し、データレベルで積極的にアプローチする革新的な手法です。ここでは、Data Rewritingの利点と限界をまとめ、今後の研究の方向性について考察します。

Data Rewritingの利点

• ポリシーギャップの低減：オフポリシー学習におけるポリシーギャップを効果的に低減し、学習の安定化と性能向上に貢献します。
• 既存手法との組み合わせ：既存のSFTやDFTなどの手法と組み合わせることで、さらなる性能向上が期待できます。
• 汎化性能の向上：モデルが自身の言葉で問題を解決する能力を促進し、汎化性能の向上に繋がります。

Data Rewritingの限界

• 評価対象の限定性：評価は、限られたモデルセット（主に中程度のパラメータースケール）に限定されています。
• タスクの限定性：数学的推論ベンチマークに焦点を当てています。
• オフライン処理：単一ラウンドのオフライン書き換え戦略を採用しており、オンラインでの適応は考慮されていません。

今後の研究の方向性

Data Rewritingはまだ発展途上の技術であり、今後の研究によって、その可能性はさらに広がると考えられます。以下に、今後の研究の方向性を示します。

• より大規模で多様なモデルへの適用：より大規模で多様なモデルへの適用性を評価し、スケーラビリティを確認する必要があります。
• 他のドメインへの拡張：ヘルスケアや金融など、数学的推論以外のドメインへの応用を検討することで、Data Rewritingの汎用性を高めることができます。
• オンラインData Rewriting：トレーニング中のポリシーシフトを軽減するために、より洗練されたオンラインアプローチ（例えば、バッチごとの書き換え）を開発することで、安定性とパフォーマンスをさらに向上させることができます。
• 高度な書き換え技術の探求：より高度なモデルからの外部知識を活用するなど、より豊富な書き換え技術を開発することで、Data Rewritingの効果を最大化することができます。

Data Rewritingは、オフライン強化学習の可能性を広げる魅力的なアプローチです。今後の研究開発によって、より安定的なオフライン強化学習の実現に貢献することが期待されます。