💬 NLP

Geometric Self-Distillation for Reasoning Generalization

本論文は、標準的なオンポリシー蒸留によって引き起こされる分布外推論のドリフトを軽減するために、ヘリンジャー損失と近接フィッシャー・ラーオ罰則を自然勾配更新と組み合わせた幾何学的自己蒸留フレームワークであるGeoSDを導入し、これにより、様々なモデルスケールにおいてイン分布の性能を維持しつつ汎化性能を大幅に向上させるものである。

原著者： Josip Jukić, Ivan Titov

公開日 2026-07-09

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Josip Jukić, Ivan Titov

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

全体像：「自信過剰な家庭教師」問題

あなたは複雑な数学の問題を解く学習をしている学生だと想像してください。そこには、非常に賢いけれど、ある「秘密の特権」を持つ家庭教師（「教師」）がいます。その特権とは、あなたがまだ答えに辿り着こうと苦戦している最中に、家庭教師はすでに解答集を見ることができているということです。

AIの世界では、これは**「特権的コンテキスト自己蒸留（Privileged Context Self-Distillation）」**と呼ばれます。AI（学生）は問題を解こうとし、同じAI（教師として機能）が、問題に加えて「完全な解答」を見ながら、学生を導きます。

問題点：
解答を見ているため、教師は次のステップに対して非常に強い自信を持っています。たとえ学生がまだその論理を理解できていなくても、教師は自信満々です。

従来の方法（標準的な学習）： 学生は教師を盲目的に模倣します。もし教師が「このステップをやりなさい！」と100%の自信を持って言えば、学生はパニックになり、なぜそうなるのかという理由を理解していなくても、無理やり自分の脳を同意させようとします。
結果： 学生は「オウム」になってしまいます。練習した特定の形式の問題（分布内：In-Distribution）を解くのは非常に上手くなりますが、新しいタイプの問題（分布外：Out-of-Distribution）を与えられると、無残に失敗します。彼らは論理を学んだのではなく、解答の書き方を暗記してしまったのです。つまり、彼らは**「自信満々に間違える」**ようになってしまうのです。

解決策：GEOSD（「幾何学的」なアプローチ）

著者らは、GEOSD（Geometric Self-Distillation：幾何学的自己蒸留）を導入しています。GEOSDは単に学生に「私を真似しなさい」と言うのではなく、学習の「幾何学」を理解している賢明なコーチのように振る舞います。学生が悪習に陥るのを防ぐために、主に2つのテクニックを使用します。

テクニック1：「重なり」によるフィルター（握手）

比喩： 学生と教師が握手をしている場面を想像してください。

標準的な学習： 教師は、学生がちゃんと握れているかどうかに関わらず、教師の手で学生の手を力一杯引っ張ります。教師が強く、学生が弱い場合、学生はコースから引きずり出されてしまいます。
GEOSD： 引っ張る力は、二人がどれくらい「握手できているか（重なっているか）」に比例します。
- もし学生がすでに教師の意見に同意している場合（「重なり」がある場合）、教師は同意を強化するために優しく引きます。
- もし学生が混乱しており、教師のアイデアをほとんど支持できていない場合、教師は握る力を緩めます。学生がまだ支えることができない考えを、無理に採用させることはしません。

なぜ助けになるのか： これにより、学生がまだ理解できていないステップにおいて、教師の「過剰な自信」を盲目的にコピーしてしまうのを防ぎます。

テクテクニック2：「アンカー付きのロープ」（安全な命綱）

比喩： 学生が新しいダンスを学ぼうとして、平均台の上を歩いている場面を想像してください。

標準的な学習： 学生は教師に合わせようとして、激しく、慌てた足取りで進みます。時間が経つにつれ、学生は元のバランスから大きく逸れてしまい、音楽が変わった（新しい問題が出た）瞬間に、バランスを崩して落下してしまいます。
GEOSD： 学生は、数分前の「安全な自分自身（少し前の状態）」と、伸縮性のあるロープでつながれています。
- 教師は依然として学生を前へと引っ張ることができますが、ロープがあることで、学生が速すぎて大きく逸脱することを防ぎます。
- これにより、学生は元のバランス感覚を失うことなく、新しい動きを学ぶことができます。

