Reusing Rollouts under Policy Lag: Prefix-Normalized Policy Optimization for LLM Reinforcement Learning

TL;DR

PNPO replaces cumulative importance ratios with prefix geometric mean, improving off-policy correction in large language model RL.

cs.AI 🔴 Advanced 2026-08-03 47 views
Wenhao Zhang Yibo Xie Rui Wang Jiahua Yang Lei Jiang Zibo Yang Yawei Wang Jiali Xu jasperawang Haoyang Long Huan Xiong alantzhao
Reinforcement Learning Large Language Models Off-policy Correction Policy Optimization Replay Buffer

Key Findings

Methodology

This paper introduces Prefix-Normalized Policy Optimization (PNPO), which replaces the cumulative importance ratio with the geometric mean along causal prefixes. This preserves causal dependence while compressing the log-scale of importance weights. Experiments on long-context mathematical reasoning tasks compare 1 and 4 PPO epochs, demonstrating that PNPO outperforms GSPO especially under multiple update cycles. The approach combines position-dependent gating and response-level averaging, effectively mitigating the scale explosion of importance weights and improving off-policy robustness.

Key Results

  • In long-context reasoning benchmarks, four-epoch PNPO achieves an average accuracy (Avg@32) of 50.24%, outperforming GSPO by 3 percentage points. Under a fixed update budget, it reaches 49.66% after 150 batches, comparable to 600 batches with one epoch, indicating higher sample efficiency. The results show that PNPO maintains performance gains as the policy drifts further off-policy, especially in multi-epoch settings.
  • Compared to existing methods like GSPO and GRPO, PNPO demonstrates superior stability and efficiency, reducing training time while maintaining or improving accuracy. The experiments validate its effectiveness across multiple reasoning benchmarks, highlighting its potential for scalable RL training of large models.

Significance

This work addresses a fundamental challenge in large language model reinforcement learning: how to effectively correct for policy lag during multiple updates. By introducing prefix normalization, PNPO offers a scalable solution that controls importance weight scales, enabling more stable and efficient training under off-policy conditions. This innovation has significant implications for deploying RL in real-world NLP applications, where data reuse and multiple updates are common. It advances the theoretical understanding of importance sampling in autoregressive models and provides practical tools for improving training stability and sample efficiency in large-scale RL systems.

Technical Contribution

PNPO's core contribution is the introduction of prefix normalization, replacing the cumulative importance ratio with a position-dependent geometric mean. This approach controls the exponential growth of importance weights, ensuring numerical stability without sacrificing causal dependence. The method integrates position-dependent gating and response-level aggregation, forming a biased but effective surrogate for off-policy correction. Theoretical analysis confirms its ability to handle policy drift, and empirical results demonstrate its superiority in long-context reasoning tasks. The framework bridges token-level and sequence-level importance weighting, opening new avenues for scalable RL optimization.

Novelty

This is the first work to propose prefix geometric mean as a scale-controlled importance weight in autoregressive RL. Unlike prior methods that rely on full response or token-wise ratios, PNPO balances causal dependence with scale compression, addressing the limitations of cumulative importance ratios. Its innovative normalization mechanism provides a new perspective on off-policy correction, especially under multiple update cycles, setting a new benchmark for scalable policy optimization in large language models.

Limitations

  • The experiments are limited to a 1.5B parameter model and specific mathematical reasoning benchmarks; generalization to larger models and diverse tasks remains to be validated. The effectiveness of the acceptance gate and response aggregation parameters needs further exploration across different settings. Additionally, the approach primarily addresses bias introduced by importance weight scaling, but other sources of off-policy mismatch, such as asynchronous data collection and offline RL, are not directly tackled. Future work should extend the framework to these broader scenarios.

Future Work

Future research will explore PNPO's applicability in asynchronous, offline, and multi-task RL settings, integrating with experience replay and off-policy data buffers. Investigating adaptive gating mechanisms and dynamic normalization strategies could further enhance robustness. Extending the approach to larger models and diverse NLP tasks, including dialogue and summarization, will be critical for practical deployment. Additionally, theoretical analysis of bias-variance trade-offs and convergence guarantees under different off-policy regimes will deepen understanding and guide further improvements.

AI Executive Summary

