技術的概要：量子化 RL における適応的重要度サンプリング（AIS）

1. 問題定義

大規模言語モデル（LLM）に対する強化学習（RL）は、現在、軌道ロールアウト生成の計算コストによってボトルネックが生じている。これを緩和するため、DeepSeek V3.2 や veRL などの最近の業界慣行では、ロールアウトエンジンに低精度量子化（例：FP8）を採用しつつ、トレーナーはより高精度（BF16）を維持している。これにより、「ロールアウトとトレーニングのミスマッチ」が生じる。すなわち、量子化されたロールアウト方策（$\hat{\pi}_{\theta_{old}}$）が生成する軌道は、勾配計算に使用される全精度の学習者方策（$\pi_{\theta_{old}}$）とは分布が異なる。

本論文は、このミスマッチを「非定常的」であり「両刃の剣」として機能すると特定している。

• トレーニング初期段階: 量子化ノイズは確率的な探索ボーナスとして作用し、トレーナーが通常なら過小サンプリングする可能性のある軌道に勾配を曝す。
• トレーニング後期段階: 方策が鋭敏化するにつれ、同じノイズは不安定化させるバイアスの源となり、推論ベンチマーク（例：Qwen3-8B における AIME25 精度の 6.63% 低下）でのトレーニング崩壊を引き起こす。

標準的な重要度サンプリング（IS）はこのバイアスを補正しようとするが、長い系列において許容できないほど高い分散が生じるため、しばしば失敗する。一方、切り捨て重要度サンプリング（TIS）はトレーニングを安定させるが、ミスマッチの進化性質に適応できない固定された補正強度を使用するため、過補正（有用な探索の抑制）または過不足ない補正（不安定性の許容）のいずれかに陥る。

2. 手法：適応的重要度サンプリング（AIS）

著者は、バッチごとに重要度サンプリング補正の強度を動的に調整するフレームワークである**適応的重要度サンプリング（AIS）**を提案する。固定された閾値の代わりに、AIS は未補正の勾配と完全に重要度重み付けされた勾配の間を補間する単一の混合係数 $\alpha(x) \in [0, 1]$ を計算する。

2.1 混合係数

AIS の核心は、方策更新中に既に計算されている統計量から導出された 3 つのリアルタイム診断を組み合わせた因数分解されたゲーティング機構である。

1. 重みの信頼性（$\alpha_{ess}$）: 切り捨てられた重要度重みの有効サンプルサイズ（ESS）比率を測定する。重み分布が集中している場合（ESS が低い）、補正は信頼性が低いとみなされ、重みが下がる。
   $$\alpha_{ess} = (1 + CV^2(\bar{w}))^{-1/2}$$
2. 発散の深刻度（$\alpha_{mis}$）: ロールアウト方策とトレーニング方策間の平均絶対対数確率の不一致を測定する。このミスマッチが飽和閾値 $\delta$ を超えた場合のみ、補正が作動する。
   $$\alpha_{mis} = \min\left(1, \frac{\bar{d}}{\delta}\right)$$
3. 分散増幅（$\alpha_{var}$）: 大きな重要度重みが大きな大きさのアドバンテージと一致し、許容限度を超えて勾配分散を増幅するかどうかを監視する。分散増幅比率 $\Delta\sigma$ が許容値 $\gamma$ を超える場合、ペナルティが適用される。
   $$\alpha_{var} = \max\left(0, \frac{\Delta\sigma - \gamma}{\gamma}\right)$$

最終的な混合係数は以下のように計算される。
$$\alpha = \text{clip}(\alpha_{ess} - \beta \cdot \alpha_{var}, 0, 1) \cdot \alpha_{mis}$$

2.2 GRPO への統合

AIS は Group Relative Policy Optimization（GRPO）に統合される。この手法は、GRPO の目的関数におけるアドバンテージ項を調整する整流重み $\tilde{w}_i$ を導入する。
$$\tilde{w}_i = 1 + \alpha \cdot (\min(\rho_{rt}, C) - 1)$$
ここで、$\rho_{rt}$ はロールアウトとトレーニングの比率、$C$ は切り捨て閾値である。この定式化により、ミスマッチが無視できる場合（$\alpha \to 0$）、目的関数はバニラ GRPO と完全に一致し、後方互換性が維持される。

