mlsanityData to 14 Oct 2026

KTAE: A Model-Free Algorithm to Key-Tokens Advantage Estimation in Mathematical Reasoning

arXiv 2505.16826 ↗First posted 22 May 2025cs.AI cs.CLcited by 19

On the map: LLM reasoning → Reinforcement learning for reasoning

Abstract (arXiv)

Recent advances have demonstrated that integrating reinforcement learning with rule-based rewards can significantly enhance the reasoning capabilities of large language models, even without supervised fine-tuning. However, prevalent reinforcement learning algorithms such as GRPO and its variants like DAPO, suffer from a coarse granularity issue when computing the advantage. Specifically, they compute rollout-level advantages that assign identical values to every token within a sequence, failing to capture token-specific contributions and hindering effective learning. To address this limitation, we propose Key-token Advantage Estimation (KTAE) - a novel algorithm that estimates fine-grained, token-level advantages without introducing additional models. KTAE leverages the correctness of sampled rollouts and applies statistical analysis to quantify the importance of individual tokens within a sequence to the final outcome. This quantified token-level importance is then combined with the rollout-level advantage to obtain a more fine-grained token-level advantage estimation. Empirical results show that models trained with GRPO+KTAE and DAPO+KTAE outperform baseline methods across five mathematical reasoning benchmarks. Notably, they achieve higher accuracy with shorter responses and even surpass R1-Distill-Qwen-1.5B using the same base model.

Key facts. Cited by 19 papers in the map (data through 14 Oct 2026). Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning. Builds on 53 of its 61 references that are in the map. First posted on arXiv 22 May 2025.

Builds on

53 of 61 references.

Closest ancestors · 15

Same direction, newest first

  1. 2025Echo Chamber: RL Post-training Amplifies Behaviors Learned in Pretraining10 Apr
  2. 2025Open-Reasoner-Zero: An Open Source Approach to Scaling Up Reinforcement Learning on the Base Model31 Mar
  3. 2025Understanding R1-Zero-Like Training: A Critical Perspective26 Mar
  4. 2025SimpleRL-Zoo: Investigating and Taming Zero Reinforcement Learning for Open Base Models in the Wild24 Mar
  5. 2025Reinforcement Learning for Reasoning in Small LLMs: What Works and What Doesn't20 Mar
  6. 2025DAPO: An Open-Source LLM Reinforcement Learning System at Scale18 Mar
  7. 2025Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning10 Mar
  8. 2025An Efficient and Precise Training Data Construction Framework for Process-supervised Reward Model in Mathematical Reasoning4 Mar
  9. 2025Process Reinforcement through Implicit Rewards3 Feb
  10. 2025DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning22 Jan
  11. 2025Kimi k1.5: Scaling Reinforcement Learning with LLMs22 Jan
  12. 2024Generative Verifiers: Reward Modeling as Next-Token Prediction27 Aug
3 more
  1. 2023Let's Verify Step by Step31 May
  2. 2023Direct Preference Optimization: Your Language Model is Secretly a Reward Model29 May
  3. 2017Proximal Policy Optimization Algorithms20 Jul

Other references · 38

  1. 2024DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Modelscited by 9,379
  2. 2022Self-Consistency Improves Chain of Thought Reasoning in Language Modelscited by 5,708
  3. 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,371
  4. 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,885
  5. 2015Continuous control with deep reinforcement learningcited by 3,407
  6. 2013Playing Atari with Deep Reinforcement Learningcited by 3,274
  7. 2015Trust Region Policy Optimizationcited by 2,580
  8. 2018Addressing Function Approximation Error in Actor-Critic Methodscited by 2,417
  9. 2016Asynchronous Methods for Deep Reinforcement Learningcited by 2,258
  10. 2024HybridFlow: A Flexible and Efficient RLHF Frameworkcited by 2,057
  11. 2022Solving Quantitative Reasoning Problems with Language Modelscited by 1,392
  12. 2025s1: Simple test-time scalingcited by 1,341
