🤖 AI

TRACE: Distilling Where It Matters via Token-Routed Self On-Policy Alignment

本論文は、検証可能な報酬を用いた強化学習中にアノテータがマークした重要なスパンに対してKLダイバージェンスを選択的に適用するトークンルーティング型自己蒸留手法TRACEを提案し、これにより勾配の無駄と特権情報の漏洩を軽減し、標準的なGRPOおよび完全応答自己蒸留のベースラインと比較して長視野の数学的推論と分布外一般化を大幅に改善する。

原著者： Jiaxuan Wang, Xuan Ouyang, Zhiyu Chen, Yulan Hu, Zheng Pan, Xin Li, Lan-Zhe Guo

公開日 2026-05-12

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

CC BY 4.0

原著者： Jiaxuan Wang, Xuan Ouyang, Zhiyu Chen, Yulan Hu, Zheng Pan, Xin Li, Lan-Zhe Guo

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

複雑な数学の問題を解く方法を学生に教える場面を想像してください。あなたは「教師」（賢い AI）と「生徒」（あなたが訓練しようとしている AI）を持っています。

過去、研究者たちはオンポリシー自己蒸留と呼ばれる手法を試みました。その考え方は単純でした。生徒が回答を生成し、その後、教師が回答全体を見て、「ここは正しく、ここも正しく、ここも...」と最後まで評価します。そして生徒は教師のすべての動きを模倣しようとします。

問題点：「過剰な教師」効果
この論文は、「すべてを模倣する」というアプローチは、特に長く複雑な推論タスクにおいて、実際には有害であると主張しています。

比喩： 学生が 10 ページの試験を受けている場面を想像してください。教師はすべての単語を赤でハイライトし、「この通りに正確に行え」と言います。
結果： 学生は圧倒されてしまいます。自分で考えることをやめてしまいます。回答は短くなります（書きすぎることを恐れるため）、ランダムに推測し始めます（エントロピーが上昇します）、そして最終的に試験に失敗します。この論文はこれを「粒度の不一致」と呼びます。学生はすべての単語で修正される必要はありません。つまずいたり間違いを犯したりする可能性のある、ごく少数の重要なステップでのみ助けが必要なのです。

解決策：TRACE（重要推論のためのトークンルーティングアライメント）
著者たちは、TRACEと呼ばれる新しい手法を提案しています。教師が論文全体を修正するのではなく、TRACE は重要な瞬間だけを指摘する賢いチューターのように機能します。

以下に、簡単な比喩を用いて TRACE の仕組みを説明します。

1. 「スポットライト」（トークンルーティング）

ページ全体に光を当てるのではなく、TRACE はスポットライトを使用します。

アノテーター： 補助者（生徒自身であっても構いません）が生徒の回答を見て、「重要な区画」のみをマークします。これらは、論理が素晴らしい（キー・スパン）か、危険なほど間違っている（エラー・スパン）少数の文です。
ルール： スポットライトが照らすのは回答の約**25%**のみです。残りの回答はそのままにされます。

2. 「二つのツール」アプローチ（FKL と RKL）

スポットライトが特定の部分に当たると、TRACE はそれに応じて 2 つの異なるツールを使用します。

生徒が苦労している場合（過少割り当て）： 教師は、「ねえ、この重要なステップを無視しているよ！私を見て、このようにやれ」と言います。これを**フォワード KL（FKL）**と呼びます。これは生徒が正しい経路に注意を払うように促します。
生徒が自信過剰で間違っている場合（過信）： 生徒は「これは間違いない！」と言いますが、実際は間違っています。教師は「いや、止まれ！自分自身を過信しすぎている。引き下がれ」と言います。これを**リバース KL（RKL）**と呼びます。これは生徒の誤った方向への自信を抑制します。

3. 「薄れるヒント」（減衰）

TRACE の最も重要な点は、教師の助けが一時的であることです。

比喩： 教師が練習テスト中に学生にヒントを与えている場面を想像してください。最初はヒントは大きく明確ですが、学生が上達するにつれて、ヒントは次第に小さくなり、最終的に完全に消えます。
結果： 短い「ウォームアップ」期間（トレーニングでは約 40 ステップ）の後、教師は完全に話しなくなります。生徒は独自の論理（強化学習）を使って問題を解くことになりますが、短い指導期間から重要な習慣をすでに学んでいます。これにより、生徒が教師に依存することを防ぎます。

なぜ機能するのか（結果）
この論文は、この手法を数学の問題（方程式の解法など）と一般知識の問題（GPQA）でテストしました。

