非同期強化学習では重要度サンプリングだけでは足りない

導入

非同期強化学習の魅力は、rolloutによるサンプル生成とlearnerによるモデル更新を同時に進められる点にある。同期方式のように、一方が終わるまで他方が待つ必要がないため、GPUや推論サーバーの稼働率を高めやすい。一方で、learnerが受け取るサンプルは、現在のモデルではなく、より古い重みで生成される。

vLLM Blogが紹介したNVIDIA ResearchとInferactのREINFORCE Zeroは、この問題を単純な「古いデータ」の問題として扱わない。非同期学習には、性質の異なる二つの差があると説明している。

重みの遅延とカーネルの不一致

サンプルが重みバージョンvで生成され、バージョンkのlearnerで学習される場合、現在の方策確率をrollout時に記録された確率で割ることで、トークン単位の重要度サンプリングを行える。これは方策の変化によるオフポリシー性を補正する方法だ。

しかし、rolloutとlearnerが同じ重みを使っていても、同じlogprobを返すとは限らない。演算カーネル、数値精度、バッチ処理、実行経路の違いによって、確率計算に差が出ることがある。これがカーネルの不一致である。

両者は挙動も異なる。重みの遅延は、更新が小さい限り、方策の変化に沿って比較的滑らかに推移する。一方、カーネル差はバッチ構成や数値計算に依存し、急激な揺らぎを生みやすい。素材の実験では、重要度重みの偏差指標のピークと中央値の比が、カーネルをそろえた場合は1.50、標準カーネルでは442.54だった。

REINFORCE Zeroの構成

この手法は次の三つの要素からなる。

• 同じ重みとトークンに対して、rolloutとlearnerが同じlogprobを計算する。
• トークン単位の重要度サンプリングを、クリップや切り捨てなしで適用する。
• 価値モデルを使わず、グループ内の平均報酬との差を各応答のアドバンテージとして通常のREINFORCEを最適化する。

確率比をクリップしないのは、正確な比率を使えば期待勾配の補正を保ちやすいからだ。クリップは分散を抑える場合がある一方、一般にはバイアスを導入する。重要なのは、クリップの有無以前に、計算系の不一致による大きな外れ値を先に取り除くことである。

結果と意味

標準カーネルの非同期実験では、rolloutとlearnerのKLが上昇し、AIME24 pass@1は34.27%から13.96%へ低下した。最大勾配ノルムは38.27に達したが、カーネルをそろえたREINFORCE ZeroではKLが安定し、系列の拒否も必要なかった。Qwen3-4BとPythonツールを使った設定では、方策遅延64でも5,000更新まで学習が続き、平均報酬は0.287から約0.78へ上昇した。

16 GPU、12,000更新の比較では、整合した非同期方式の学習ステップスループットは、整合した同期方式の2.40倍だった。DeepSeek-V4.1-Flash、Kimi K3.1、AMD MI355Xにも対応範囲が広がっている。

この研究の示唆は、重要度サンプリングを強化すればすべての非同期誤差を解決できる、という考えへの注意である。まず確率計算を一致させ、その後に方策遅延を統計的に補正する。ツール利用や長いコンテキストを扱うエージェント学習では有効な設計原則になり得るが、モデルやハードウェアごとのカーネル調整コストも評価する必要がある。

出典：vLLM Blog