HyperAIHyperAI

Command Palette

Search for a command to run...

정책 지연 하에서의 롤아웃 재사용: 대규모 언어 모델 강화 학습을 위한 접두사 정규화 정책 최적화

초록

자기회귀적 롤아웃 생성은 대규모 언어 모델의 강화 학습에서 주요 계산 비용을 차지한다. 각 롤아웃 배치를 추가적인 학습자 업데이트에 재사용하면 이 비용을 분산 상각할 수 있지만, 학습자가 행동 정책에서 멀어짐에 따라 이후 업데이트는 점점 더 오프폴리시(off-policy) 상태가 된다. 특정 토큰 위치에서 정확한 오프폴리시 보정은 현재 행동과 해당 접두사에 도달할 확률을 모두 고려해야 한다. 누적 중요도 비율이 이러한 보정을 제공하지만, 그 곱셈 형태는 다루기 어려운 동적 범위를 만들어낼 수 있다. 본 논문은 누적 비율을 각 인과적 접두사를 따라 계산된 우도비의 기하 평균으로 대체하는 접두사 정규화 정책 최적화(PNPO)를 연구한다. 이 방법은 각 위치에서 인과적 접두사 의존성을 보존하면서 로그 가중치 척도를 압축한다. 통제된 장문 맥락 수학 추론 실험에서 롤아웃 배치당 하나 또는 네 번의 정책 업데이트 에폭을 사용하여 두 가지 오프폴리시 환경을 유도했다. PNPO는 한 에폭에서는 GSPO를 지속적으로 능가하지 못했다. 네 에폭에서는 각 벤치마크에서 가장 높은 Avg@32를 달성했으며, 독립적으로 선택된 세 벤치마크 최고치의 비가중 평균은 50.24로 GSPO보다 3.00퍼센트 포인트 높았다. 2,400회 업데이트 예산 하에서 네 에폭 PNPO는 150개의 롤아웃 배치 후 최종 매크로 Avg@32 49.66에 도달했으며, 이는 한 에폭으로 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=πθ(atst)/πβ(atst)\rho_t = \pi_\theta(a_t|s_t)/\pi_\beta(a_t|s_t)ρt=πθ(atst)/πβ(atst) 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=1tρk.C_t = \prod_{k=1}^{t} \rho_k.Ct=k=1tρk.

This cumulative ratio, when paired with the appropriate advantage, can yield an unbiased gradient estimate. However, using CtC_tCt 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βA_t^\betaAtβ while using the current-token score zt=θlogπθ(atst)z_t = \nabla_\theta \log \pi_\theta(a_t|s_t)zt=θlogπθ(atst) would be biased; to pair AtβA_t^\betaAtβ with an unbiased score, one must use the cumulative prefix score k=1tzk\sum_{k=1}^t z_kk=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 CtC_tCt:

wi,tPN=Ci,t1/t=exp ⁣(1tk=1tlogρi,k).w_{i,t}^{\mathrm{PN}} = C_{i,t}^{1/t} = \exp\!\left(\frac{1}{t}\sum_{k=1}^{t} \log \rho_{i,k}\right).wi,tPN=Ci,t1/t=exp(t1k=1tlogρi,k).

This transform preserves the sign of logCt\log C_tlogCt and the order across responses, while reducing the dynamic range. At t=1t=1t=1, it recovers the local action ratio; at t=Lit=L_it=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/th(t,L_i) = \sqrt{L_i/t}h(t,Li)=Li/t, which widens the bounds at early positions and shrinks them to the base interval [1ϵ,1+ϵ+][1-\epsilon_-, 1+\epsilon_+][1ϵ,1+ϵ+] at the final token. The mask Mi,tM_{i,t}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 GGG responses from the behavior policy, compute a group-relative advantage A^i,t\widehat{A}_{i,t}Ai,t (using the same normalization as GRPO), and then maximize

JPNPO(θ)=E ⁣[1Gi=1G1Lit=1Lisg ⁣[Mi,twi,tPNA^i,t]logπθ(yi,tx,yi,<t)].\mathcal{J}_{\mathrm{PNPO}}(\theta) = \mathbb{E}\!\left[\frac{1}{G}\sum_{i=1}^{G}\frac{1}{L_i}\sum_{t=1}^{L_i} \operatorname{sg}\!\left[M_{i,t} w_{i,t}^{\mathrm{PN}} \widehat{A}_{i,t}\right] \log \pi_\theta(y_{i,t}\mid x,y_{i,<t})\right].JPNPO(θ)=E[G1i=1GLi1t=1Lisg[Mi,twi,tPNAi,t]logπθ(yi,tx,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.


AI로 AI 구축

아이디어에서 출시까지 — 무료 AI 코코딩, 즉시 사용 가능한 환경, 최적의 GPU 가격으로 AI 개발을 가속화하세요.

AI 협업 코딩
바로 사용 가능한 GPU
최적의 가격

HyperAI Newsletters

최신 정보 구독하기
한국 시간 매주 월요일 오전 9시 에 이번 주의 최신 업데이트를 메일로 발송합니다
이메일 서비스 제공: MailChimp