なぜ助けになるのか： これにより「ドリフト（漂流）」を防ぎます。学生は、自分自身の核となる感覚を完全に忘れることなく、元の自分を維持しながら向上し、未知の問題に対処できる柔軟性を保ちます。

「秘伝のソース」：「幾何学」対「直線」

論文では、「ヘリンジャー損失（Hellinger loss）」や「フィッシャー・ラオ距離（Fisher–Rao distance）」といった高度な数学用語が使われています。これを簡単に説明すると以下の通りです。

通常の学習（ユークリッド幾何学）： 学習を平坦なグリッド上の歩行として捉えます。「左」へ動くときは一定量動きます。しかしAIの世界では、「左」へ少し動くだけでモデルの挙動が劇的に変わってしまうこともあれば、「右」へ大きく動いても何も変わらないこともあります。
GEOSD（幾何学的）： 学習を**「球体」**の上を歩くこととして捉えます。
- AIの予測を、球体上の点だと想像してください。
- GEOSDは、ある点から別の点へ移動するために、**「球の表面に沿って」**どれくらい歩かなければならないかで距離を測定します。
- これにより、AIが取る一歩一歩が、単にコンピュータ上のコードがどれくらい変わったかではなく、**「AIの挙動が実際にどれくらい変化したか」**によって測定されるようになります。

結果：何が起きたのか？

著者らは、さまざまなサイズのAIモデル（小型から巨大なものまで）を用い、数学の問題（AIMEやAMCコンテストなど）でこのテストを行いました。

従来の方法： モデルは練習問題については上手くなりましたが、新しい難問については性能が悪化しました。彼らは自信過剰で硬直的になってしまったのです。
GEOSD： モデルは練習問題が上手くなっただけでなく、新しい難問に対しても大幅に性能が向上しました（平均で5.7%から8.6%の向上）。
「なぜか」： 従来の方法が失敗したとき、彼らは代替となる回答から「質量を奪って」いました。つまり、一つの誤った回答に対して100%の自信を持つよう強制していたのです。GEOSDは他の選択肢を「生きたまま」にしておいたため、AIが不確実性を保ち、柔軟であり続けることができました。これが新しい問題を解く上で極めて重要です。

まとめとしての比喩

標準的な学習は、台本を暗記している学生のようなものです。場面が変わると、彼らは固まってしまいます。
GEOSDは、演技の「原理」を学んでいる学生のようなものです。彼らはディレクター（教師）の指示を聞きますが、その指示を自分のパフォーマンスとして受け入れるのは、それが自分の中で実感できる場合に限られます。彼らは自分のスタイル（アンカー）に根ざしているため、どんな新しい場面にも適応できるのです。

論文は、AIの推論能力を高めるためには、単に「解答の鍵」をコピーさせるだけでは不十分であると結論付けています。現在の理解を尊重し、核となる論理にしっかりと根ざしたまま、優しく導いていく必要があるのです。

技術要約：推論の汎化のための幾何学的自己蒸留（Geometric Self-Distillation）

1. 問題提起

オンポリシー蒸留（On-policy Distillation, OPD）は、学生モデル自身が生成した軌跡に対して、トークンレベルの密な教師あり学習を提供する、大規模言語モデル（LLM）の実用的なポストトレーニング戦略として台頭している。特定の効果的なバリアントである**特権コンテキスト自己蒸留（Privileged-context Self-Distillation, OPSD）**は、「教師」（同一のモデル）に対して、完全な解法トレースやヒントなどの追加情報を条件付け、一方で学生は元の問題コンテキストのみから推論を行う。