「全トークン」手法は失敗した： 他の手法が回答全体を修正しようとしたとき、AI のパフォーマンスは急落し、より短く創造性の低い回答を与えるようになりました。
TRACE は成功した： 重要な部分のみを修正し、その後姿を消すことで、TRACE は標準的な手法と比較して AI の数学スコアを大幅に改善しました（平均約 2.76%）。
「自己学習」ボーナス： AI が自分自身をチューターとして行動する場合（人間やより強力な AI が助けない場合）でも、TRACE はうまく機能します。これは、この手法が単に超賢い教師の知識を「不正に持ち込む」ものではなく、AI により良く学ぶ方法を実際に教えていることを証明しています。

まとめ：
TRACE は、ゲーム中のすべてのミスに対して選手に怒鳴りつけるようなコーチではありません。代わりに、コーチは重要な瞬間を待ち、鋭く具体的な修正を与え、その後、選手が自分でゲームをプレイできるように一歩引きます。これにより、選手は集中し、自信を持ち、燃え尽きることなく長く困難な課題に対処できるようになります。

技術的サマリー：TRACE - トークンルーティングによるオンポリシーアライメントを通じた重要な部分の蒸留

1. 問題定義

本論文は、数学的推論タスクにおいて特に、検証可能な報酬を伴う長期の強化学習（RLVR）に適用される際の**オンポリシー自己蒸留（Self-OPD）**の限界に取り組む。Self-OPD 手法（SDPO、SRPO など）は、教師モデルとしてポリシーのスナップショットを用いてスパースな報酬信号を高密度化させるが、著者らは長期トレーニングにおいて以下の決定的な失敗モードを特定した：

粒度の不一致： 既存の手法は、応答内のすべてのトークンに知識蒸留（KL）損失を適用する。しかし、分析によると、正しいロールアウトにおける確率質量の約 97% は学生と教師の間ですでに共有されており、つまり KL 勾配の大部分は冗長である。
「均一蒸留税」： すべてのトークンに KL を適用することは、修正を必要としない位置に対してモデルが勾配を費やすことを強制する。これにより以下の問題が生じる：
- 特権情報の漏洩： 学生は推論時には存在しない特権コンテキスト（例：検証済みトレース）に含まれる相関を吸収し、分布外（OOD）での性能低下を引き起こす。
- エントロピーの上昇と推論の崩壊： 著者らは長期トレーニングにおいて「三症状の崩壊」を観測した：応答長が劇的に短縮し、トークンごとのアクターエントロピーが上昇（しばしばベースラインの 4 倍以上）、検証精度が著しく低下する。
核心的な問い： 適切な蒸留単位は応答全体ではなく、学生の分布が実際に修正を必要とする決定的な推論トークンの小さな部分集合である。

2. 手法：TRACE

著者はTRACE（Critical rEasoning のためのトークンルーティングアライメント）を提案する。これは、均一に適用するのではなく、特定のトークンスパンにのみ蒸留信号をルーティングする手法である。

中核コンポーネント

トークンルーティング型自己 OPD： グローバルな KL 損失の代わりに、TRACE は特定のロールアウト内の特定スパンを識別するために特権アノテーター（ $\pi_A$ ）を使用する：
- キースパン（ $K_y$ ）： 報酬=1 の正しいロールアウトにおける決定的な推論ステップ。
- エラースパン（ $E_y$ ）： 報酬=0 の誤ったロールアウトにおける根本原因となる推論エラー。
- 非スパン（ $N_y$ ）： 残りのすべてのトークン。
- 制約： 特権情報の露出を制限するため、スパンのカバレッジ合計は $\alpha = 0.25$ （トークンの 25%）に上限設定される。
粗い診断ルーティング：
- アノテーターは粗いタイプラベル（例：case_split、missed_case）を出力するが、スパンのテキストや完全なロールアウトの内容は出力しない。
- 教師（ $\pi_T$ ）は、問題とこの粗いラベルをプライベートな診断プレフィックスとして受け取る。学生が生成したプレフィックスに対して因果的にログオッズを計算するが、特定のスパンの位置や内容は視認しない。
- 学生（ $\pi_S$ ）は、KL 項がマスクされたスパンにのみ適用されるルーティング損失を用いてトレーニングされる。
非対称発散ルーティング：
- キースパンにおける前方 KL（FKL）： 正しいロールアウトに適用される。FKL（ $KL(\pi_T \parallel \pi_S)$ ）は、学生の現在の確率質量に関わらず、学生が過少割り当てしているトークンに「持ち上げ」効果をもたらす。
- エラースパンにおける逆 KL（RKL）（オプション）： 誤ったロールアウトに適用される。RKL（ $KL(\pi_S \parallel \pi_T)$ ）は、学生が過信しているが教師が好まないトークンを抑制する。
- 非スパンにおける KL なし： これらのトークンは標準的な GRPO（Group Relative Policy Optimization）損失のみに依存する。
有限露出スケジュール：
- KL 重み（ $\lambda_k$ ）は、短いウォームアップ（例：40 ステップ）後にゼロにアニーリングされる。
- これにより、トレーニング期間全体を通じて累積的な特権勾配露出が有限に保たれ、OOD 性能を低下させる相関の長期的な蓄積が防止される。

