Beyond Uniform Token-Level Trust Region in LLM Reinforcement Learning

2026-06-09

AI Summary

This work addresses the limitations of existing Proximal Policy Optimization (PPO) methods in reinforcement learning for large language models, which employ position-agnostic trust region thresholds and thereby neglect the compounding effect of early-token deviations and prefix drift inherent in autoregressive generation. To overcome this, the authors propose Context-aware PPO (CPPO), the first approach to incorporate positional sensitivity and prefix drift into token-level trust region design. CPPO introduces two key mechanisms—position-weighted thresholds and a cumulative prefix deviation budget—to dynamically adjust policy update constraints in alignment with finite-horizon policy improvement bounds. This enables stringent control over early-token updates while permitting flexible exploration in later positions. Experimental results demonstrate that CPPO substantially enhances training stability and significantly improves reasoning accuracy across multiple model scales.

Abstract

Reinforcement learning with verifiable rewards (RLVR) has become standard for improving LLM reasoning. However, existing PPO-style trust-region mechanisms remain position-agnostic by enforcing uniform thresholds across all tokens independently. This pointwise treatment conflicts with autoregressive generation in two critical ways. First, uniform thresholds ignore autoregressive asymmetry. Early-stage deviations produce compounding sequence-level drift, causing static thresholds to under-regulate early divergence and excessively constrain late-stage exploration. Second, evaluating token-level divergence in isolation overlooks cumulative prefix drift, granting the same divergence allowance regardless of how far the conditioning history has already deviated from the rollout policy. To address this limitation, we propose CPPO (Cumulative Prefix-divergence Policy Optimization), a token-level masking rule that aligns updates with a finite-horizon policy-improvement bound via two coupled mechanisms. First, a position-weighted threshold imposes stricter limits at early positions whose effects persist longer, relaxing constraints for late-stage tokens. Second, a cumulative prefix budget tracks historical deviations, dynamically restricting further token-level deviation to prevent compounding errors along the prefix. Empirically, CPPO enhances training stability and significantly improves reasoning accuracy across various model scales.

Problem

• trust region
• autoregressive generation
• token-level divergence
• prefix drift
• reinforcement learning

Innovation

• trust region
• autoregressive generation
• policy optimization
• cumulative divergence
• reinforcement learning