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.
College of Computer Science and Technology, Jilin University · Key Laboratory of Symbolic Computation and Knowledge Engineering of MOE, Jilin University · Baidu Inc. +2