Title: Multi-Token Policy Gradient Optimization
Primary area: reinforcement learning
Keywords: reinforcement learning, large language model, reasoning

Abstract: Policy-gradient optimization methods like PPO typically operate at the token level, estimating action probabilities for each next-token prediction. While effective, this formulation often mismatches the sequence-level nature of rewards in reasoning tasks, such as math equations and multi-step derivations. To address this issue, we propose Multi-Token Policy Gradient Optimization (MPO), a framework that treats blocks of K future tokens as atomic actions, directly aligning the optimization granularity with semantic segments of text. This block-level view better captures the structure of reasoning trajectories and provides a more faithful connection between policy optimization and sequence-level objectives. To stabilize training, we additionally adopt a log-space approximation of importance ratios, which mitigates numerical instability during optimization. Experiments on GSM8K and MATH show that MPO consistently outperforms standard token-level policy gradient baselines, demonstrating the effectiveness of modeling multi-token actions for large-scale language model post-training.