🤖 machine learning

Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration

本論文は、勾配相互作用の最適化分析に由来する適応的スキームである最適係数較正（OCC）を提案し、これによりマルチトークン予測と検証可能報酬からの強化学習の効果的な共同訓練を可能にし、それによって従来の性能低下を克服し、数学的推論ベンチマークにおいて優れた結果を達成するものである。

原著者： Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, Guojun Yin

公開日 2026-05-28

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

CC BY 4.0

原著者： Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, Guojun Yin

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

この論文を、平易な言葉と創造的な比喩を用いて解説します。

全体像：2 人のコーチ、1 人のアスリート

巨大で超賢いロボット（大規模言語モデル）を、難しい数学の問題を解くように訓練している状況を想像してください。

メインコーチ（RL）： このコーチはロボットに勝利する方法を教えます。彼らはロボットの回答を見て、正解には「サムズアップ（親指を立てる）」を、不正解には「サムズダウン（親指を下げる）」を与えます。ロボットは、うまくいったことを繰り返し、うまくいかないことをやめるように学習します。これを**強化学習（RL）**と呼びます。
アシスタントコーチ（MTP）： このコーチは、ロボットが次に言う単語だけでなく、その「次の数語」を予測しようとする特別なモジュールです。これを**マルチトークン予測（MTP）**と呼びます。まるで、ロボットが事前に文を計画するのを手助けするコーチのようです。

問題点：
以前、研究者たちは両方のコーチが同時にロボットを訓練しようと試みました。しかし、何かがうまくいきませんでした。ロボットは、メインコーチだけがいる場合よりもパフォーマンスが低下し始めたのです。

このため、ほとんどのチームは「勝利」を目指す訓練段階において、アシスタントコーチを「解雇」することにしました。彼らはアシスタントコーチに見守らせていましたが、その助言がロボットの脳（モデル）を変えることは許しませんでした。彼らはアシスタントコーチを「切り離（detached）」したのです。

問い：
この論文の著者たちは、次のように問いかけました：なぜ両方のコーチがいることがロボットに悪影響を与えるのか？両方のコーチが訓練を台無しにすることなく協力できるように、それを修正することはできるのか？

診断：「押し引き」の問題

著者たちは訓練の背後にある数学を調べ、答えを見つけました。彼らは、アシスタントコーチが助言を与えようとすると、ロボットの脳に 2 つの対立する力が生じることに気づいたのです。

有益な押し（相関）： 時々、アシスタントコーチの助言はメインコーチと同じ方向を指します。これは素晴らしいことです！まるで 2 人が同じ方向に車を押しているようなものです。
厄介な揺らぎ（ペナルティ）： 時々、アシスタントコーチの助言は単なる「ノイズ」です。それはロボットを勝つのに役立たない無作為な方向へ押しやります。まるで、運転しようとしている間に誰かが車を揺さぶっているようなものです。

なぜ従来の手法は失敗したのか：

手法 A（切り離し）： アシスタントコーチの押しを無効化しました。安全ですが、「有益な押し」を見逃してしまいました。
手法 B（交差エントロピー）： アシスタントコーチに押し続けさせましたが、アシスタントコーチは間違ったものさえ含めて、すべての単語を均等に予測するようにロボットに教えようとしていました。これにより、「有益な押し」はほとんどなく、「厄介な揺らぎ」が大量に発生しました。ロボットは混乱し、パフォーマンスが悪化しました。
手法 C（ポリシー損失）： アシスタントコーチにメインコーチと同じ「勝利」のルールを使うよう指示しました。最初は非常にうまくいきました！「押し」は強力でした。しかし、ロボットが上達するにつれて、2 人のコーチはわずかに意見が食い違い始めました。「揺らぎ」は同じままですが、「有益な押し」は弱まりました。最終的に揺らぎが支配的になり、ロボットの性能は急落しました。

比喩：
あなたが綱渡り（RL 訓練）をしている状況を想像してください。

切り離しは、一人で歩くことです。安全ですが、遅いです。
従来の共同訓練は、酔っ払った友人に手を取ってもらうことです。彼らはあなたをロープから引きずり落とします。
「上昇と下降」の問題は、最初は助けてくれる清醒な友人がいるものの、あなたが疲れてくると間違った方向に引き始め、転落させるようなものです。

解決策：「スマート調整器（OCC）」

著者たちは、問題がアシスタントコーチ自体が悪いからではなく、彼らの声の音量が固定されたレベルに設定されていたからだと気づきました。

訓練の初期段階： アシスタントコーチは非常に有益です。彼らの音量を上げるべきです。
訓練の後期段階： アシスタントコーチは邪魔になり始めます。彼らの音量を下げるべきです。

彼らは**最適係数較正（OCC）**と呼ばれる新しいシステムを作成しました。