この設定は豊富な教師信号を提供する一方で、本論文は決定的な失敗モードを特定している。それは**予測ドリフト（Predictive Drift）**である。教師は特権コンテキストを保持しているため、解法に基づけば明白な継続（continuation）を自信を持ってコミットできるが、学生の視点からはそれが不確実であったり、根拠がなかったりする場合がある。標準的な蒸留目的関数（KLダイバージェンスなど）は、教師自身の確率質量によって教師の好みを重み付けする。その結果、教師が高度に確信している一方で学生が低い確率しか割り当てていない状態（オーバーラップが低い状態）では、勾配信号が強くなる。多くの更新を経て、これらの強い引き寄せは、学生自身の情報がそれを正当化する前に、学生に教師の自信に満せった継続を採用させる。このような不一致のある更新の蓄積は、ドリフトを引き起こす。つまり、モデルが不確実性を抑制し、**分布外（OOD）**の推論性能を低下させ、「誤った答えに対する自信に満ちた合意」という結果をもたらすのである。

既存のアプローチは、信頼できない教師信号をフィルタリングしたり、ダイバージェンスの重みを再形成したり（例：逆KLダイバージェンスやフィルタリング機構の使用）することで、これを軽減しようとしている。しかし、著者らはこれらの手法が核心的な問題に対処できていないと主張する。すなわち、**「信頼できる信号が、学生をどれほど遠くまで動かすべきか」**という問題である。選択（Selection）だけでは更新の大きさを制御できず、チェックのない移動はドリフトへと蓄積される。

2. 手法：GEOSD

著者らは、次トークン分布の特定の幾何学的空間内での移動として予測行動を扱うことで、このドリフトに対抗するように設計された目的関数、**GEOSD（Geometric Self-Distillation）**を導入する。この手法は、2つの補完的な方法で機能する。

A. 局所的な圧力減衰（Hellinger Loss）

GEOSDは、標準的なKLダイバージェンスの代わりに、各訪問状態における教師の引き寄せを調整するためにHellinger損失を採用する。

メカニズム: Hellinger距離は、確率の平方根埋め込み（ $p \to \sqrt{p}$ ）の下での確率単体（probability simplex）の幾何学、すなわち分布がハイパースフィア（超球面）上に位置する性質から導出される。Hellinger損失の勾配は、各教師の好みを、学生と教師の確率の幾像平均（ $\sqrt{p_i q_i}$ ）によってスケーリングする。
効果: 教師があるトークンに高い確率を割り当てているが、学生がほぼゼロに近い確率しか割り当てていない場合、オーバーラップ項が消失し、勾配による引き寄せが減衰する。逆に、共通の好みが存在する場合は強く転移する。これにより、学生が教師の自信はあるが根拠のない継続を即座に採用することを強制するのではなく、学生と教師が合意に近づくにつれて、指導が段階的に強化されるようになる。

B. グローバルなドリフト制御（Fisher–Rao 近似項）

オーバーラップによる重み付けは個々の状態における引き寄せを制限するが、小さな妥当な更新が訓練を通じて大きなグローバル・ドリフトへと蓄積することを防ぐことはできない。

メカニズム: GEOSDは、現在の学生の分布と直近のチェックポイント（ $\theta_{ckpt}$ ）との間の距離を罰する近似項を追加する。この距離は、予測分布の幾何学における弧長（arc length）であるFisher–Rao距離を用いて測定され、累積された変位を追跡する。
効果: これは「移動する信頼領域（moving trust region）」として機能し、急激な変位を抑制しながら、教師の信号が一貫している場合には着実な移動を許容する。初期モデルへのアンカリング（これは適応を禁じることになる）とは異なり、直近のチェックポイントへのアンカリングは、モデルが適応することを許容しつつ、直近の安定した挙動から離れすぎるのを防ぐ。

C. 自然勾配更新（Natural-Gradient Update）

Hellinger損失とFisher–Rao近似項は、どちらも次トークン分布の同一の幾何学における距離である。この目的関数を最適化するために、GEOSDは自然勾配更新を使用する。

メカニズム: パラメータをユークリッド空間で更新する代わりに、更新はフィッシャー情報行列（効率化のためにK-FACを用いて近似）によって事前条件付けされる。
効果: これにより、各ステップがパラメータ空間において固定のサイズで移動するのではなく、目的関数を改善するための、学生の予測に対する最小の変化を実現する。近似項は、この幾何学における「チェックポイントへの引き戻し」として機能し、パラメータの方向の行動的顕著性（behavioral salience）に基づいて、ポリシーをチェックポイントへと戻す。

