💬 NLP

Beyond Token-Level Policy Gradients for Complex Reasoning with Large Language Models

The paper proposes Multi-token Policy Gradient Optimization (MPO), a framework that treats sequences of K consecutive tokens as unified semantic actions to better align policy gradient optimization with the block-level structure of complex reasoning tasks, demonstrating superior performance over standard token-level methods on mathematical and coding benchmarks.

Original authors: Mufan Xu, Kehai Chen, Xuefeng Bai, Zhengyu Niu, Muyun Yang, Tiejun Zhao, Min Zhang

Published 2026-02-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mufan Xu, Kehai Chen, Xuefeng Bai, Zhengyu Niu, Muyun Yang, Tiejun Zhao, Min Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a very smart but very literal robot how to write a story or solve a math problem.

The Old Way: The "One-Step" Robot
Currently, most AI training methods (like PPO) teach the robot one word at a time.

The Analogy: Imagine you are teaching a child to build a Lego castle. The old method says, "Pick up one brick. Good job! Now pick up the next brick. Good job!"
The Problem: The robot doesn't understand that a "wall" is made of 10 bricks working together. It treats every single brick as an isolated decision. If the robot needs to write the phrase "if x equals 5," it might get the "if" right, but then forget the logic for "x equals 5" because it's only looking at the very next word. It's like trying to understand a sentence by only looking at one letter at a time.

The New Idea: The "Block" Robot (MPO)
This paper proposes a new method called Multi-Token Policy Gradient Optimization (MPO). Instead of looking at one word, the robot learns to look at a chunk of words (a "block") as a single unit of thought.

The Analogy: Now, you tell the child, "Build a whole wall section (about 10 bricks) as one single move."
How it works: When the robot decides to define a variable (like a = 5) or write a math equation, it treats the entire phrase a = 5 as one single action. It doesn't just guess the next letter; it plans the whole semantic "block" at once.

Why is this a big deal?

1. It stops the robot from getting "tunnel vision."
In complex reasoning (like math or coding), the meaning often spans several words.

Old Way: The robot sees "if" and thinks, "Okay, what comes next?" It might pick a word that makes sense grammatically but breaks the logic of the equation.
New Way: The robot sees the whole block "if a equals 5" and thinks, "This is a complete logical step." It ensures the whole chunk makes sense together before moving on.

2. It's like learning to drive in "chunks" instead of "pixels."
Imagine driving a car.

Token-level (Old): You look at the road one pixel at a time. "Pixel is red. Pixel is red. Pixel is green." You might miss that the whole traffic light is red.
Block-level (New): You look at the whole traffic light as one object. You see "Red Light" and stop. You understand the intent of the signal, not just the individual colors.

3. It makes the robot smarter at math and coding.
The researchers tested this on hard math problems (like solving algebra) and coding tasks.

The Result: The "Block" robot (MPO) got significantly better scores than the "One-Step" robot. It made fewer mistakes because it stopped trying to guess the next word in isolation and started planning the next thought.

The "Secret Sauce" (How they did it)

The researchers didn't just tell the robot to think bigger; they gave it a special training tool called MTP (Multi-Token Prediction).

Think of this as giving the robot goggles that let it see 5 steps ahead while it's still standing at step 1.
During training, the robot practices predicting a whole sentence chunk at once. Once it gets good at seeing the whole picture, they use that skill to improve its decision-making.

The Bottom Line

This paper suggests that to make AI truly good at reasoning, we need to stop treating it like a machine that just predicts the next letter. Instead, we should train it to think in meaningful chunks, just like humans do when we solve a problem. We don't think word-by-word; we think in ideas, equations, and code blocks. MPO teaches the AI to do the same.

1. Problem Statement

Current policy gradient methods for fine-tuning Large Language Models (LLMs), such as Proximal Policy Optimization (PPO), Group Relative Policy Optimization (GRPO), and DAPO, operate at the token level. They treat the generation of each subsequent token as an isolated action.

The authors argue that this granularity mismatch is detrimental to complex reasoning tasks (e.g., mathematical problem solving, code generation). In these tasks, semantic decisions (like defining a variable, composing an equation, or writing a function) are inherently block-level structures realized across multiple tokens.

Limitation: Decomposing a coherent reasoning step into a series of local token predictions disrupts the semantic coherence of the action.
Consequence: Token-level optimization fails to capture long-range dependencies and structural integrity, leading to suboptimal reasoning performance and unstable training due to high variance in importance sampling ratios.

2. Methodology: Multi-token Policy Gradient Optimization (MPO)

