Existing policy-gradient methods for auto-regressive language models typically select subsequent tokens one at a time as actions in the policy. While effective for many generation tasks, such an approach may not fully capture the structure of complex reasoning tasks, where a single semantic decision is often realized across multiple tokens--for example, when defining variables or composing equations. This introduces a potential mismatch between token-level optimization and the inherently block-level nature of reasoning in these settings. To bridge this gap, we propose Multi-token Policy Gradient Optimization (MPO), a framework that treats sequences of K consecutive tokens as unified semantic actions. This block-level perspective enables our method to capture the compositional structure of reasoning trajectories and supports optimization over coherent, higher-level objectives. Experiments on mathematical reasoning and coding benchmarks show that MPO outperforms standard token-level policy gradient baselines, highlight the limitations of token-level policy gradients for complex reasoning, motivating future research to look beyond token-level granularity for reasoning-intensive language tasks.
Figures & tables
Figure 1: Left : In reasoning tasks such as mathematical problem-solving or code generation, the model’s decision process often spans across blocks of tokens—such as equations or functions—rather than being determined by each token independently; Right : illustration of token-level (NTP/PPO) vs. block-level (MTP/MPO) optimization. MPO aggregates K tokens as a semantically meaningful block for prediction and optimization, thereby better capturing sequence structure and long-range dependencies.
Figure 2: (a) Demonstration of the implementation of MPO warm-up and training process; (b) illustration of the united importance sampling ratio proposed in MPO method; (c) Comparison between single-token and multi-token optimization, multi-token optimization jointly models contiguous tokens as a structural reasoning action.
Figure 3: The performance of proposed MPO and baseline methods. MPO outperforms the baselines in most scenarios, demonstrating the effectiveness of aggregating block-wise information.
Model
Zero-Shot
GRPO
DAPO
PPO
MPO (Ours)
Llama3.2 - 1B - Instruct
0.354
0.372
0.396
0.390
0.403
DeepSeek - Qwen2.5 - 1.5B
0.451
0.591
0.603
0.598
0.640
DeepSeek - Qwen2.5 - 7B
0.689
0.811
0.817
0.805
0.841
Table 1: The performance of proposed MPO and baseline methods on coding task HumanEval, MPO consistently outperforms the baselines, demonstrating the effectiveness of aggregating block-wise actions during optimization.
Figure 4: Comparison of the variance of importance sampling ratios and clip fraction during training.
Figure 5: Effect of MTP block size K and decay rate λ on training stability and performance. Extending the block size to K=5 and applying a moderate decay λ=0.8 produces the most stable and effective result.
Figure 6: The effect of varying the proportion of weights incorporated from the MTP modules.
Figure 7: Reliability analysis of the multi-token information. (a) Larger prediction errors correlate with higher gradient magnitudes. (b) Progressive improvement in the MTP module’s next-K-token prediction capability.
Training Speed (Second per iteration)
Model
PPO
K=3
K=5
Deepseek-Qwen2.5-1.5b
25.6
30.3
33.3
Llama3.2-1b
19.1
25.0
28.3
Table 2: The average training time with block size K.
Method
Memory Usage
Time Per Step
PPO
22.1%
25.6 s
GRPO
30.6%
43.7 s
MPO
31.0%
33.3 s
Table 3: Memory and time efficiency.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: (a) The implementation of the MTP module used in MPO, the structure preserves causal consistency following Deepseek-V3; (b) Demonstration of the MPO warmup stage to initialize MTP Modules before the RL stage.
Dataset
# Train Samples
# Test Samples
Avg. Sample Length
GSM8K
7,473
1,319
∼ 80 tokens
MATH
7,500
5,000
∼ 200 tokens
HumanEval
1399 (MBPP)
164
∼ 130 tokens
Appendix
Table 4: Statistics of datasets used for training and evaluation.
Parameter
Value
Train batch size
256
PPO mini batch size
128
PPO micro batch size / GPU
4
Number of GPUs
8
Number of Nodes
1
Max prompt length
1024
Appendix
Table 5: Key hyperparameters for MPO training.
Settings
GSM8K
MATH
K=2
K=3
K=5
K=2
K=3
K=5
MPO
0.871
0.875
0.882
0.771
0.753
0.789
+ VMPO
0.879
0.875
0.882
0.773
0.756
0.762
Appendix
Table 6: MPO performance with different token block size K and ablation study of multi-token value estimation.
Figure 9: The implementation of the value function on the original auto-regressive backbone model (left) and those with the MTP module (right). For multi-token prediction models, we try a weighted sum of the output from the shared value head.
Group Relative Policy Optimization (GRPO) has significantly advanced the reasoning ability of large language models (LLMs), particularly in their mathemat ical reasoning performance. However, GRPO and related entropy regularization methods still struggle with token-level sparse-rewards, which is an inherent chal lenge in chain-of-thought (CoT) reasoning. These approaches often rely on undifferen tiated token-level entropy regularization, which easily leads to entropy collapse or model degradation under sparse token rewards. In this work, we propose TEPO, a novel token-level framework that (1) leverages sequence-level likelihood to link group-level rewards with individual tokens via token-level aggregation, and (2) introduces a token-level KL-Divergence mask constraint that targets tokens with positive advantages and decreasing entropy to mitigate abrupt policy updates. Experiments demonstrate that TEPO not only achieves state-of-the-art performance on mathematical reasoning benchmarks but also markedly enhances training stability, reducing convergence time by 50% compared with GRPO/DAPO.
Xingyu Lin, Yilin Wen, Du Su +5
College of Computer Science and Technology, Jilin University · Key Laboratory of Symbolic Computation and Knowledge Engineering of MOE, Jilin University · Baidu Inc. +2
Reinforcement learning from verifiable rewards has significantly advanced the reasoning capabilities of large language models. However, Group Relative Policy Optimization (GRPO) typically assigns a uniform, sequence-level advantage to all tokens, thereby overlooking the intrinsic information heterogeneity along reasoning chains. We show that this coarse-grained credit assignment leads to premature entropy collapse and encourages the model to generate redundant, low-quality reasoning paths. Through systematic empirical analysis, we identify Critical Decision Pivots (CDPs): transient high-entropy states where the policy's trajectory is most sensitive to perturbations. These pivots represent the "forks in the road" where effective multi-path exploration is most crucial yet often suppressed by uniform advantage signals. Building on these insights, we propose Entropy-Regulated Policy Optimization (ERPO), which transitions the optimization focus from coarse sequences to fine-grained token dynamics. ERPO introduces three synergistic components: (i) Entropy-aware Gating, which adaptively amplifies exploration at CDPs to facilitate diverse path discovery; (ii) Bucket-based Implicit Normalization, which mitigates difficulty bias by aligning token progress windows; and (iii) Result-anchored Advantage Synthesis, which re-weights token-level signals via outcome-driven anchors. Extensive experiments on competitive mathematical benchmarks demonstrate that ERPO significantly outperforms GRPO. Notably, ERPO not only boosts reasoning accuracy but also yields significantly more concise and robust derivation paths, while achieving performance comparable to large models with orders of magnitude more parameters.
Song Yu, Li Li, Wenwen Zhao +1
School of Computer and Information Science, Southwest University, Chongqing, 400715, China
Existing reinforcement learning approaches for Large Language Models typically perform policy optimization at the granularity of individual tokens or entire response sequences. However, such formulations often misalign with the natural step-wise structure of reasoning processes, leading to suboptimal credit assignment and unstable training in multi-modal reasoning tasks. To bridge this gap, we propose Segment-Aligned Policy Optimization (SAPO), a novel reinforcement learning paradigm that treats coherent reasoning steps, rather than tokens or full sequences as fundamental units of policy update. SAPO introduces a step-wise Markov decision process abstraction over reasoning segments, accompanied by segment-level value estimation, advantage computation, and importance sampling mechanisms that are semantically aligned with reasoning boundaries. Experiments on representative reasoning benchmarks demonstrate that SAPO consistently outperforms token-level and sequence-level policy optimization methods, achieving significant accuracy improvements while exhibiting better training stability and value estimation consistency. Our work underscores the importance of aligning reinforcement learning updates with the intrinsic structure of reasoning, paving the way for more efficient and semantically grounded policy optimization in complex reasoning tasks. Codes and models will be released to ensure full reproducibility.
Lei Gao, Zhuoming Li, Mengxi Jia +4
Fudan University · China Telecom Artificial Intelligence Technology (Beijing) Co., Ltd. · Southeast University +1