3. 主要な貢献

• 非定常性の特性評価: 本論文は、ロールアウトとトレーニングのミスマッチが均一的に有害ではないことを確立している。トレーニング初期には探索ボーナスを提供するが、後期にはバイアスの源となる。これは、量子化ミスマッチを単なる敵対者として見る見方に挑戦するものである。
• 適応的補正フレームワーク: AIS は、静的な閾値を使用するのではなく、リアルタイム診断（信頼性、発散、分散）に基づいて重要度サンプリング補正の強度を動的に調整する最初の手法である。
• 拡散モデルへの応用: この研究は、拡散ベースの LLM（具体的には LLaDA-8B-Instruct）に対する RL における低精度ロールアウトの最初の研究であり、ミスマッチ問題と AIS 解決策が自己回帰アーキテクチャを超えて一般化することを示している。
• 効率性と精度のトレードオフ: この手法は、全精度 BF16 トレーニングの精度を回復しつつ、FP8 ロールアウトのスループットとメモリ上の利点を維持するバランスを達成する。

4. 実験結果

著者は、数学的推論（GSM8K、MATH500、AIME25、AMC、Olympiad）および計画（Sudoku、Countdown）ベンチマークにおいて、Qwen3-8B、Qwen3.5-9B（自己回帰）、およびLLaDA-8B-Instruct（拡散）に対して AIS を評価した。

• 精度回復:
  • 自己回帰モデル: Qwen3-8B において、単純な FP8 ロールアウトは AIME25 精度を BF16 の 36.70% から 26.81% まで低下させた。AIS はこれを**43.33%**まで回復させ、BF16 ベースラインを上回った。Qwen3.5-9B でも同様の改善が見られ、AIS は 12 組のベンチマークペアのうち 11 組で BF16 と同等かそれ以上の性能を示した。
  • 拡散モデル: LLaDA-8B において、FP8 ロールアウトは大幅な劣化を引き起こした（例：512 トークンでの Countdown で -19.00%）。AIS は、さまざまな系列長および微調整手法（フル FT と LoRA）において、BF16 レベルに匹敵または上回る性能まで回復させた。
• 汎用能力: AIS は、MMLU、PIQA、ARC などの汎用ベンチマークでの性能を維持しており、RL 報酬信号への過学習がないことを示している。
• 効率性の向上:
  • 高速化: FP8 ロールアウトは、バッチサイズとモデルアーキテクチャに応じて、エンドツーエンドのロールアウト時間を1.5 倍から 2.76 倍高速化した（拡散モデルではテンソルコアの活用により、より高い向上が観測された）。
  • メモリ: FP8 ロールアウトはロールアウトモデルのフットプリントを半減させた（例：LLaDA-8B で約 7.95 GB 対 約 15.9 GB）。これにより、BF16 下ではメモリ不足（OOM）となっていたより大きなバッチサイズが可能になった。

5. 意義と主張

本論文は、AIS が拡散 LLM に対する RL における低精度ロールアウトの最初の研究であると主張している。その主な意義は、量子化の効率性向上を犠牲にすることなく「ロールアウトとトレーニングのミスマッチ」を解決することにある。

著者は以下を主張している。

1. 排除ではなく管理: ミスマッチは排除されるべきではなく、管理されるべきである。AIS は、トレーニング後期の不安定化バイアスを効果的に抑制しつつ、初期段階における量子化ノイズの探索的利点を維持することに成功している。
2. スケーラビリティ: この手法は計算効率が良く、診断が既存の統計量に対するスカラー縮約であるため、オーバーヘッドは無視できる。
3. ロバスト性: AIS は、異なるモデルファミリー（自己回帰対拡散）およびトレーニングパラダイム（フル FT 対 LoRA）にわたって一般化する。

本論文は、AIS が高スループットかつ低精度の RL トレーニングを可能にし、全精度ベースラインに匹敵または上回る性能を発揮するため、リソース制約のある環境における RL のスケーリングに対する実用的な解決策となると結論付けている。