3. 理論的解析

本論文は、ルーティング戦略の理論的根拠を提供する：

非対称ポイントごとの圧力（系 2）：
- FKLは、過少割り当てされたトークン（学生確率 $p \ll$ 教師確率 $q$ ）に対して有効である。なぜなら、勾配圧力は差 $(q-p)$ に比例し、 $p \to 0$ になっても消えないからである。
- RKLは、過信して誤ったトークン（ $p \gg q$ ）に対して有効である。なぜなら、勾配は学生の過信度（ $p$ ）に比例するからである。
累積露出の上限（命題 3）： スパンマスク（帯域幅を $\alpha$ に制限）と KL 減衰（期間を制限）を組み合わせることで、TRACE は累積的な特権勾配露出を $O(\alpha \Lambda^2)$ に制限し、永続的な全トークン KL ベースラインで見られる無制限の成長を防ぐ。
信号対漏洩： この手法は、キースパンにおけるアライメント信号を最大化しつつ、偽の相関を引き起こす特権情報の漏洩リスクを最小化する。

4. 主要な貢献

診断： 全トークン型 Self-OPD における「均一蒸留税」の特定。エントロピー上昇、認識的抑制、特権勾配漏洩を粒度の不一致と関連づけた。
手法と理論： トークンクラス（キー、エラー、非スパン）に基づいて $\{FKL, RKL, \emptyset\}$ をルーティングし、特権チャネルを減衰させる TRACE の導入。強いベースモデル領域では FKL がキースパンで支配的であり、弱いベースモデル領域では最適なコーナーがエラースパンにおける RKL にシフトする理由を理論的に説明した。
実証的証拠：
- 性能： Qwen3-8B において、TRACE は 5 つのベンチマークの平均を GRPO より2.76 パーセントポイント改善した。
- 安定性： SDPO/SRPO と異なり、TRACE はベースラインが劣化する GPQA-Diamond においてベースモデルの OOD 性能を維持した。
- 自己アノテーション： この手法はオンライン自己アノテーター（学生ポリシー自身）で機能し、強力な外部 API アノテーターからの利益の約 69% を回復した。これは利益が外部能力ではなく、ルーティングメカニズムに由来することを証明する。
- クロススケール適応： 弱いベース（Qwen3-1.7B）において、支配的なルーティングのコーナーは自然に FKL-on-Key から RKL-on-Error に反転し、理論的予測と一致した。

5. 結果

Qwen3-8B（強いベース）： TRACE（キースパンにおける FKL）は、MATH-500、AIME 24/25、AMC 23、GPQA-Diamond の 5 つのベンチマークで平均**81.51%**の精度を達成し、GRPO（78.75%）および最も強力な既存の Self-OPD ベースラインである RLSD（78.99%）を上回った。重要なのは、GRPO や全トークンベースラインが大幅に低下するのに対し、ベースモデルの GPQA-Diamond スコア（58.33% vs ベース 58.27%）を維持している点である。
Qwen3-1.7B（弱いベース）： TRACE（エラースパンにおける RKL）は平均**60.16%**を達成し、ベースモデル（58.77%）を上回った唯一のトレーニング済み手法となった。全トークンベースライン（SDPO/SRPO）は崩壊し、GRPO より 11〜19 ポイント低下した。
アブレーション研究：
- マスクの局所化： TRACE の特定のスパンマスクをランダムな 25% マスクや全トークン KL に置き換えると性能が低下し、どのトークンが選択されるかが量よりも重要であることを確認した。
- アノテーターの品質： 能動的にトレーニング中の学生を自身のアノテーターとして使用することで +1.90 ポイントの向上が見られ、この手法が外部の「超人的」推論能力の導入に依存していないことを示した。

6. 意義と主張

本論文は、TRACE がドメイン固有の修正ではなく、長期 Self-OPD の不安定性に対する損失幾何学的な解決策を提供すると主張する。その意義は以下の点にある：

蒸留と不安定性の分離： 「特権」信号を重要なスパンに制限し、それを減衰させることで、TRACE は SDPO や SRPO で観測されたエントロピー上昇や推論の崩壊を回避する。
適応的ルーティング： 最適な蒸留方向（FKL 対 RKL）は固定されたハイパーパラメータではなく、ベースモデルのエラープロファイル（過少割り当て対過信）に依存することを示した。
実用性： 外部教師や高コストなクエリごとの評価基準なしに、高品質な蒸留を達成可能であることを証明した。代わりに、自己生成可能な軽量で粗い診断アノテーションに依存する。

著者らは、TRACE が固定されたレシピではなく「コーナールーティングフレームワーク」であると強調しており、特定のルーティング動作（FKL 対 RKL）はベースモデルの残差エラーモードに基づいて選択されるべきである。