26 more
  1. 2024Qwen2.5-Math Technical Report: Toward Mathematical Expert Model via Self-Improvementcited by 1,090
  2. 2024Large Language Monkeys: Scaling Inference Compute with Repeated Samplingcited by 991
  3. 2024Do NOT Think That Much for 2+3=? On the Overthinking of o1-Like LLMscited by 570
  4. 2025LIMO: Less is More for Reasoningcited by 498
  5. 2025Stop Overthinking: A Survey on Efficient Reasoning for Large Language Modelscited by 464
  6. 2025L1: Controlling How Long A Reasoning Model Thinks With Reinforcement Learningcited by 384
  7. 2025From System 1 to System 2: A Survey of Reasoning Large Language Modelscited by 320
  8. 2025Chain of Draft: Thinking Faster by Writing Lesscited by 297
  9. 2025O1-Pruner: Length-Harmonizing Fine-Tuning for O1-Like Reasoning Pruningcited by 279
  10. 2024Token-Budget-Aware LLM Reasoningcited by 247
  11. 2025ThinkPrune: Pruning Long Chain-of-Thought of LLMs via Reinforcement Learningcited by 192
  12. 2023ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Modelscited by 187
  13. 2025When More is Less: Understanding Chain-of-Thought Length in LLMscited by 183
  14. 2025Light-R1: Curriculum SFT, DPO and RL for Long COT from Scratch and Beyondcited by 175
  15. 2017Equivalence Between Policy Gradients and Soft Q-Learningcited by 173
  16. 2025Phi-4-reasoning Technical Reportcited by 151
  17. 2017Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximationcited by 121
  18. 2025SoftCoT: Soft Chain-of-Thought for Efficient Reasoning with LLMscited by 116
  19. 2024Escape Sky-high Cost: Early-stopping Self-Consistency for Multi-step Reasoningcited by 108
  20. 2025Towards Scientific Intelligence: A Survey of LLM-based Scientific Agentscited by 98
  21. 2025OverThink: Slowdown Attacks on Reasoning LLMscited by 96
  22. 2024OlympicArena: Benchmarking Multi-discipline Cognitive Reasoning for Superintelligent AIcited by 71
  23. 2024Smaller, Weaker, Yet Better: Training LLM Reasoners via Compute-Optimal Samplingcited by 61
  24. 2024Can Language Models Learn to Skip Steps?cited by 59
  25. 2025InftyThink: Breaking the Length Limits of Long-Context Reasoning in Large Language Modelscited by 54
  26. 2023ChineseWebText: Large-scale High-quality Chinese Web Text Extracted with Effective Evaluation Modelcited by 16

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 7; 2026: 12.

19 papers.

Sort:

By area:

  1. 2026Self-Distilled RLVRcited by 131
  2. 2025Towards a Unified View of Large Language Model Post-Trainingcited by 36
  3. 2025Unlocking Exploration in RLVR: Uncertainty-aware Advantage Shaping for Deeper Reasoningcited by 34
  4. 2026Your Group-Relative Advantage Is Biasedcited by 34
  5. 2026Outcome-Grounded Advantage Reshaping for Fine-Grained Credit Assignment in Mathematical Reasoningcited by 19
  6. 2025Look Again, Think Slowly: Enhancing Visual Reflection in Vision-Language Modelscited by 16
  7. 2025CROP: Contextual Region-Oriented Visual Token Pruningcited by 14
  8. 2026What and When to Distill: Selective Hindsight Distillation for Multi-Turn Agentscited by 13
  9. 2025Teaching Vision-Language Models to Ask: Resolving Ambiguity in Visual Questionscited by 11
  10. 2026Reinforcement Learning from Rich Feedback with Distributional DAggercited by 10
  11. 2025VRPO: Rethinking Value Modeling for Robust RL under Noisy Supervision in LLM Post-Trainingcited by 9
  12. 2026Beyond Uniform Credit Assignment: Selective Eligibility Traces for RLVRcited by 5
7 more
  1. 2025Parallel Scaling Law: Unveiling Reasoning Generalization through A Cross-Linguistic Perspectivecited by 4
  2. 2026Know What You Know: Metacognitive Entropy Calibration for Verifiable RL Reasoningcited by 3
  3. 2026Policy Improvement Reinforcement Learningcited by 2
  4. 2026Beyond Entropy: Learning from Token-Level Distributional Deviations for LLM Reasoningcited by 1
  5. 2026DFPO: Scaling Value Modeling via Distributional Flow towards Robust and Generalizable LLM Post-Trainingcited by 0
  6. 2026Guidance Contrastive Token Credit Assignment for Discrete Policy Optimizationcited by 0
  7. 2026BALTO: Balanced Token-Level Policy Optimization for Hallucination Mitigationcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2026Step Potential Advantage Estimation: Harnessing Intermediate Confidence and Correctness for Efficient Mathematical Reasoning
  2. 2025Conditional Advantage Estimation for Reinforcement Learning in Large Reasoning Models
  3. 2026Selective-Advantage Entropy-Adaptive Horizon GRPO: Asymmetric Token-Level Discounting for Efficient Reinforcement Learning of Language Models
  4. 2026GRAIL: Gradient-Reweighted Advantages for Reinforcement Learning with Verifiable Rewards
  5. 2026DISPO: Enhancing Training Efficiency and Stability in Reinforcement Learning for Large Language Model Mathematical Reasoning
  6. 2026BASIS: Batchwise Advantage Estimation from Single-Rollout Information Sharing for LLM Reasoning
  7. 2025Act Only When It Pays: Efficient Reinforcement Learning for LLM Reasoning via Selective Rollouts
  8. 2025Tree-OPO: Off-policy Monte Carlo Tree-Guided Advantage Optimization for Multistep Reasoning