仕組み（創造的な比喩）：
OCC を訓練プロセスのためのスマートなサーモスタットだと考えてください。

温度（訓練の重み）を固定された数値に設定するのではなく、そのサーモスタットは部屋を常時チェックします。
アシスタントコーチが現在役立っているか、害を及ぼしているかを確認するために、「プロキシ（簡易で安価なセンサー）」を使用します。
アシスタントコーチが正しい方向に押している場合、サーモスタットは音量を上げます。
アシスタントコーチが漂移し始め、ノイズを引き起こす場合、サーモスタットは音量を下げます。

これは、すべてを遅くする複雑な数学を計算して停止する必要なく、ステップごとに自動的に起こります。

結果：ハッピーエンド

著者たちは、この新しい「スマート調整器」を 6 つの難易度の高い数学コンテスト（AIME やオリンピックなど）でテストしました。

従来の手法を凌駕： 「スマート調整器（OCC）」で訓練されたロボットは、「切り離し（メインコーチのみ）」で訓練されたロボットよりも多くの問題を解きました。
急落を修正： 最初は強くても後に失敗した「ポリシー損失」手法とは異なり、OCC 手法は最初のステップから最後のステップまで一貫して強さを保ちました。
高速： 「スマート調整器」は訓練を遅くしませんでした。重い計算を行う代わりに、音量をチェックするための巧妙なショートカットを使用したため、従来の「切り離し」手法と同じ速度でした。

まとめ

この論文は、メインコーチとアシスタントコーチの両方でロボットを訓練することは可能だが、固定された設定のままにしておくことはできないことを証明しています。訓練プロセスに耳を傾け、リアルタイムでアシスタントコーチの影響を調整する動的システムが必要です。そうすることで、ロボットはこれまで以上に効率的かつ高速に学習できるようになります。

技術的サマリー：最適係数較正による強化学習におけるマルチトークン予測の共同訓練

問題定義

検証可能な報酬からの強化学習（RLVR）は、大規模言語モデル（LLM）の推論能力を強化するための標準的なパラダイムとなっている。同時に、事前学習中に複数の将来トークンを予測するために訓練される補助目的であるマルチトークン予測（MTP）は、サンプル効率の向上や、推測的デコーディングのためのドラフトモデルとしての機能において、顕著な利点を示している。

直感的には、事後学習中にMTPとRL目的を共同で最適化することで、モデルは推論能力を獲得しつつマルチトークン監督の恩恵を受けることができるはずである。しかし、実証的証拠は、MTP-RLの共同訓練がしばしば主要モデルの性能を著しく低下させることを示している。その結果、現在の慣行（veRL や slime などのフレームワークにおけるものなど）では、通常、RL 訓練中に主要モデルから MTP の勾配を切り離す（detach）。この切り離しにより、MTP パラメータは独立した事後微調整を通じてのみ適応することを余儀なくされ、同時マルチトークン監督からの潜在的な利益が犠牲になっている。ここで扱われる核心的な問題は以下の通りである：なぜ MTP-RL の共同訓練は失敗するのか、そして効果的な共同訓練を可能にする原理的な戦略を設計できるか？

手法

理論的解析

著者らは最適化の観点からこの問題にアプローチし、RL 目的関数 $J(\theta)$ に対する MTP のステップごとの影響を解析する。 $L$ -滑らかな目的関数を仮定し、ステップごとの改善の下限を導出する。この分解により、MTP の影響は以下の 2 つの競合する項によって支配されることが明らかになる：

一次相関項：RL 勾配（ $g_{RL}$ ）と MTP 勾配（ $g_{MTP}$ ）間の方向的な一致（内積）によって決定される。
二次摂動ペナルティ：MTP 勾配の大きさによって引き起こされる摂動を反映する、非正の項。

性能の向上が生じるのは、相関項がペナルティ項を上回る場合に限られる。

既存の訓練レジームの解析

この分解を用いて、本論文は 3 つの既存の訓練レジームを統合し、説明する：

Detach：主要モデルに対する $g_{MTP} = 0$ 。 $\Delta_{MTP} = 0$ 。これは安全だが、潜在的な利益を放棄する。
交差エントロピー（CE）損失：RL は高報酬サンプルの重みを増やすのに対し、CE はすべてのサンプルを均等に扱うため、期待値において相関項はほぼゼロになる。その結果、負の二次ペナルティが支配的となり、劣化を引き起こす。
ポリシー損失：MTP に対して同じ RL 目的関数形式を使用する。初期段階では勾配が整合しており、相関項が支配的となって利益が得られる。しかし、訓練が進みモデルが損失地形の平坦な領域に入ると、相関は減衰する一方で勾配ノルム（したがってペナルティ）は持続する。これは、固定係数がフェーズ遷移を追跡できない「上昇後、低下」の性能曲線をもたらす。

