Reinforcement learning improves answer quality by focusing on hesitation points
Margin-Guided Token Credit for Reinforcement Learning with Verifiable Rewards
Machine LearningArtificial Intelligence
Summary
When teaching language models to solve math problems, it helps to give them credit not just for correct answers, but exactly where they struggled when writing those answers. The authors created a new way to measure these hesitation points using the model’s uncertainty between its top guesses. By giving more learning reward to tokens where the model was unsure, their method leads to better performance on math challenges and longer, more detailed answers. This method works without needing extra computation and improves existing approaches consistently.
What this means in practice
- •For ai model developers: Enhance training methods for language models solving complex problems by focusing reward signals on uncertain parts of their generated answers.
- •For math competition platform teams: Improve automated solution generation quality and reliability for math problems by adopting margin-guided credit assignment in reinforcement learning training.
Authors
Chi Liu, Xin Chen, Yuchuan Weng, Srinivasan Manoharan
Abstract
In reinforcement learning with verifiable rewards, a language model is rewarded for a correct final answer, and GRPO passes that reward to every token of the solution in equal measure. We instead place the credit where the model hesitated, in a method we call MCPO (Margin-Credit Policy Optimization). We measure hesitation by the top-2 margin $M=p_1-p_2$ of the policy's next-token distribution, map it to a score with two fixed thresholds, and spread each solution's advantage over its tokens so that the mean stays unchanged. The margin is read from a forward pass the trainer already performs, so the method needs no extra model evaluation and no batch statistics. We train MiniCPM5-2B-Base with LoRA for 200 steps on 2,048 DAPO-Math problems and evaluate on six competition-math benchmarks. MCPO beats an unmodified GRPO baseline for all five data seeds, by 6.55 points of macro avg@32 on average (95\% CI $[4.1, 9.0]$), and a Qwen3-4B-Base replication improves for all three seeds ($+4.75$). The gain is largest on problems that GRPO solves only some of the time and grows with $k$ in pass@$k$, reaching 13.8 points at pass@32. MCPO models also write two to three times longer answers, and runs that write longer score higher; our experiments do not separate the two effects.