3. 主な貢献

GEOSD目的関数: 教師の好みを（Hellinger損失を介して）教師と学生のオーバーラップによって重み付けし、（Fisher–Rao距離を介して）直近のチェックポイントからの予測ドリフトを罰し、統一された幾何学の中で自然勾配によって最適化される自己蒸留フレームワーク。
実証的な改善: GEOSDは、自己蒸留の分布内（ID）での利点を維持しつつ、OOD推論性能を大幅に向上させる。3つのモデルファミリー（Qwen3, Olmo, DeepSeek-R1）および5つのスケール（1.7Bから32B）にわたり、標準的なKL目的関数がOOD性能を低下させるのに対し、GEOSDはベースモデルに対して平均で5.7～8.6ポイントのOOD精度向上を実現した。
メカニズム分析: 本論文は、なぜ標準的なマッチングがOODで失敗するのかを特徴づけている。それは、高エントロピー状態において代替案を「枯渇」させることで、教師のトップの選択肢に確率質量を集中させてしまうためであり、これが自信に満ちた誤答への合意を招く。GEOSDは、トップの選択肢の優位性を維持しながら、代替となる継続の可能性を「手の届く範囲」に留め、早期の崩壊を防ぐ。

4. 実験結果

著者らは、数学的推論ベンチマーク（訓練にはDAPO-Math、OODテストにはAIME 2024/2025, AMC 2023, MATH-500）を用いてGEOSDを評価した。

パフォーマンス: GEOSDは、テストされたすべてのベンチマークとモデルファミリーにおいて最高のOOD精度を達成した。例えば、Qwen3-8Bにおいて、GEOSDはForward KLによる8.1ポイントの低下に対し、ベースモデルから**+8.6ポイント**のOOD精度向上を実現した。
堅牢性: 獲得された利点は、様々な量の特権コンテキスト（短いプレフィックスから完全な解法まで）および様々なモデルスケール（1.7Bから32B）において維持された。
アブレーション研究: Hellingerによるオーバーラップ重み付け、Fisher–Rao近似項、または自然勾配更新のいずれかを削除すると、OOD性能が大幅に低下した。これは、これら3つの要素すべてが必要であることを裏付けている。
偽の合意（False Consensus）: 分析によれば、標準的な手法（Forward KLなど）は、高エントロピー状態において「誤った答えに対して強く合意する」高い割合の「偽の合意」を示す。GEOSDはこの割合をベースモデルよりも低く抑え、情報量の多い自己一貫性を維持した。

5. 意義と主張

本論文は、特権コンテキストを用いた自己蒸留におけるOOD推論の劣化は、単に信頼できない信号を使用している結果ではなく、それらの信号がモデルをどのように動かすかという結果であると主張している。標準的な目的関数は、教師の確信度によって学生を狭いモードへと強制的に押し込み、実行可能な代替案の確率質量を枯渇させることを許してしまう。

GEOSDは、効果的な推論蒸留には、どの教師信号を信頼するかだけでなく、それらの信号が学生をどれほど遠くまで動かすことが許されるかに依存することを示している。これらの移動を幾何学的に制御すること（学生の支持がない場所では引き寄せを減衰させ、直近の安定した状態に軌跡を固定すること）により、GEOSDは分布内パフォーマンスと分布外汎化の優れたトレードオフを達成する。著者らは、推論タスクにおける自己蒸公式のスケールアップにおいて、予測ドリフトの制御が不可欠であると結論付けている。

注記された限界: 著者らは、彼らの知見が現在は数学的推論においてのみ確立されていることを認めている。高エントロピー状態が異なる構造を持つ可能性のある他のドメイン（コード、科学的推論など）に対する、オーバーラップ重み付け転移およびFisher–Raoドリフト制御の適用可能性は、まだテストされていない。さらに、この手法はベースとなるポリシーが広く有能であることを前提としている。非常に弱いベースモデルの場合、このメカニズムは、大きな行動変化を必要とする有用な信号さえも抑制してしまう可能性がある。