Command Palette
Search for a command to run...
ポリシーラグ下でのロールアウト再利用:LLM強化学習のためのプレフィックス正規化方策最適化
ポリシーラグ下でのロールアウト再利用:LLM強化学習のためのプレフィックス正規化方策最適化
概要
自己回帰的なロールアウト生成は、大規模言語モデルの強化学習における主要な計算コストである。各ロールアウトバッチを学習器の追加更新に再利用することでこのコストを償却できるが、学習器が行動方策から離れるにつれて、後の更新はますますオフポリシーとなる。あるトークン位置において、正確なオフポリシー補正には、現在の行動とそのプレフィックスに到達する確率の両方を考慮しなければならない。累積重点比はこの補正を提供するが、その積形式は扱いにくいダイナミックレンジを生じさせる可能性がある。我々は、累積比を各因果的プレフィックスに沿った尤度比の幾何平均で置き換え、各位置での因果的プレフィックス依存性を保持しつつ対数重みのスケールを圧縮する、プレフィックス正規化方策最適化(PNPO)を研究する。制御された長文脈の数学的推論実験において、ロールアウトバッチあたり1回または4回の方策更新エポックを用いることで、2つのオフポリシー状況を誘発する。PNPOは、1エポックではGSPOを一貫して上回ることはない。4エポックでは、各ベンチマークで観測された最高のAvg@32を達成し、独立に選択された3つのベンチマークピークの重み付けなし平均は50.24であり、GSPOを3.00パーセントポイント上回る。一致させた2400更新の予算の下で、4エポックのPNPOは150ロールアウトバッチ後に最終的なマクロAvg@32が49.66に達し、これは1エポックで600バッチ後に達成された49.56に匹敵する。これらの結果は、学習がよりオフポリシーに進むにつれてPNPOが有利になりうるという予備的な証拠を提供する。
One-sentence Summary
Tencent et al. propose Prefix-Normalized Policy Optimization (PNPO), which replaces the cumulative importance ratio with the geometric mean of likelihood ratios along each causal prefix to compress log-weight scale while preserving causal-prefix dependence, and in long-context mathematical reasoning experiments, PNPO attains a 3 percentage point higher Avg@32 than GSPO under four-epoch updates, demonstrating efficient off-policy LLM reinforcement learning.
Key Contributions
- Prefix-Normalized Policy Optimization (PNPO) replaces the exact cumulative importance ratio with the geometric mean of token likelihood ratios over each causal prefix, preserving prefix-conditioned dependence while compressing the log-weight dynamic range.
- In controlled mathematical reasoning experiments, PNPO attains a 3.00 percentage point higher unweighted mean peak Avg@32 than GSPO in a four-epoch off-policy regime, with the advantage becoming more pronounced as learner–behavior mismatch increases.
- Under a matched 2,400-update budget, four-epoch PNPO reaches a final macro Avg@32 of 49.66 after 150 rollout batches, comparable to the 49.56 reached after 600 batches with one epoch, providing preliminary evidence of more effective rollout reuse under greater off-policy mismatch.
Introduction
Policy-gradient methods are a standard approach for post-training language models, but collecting rollout batches through autoregressive generation is expensive. To amortize this cost, methods like PPO reuse collected rollouts for multiple learner updates, which introduces off-policy drift as the learner changes while the behavior data remains fixed. Proximal policy methods address this mismatch with local-ratio surrogates, but these omit a full state–action correction and weaken as reuse increases. The authors propose Prefix-Normalized Policy Optimization (PNPO), which uses the geometric mean of likelihood ratios along the causal prefix to correct for prefix visitation while controlling the scale of cumulative log-ratios. Their main contribution is formulating this prefix-normalized weight and evaluating it in controlled off-policy regimes, where PNPO shows a pronounced advantage over sequence-shared weighting when repeated updates induce greater learner–behavior mismatch.
Method
The authors propose Prefix-Normalized Policy Optimization (PNPO), a method for reinforcement learning fine-tuning of large language models that addresses the state-distribution shift inherent in standard proximal policy optimization (PPO) and its variants such as GRPO. These methods typically rely on a local surrogate that freezes the state occupancy at the behavior policy, using only the token-level action ratio ρt=πθ(at∣st)/πβ(at∣st) and a clipped advantage. While the resulting policy gradient is unbiased at the behavior policy, it omits the prefix-state correction required for exact off-policy estimation when the learner drifts away from the behavior policy.
In the autoregressive setting, the prefix-state ratio factorizes as a product of past action ratios, leading to the exact joint state–action change-of-measure ratio
Ct=k=1∏tρk.This cumulative ratio, when paired with the appropriate advantage, can yield an unbiased gradient estimate. However, using Ct directly introduces a scale problem: its log-variance grows with prefix length and policy drift, making gradient estimates unstable. Furthermore, retaining the behavior-advantage Atβ while using the current-token score zt=∇θlogπθ(at∣st) would be biased; to pair Atβ with an unbiased score, one must use the cumulative prefix score ∑k=1tzk, not the current-token score alone. The authors therefore develop a practical surrogate that interpolates between exact correction and scale control.
The core of PNPO is the prefix-normalized policy weight, which compresses the cumulative log-ratio by taking the t-th root of Ct:
wi,tPN=Ci,t1/t=exp(t1k=1∑tlogρi,k).This transform preserves the sign of logCt and the order across responses, while reducing the dynamic range. At t=1, it recovers the local action ratio; at t=Li, it equals the geometric mean of the full-response ratio. For intermediate positions it depends only on the prefix, excluding future-suffix likelihood shifts, and thus is not a true density ratio but a monotone approximation.
To further control the influence of tokens with extreme weight values, PNPO employs an acceptance gate that hard-rejects score terms at positions where the weight falls outside a position-dependent tolerance interval. The gate uses a heuristic scaling factor h(t,Li)=Li/t, which widens the bounds at early positions and shrinks them to the base interval [1−ϵ−,1+ϵ+] at the final token. The mask Mi,t is a binary indicator; if the weight is outside the bounds, the corresponding token’s score term is simply omitted from the objective, while later positions remain eligible.
The complete PNPO objective is a detached score-function surrogate that treats the gate, prefix-normalized weight, and advantage as constant coefficients via stop-gradient. For a batch of prompts, the authors sample a group of G responses from the behavior policy, compute a group-relative advantage Ai,t (using the same normalization as GRPO), and then maximize
JPNPO(θ)=E[G1i=1∑GLi1t=1∑Lisg[Mi,twi,tPNAi,t]logπθ(yi,t∣x,yi,<t)].The per-response averaging first normalizes over valid tokens (rejected tokens contribute zero), then averages the per-response means across responses. The objective thus retains the current-token-score structure of PPO/GRPO while incorporating a prefix-aware weight that approximates the state-distribution correction.
During training, PNPO reuses a fixed rollout batch for multiple optimizer epochs. The behavior-policy log probabilities, group-derived advantages, and response lengths are held constant. In each minibatch update, the learner evaluates the current numerator log probabilities, recomputes the cumulative ratios, the prefix-normalized weight, and the acceptance mask, and applies the objective. This allows the weight to adapt as the policy evolves, while the advantage and behavior-policy denominator remain fixed. The overall design trades exact off-policy correction for stable scale control, aiming to mitigate the state-distribution drift that can accumulate when a surrogate is repeatedly optimized over the same stale trajectories.
Experiment
The experiments evaluate PNPO on long-context mathematical reasoning using DeepSeek-R1-Distill-Qwen-1.5B fine-tuned on DAPO-Math-17k, comparing against GSPO and GRPO on AMC 2023, AIME 2024, and AIME 2025. Under four PPO epochs per rollout batch, PNPO achieves the best performance on all benchmarks, with a macro average 3.00 points above GSPO, and the advantage persists across the entire training trajectory rather than only at peak scores. Four-epoch PNPO matches the final performance of the one-epoch setting while using only a quarter of the newly generated responses, demonstrating more effective rollout reuse and reaching a given reward threshold substantially earlier than GSPO. GRPO fails to maintain early gains over long horizons, and while GSPO remains stable, a gap in evaluation metrics remains, consistent with the benefit of PNPO's token-level weighting granularity.
Under four PPO epochs, PNPO yields the highest average evaluation scores across all benchmarks, with a 3 percentage point advantage over GSPO. This lead persists across the evaluation trajectory, and PNPO achieves comparable final performance to its one-epoch counterpart using only a quarter of the newly generated responses, indicating more efficient rollout reuse. PNPO also reaches a fixed reward threshold hours earlier than GSPO in the four-epoch setting. With four PPO epochs, PNPO attains the best observed Avg@32 on all three benchmarks, averaging 50.24 compared to GSPO's 47.24. In the one-epoch setting, PNPO and GSPO remain close and exchange the lead, while in the four-epoch setting PNPO leads at 14 of 15 evaluations and finishes 2.66 percentage points ahead. At the same total optimizer updates, four-epoch PNPO reaches a final macro Avg@32 of 49.66 after 150 rollout batches, comparable to the 49.56 reached after 600 one-epoch batches, showing more effective rollout reuse. Four-epoch PNPO reaches a centered reward threshold of 0.25 in 16.0 hours, 6.4 hours earlier than GSPO and much earlier than the one-epoch runs. GRPO's four-epoch evaluation declines after peaking, and GSPO's stable reward does not close the evaluation gap, consistent with PNPO's position-dependent prefix statistic providing an advantage.
In experiments comparing PNPO, GSPO, and GRPO across benchmarks under one and four PPO epochs, PNPO consistently achieves the highest evaluation scores, particularly in the four-epoch setting where it maintains a clear lead and reaches reward thresholds faster. PNPO demonstrates more efficient reuse of rollouts, matching the final performance of one-epoch training with far fewer new responses, while GRPO's performance degrades after peaking and GSPO fails to close the gap. The advantage is attributed to PNPO's position-dependent prefix statistic.