On-policy のはずが Off-policy になる：LLM 強化学習の rollout mismatchと対策(rollout correction)

ABEJA でデータサイエンティストをしている服部です。 今回は GRPO などの LLM の強化学習の中で実際に起こりうる少しマイナーだけど重要かもしれない強化学習の Rollout（推論実行）時に発生するトピックについての記事です。 ![image](https://cdn-ak.f.st-hatena.com/images/fotolife/g/go5paopao/20260603/20260603004017.png) 過去にも AgenticRL に関しての記事も書いてますので、こちらも興味ある方いましたら読んでいただけると幸いです。 [tech-blog.abeja.asia](https://tech-blog.abeja.asia/entry/geniac3-agentic-rl-process) 本記事は何となくでも GRPO などの手法で何をしているか知っている前提の記事になっています。 上記記事をベースにした内容を W&B 様のホワイトペーパーにも寄稿しており、こちらには強化学習の基礎から書かれているのでこちらももし良ければご確認いただければです。 [wandb.ai](https://wandb.ai/site/ja/resources/whitepapers/agentic-rl/) - 学習と Rollout（推論実行）でのズレ - 実際にズレが起こりうるかを見てみる。 同じ token でも、確率は一致しない - ズレの大きさの分布を見てみる - 今回の範囲で言えそうなこと - 分布の mismatch（不一致） - mismatch（不一致）への対策 token-level IS（重要性サンプリング） - truncated IS（切り捨てられた重要性サンプリング） - 幾何平均 - rejection sampling（拒否サンプリング） - 実ライブラリでの実装 NeMo-RL - verl - まとめ - We Are Hiring! GRPO を始めとする最近の LLM（大規模言語モデル）での強化学習では、モデルの Policy（方策）に従って Rollout（推論実行）を行い、その Rollout 結果を元に与えられた報酬を用いてモデルを更新します。 この Rollout（推論実行）というのは学習データ（指示などのインプット）に対しての推論を行うわけですが - GRPO 系では同一 sample（サンプル）に対して Group 内での出力数分、推論を行う - 一つの Rollout 出力が、Reasoning（推論過程）も含めると時間がかかる - Tool Use（ツール利用）などの処理時間、multi-step 化により時間がかかる といった背景から、学習時間より Rollout（推論実行）の時間がボトルネックになりがちです。 そのため実際のシステムでは、 - 学習：Megatron / FSDP / DeepSpeed - Rollout（推論実行）：vLLM / SG-Lang のように、学習 backend（バックエンド）と Rollout（推論実行）backend（バックエンド）を分離する構成が一般的です。 しかし、**Rollout 時の Policy（方策）と、学習時に仮定している Policy（方策）が一致しないことがある**という課題があります。 このズレは例えば、 - FP16 / BF16 / FP8 の違い - attention 実装の差 - backend 実装の差 - sampling 実装の差 - CUDA kernel の差 などによって、同じ重み（weight）を使っていてもトークンの出力確率が完全には一致しないことがあります。 GRPO/PPO シリーズの学習では、Rollout で得られたトークン列に対して、次のような確率比を使ってポリシーを更新します。 ここで重要なのは、分母の $\pi_{old}$ が単なる正規化項ではなく、そのトークン列を実際に生成したポリシーであるという前提を持っている点です。 つまり PPO/GRPO では、Rollout で得られたアクション（action）$a_t$ は からサンプリングされたものとして扱われます。 しかし、実際のシステムでは Rollout のバックエンドと学習のバックエンドが分かれていることが多く、Rollout を生成したポリシーを $\pi_{rollout}$ 、学習側でオールド・ポリシー（old policy）として再計算したポリシーを $\pi_{learn}$ とすると、 となることがあります。 このとき、**サンプルは $\pi_{rollout}$ から来ているのに、損失関数上では $\pi_{learn}$ から来たものとして扱ってしまう**ことになってしまいます。 例えば、あるトークンについて Rollout バックエンドでは だったのに、学習バックエンドで再計算すると だったとします。 このトークンは、Rollout バックエンド上では比較的出やすいトークンとしてサンプリングされています。 しかし学習バックエンド側で見ると、同じトークンが「オールド・ポリシーではかなり出にくいトークン」として扱われます。 例えば現在のポリシーがそのトークンに確率 $p_{rollout}$ を割り当てていた場合、本来 Rollout 時の確率 $p_{rollout}$ を基準にすれば比率（ratio）は ですが、学習側の $p_{learn}$ を分母に使うと になります。 つまり、ポリシーが Rollout 時から変わっていないようなケースでも、バックエンド間のミスマッチによって「確率が 2 倍に増えた」ように見えてしまいます。このように、バックエンド間の差が PPO/GRPO の比率（ratio）に入ると、損失関数上の更新量やクリップ（clip）のかかり方に影響が出ます。 ## 実際にズレが起こりうるかを見てみる。 ここまで「Rollout 時のポリシーと学習時のポリシーがズレることがある」という話をしてきましたが、実際どのくらいズレるのか、手元で測ってみます。 今回は Qwen3-4B を使って、次のような実験をしました。 - まず vLLM でいくつかのプロンプトに対して生成（Rollout）を行い、各トークンについて「vLLM がそのトークンを選んだ確率」を記録します。これが先ほどの $p_{rollout}$ にあたります。 - 次に、生成された まったく同じトークン列 を transformers (HuggingFace) 側にもう一度通して、同じトークンの確率を測り直します。これが学習側が見る $p_{learn}$ にあたります。 - この 2 つを トークンごとに並べて比較します。 重み（weight）は同じものを使い、データ型（dtype）も両方 bf16 に揃えています。つまりここで出てくる差は、モデルや精度の違いではなく、純粋に vLLM と transformers の実装（attention や kernel など）の違い から来るもの、ということになります。 実験は A6000 上で行っています。 まず、vLLMが出した確率と transformers が出した確率を、tokenごとに散布図にしてみます。 ![image](https://cdn-ak.f.st-hatena.com/images/fotolife/g/go5paopao/20260529/20260529232023.png) 対角線（赤い点線）の上に乗っていれば「2つの backend(バックエンド) が完全に一致している」状態です。見てのとおり、多くの点は対角線の近くに集まっていますが、そこそこの数の点が対角線から外れています。同じ weight(重み) で同じ token を見ているはずなのに、vLLM と transformers で「この token を出す確率」がそれなりに違う、というのがまず確認できます。 具体的にどれくらい違うのか、ズレの大きかった token をいくつか挙げてみます。 token vLLM が出した確率 transformers が出した確率 Sydney 0.59 0.38 the 0.67 0.50 is 0.57 0.44 We 0.38 0.50 上記の例は「オーストラリアの首都はどこか、なぜ Sydney ではないのか（原文: What is the capital of Australia, and why is it not Sydney?）」という質問に対して、 The capital of Australia is **Canberra**. **Why is it not Sydney?** While Sydney is the largest and most well-known city in Australia, it is not the capital. The capital was chosen to be **Canberra** for several reasons: 1. **Balanced Geographical Position**: Sydney .... といった文章を生成しており、"1. **Balanced Geographical Position**: Sydney" の Sydney という token に対して、vLLM は 0.59 の確率を割り当てているのに、transformers では 0.38 になっています。0.2 以上の差があり、この 1 token の出力だけでも計算上影響がでそうに思えます。 ## ズレの大きさの分布を見てみる とはいえ、全部の token がこんなに大きくズレているわけではありません。むしろ逆で、大半の token はほとんど一致しています。 ズレの大きさ（確率の差 |Δp|）で分布を見ると、今回のデータ（約 3900 token）では次のようになりました。 - 確率の差が 0.001 未満（ほぼ完全一致）: 81.1% - 確率の差が 0.01 以上: 10.9% - 確率の差が 0.05 以上: 2.0% - 確率の差が 0.10 以上: 0.3% - 最大のズレ: 0.215 ![image](https://cdn-ak.f.st-hatena.com/images/fotolife/g/go5paopao/20260529/20260529232028.png) ヒストグラム（縦軸は対数目盛にしています）を見ると、ゼロ付近に鋭いピークがあって、そこから両側に細く裾が伸びている形になっています。8割以上のtokenはほぼほぼ一致していて、ズレているのは一部のtokenだけです。ただしその一部は、上記の例のように 0.2 近くズレることもある、という分布です。 また、このズレが系列の中で均一に散らばっているわけではなく、特定の位置に固まって現れるのも興味深いですね。1つの出力について、tokenごとのズレを生成順に並べてみると、こんな感じになりました。 ![image](https://cdn-ak.f.st-hatena.com/images/fotolife/g/go5paopao/20260529/20260529232042.png) ほとんどの位置はゼロ付近に張り付いていて、ところどころ大きく飛び出している箇所があるように見えてます。 ## 今回の範囲で言えそうなこと 今回のデータで見えたのは次のような傾向でした。 - vLLM と transformers のズレは、確かに存在する - ただし全tokenが一律に少しずつズレるわけではなく、大半は一致していて、一部のtokenだけが大きくズレる - 平均的なズレ（log-ratio の平均）はほぼ 0 に近く、今回の範囲では効いてきそうなのはむしろ「たまに大きくズレるtoken」の方だった 冒頭で触れたような「各tokenが少しずつズレて、それが sequence 全体で積み重なっていく」というイメージとは少し違って、今回見た限りでは「ほとんどのtokenは一致するけれど、たまに大きく外れるtokenが混じる」という印象でした。 ただし、これはあくまで今回試した条件（Qwen3-4B、temperature=1.0、bf16、vLLMとtransformersの組み合わせ）での結果です。 モデルやサンプリング設定、backendのバージョンによって傾向は変わりうるので、あくまで一例と思ってもらえるとです。 （ハイパラなどのデフォルト値などで予期せぬ差分が発生してる可能性なども捨てきれません） ## 分布のmismatch 実際にはrolloutとtrainのbackendの差分以外でもこういった分布差は問題として存在しています。 主なところでいうと、今回のもの含めて以下の3つがあります。 - 推論ライブラリ（今回） - 非同期RL の stale policy - classical off-policy RL 2つ目の非同期RLというのは、RLの推論に時間がかかりすぎる問題への対策として最近ホットになってきているものです。 通常の同期RLでは、Rollout→Policy更新→ Rolloutという順で更新後の最新のPolicyでRolloutを行います。 これだとPolicy更新中Rolloutを止めることになりGPU効率が落ちるため、非同期RLでは推論効率重視でRolloutするPolicyが最新である保証を捨ててRolloutを行います。（今回の記事の本題ではないためざっくりとした説明です） この場合も、実際にRolloutを生成したPolicyと学習側が仮定しているPolicyの間に違いが生じます。 3 つ目の古典的オフポリシー RL とは、従来のオフポリシーを前提とした RL 全般のことを指します。 LLM の文脈では、DPO やその派生手法のように、別モデルや過去モデルから収集したデータを用いて学習するアプローチがこれに該当します。オフポリシー RL は古くから研究されており、補正手法が提案されてきました。こうした古典的なオフポリシーの枠組みも切り離せないものでしょう。 ## mismatch への対策 この補正のために使われる代表的な手法が Importance Sampling (IS) です。 IS は、ある分布から得られたサンプルを、別の分布から得られたサンプルであるかのように扱うための補正です。今回の文脈では、Rollout backend で実際にサンプリングに使った Policy を π とし、学習 backend 側で old policy として扱う Policy を μ とすると、π から得られた token 列を μ から得られたものとして補正したい、ということになります。 sequence 全体で見ると、Rollout で得られる token 列は π によって生成されています。そのため、理論上の重みは sequence 全体の生成確率の比になります。 自己回帰モデルでは sequence の生成確率は token ごとの条件付き確率の積で書けるため、これは次のように分解できます。 ここで、各 token に対応する確率比 が、sequence 全体の重みを構成する因子になっています。対数で見ると、sequence 全体の重みは token 単位の log-ratio の和になります。 ここで問題になるのが、この sequence 全体の重みをそのまま使うと非常に不安定になりやすいという点です。 先ほどの実験では、token 単位の log-ratio の平均はほぼ 0 に近い値でした。つまり今回の結果に限れば、「各 token が少しずつ同じ方向にズレて、それがのように膨らむ」というより、ほとんどの token は一致しているが、一部の token だけ大きくズレるという形でした。 このような分布では、sequence 全体で log-ratio を足し合わせたとき、平均よりも分散や外れ値の影響が問題になります。平均が 0 付近であっても、足し合わせる token 数が増えるほど sequence ごとのばらつきは大きくなります。また、log-ratio が大きい token が数個混ざるだけで、は数倍、あるいはその逆方向に大きく振れる可能性があります。 その結果、sequence 全体の重みはサンプルごとに大きく上下する高分散な量になり、loss にそのまま掛けると学習が不安定になりやすくなります。 そのため実際には、素朴な sequence-level IS をそのまま使うのではなく、 - token-level IS - truncated IS - 幾何平均 - rejection sampling - IcePOP などの安定化手法が使われます。 ## token-level IS 最もシンプルかつ広く使われているアプローチです。 sequence 全体でまとめて一つの重みを掛けるのではなく、token 単位で重みを計算して loss に掛けます。 ここで は学習 backend 側で計算し直した logprob、は Rollout backend が実際にサンプリングに使った logprob です。 実際の実装では、ロールアウト（rollout）で得られたサンプルに対して、学習用バックエンド側でフォワードパスを一度実行して の対数尤度を取得し、ロールアウト用バックエンドが返した対数尤度 との差分から重みを作成します。 これを PPO/GRPO の損失項に掛ける形で、] のように補正します。 シークエンス全体の積を取らないため、先ほどのような指数的な発散は起こりません。verl や NeMo-RL などの主要な強化学習（RL）ライブラリでは、この形のトークンレベルの補正が標準的に実装されています。 ## トランクテッド IS トークンレベルの重要性サンプリング（IS）であっても、特定のトークンで の確率が極端に小さい場合、重み が巨大になり学習が不安定になることがあります。 これを防ぐために、重みに上限を設けるのがトランクテッド IS です。 ここで は事前に決めた閾値です（例えば など）。 トランケーションを入れると厳密にはバイアスを含んだ推定量になりますが、分散が大幅に下がるため実用上は安定して学習が進むケースが多いです。 バックエンド間のズレは通常そこまで大きくないため、ほとんどのトークンではトランケーションは発動せず、外れ値的に大きく乖離したトークンだけを抑え込むイメージになります。 なお、これは PPO/GRPO で元々ある へのクリップ（ポリシードリフト補正側のクリップ）とは別物で、あくまでバックエンドミスマッチ補正用の重み に対する処理です。実装上は両方が同時に効いている形になります。 verl の Rollout Correction や、TIS (Truncated Importance Sampling) として論文化されているのもこのアプローチです。 ## 幾何平均 トークンレベルの重要性サンプリング（IS）は安定性が高い一方で、トークン単位で独立に重みを掛けるため「シークエンスとしての一貫した重み付け」にはなっていません。 本来、シークエンスとして生成されたサンプルに対する重みは ですが、これをそのまま使うとシークエンス長 に対して指数的に発散します。 そこで、積の代わりに幾何平均を取って一つのシークエンスレベルの重みにするアプローチがあります。 対数の世界で見ると、 つまり、トークン単位の対数比の平均を取っていることになります。 こうすることで、シークエンス長 による指数的なスケーリングを抑えつつ、シークエンス全体としての と のズレを一つの重みに集約できます。 なお、類似の幾何平均をシークエンスレベルの比率に使うアプローチとして GSPO (Group Sequence Policy Optimization) や GMPO (Geometric-Mean Policy Optimization) が知られていますが、これらはここで議論しているバックエンドミスマッチの補正ではなく、学習中の と old policy の比をシークエンスレベルで取るための安定化を目的としたものです。問題設定は別ですが、「トークン単位の対数比を平均する」という手法のアイデアは共通しています。 これは重みで補正するのではなく、「 と のズレが大きすぎるサンプルは学習から外す」という割り切ったアプローチです。 具体的には、sample（あるいはtoken）ごとに を計算し、これが事前に決めた範囲 ] から外れる場合には、そのサンプルをloss計算から除外します。 メリットとしては、 - 実装が単純 - 巨大な重みによる勾配の暴走が原理的に起きない - 残ったサンプルだけで見ると、Rollout backendと学習backendがほぼ一致している状況で学習できる といった点があります。 一方で、 - 一定割合のサンプルを捨てることになるためデータ効率が落ちる - 特定のtoken（例えば低確率token、長いsequenceの後半など）でズレが系統的に大きくなる場合、それらが選択的に捨てられることでデータ分布に偏りが生じる可能性がある というデメリットもあります。 verlの Rollout Correction でも、Importance Sampling weight と並んで Rejection Sampling が用意されており、両者を組み合わせて使うことができます。 ## 実ライブラリでの実装 実際にRLを行えるライブラリで、このあたりどう対応されているかを見ていきます。 ## NeMo-RL NeMo-RLでは、loss関数（ClippedPGLossConfig）の設定で対応がされています。（ブログ執筆時点） [github.com](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/loss/loss_functions.py) # --- 重要度サンプリング補正 --- # 非同期 GRPO では、重要度サンプリング補正を有効にする必要があります # async_grpo.enabled が true の場合に true に設定します use_importance_sampling_correction: bool = False # --- 截断重要度サンプリング --- # 截断重要度サンプリングのタイプ: # "tis" – IS 重みを最大値にクランプする # "icepop" – IS 重みが [min, max] の範囲外にあるトークンをゼロアウトする # "seq-mask-tis" – 幾何平均の IS 比率に基づいてシーケンス全体をゼロアウトし、截断されていないトークンの IS 補正を行う truncated_importance_sampling_type: Optional[str] = None truncated_importance_sampling_ratio: Optional[float] = None # ICE-POP / seq-mask-tis のフィルタリングにおける下限値 truncated_importance_sampling_ratio_min: Optional[float] = None use_importance_sampling_correction というフラグが用意されています。 また、truncated_importance_sampling_type や truncated_importance_sampling_ratio といったパラメータも存在しており、 truncated_importance_sampling_type には、tis（截断重要度サンプリング）、icepop、seq-mask-tis（幾何平均によるシーケンス単位の IS）が用意されているようです。 ## verl verl では、この一連の補正が「Rollout Correction」という機能としてまとめられており、重要度サンプリング (IS) とリジェクト・サンプリング (RS) を組み合わせて使用できるようになっています。（ブログ執筆時点） [verl.readthedocs.io](https://verl.readthedocs.io/en/latest/algo/rollout_corr.html) 主な設定パラメータは以下の通りです。 - rollout_is : IS 補正のレベル（null / token / sequence） - rollout_is_threshold : IS 重みの閾値（デフォルト 2.0。"0.5_5.0" のように指定すると IcePop） - rollout_rs : リジェクトサンプリングのモード（token_k1 / seq_sum_k1 / seq_mean_k1 など） - rollout_rs_threshold : リジェクトサンプリングの閾値 - bypass_mode : πold の計算を省略し、πold = π_rollout とする rollout_is でトークンレベルかシーケンスレベルかを選び、rollout_rs を併用するかどうかで、truncated IS（逐次重要性推定）だけにするかリジェクトサンプリングも効かせるかを切り替えます。よく使う組み合わせは、decoupled_token_is（トークンレベルの IS）や decoupled_seq_is（対数比の平均で集約するシーケンスレベルの IS）などの preset として用意されており、ここまで紹介した「トークンレベルの IS」「truncated IS」「幾何平均的なシーケンスレベル補正」「リジェクトサンプリング」が設定の組み合わせで一通り表現できます。 ## まとめ というわけで今回は rollout（ロールアウト）でのズレによる問題をとりあげました。 バックエンドの違いによる差分という直近の LLM 学習の課題のところから、On-policy/Off-policy という従来の課題まで繋がっている感じがして個人的には調べながら面白かったです。 これ自体がどこまで学習に影響を与えるかまでは調べられていませんが、今後気にしていきたいと思っています。 ## We Are Hiring! ABEJA は、テクノロジーの社会実装に取り組んでいます。技術はもちろん、技術をどのようにして社会やビジネスに組み込んでいくかを考えるのが好きな方は、下記採用ページからエントリーください！（新卒の方やインターンシップのエントリーもお待ちしております！） [careers.abejainc.com](http://careers.abejainc.com/) 特に下記ポジションの募集を強化しています！ぜひ御覧ください！ [トランスフォーメーション領域：データサイエンティスト](https://hrmos.co/pages/abeja/jobs/dev-ds41-sol11) [トランスフォーメーション領域：データサイエンティスト（ミドル）](https://hrmos.co/pages/abeja/jobs/dev-ds31-41-sol11) [トランスフォーメーション領域：データサイエンティスト（シニア）](https://hrmos.co/pages/abeja/jobs/dev-dse-41-sol11)

On-policy のはずが Off-policy になる：LLM 強化学習の rollout mismatchと対策(rollout correction)

背景や根拠まで確認しますか？

調べる

選ぶ

サイト