LLM（大規模言語モデル）の強化学習では、rollout（学習用の回答例をモデルに生成させる工程）が計算時間の大部分を占めます。
ここでいう方策（policy）とは、モデルのアーキテクチャや勾配そのものではありません。「現在の文脈に対して、次のtoken（文章を分割した処理単位）をどの確率で選ぶか」を決めるルールです。LLMでは、その時点のモデル重みと、そこから計算されるtokenの確率分布が方策に相当します。
「方策がデータを集める」とは、モデルがこの確率分布に従って回答例を生成することです。本記事では、初出以降も意味が曖昧になりそうな箇所では「モデルの確率分布」と言い換えます。
そこでrolloutだけをFP8へ量子化し、学習はBF16で行うと高速化できます。しかし、回答例を生成したFP8モデルと、その回答例を使って勾配（モデルの重みをどちらへ変えるかを表す値）を計算するBF16モデルでは、tokenの確率分布が少し異なります。このずれが「rollout-training mismatch」です。
本記事では、なぜ同じモデルでもFP8とBF16でtoken確率が変わるのか、量子化誤差が学習初期には探索を助け、後期には有害なbias（勾配の系統的なずれ）へ変わる理由、そして提案手法AISが3つの指標から補正強度を決める仕組みを、具体的な数値例と図で解説します。
3文要約
FP8で回答を生成し、BF16で勾配を計算すると、丸め誤差などによって生成時と学習時のtoken確率が一致しません。その結果、「回答例を生成したモデルと、学習対象のモデルが同じ確率分布で動く」というon-policy（オンポリシー）の前提が崩れます。
論文では、この不一致が常に悪いわけではなく、学習初期には回答の多様性を増やす探索ノイズとして働く一方、少数のtokenへ確率が集中する学習後期には、勾配を不安定にするbias（偏り）へ変わると示しています。
AIS（Adaptive Importance Sampling）は、重みの信頼性、分布差の深刻度、勾配分散の増幅という3指標をバッチごとに測り、FP8の探索効果を残しながら必要なときだけ重点サンプリング補正を強めます。
論文情報
| 項目 | 内容 |
|---|---|
| 論文タイトル | AIS: Adaptive Importance Sampling for Quantized RL |
| 著者 | Jiajun Zhou, Wei Shao, Lingchao Zheng, Yuwei Fan, Ngai Wong |
| 研究機関 | Huawei、The University of Hong Kong |
| 発表年 | 2026年 |
| arXiv提出 | 2026年5月13日（v1） |
| 論文リンク | AIS: Adaptive Importance Sampling for Quantized RL |
| arXiv PDF | |
| 公式実装 | 論文では公開コードベースに言及していますが、v1のarXivページから公式URLは確認できませんでした |
本記事はarXiv v1を対象にします。
論文の実験結果と、本記事で理解のために作った数値例・概念図は明確に分けます。
本記事の目的
本記事の目的は、「FP8は精度が低いので誤差が出る」という一言で終わらせず、その誤差がtoken確率を変え、LLM強化学習の勾配計算へどのように影響するかを順に理解することです。
特に、次の問いへ答えます。
- FP8 rolloutとBF16 trainerの間で、なぜ確率分布がずれるのか
- 量子化誤差が、なぜ学習初期には役立ち、後期には有害になるのか
- 重点サンプリングの確率比は、どのように勾配を補正するのか
- AISの3指標は何を測り、それぞれどのような効果を持つのか
- 論文の実験で、精度と速度はどの程度改善したのか
背景：LLM強化学習ではrolloutがボトルネックになる
LLMの強化学習では、大きく次の処理を繰り返します。
- promptに対して複数の回答を生成する
- 正解判定などからreward（望ましさを表す報酬）を計算する
- 良かった回答の確率を上げ、悪かった回答の確率を下げるように勾配を計算する
- モデルの重みを更新する
このうち、長い文章をtokenずつ生成するrolloutが重い処理です。論文では、rolloutがend-to-end latency（1学習step全体の所要時間）の最大70%を占める場合があると説明しています。
rolloutをFP8で実行すれば、論文の実験では1.5〜2.76倍のrollout高速化と、およそ50%のメモリ削減が得られました。
一方、学習側は小さな勾配や更新量を扱うため、数値的な安定性を保ちやすいBF16を使います。
| 処理 | 主な役割 | 精度 | 利点 | 問題 |
|---|---|---|---|---|
| rollout engine | 回答を大量に生成する | FP8 | 高速、省メモリ | 量子化誤差で生成確率が変わる |
| trainer | log確率、勾配、重み更新を計算する | BF16 | 数値的に安定 | FP8が集めたデータ分布と一致しない |
速度だけを見ると合理的な分業です。しかし強化学習では、「どのモデルの確率分布から回答例を生成したか」が、勾配計算の正しさに直結します。
rollout-training mismatchはなぜ起きるのか
FP8とBF16は、表現できる数の細かさが違う
浮動小数点数は、符号、指数部、仮数部で表します。
この論文で使うFP8 E4M3は、1bitの符号、4bitの指数、3bitの仮数を持ちます。BF16は、1bitの符号、8bitの指数、7bitの仮数を持ちます。
| 形式 | 合計bit | 指数部 | 仮数部 | おおまかな特徴 |
|---|---|---|---|---|
| FP8 E4M3 | 8 | 4 | 3 | 省メモリで高速だが、近い数を区別する刻みが粗い |
| BF16 | 16 | 8 | 7 | 広い範囲を保ち、FP8より細かい値を表現できる |
| FP16 | 16 | 5 | 10 | BF16より狭い範囲だが、仮数は細かい |
ユーザー指定ではFP8とFP16の不一致が挙げられていますが、本論文のtrainerは厳密にはBF16です。したがって本記事では、論文の条件である「FP8 rollout対BF16 trainer」を中心に説明し、FP16は16bit高精度側の比較として扱います。
もしtrainerがFP16であっても、不一致が起きる基本原理は同じです。FP16の10bit仮数に対してFP8 E4M3の仮数は3bitしかないため、FP8側の重み・活性化はより粗く丸められます。違いは高精度側の指数範囲や仮数精度であり、「FP8で生成した回答の分布と16bit側で勾配を計算する分布が一致しない」という問題構造は変わりません。
量子化は、元の実数 \(X\) を表現可能な離散値へ写します。論文が基礎説明に使う対称線形量子化は、次の形です。
\[
Q(X)
=
\mathrm{clamp}\left(
\mathrm{round}\left(\frac{X}{s}\right),
-2^{b-1},
2^{b-1}-1
\right)s,
\qquad
s=\frac{\max |X|}{2^{b-1}-1}
\]
量子化誤差は次の差です。
\[
\epsilon = Q(X)-X
\]
FP8は整数量子化とは異なる浮動小数点形式ですが、表現可能な値の間にある数を近い値へ丸めるため、同じ考え方で誤差が生じます。
量子化誤差は、丸めだけでなくscaleと外れ値からも生じる
量子化誤差の主な原因は次のとおりです。
| 原因 | 何が起きるか | LLMへの影響 |
|---|---|---|
| 丸め誤差 | 表現できない値を近い格子点へ丸める | 各行列積に小さな差が入る |
| 仮数bit不足 | 近い2値を区別できない | attention scoreやlogitの順位が変わり得る |
| dynamic scaling | tensorの最大値に合わせてscaleを変える | 外れ値があると多数の通常値へ使える解像度が減る |
| overflow / underflow | 大きすぎる値、小さすぎる値を十分に表せない | 飽和や0への丸めが起きる |
| 誤差伝播 | 層ごとの小さな誤差が残差接続や非線形関数を通る | 最終logitで差が拡大する場合がある |
| 実装差 | kernel、演算順序、decoding設定が異なる | 精度形式以外の差も分布ずれへ加わる |
論文の実験では、rolloutの重みと活性化をFP8 E4M3へper-tensor dynamic scaling（tensor全体で動的scaleを決める方式）し、log確率評価と逆伝播はBF16のままにしています。
小さなlogit差がtoken確率と生成文を変える
次は説明用の単純な例です。実験から抜き出した値ではありません。
BF16側で3 tokenのlogitが \([1.06, 1.00, 0.94]\) だったとします。理解しやすいように、FP8側では丸め後に \([1.00, 1.00, 0.9375]\) になったと仮定します。
softmax後の確率は次のように変わります。
| token | BF16 logit | BF16確率 | FP8 logit | FP8確率 | 差 |
|---|---|---|---|---|---|
| A | 1.0600 | 35.35% | 1.0000 | 34.02% | -1.33pt |
| B | 1.0000 | 33.29% | 1.0000 | 34.02% | +0.73pt |
| C | 0.9400 | 31.35% | 0.9375 | 31.96% | +0.61pt |
一回の選択では小さな差に見えます。しかし、token AではBF16が最有力なのに、FP8ではAとBが同率になります。サンプリング（確率的なtoken選択）では、ここから別のtokenが選ばれ、その後の文脈自体が分岐します。
さらに、系列 \(x=(x_1,\dots,x_T)\) の確率はtoken確率の積です。
\[
\pi(x)=\prod_{t=1}^{T}\pi(x_t\mid x_{\lt t})
\]
各tokenでFP8側の確率がBF16側の0.99倍だったとしても、1000 tokenでは \(0.99^{1000}\approx 4.3\times10^{-5}\) です。小さな差が長い推論過程で累積する理由がここにあります。
mismatchは「別の確率分布で生成した回答を、補正せずに学ぶ」問題
以降、FP8モデルが回答を生成するときの確率分布を \(\pi_{\mathrm{rollout}}\)、BF16モデルが勾配を計算するときの確率分布を \(\pi_{\mathrm{train}}\) と表します。これらを専門用語では、それぞれrollout方策、学習方策と呼びます。
実際のデータは \(x\sim\pi_{\mathrm{rollout}}\) から得られるのに、補正なしの勾配は次の形で計算されます。
\[
\widehat{g}_{\mathrm{uncorrected}}
=
\mathbb{E}_{x\sim\pi_{\mathrm{rollout}}}
\left[
A(x)\nabla_\theta\log\pi_{\mathrm{train}}(x)
\right]
\]
\(A(x)\) はadvantage（その回答が平均よりどれだけ良かったか）です。
本来求めたいon-policy勾配は、BF16モデル自身の確率分布 \(x\sim\pi_{\mathrm{train}}\) から回答を生成した場合の平均です。実際にはFP8側から回答を生成しているため、その違いを補正しなければ、推定した勾配にbias（系統的な偏り）が入ります。
これは、別の人が偏った方法で集めたアンケートを、母集団から均等に集めたデータだと思って平均する状況に似ています。
量子化誤差は常に有害ではない：学習時期で役割が変わる
この論文の重要な発見は、量子化誤差を単なる敵として扱わないことです。
同じFP8ノイズでも、モデルがまだ多くのtoken候補へ確率を割り当てる学習初期と、正解らしい少数の出力へ確率が集中した学習後期では、意味が変わります。
学習初期：誤差が探索のきっかけになる
学習初期のモデルは、複数のtokenや推論経路へ比較的広く確率を割り当てています。
FP8の丸めが確率順位を少し入れ替えると、BF16モデルでは選ばれにくかった回答も生成されます。その回答が高いrewardを得れば、trainerは「この推論経路の確率を上げるべきだ」という新しい学習信号を得られます。
これは、地図がまだ曖昧な探索初期に、少しランダムな寄り道をすることで近道を発見する効果に近いです。
論文はこれをstochastic exploration bonus（確率的な探索ボーナス）と捉えています。量子化誤差そのものがrewardを加算するわけではなく、観測するtrajectory（生成系列）の範囲を広げることが利益です。
学習後期：同じ誤差がbiasと分散を増やす
学習が進むと、モデルは高rewardのtokenへ確率を集中させます。このように少数候補へ確率が偏ると、わずかなlogit差でも重要tokenの相対確率が大きく変わり得ます。
この段階でFP8 rolloutが別の経路を頻繁に選ぶと、それは有益な未知探索よりも、ほぼ収束したモデルの確率分布からの系統的なずれになります。
さらに、後述する確率比で完全補正しようとすると、一部のサンプルの重みが極端に大きくなり、勾配分散も増えます。
| 学習時期 | token確率分布 | FP8ノイズの主な意味 | 望ましい対応 |
|---|---|---|---|
| 初期 | 広く、まだ不確か | 多様なtrajectoryを見つける探索 | 補正を緩め、探索効果を残す |
| 中期 | 有望な経路へ集中し始める | 探索とbiasが混在 | バッチ状態に応じて適応補正 |
| 後期 | 少数のtokenへ集中 | on-policy勾配を崩す不安定なbias | 信頼できる範囲で補正を強める |
論文のグラフと数値が示すこと
論文Figure 4では、補正なしFP8 rolloutのESS ratio（実効サンプルサイズ比）が学習初期の0.95超から後期の0.70未満まで低下しています。
ESS ratioが1に近ければ、各サンプルが比較的均等に勾配へ寄与します。0.70未満へ下がることは、少数のサンプルへ補正重みが集中し、勾配推定が不安定になっていることを示します。
また、Qwen3-8BのAIME25では、BF16の36.70%に対し、補正なしFP8 rolloutは26.81%でした。差は9.89ポイントです。一方、AISは43.33%まで回復しました。
上の棒グラフは、論文Table 1のQwen3-8Bの値を使って本記事で再作図したものです。学習時期ごとのreward曲線そのものではありませんが、後期までmismatchを放置した最終性能への影響を、具体的な数値で比較できます。
重要なのは、「FP8が常にBF16より悪い」でも「ノイズを残せば常に良い」でもないことです。必要なのは、変化するbiasとvariance（推定値のばらつき）の釣り合いを学習中に監視することです。
提案手法AIS：補正を0か1かではなく連続的に混ぜる
重点サンプリングは「集め方の偏り」を確率比で直す
FP8モデルが生成したサンプルを、BF16モデルから生成したかのように評価する基本的な補正がimportance sampling（重点サンプリング）です。
trajectory全体の重みは次の確率比です。
\[
w(x)
=
\frac{\pi_{\mathrm{train}}(x)}{\pi_{\mathrm{rollout}}(x)}
=
\prod_{t=1}^{T}
\underbrace{
\frac{\pi_{\mathrm{train}}(x_t\mid x_{\lt t})}
{\pi_{\mathrm{rollout}}(x_t\mid x_{\lt t})}
}_{\rho_t}
\]
たとえば、あるtokenの確率がFP8 rolloutでは0.20、BF16 trainerでは0.30なら、確率比は \(\rho=0.30/0.20=1.5\) です。そのtokenはBF16モデルでは本来1.5倍現れやすいので、勾配への寄与を1.5倍へ増やします。
逆に、rolloutで0.40、trainerで0.20なら \(\rho=0.5\) で、寄与を半分にします。
ただしtrajectory全体では比の積になるため、長い系列では重みが極端に大きくなるか、ほぼ0になります。そこで論文は、per-tokenの重みを上限 \(C\) で切るtruncated importance sampling（打ち切り重点サンプリング）を土台にします。
\[
\bar{w}(x)=\min(w(x),C)
\]
実験では \(C=5\) です。
AISの混合勾配
AISは、補正なしを \(\alpha=0\)、完全な重点サンプリング補正を \(\alpha=1\) とし、その間を連続的に混ぜます。
\[
\widehat{g}
=
\mathbb{E}_{x\sim\pi_{\mathrm{rollout}}}
\left[
\left(
1-\alpha(x)+\alpha(x)\bar{w}(x)
\right)
A(x)\nabla_\theta\log\pi_{\mathrm{train}}(x)
\right]
\]
重み部分は次のようにも書けます。
\[
\widetilde{w}(x)
=1+\alpha\left(\bar{w}(x)-1\right)
\]
| \(\alpha\) | \(\widetilde{w}\) | 意味 |
|---|---|---|
| 0 | 1 | mismatchを補正しない |
| 0.5 | \(0.5+0.5\bar{w}\) | 補正を半分だけ反映する |
| 1 | \(\bar{w}\) | 打ち切り重点サンプリングを完全に反映する |
\(w\to1\)、つまり2つの確率分布が一致する場合は、\(\alpha\) が何であっても \(\widetilde{w}=1\) です。AISはmismatchがなければ通常の学習へ戻ります。
論文は、補正なしのbiasの大きさを \(\|b_0\|^2\)、完全補正側の分散を \(\sigma_1^2\) とすると、近似的な最適係数を次のように導きます。
\[
\alpha^*
=
\frac{\|b_0\|^2}
{\|b_0\|^2+\sigma_1^2}
\]
bias（偏り）の影響が大きければ \(\alpha^*\) は1へ近づき、補正を強めます。補正によるvariance（ばらつき）の影響が大きければ0へ近づきます。
実際には真のbiasも真の分散も直接は分からないため、AISは次の3指標で近似します。
AISの補正強度を決める3つの指標
1. Weight Reliability：補正重みを信用できるか
重点サンプリング重みが全サンプルでほぼ同じなら、多くのサンプルが勾配へ寄与します。一方、1個だけ極端に大きいと、ミニバッチ全体の更新方向がそのサンプルに支配されます。
AISは、重みのCV（coefficient of variation、平均に対する標準偏差の比）からESS ratioを計算します。
\[
\alpha_{\mathrm{ess}}
=
\sqrt{\mathrm{ESS}_{\mathrm{ratio}}}
=
\left(1+\mathrm{CV}^2(\bar{w})\right)^{-1/2}
\]
| 重みの状態 | CV | \(\alpha_{\mathrm{ess}}\) | 効果 |
|---|---|---|---|
| ほぼ均一 | 小さい | 1に近い | 補正値を信用しやすい |
| 一部へ集中 | 大きい | 0へ近づく | 極端なサンプルによる更新を弱める |
この指標は「補正が必要か」ではなく、「計算した補正を統計的に信用できるか」を測ります。
2. Divergence Severity：そもそも補正が必要か
重みが均一でも、rolloutとtrainerがすでに一致していれば補正は不要です。
AISは、ミニバッチ \(B\) の全tokenについて平均絶対log確率差を測ります。
\[
\bar{d}
=
\frac{1}{|B|}
\sum_{x_t\in B}
\left|
\log\pi_{\mathrm{train}}(x_t\mid x_{\lt t})
–
\log\pi_{\mathrm{rollout}}(x_t\mid x_{\lt t})
\right|
\]
\[
\alpha_{\mathrm{mis}}
=
\min\left(1,\frac{\bar{d}}{\delta}\right)
\]
論文の実験では \(\delta=0.02\) です。
| \(\bar d\) の例 | \(\alpha_{\mathrm{mis}}\) | 効果 |
|---|---|---|
| 0.002 | 0.10 | ほぼ一致しているため補正を弱くする |
| 0.010 | 0.50 | mismatchに応じて半分だけ必要性を認める |
| 0.020以上 | 1.00 | mismatchが十分大きく、補正を有効にする |
この指標はnecessity gate（補正の必要性を決める門）です。
3. Variance Amplification：補正が勾配を荒らしていないか
重みの分布だけを見ても不十分です。大きな重要度重みが、大きな絶対値のadvantageと偶然重なると、実際の勾配寄与 \(A\bar w\) のばらつきが急増します。
AISは次の分散増幅比を測ります。
\[
\Delta_\sigma
=
\frac{\mathrm{std}(A(x)\bar{w}(x))}
{\mathrm{std}(A(x))+\epsilon}
\]
許容値 \(\gamma\) を超えた分をpenaltyにします。
\[
\alpha_{\mathrm{var}}
=
\max\left(0,\frac{\Delta_\sigma-\gamma}{\gamma}\right)
\]
論文の実験では \(\gamma=1.2\)、\(\epsilon=10^{-6}\) です。
| \(\Delta_\sigma\) | \(\alpha_{\mathrm{var}}\) | 効果 |
|---|---|---|
| 1.0 | 0 | 補正後も分散が増えていないためpenaltyなし |
| 1.2 | 0 | 許容境界 |
| 1.5 | 0.25 | 補正強度からpenaltyを差し引く |
| 2.4 | 1.00 | 勾配が大きく荒れるため補正を強く抑える |
この指標の役割は、bias（偏り）を直そうとしてvariance（ばらつき）を急増させる失敗を防ぐことです。
3指標を1つの補正係数へ統合する
最終係数は次の式です。
\[
\alpha
=
\underbrace{
\mathrm{clip}
\left(
\alpha_{\mathrm{ess}}-\beta_{\mathrm{var}}\alpha_{\mathrm{var}},
0,1
\right)
}_{\text{Reliability}}
\cdot
\underbrace{\alpha_{\mathrm{mis}}}_{\text{Necessity}}
\]
掛け算になっているため、mismatchが小さければ補正は0へ近づきます。また、mismatchが大きくても重みが信頼できない、または勾配分散が危険なら補正を抑えます。
| 状態 | \(\alpha_{\mathrm{ess}}\) | \(\alpha_{\mathrm{mis}}\) | \(\alpha_{\mathrm{var}}\) | 最終判断 |
|---|---|---|---|---|
| 一致している | 高 | 低 | 低 | 不要なので補正しない |
| mismatch大、重み安定 | 高 | 高 | 低 | 補正を強める |
| mismatch大、重み集中 | 低 | 高 | 中 | 信頼できないため弱める |
| mismatch大、勾配分散大 | 中 | 高 | 高 | 学習を荒らすため弱める |
具体例：AISはtokenごとの勾配をどう変えるのか
次はAISの計算を追うための説明用ミニバッチです。論文の実測値ではありません。
4 tokenの打ち切り済み重要度重みを \(\bar w=[0.8,1.0,1.2,2.0]\)、advantageを \(A=[1.2,0.4,-0.2,-0.8]\) とします。
このとき、重みの \(\mathrm{CV}=0.421\)、\(\alpha_{\mathrm{ess}}=0.922\) です。
\(\mathrm{std}(A\bar w)/\mathrm{std}(A)=1.290\) なので、\(\gamma=1.2\) では \(\alpha_{\mathrm{var}}=0.075\) です。mismatchは閾値以上で \(\alpha_{\mathrm{mis}}=1\) とします。
説明を単純にするため \(\beta_{\mathrm{var}}=1\) と置くと、最終係数は次の値です。
\[
\alpha
=
(0.922-1\times0.075)\times1
=0.847
\]
各tokenの補正重みは \(\widetilde w=1+0.847(\bar w-1)\) です。
| token | \(\bar w\) | 元の \(A\) | AIS重み \(\widetilde w\) | 補正後 \(\widetilde A=\widetilde wA\) | 学習への意味 |
|---|---|---|---|---|---|
| 1 | 0.8 | 1.2 | 0.831 | 0.997 | rolloutで過剰だった正例を少し弱める |
| 2 | 1.0 | 0.4 | 1.000 | 0.400 | 分布が一致するため変更しない |
| 3 | 1.2 | -0.2 | 1.169 | -0.234 | trainerで重要な負例を少し強める |
| 4 | 2.0 | -0.8 | 1.847 | -1.478 | BF16モデル側で出現確率が高い強い負例を大きく反映する |
実際のpolicy gradientでは、各tokenについて次の方向へ更新します。
\[
\widetilde A_t\nabla_\theta
\log\pi_{\mathrm{train}}(x_t\mid x_{\lt t})
\]
\(\widetilde A_t>0\) ならそのtokenのlog確率を上げる方向、\(\widetilde A_t<0\) なら下げる方向です。AISは勾配の向きを新しく作るのではなく、確率比とバッチ診断から各tokenのadvantageの強さを調整します。
なお、論文は \(\alpha\) の3指標をdetached tensor（勾配追跡から切り離した値）で計算します。診断値そのものへ逆伝播し、意図しない二階的な相互作用が生じるのを避けるためです。
AISをGRPOへ組み込む学習フロー
AIS-GRPOでは、2種類の確率比を区別することが重要です。
| 確率比 | 比較する方策 | 直す問題 |
|---|---|---|
| rollout-training ratio \(\rho_{\mathrm{rt}}\) | 同じcheckpointのFP8 rolloutとBF16 trainer | 量子化・実装差による分布ずれ |
| old-new ratio \(\rho_{\mathrm{on}}\) | 更新前方策と現在の学習方策 | PPO/GRPOの通常のpolicy更新幅 |
処理の流れは次のとおりです。
- 現在の重みをFP8 rollout engineへ同期する
- FP8モデルで各promptから複数回答を生成する
- rewardとgroup-relative advantageを計算する
- 同じ生成tokenのlog確率をBF16 trainerでも計算する
- \(\rho_{\mathrm{rt}}=\exp(\log\pi_{\mathrm{train}}-\log\pi_{\mathrm{rollout}})\) を求める
- \(\alpha_{\mathrm{ess}}\)、\(\alpha_{\mathrm{mis}}\)、\(\alpha_{\mathrm{var}}\) から \(\alpha\) を求める
- \(\widetilde w=1+\alpha(\min(\rho_{\mathrm{rt}},C)-1)\) でadvantageを補正する
- 通常のGRPO clippingとKL penaltyを使ってBF16で更新する
AISは既存のGRPO目的関数を置き換えるのではなく、量子化由来のずれを補正したadvantageを渡します。3指標はいずれも学習時にすでに計算する統計量のscalar reduction（多数の値を平均や標準偏差へ集約する処理）なので、論文は追加overheadを無視できる程度としています。
実験結果：AISはFP8の速度を残して精度を回復したか
論文は、autoregressive（左から右へtoken生成する）Qwen3-8B、Qwen3.5-9Bと、diffusion型LLMのLLaDA-8B-Instructを評価しました。
Qwen3-8Bの主結果
| 手法 | GSM8K | MATH500 | AIME25 | AMC | Olympiad | Sudoku |
|---|---|---|---|---|---|---|
| BF16 | 89.84 | 76.80 | 36.70 | 68.50 | 48.80 | 38.28 |
| FP8 rollout | 87.65 | 75.80 | 26.81 | 61.25 | 42.10 | 36.33 |
| TIS | 88.63 | 76.20 | 34.23 | 63.40 | 46.50 | 37.50 |
| AIS | 89.99 | 77.00 | 43.33 | 69.75 | 48.20 | 39.10 |
補正なしFP8は、難しいAIME25で9.89ポイント低下しました。固定補正のTISは34.23%まで回復しますが、BF16に届きません。AISは43.33%でBF16を6.63ポイント上回りました。
ただし、これは「FP8が本質的にBF16より高精度」という意味ではありません。学習中の探索と適応補正の組み合わせが、今回の設定でBF16 baselineより良い最終checkpointへ到達した結果です。
Qwen3とQwen3.5を合わせた12個の推論benchmark中、AISは11個でBF16と同等以上でした。例外はQwen3-8BのOlympiadとQwen3.5-9BのSudokuです。
LLaDAと長い推論系列
LLaDA-8BのCountdownでは、生成長512 tokenでBF16が52.20%、補正なしFP8が33.20%でした。差は19.00ポイントです。AISは54.69%まで回復しました。
これは、系列が長いほどtokenごとの分布差が累積しやすいという説明と整合します。
効率
| 項目 | 論文の報告 |
|---|---|
| FP8 rollout高速化 | 1.5〜2.76倍（abstractの範囲） |
| rolloutメモリ | およそ50%削減 |
| AIS追加計算 | 既存統計のscalar reduction中心で無視できる程度 |
| trainer / optimizer | BF16のまま |
速度向上幅はmodel、batch size、生成方式で変わります。AISが速くするのはFP8 kernelそのものではなく、FP8 rolloutを精度低下で諦めずに使えるようにする部分です。
どのようなサービス・技術に使えるか
AISが特に有効になり得るのは、rolloutが長く、大量のサンプルを必要とするRLVR（検証可能な報酬を使う強化学習）です。
| 適用候補 | 期待できる効果 | 実装上の注意 |
|---|---|---|
| 数学推論LLM | 長い思考列のrolloutを高速化 | reward verifierの品質が必要 |
| code生成RL | unit testを報酬に大量sampling | 本論文ではcode benchmark未評価 |
| planning agent | 複数stepの試行を省メモリ化 | 環境interactionの分布差も考慮が必要 |
| diffusion LLM | 反復denoising型生成のFP8化 | architecture固有のlog確率計算が必要 |
| 大規模RLHF基盤 | rollout workerのthroughput向上 | 人手選好rewardでは再検証が必要 |
画像高画質化タスクへ直接適用する手法ではありません。ただし、画像生成・復元modelをRLでfine-tuningし、生成候補の評価をrewardへ変換するpipelineでは、低精度samplingと高精度trainerの分布差を管理する考え方が応用できる可能性があります。
その場合、画像tokenやlatent trajectoryは文章tokenと統計が異なり、per-token確率比、advantage、長さ正規化をそのまま流用できるとは限りません。実験による確認が必要です。
限界と実装時の注意点
- 評価は数学、planning、一般能力が中心で、code生成、対話、人手選好RLHFは未評価です。
- 実験のAIS hyperparameterはGSM8K validationで調整し、他taskへ固定しています。異なるmodel規模やFP4では再調整が必要な可能性があります。
- 確率比を計算するには、rollout時のlog確率を正しく保存し、同じtokenをBF16 trainerでも再評価する必要があります。
- \(C=5\) の打ち切りは分散を抑えますが、完全なunbiased estimatorではなく、truncation biasが残ります。
- 「量子化ノイズが探索に有益」という観察は、無制限にノイズを増やしてよいことを意味しません。
- 論文v1では公式実装URLをarXivページ上で確認できず、細かな再現条件は追加確認が必要です。
- 論文の理論上のoracle係数は近似です。実用式の3指標は、真のbiasとvarianceを直接測るものではありません。
特に、Equation 17の \(\beta_{\mathrm{var}}\) はvariance penaltyへの感度を調整しますが、論文v1のhyperparameter一覧では具体値が明記されていません。実装時には公開コードや著者設定の確認が必要です。
よくある誤解
| 誤解 | 正確な理解 |
|---|---|
| FP8の重みが同じ元modelから作られれば、出力確率も同じ | 丸めとscaleによりforward結果とlogitが変わる |
| 量子化誤差は常に学習を悪化させる | 初期にはtrajectoryの多様性を増やす場合がある |
| ノイズが探索に良いなら補正しない方がよい | 後期にはbiasと分散が増え、崩壊を招き得る |
| importance samplingで完全に直せる | 長系列では確率比の積が高分散になりやすい |
| AISは常に補正を強める | 必要性と信頼性が低いbatchでは補正を弱める |
| \(\alpha_{\mathrm{ess}}\) が高ければ補正すべき | 重みを信用できるだけで、必要性は \(\alpha_{\mathrm{mis}}\) が決める |
| AISがFP8演算を高速化する | 速度はFP8 kernel由来で、AISは精度・安定性を回復する補正 |
| 論文はFP16 trainerを使う | 実験のtrainerとbackpropagationはBF16 |
関連技術
| 技術 | AISとの関係 | 主な違い |
|---|---|---|
| GRPO | AISを統合するpolicy optimization手法 | group内reward平均との差でadvantageを作る |
| PPO | GRPOの基礎となるtrust-region型手法 | value modelを使う構成が一般的 |
| TIS | per-token確率比を固定閾値で打ち切るbaseline | batchごとの適応係数を持たない |
| FP8-RL | FP8 rolloutを実用化する低精度stack | AISは補正強度の非定常性へ焦点を当てる |
| Jet-RL | trainingとrolloutのprecision flow統一を目指す | mismatchを補正するより発生源を揃える発想 |
| exploration noise | 未知の行動を試すための摂動 | AISでは量子化誤差から副次的に生じる |
関連する一次資料として、FP8-RLとJet-RLも、低精度rolloutと学習安定性の別アプローチを比較するうえで有用です。
まとめ
FP8 rolloutとBF16 trainerの組み合わせでは、量子化された重みと活性化の丸め、dynamic scaling、層間の誤差伝播、実装差により、同じtokenの生成確率が変わります。
その結果、FP8モデルが生成したtrajectoryを、BF16モデル自身が生成したサンプルとして補正せずに扱うと、policy gradientにbiasが入ります。長い推論系列ではtokenごとの差が累積するため、難しいreasoning taskほど影響が大きくなり得ます。
一方、量子化誤差は学習初期には探索範囲を広げる場合があります。問題はノイズの存在そのものではなく、モデルの確率が少数tokenへ集中した後期にも同じ扱いを続けることです。
AISは、次の3つを分けて測ります。
- \(\alpha_{\mathrm{ess}}\)：補正重みを信用できるか
- \(\alpha_{\mathrm{mis}}\)：補正が必要なほど分布がずれているか
- \(\alpha_{\mathrm{var}}\)：補正が勾配分散を増幅していないか
この3指標から補正係数 \(\alpha\) をbatchごとに決めることで、FP8の探索効果を残しながら、後期の不安定なbiasを抑えます。
論文では、Qwen3-8BのAIME25で補正なしFP8の26.81%をAISで43.33%へ改善し、FP8 rolloutの1.5〜2.76倍の高速化とおよそ50%のメモリ削減を維持したと報告しています。
次に読むべき記事
- FP8とは？E4M3・E5M2・BF16の違いとdynamic scaling
- GRPOとは？group-relative advantageとPPOとの違い
- Importance Samplingとは？biasとvarianceを数式で理解する
- FP8-RLとJet-RLの違い：補正するかprecisionを統一するか
- 長いChain-of-Thoughtで確率比が不安定になる理由
コメント