Reinforcement learning for large language models faces a significant challenge: how to efficiently reuse data while correcting for policy lag. Traditional importance sampling methods, such as cumulative importance ratios, tend to explode in scale as the prefix length increases, especially under multiple policy updates. This problem hampers training stability and sample efficiency, limiting the scalability of RL in NLP applications.

This paper introduces Prefix-Normalized Policy Optimization (PNPO), a novel approach that replaces the cumulative importance ratio with a prefix-based geometric mean. This normalization preserves the causal dependence of token generation while controlling the scale of importance weights. By doing so, PNPO effectively mitigates the scale explosion problem, enabling more stable and efficient off-policy training.

The authors validate PNPO through controlled experiments on long-context mathematical reasoning tasks, comparing 1 and 4 PPO epochs per batch. Results show that, under multiple update cycles, PNPO consistently outperforms existing methods like GSPO and GRPO, achieving higher accuracy with fewer responses and less training time. Specifically, in a fixed update budget, PNPO reaches near-peak performance with only a quarter of the responses needed by traditional methods, demonstrating superior sample efficiency.

The core innovation lies in the prefix normalization mechanism, which balances bias and scale control, maintaining causal dependence while compressing importance weight ranges. This approach addresses a longstanding bottleneck in autoregressive RL, opening new avenues for scalable, stable training of large models. Its implications extend to real-world NLP applications such as dialogue systems, automated reasoning, and multi-task learning, where data reuse and multiple updates are essential.

While promising, the study is limited to a 1.5B parameter model and specific reasoning benchmarks. Future work will focus on extending PNPO to larger models, diverse tasks, and more complex RL scenarios, including asynchronous and offline settings, to fully realize its potential in industrial-scale NLP systems.

Deep Analysis

Background

近年来,随着大规模预训练模型的广泛应用,强化学习(RL)在提升模型生成质量和任务适应性方面逐渐成为研究热点。早期代表性工作如PPO、TRPO等通过策略梯度方法实现了较好的性能,但在大模型训练中,回滚生成(rollout)成本极高,限制了训练效率。为了缓解这一问题,研究者提出多轮数据重用(多次epoch),但带来了偏离行为策略的逐步积累,导致梯度偏差和训练不稳定。近年来,偏离校正技术如GSPO、GRPO等被提出,试图在保持偏离校正的同时控制数值范围,但累积比率的指数爆炸仍是难题。本文在此背景下,提出前缀归一化机制,旨在通过压缩偏离校正的尺度,提升多轮偏离环境下的训练稳定性和效率。

Core Problem

核心问题在于多轮策略更新引起的偏离行为策略,导致梯度偏差累积,训练不稳定。传统偏离校正方法如GSPO在偏离较深时效果减弱,累积比率的数值范围难以控制,限制了偏离缓冲的效果。如何在保持因果前缀依赖的基础上,设计一种压缩偏离尺度的机制,成为提升大模型RL训练效率的关键。解决这一问题对于实现更稳定、更高效的偏离校正具有重要意义,尤其在长上下文推理任务中。

Innovation

PNPO的创新在于引入前缀归一化机制,将累积比率的对数值通过几何平均进行压缩,既保持因果前缀依赖,又避免数值爆炸。具体做法是用位置依赖门控和响应级平均,替代传统的全响应或逐步比率,形成偏离校正的中间机制。该方法在保持偏离校正的因果关系的同时,有效控制偏离尺度,特别适用于多轮偏离环境。实验验证显示,PNPO在长上下文推理任务中表现优越,显著优于GSPO和GRPO,彰显其理论创新和实用价值。

Methodology

  • �� 以标准自回归生成模型为基础,定义状态为前缀,动作为当前Token。• 引入局部比率ρt,校正当前Token的策略偏差。• 计算累积比率Ct,反映整个前缀的偏离程度。• 提出前缀归一化权重wPNt,将log Ct除以前缀长度t的倒数,压缩尺度。• 设计位置依赖门控,根据wPNt动态调整采样接受范围。• 结合偏离优势和响应级平均,优化偏离缓冲效果。• 在长上下文数学推理任务中,采用多轮策略更新,验证PNPO在偏离环境中的表现。

Experiments