The paper proposes MPO, a framework that shifts the optimization unit from single tokens to blocks of $K$ consecutive tokens, treating them as unified semantic actions.

Core Components:

Multi-Token Prediction (MTP) Modules:
- The model is augmented with $K$ additional prediction heads (MTP modules) attached to the backbone LLM.
- Unlike standard autoregressive models that predict $o_{t+1}$ , MTP modules predict a sequence $o_{t+1}, \dots, o_{t+K}$ conditioned on the prefix $o_{1:t}$ .
- These modules are initialized from the backbone's final layer and undergo a warm-up phase (supervised fine-tuning) before the reinforcement learning stage to ensure prediction quality.
Block-Level Importance Sampling:
- Instead of calculating the importance ratio $r_t$ for a single token, MPO aggregates the ratio over a span of $K$ tokens.
- Variance Control: A simple product of ratios ( $\prod r_{t+n}$ ) leads to high variance. To mitigate this, MPO adopts a weighted log-sum formulation inspired by Log-COP-TD:
  $\tilde{R}^{(K)}_{i,t}(\theta) = \exp \left( \sum_{n=1}^{K} \beta_n \log r_{i,t+n}(\theta) \right)$
  where $\beta_n$ are decaying weights ( $\beta_k = \beta_2 \times \lambda^{k-2}$ ) that prioritize immediate predictions while incorporating future context with diminishing influence.
Policy Gradient Objective:
- The surrogate objective function replaces the single-token ratio with the aggregated block ratio $\tilde{R}^{(K)}_{i,t}$ .
- The loss function encourages the policy to optimize for coherent multi-token patterns rather than isolated next-token moves.
- Compatibility: MPO is designed to be compatible with existing frameworks (like PPO). It only modifies the calculation of the importance sampling ratio and does not require changing the value network architecture (though experiments with multi-token value heads were conducted).

3. Key Contributions

Novel Framework: Introduction of MPO, the first policy gradient method to explicitly optimize over structured, block-level semantic actions rather than isolated tokens.
Theoretical Insight: Demonstrates that token-level optimization introduces a granularity mismatch for reasoning tasks. MPO aligns the optimization signal with the natural structure of reasoning (e.g., equations, code blocks).
Empirical Validation: Extensive experiments showing MPO outperforms state-of-the-art baselines (PPO, GRPO, DAPO) across different model scales and reasoning domains.
Stability Analysis: Proves that aggregating multi-token information reduces the variance of importance sampling ratios and the frequency of gradient clipping, leading to more stable training.

4. Experimental Results

The authors evaluated MPO on three benchmarks using models ranging from 1B to 7B parameters (Llama3.2, DeepSeek-Distilled-Qwen2.5):

Mathematical Reasoning (GSM8K & MATH):
- MPO consistently outperformed baselines.
- On GSM8K, MPO achieved 0.882 accuracy (vs. 0.866 for PPO) on the 1.5B model and 0.941 on the 7B model.
- On the harder MATH dataset, MPO showed significant gains, reaching 0.789 (vs. 0.769 for PPO) on the 1.5B model.
Code Generation (HumanEval):
- MPO achieved the highest pass@1 scores across all models, reaching 0.640 on the 1.5B model (vs. 0.598 for PPO) and 0.841 on the 7B model.
Training Efficiency & Stability:
- Variance Reduction: MPO significantly lowered the variance of importance sampling ratios and reduced the "clip fraction" (the percentage of updates discarded due to clipping), indicating a more stable optimization landscape.
- Efficiency: While MPO adds a small computational overhead (approx. 30% slower per iteration than PPO due to extra forward passes for MTP), it offers substantial performance gains. The memory usage is comparable to GRPO.
- Ablation: Optimal performance was found with a block size of $K=5$ and a decay rate $\lambda=0.8$ . Integrating too much multi-token information (>30% weight) introduced noise and degraded performance.

5. Significance and Future Directions

Paradigm Shift: The paper challenges the prevailing assumption that token-level granularity is sufficient for LLM fine-tuning. It suggests that for reasoning-intensive tasks, semantic block-level optimization is necessary to capture structural dependencies.
Practicality: MPO is a lightweight modification that can be integrated into existing post-training pipelines without retraining the entire backbone from scratch.
Future Work: The authors identify open questions regarding scaling the block size ( $K$ ), optimizing warm-up strategies for models not pre-trained with MTP, and integrating MPO with other objective functions (e.g., GRPO variants).

In conclusion, MPO bridges the gap between the discrete nature of token generation and the continuous, structured nature of human reasoning, offering a robust method to enhance the reasoning capabilities of LLMs during the reinforcement learning phase.