Not all tokens are needed(NAT): token efficient reinforcement learning

この論文は、大規模言語モデルの強化学習において、Horvitz-Thompson 再重み付けを用いて生成トークンの一部のみで偏りのない勾配を推定する「NAT（Not All Tokens Are Needed）」フレームワークを提案し、計算コストやメモリ使用量を削減しながらフルトークン学習と同等の性能を達成することを示しています。

要約

この論文は、人工知能（AI）が「長い思考プロセス」を学ぶ際の**「もったいないな」を解消する**画期的な方法について書かれています。

タイトルは**「NAT（Not All Tokens are Needed）」、つまり「すべての単語（トークン）が必要じゃない」**というものです。

以下に、難しい専門用語を避け、日常の例え話を使って簡単に解説します。

---

🧠 背景：AI は「考えすぎ」で疲れている

最近の AI（大規模言語モデル）は、数学の問題や複雑なタスクを解くとき、人間のように「ステップバイステップ」で考えることができます（これを「思考の連鎖」と呼びます）。

しかし、ここで大きな問題が起きます。
AI が長い文章を生成して学習する際、「生成したすべての単語」に対して、AI は「正解だったか間違っていたか」を振り返り、脳（メモリ）をフル回転させて修正しようとするのです。

• 例え話：
  料理のレシピを覚えるために、料理人が「玉ねぎを切る」「炒める」「塩を振る」「火を止める」「盛り付ける」という 100 の手順をすべて記録しました。
  味見をした結果、「塩が少し足りなかった」という評価が返ってきました。
  従来の AI は、「玉ねぎを切った瞬間」から「盛り付け」まで、すべての 100 の手順を頭の中で 1 回ずつやり直して、どこを直せばいいか計算します。

  でも、よく考えてみてください。「塩を振る」直前の手順や、盛り付けの最後の数歩は、味（正解）に直接関係ないかもしれません。なのに、AI はそれらすべてを計算し直しているのです。これでは、計算リソース（GPU メモリや時間）が大量に無駄遣いされてしまいます。

💡 解決策：NAT（Not All Tokens are Needed）

この論文が提案する「NAT」は、**「すべての手順を振り返る必要はないよ！」**という考え方です。

AI が生成した長い文章の中から、「本当に学習に必要そうな部分」だけをランダムに選び出し、その部分だけを使って学習（修正）するという方法です。

🎯 2 つの新しい「選び方」

論文では、この「選び方」を 2 つの簡単な方法で実現しました。

1. ランダムな抜き取り（URS）：

   • 文章全体から、サイコロを振って「ここは勉強する」「ここは飛ばす」を決めます。
   • 効果： 計算量は減りますが、AI が「前の文脈」を読み飛ばせないため、メモリの節約効果は限定的です。

2. ランダムな「途中まで」切り取り（RPC）：

   • これが今回の主役です。
   • 「この文章は、前半の 50% までしか読まないで勉強しよう！」と、先頭から連続した部分だけをランダムに切り取って学習します。
   • 例え話： 料理のレシピを学ぶとき、「玉ねぎを切る」から「炒める」までだけを見て、「塩を振る」以降は最初から読まないことにします。
   • すごい点： これなら、AI は「後半の文」を頭の中に一度も展開（メモリに展開）する必要がなくなります。つまり、メモリの圧迫が劇的に減ります。

⚖️ 重要なポイント：「偏り」をなくす魔法

「後半を捨てちゃっていいの？AI が変な方向に学習しない？」と心配になるかもしれません。
ここで、この論文の**「魔法の重み付け（ホーヴィッツ＝トンプソン推定）」**が登場します。

• 仕組み：
  「後半を捨てた分、前半の学習結果を少しだけ『重く』して評価する」という計算を行います。
  • 例：「100 回中 50 回しか後半を見ていないなら、その 50 回の学習結果を 2 倍の重みで評価する」
• 結果：
  数学的に証明されている通り、「すべてを見た場合」と「一部だけを選んで重み付けした場合」の学習結果は、長期的には全く同じになります。
  つまり、**「質を落とさずに、コストだけ下げる」**ことが可能になったのです。

📊 実際の効果：どれくらい速くなる？

実験結果（Qwen3-8B という AI モデルを使った場合）によると：

• メモリ使用量： 約 18% 削減（ピーク時のメモリが圧迫されにくくなり、より大きなモデルや長い文章を扱えるようになります）。
• 学習時間： 約 29% 短縮（1 ステップあたりの計算が速くなりました）。
• 性能： 数学の問題を解く正解率は、従来の「すべて見る方法」と全く同じレベルを維持しました。

🌟 まとめ

この論文が伝えたかったことはシンプルです。

「AI に長い思考プロセスを学ばせる際、すべての単語を丁寧に振り返る必要はありません。ランダムに『前半だけ』を切り取って、少し計算を工夫すれば、同じように賢くなりながら、計算コストとメモリを大幅に節約できます。」

これは、AI がもっと複雑で長い思考（例：数時間かかるような研究やプログラミング）をできるようになるための、**「賢い節約術」**と言えます。これにより、より高性能な AI を、より安く、速く作れるようになるでしょう。

技術概要

以下は、提示された論文「Not All Tokens are Needed: Token-Efficient Reinforcement Learning (NAT)」の技術的な詳細な要約です。

論文要約：Not All Tokens are Needed (NAT)

1. 背景と課題 (Problem)

大規模言語モデル（LLM）における強化学習（RL）、特に検証可能な報酬を用いた強化学習（RLVR）は、数学的推論やコード生成などのタスクで劇的な性能向上をもたらしています。しかし、推論プロセスが長くなる（Chain-of-Thought: CoT）につれ、スケーリングには重大なボトルネックが生じています。