提案手法：最適係数較正（OCC）

フェーズ遷移に対処するため、著者らは 最適係数較正（OCC） を提案する。これは、最適な MTP 係数（ $\lambda$ ）をオンラインで追跡する適応的スキームである。

閉形式の最適解：解析により、各ステップにおける最適な $\lambda^*$ の閉形式表現が得られ、これは勾配相関（ $c$ ）と勾配ノルムの二乗（ $v^2$ ）の比率に比例する。
対数確率プロキシ：大規模分散システムにおいて $c$ $c$ と $v^2$ $v^{2}$ のためのフルモデル勾配を計算することは著しく高コストである。OCC は、小ステップ更新下での対数確率の変化から導出される対数確率プロキシ（ $\delta \approx \log \pi_\theta - \log \pi_{old}$ $δ \approx lo g π_{θ} - lo g π_{o l d}$ ）を利用する。
- 整合性プロキシ（ $\hat{c}_t$ ）：内積を推定する。
- 分散プロキシ（ $\hat{v}^2_t$ ）：ノルムの二乗を推定する。
適応的更新：係数は $\lambda_t = \lambda_+ \cdot \frac{\hat{c}_t}{\hat{v}^2_t + \epsilon}$ として更新される。これにより、OCC はフル勾配を計算することなく、訓練全体を通じて正の利益を維持しつつ、MTP の重みを動的に調整できる。

主要な貢献

理論的統合：本論文は、Detach、交差エントロピー、およびポリシー損失の各レジームを統合する理論的解析を提供し、一次相関と二次摂動の相互作用に基づいてそれらの成功または失敗を説明する。
OCC アルゴリズム：対数確率プロキシを用いてオンラインで MTP 係数を較正する適応的手順である最適係数較正（OCC）を提案する。この手法は計算オーバーヘッドをほぼ必要とせず、フルモデル勾配計算を回避する。
実証的検証：コンペティションレベルの数学的推論ベンチマークにおいて、複数のモデル（MiMo-7B、GLM-4.5-Air）およびアルゴリズム（DAPO、GSPO）にわたる広範な実験により、OCC が Detach ベースラインと一貫して同等かそれ以上の性能を発揮することが示された。

実験結果

著者らは、AIME24、AIME25、AMC、MATH、Minerva、OlympiadBench の 6 つのベンチマークにおいて OCC を評価した。

性能：すべてのベンチマークおよびモデル規模（7B 密および 106B MoE）において、OCC は Detach ベースラインおよび固定係数のポリシー損失を一貫して上回った。
- DAPO を用いた MiMo-7B において、OCC は平均精度 61.7% を達成し、Detach（58.9%）およびポリシー損失（57.4%）を上回った。
- GLM-4.5-Air において、OCC は平均精度 67.6% を達成し、Detach の 65.9% を上回った。
訓練ダイナミクス：ポリシー損失が「上昇後、低下」の傾向（後期段階での劣化）を示したのに対し、OCC は訓練全体を通じて安定した改善を維持した。CE 損失は Detach よりも一貫して低い性能を示した。
効率性：OCC は無視できるオーバーヘッドしか導入しなかった。OCC のステップあたりのウォールクロック訓練時間（8.07 秒）は Detach（8.11 秒）と同等であったのに対し、フルモデルの逆伝播には 45.06 秒を要し（5.6 倍の遅延）、であった。
アブレーション：感度分析により、事前定義されたスケーリング因子 $\lambda_+$ は広範な値の範囲でロバストであり、 $\lambda_+ = 1.0$ が最高の平均性能をもたらすことが示された。

意義と主張

本論文は、MTP-RL 共同訓練の失敗は本質的なものではなく、静的な係数と訓練中の勾配整合の動的なドリフトとの間のミスマッチに起因すると主張する。最適化地形の理論的分解を提供し、OCC を提案することで、著者らは MTP を RL 事後訓練に安全に再導入できることを実証する。

その意義は、モデルが RL を通じて推論能力を獲得すると同時に、性能を犠牲にしたり著しい計算コストを伴ったりすることなくマルチトークン表現を活用することを可能にする点にある。この研究は、非凸深層学習地形における共同訓練のために適応的較正が不可欠であることを示唆し、現在の業界標準である勾配の切り離しに対する原理的な代替案を提供する。

限界

著者らは、理論的解析が非凸地形では正確に成り立たない可能性のある $L$ -滑らかさの仮定に依存しており、定量的というよりは定性的なガイダンスを提供している点を指摘している。対数確率プロキシは一次近似であり、大きな学習率下では忠実度を失う可能性がある。さらに、実験は検証可能な報酬を伴う数学的推論に焦点を当てており、学習された報酬モデルを伴う RLHF やオープンエンド生成などの設定への一般化は、まだ検証されていない。