采用DeepSeek-R1-Distill-Qwen-1.5B模型,在DAPO-Math-17k数据集上训练,评估AMC 2023、AIME 2024和AIME 2025。实验包括每轮采样256提示,每提示8响应,总2048响应,最大长度1024和15360。训练采用学习率10^-6,批次64,优化4次。比较方法包括GSPO、GRPO和PNPO,重点验证偏离环境下的性能差异。通过不同策略更新轮次(1轮和4轮)验证PNPO的效果。

Results

在长上下文推理任务中,四轮策略更新下,PNPO在三个基准中的平均准确率达50.24%,比GSPO高3个百分点。有限预算下,150轮后,PNPO达49.66%的宏平均,接近600轮单轮的49.56%。实验显示,PNPO在偏离更严重时,能持续保持性能优势,验证其高效的样本利用能力。

Plain Language Accessible to non-experts

想象你在做一道非常复杂的数学题,每次尝试都用不同的方法。有时候,你用的方法会逐渐偏离正确答案,就像走错了路。传统方法就像每次都用同一张地图,偏离会越来越远。而PNPO像给你一副特殊的放大镜,可以自动调整偏差,把偏离范围变得更小。这样,你每次尝试都能更接近正确答案,不会因为反复试错而偏离太远。它用一种聪明的数学技巧,把偏差压缩到合理范围,让你在不断尝试中变得更稳、更准,就像用放大镜看东西,能看得更清楚,走得更稳,最终找到最正确的路。

ELI14 Explained like you're 14

想象你在玩一个超级难的游戏,每次都要做很多决定。有时候,你用的策略会偏离最好的路线,因为你不断调整。以前的方法就像每次都用同一张地图,结果偏离太远了。PNPO就像给你一副神奇的眼镜,每次看地图时都能自动调整偏差,把偏离变得更小。这样,你每次走的路都更接近正确的方向,不会因为反复试错而偏离太远。它用一种聪明的数学方法,把偏差压缩到合理范围,让你在不断尝试中变得更厉害、更准,最终找到最好的路线。

Glossary

Importance Sampling(重要性采样)

一种调整样本权重以校正偏离行为策略的技术,确保梯度估计的无偏性。技术上通过比率校正策略偏差。

用于校正偏离行为策略,确保训练稳定性。

Cumulative Importance Ratio(累积重要性比率)

沿前缀连续乘积的偏离校正因子,反映从行为策略到目标策略的整体偏离。

衡量偏离程度的核心指标,但数值范围难控。

Prefix-Normalized Policy(前缀归一化策略)

将累积比率的对数平均化,压缩偏离尺度,保持因果前缀依赖。

本文提出的机制,用于缓解偏离比率爆炸问题。

Geometric Mean(几何平均)

一组数的乘积的n次方根,用于压缩偏离比率的尺度。

PNPO中用以替代累积比率,控制偏离尺度。

Open Questions Unanswered questions from this research

  • 1 在更大模型(如百亿参数)中,偏离校正的稳定性和效果仍需验证。
  • 2 多任务、多模态环境下PNPO的适应性和效果尚未充分探索。
  • 3 异步采样和离线数据引入偏离机制与PNPO结合的研究仍待深入。

Applications

Immediate Applications

自动推理系统

利用PNPO提升长上下文推理中的准确性,减少偏离引起的不稳定。

Long-term Vision

高效大模型训练

结合偏离缓冲、多轮优化,推动大模型在多任务、多场景中的快速适应和部署。

Abstract

Autoregressive rollout generation is a major computational cost in reinforcement learning for large language models. Reusing each rollout batch for additional learner updates amortizes this cost, but later updates become increasingly off-policy as the learner departs from the behavior policy. At a token position, exact off-policy correction must account for both the current action and the probability of reaching its prefix. The cumulative importance ratio provides this correction, but its product form can produce an unwieldy dynamic range. We study Prefix-Normalized Policy Optimization (PNPO), which replaces the cumulative ratio with the geometric mean of likelihood ratios along each causal prefix, preserving causal-prefix dependence at each position while compressing the log-weight scale. In controlled long-context mathematical reasoning experiments, we induce two off-policy regimes by using one or four policy-update epochs per rollout batch. PNPO does not consistently outperform GSPO with one epoch. With four epochs, it attains the highest observed Avg@32 on each benchmark; the unweighted mean of the three independently selected benchmark peaks is 50.24, 3.00 percentage points above GSPO. 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. These results provide preliminary evidence that PNPO can be advantageous as training moves further off-policy.

cs.AI cs.LG