• 既存の課題: 標準的な RLVR パイプライン（例：GRPO）では、生成されたすべてのトークンに対して方策勾配（policy gradient）を計算し、バックプロパゲーションを行います。
• コストの増大: 長い CoT 軌道（trajectory）は、活性化メモリ（activation memory）の爆発的な増加、1 更新あたりの FLOPs の増大、および OOM（Out-of-Memory）エラーや不安定な最適化ダイナミクスを引き起こします。
• 非効率性: 最近の研究は推論（ロールアウト生成）の高速化に焦点を当てていますが、学習フェーズ（フォワード/バックワードパス）は依然としてメモリ制約と計算集約的な課題を抱えており、生成速度の向上を学習効率の向上に変換できていません。
• 核心的な問い: 「強力な RL 推論者を訓練するために、本当にすべてのトークンが必要なのか？」

2. 提案手法 (Methodology)

著者は、NAT (Not All Tokens are Needed) という新しいフレームワークを提案しました。これは、報酬計算には完全な応答を使用しつつ、方策更新には選択されたトークンのサブセットのみを使用するアプローチです。

2.1 核となる技術：Horvitz-Thompson (HT) 推定

NAT の核心は、トークン選択によるバイアスを補正するためのHorvitz-Thompson (HT) 再重み付けです。

• 仕組み: 各トークン $t$ に「選択確率（inclusion probability）$p_{i,t}$」を割り当てます。バックプロパゲーションには選択されたトークンのみを使用しますが、勾配を $1/p_{i,t}$ で重み付けします。
• 理論的保証: この手法により、部分トークンでの更新が、完全なシーケンスでの RLVR 勾配の**不偏推定量（unbiased estimator）**となることが証明されています。つまり、統計的に正しい勾配を維持しつつ、計算コストを削減できます。

2.2 具体的な実装手法

NAT は、2 つのプラグアンドプレイなトークン選択スキームを実装しています。

1. 一様ランダムサンプリング (URS: Uniform Random Sampling)

   • 各トークンを独立して確率 $p$ で選択します。
   • 効果: バックワードパスの計算量とメモリを削減します。
   • 限界: 因果的アテンション（causal self-attention）の性質上、フォワードパスではすべての先行トークンを処理する必要があるため、フォワード計算コストの削減には寄与しません。

2. ランダムプレフィックスカット (RPC: Random Prefix Cutting)

   • 各応答から、連続したプレフィックス（先頭部分）のみをランダムに選択します（例：長さ $L_i$ のプレフィックス）。
   • 効果:
     • フォワードパス: 処理するシーケンス長が短くなるため、アテンション計算量（$O(T^2) \to O(L^2)$）と活性化メモリが大幅に削減されます。
     • バックワードパス: URS と同様に HT 補正により不偏性を保ちます。
     • バイアスの回避: 決定論的な先頭カット（常に前半だけを使う）とは異なり、後続のトークンも確率的に選択されるため、学習信号の系統的バイアスを防ぎます。

3. 主要な貢献 (Key Contributions)

1. 統一フレームワークの提案: 報酬評価と方策最適化を分離し、HT 推定を用いてトークン効率化を実現する NAT フレームワークを提案しました。
2. 不偏性の証明: 任意の正の選択確率に対して、HT 再重み付けが完全なトークン勾配の不偏推定量となることを理論的に示しました。
3. 実用的な高速化と品質維持: 数学推論ベンチマークにおいて、バックプロパゲーションするトークンを最大 50% に削減しながら、フルトークンの GRPO と同等の性能を達成しました。特に RPC は、ピーク GPU メモリと学習時間を大幅に削減します。

4. 実験結果 (Results)

DAPO-Math-17K データセットを用いて、Qwen2.5-Math-7B と Qwen3-8B モデルで評価を行いました。

• 精度 (Accuracy):
  • RPC と URS: MATH、AIME24、AIME25 などのベンチマークにおいて、フルトークンの GRPO と統計的に有意差のない（重なり合う 95% 信頼区間を持つ）性能を達成しました。
  • Deterministic Truncation (決定論的カット): 後半のトークンを常に削除する方法は、学習信号の歪みにより精度が著しく低下し、不安定になりました。
• メモリ効率:
  • RPC: Qwen3-8B において、ピーク GPU メモリを約 18% 削減（47.72GB → 39.23GB）しました。URS はメモリ削減効果がほとんどありませんでした（フォワードパスの制約のため）。
• 学習時間:
  • RPC: Qwen3-8B において、推論を除く学習ステップあたりの時間を約 29% 削減（311.50 秒 → 220.32 秒）しました。エンドツーエンドの時間削減も約 36% でした。
• エントロピー: RPC と URS は、GRPO と同様にエントロピーが収束し、学習が安定していることを示しました。

5. 意義と結論 (Significance & Conclusion)

• パラダイムシフト: 「長い推論タスクにはすべてのトークンのバックプロパゲーションが必要である」という従来の仮定に挑戦し、統計的に正当化された部分更新が可能であることを示しました。
• スケーラビリティ: NAT、特に RPC は、長い CoT 軌道によるメモリと計算のボトルネックを解決する直交的な（orthogonal）アプローチを提供します。これはロールアウト生成の最適化とは別に機能するため、既存の高速化技術と組み合わせ可能です。
• 将来展望: 現在の手法は情報非依存（ランダム選択）ですが、将来的には勾配の大きさやトークンの不確実性に基づいた「情報意識型（information-aware）」の選択確率の学習や、システムカーネルとの共同設計が有望視されています。

結論として、NAT は最先端の AI システムを効率的にスケーリングするための、理論的根拠に基づき実用的なツールとして位置づけられています。