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

TL;DR

KTAE addresses the limited granularity of rollout-level advantage signals in GRPO and DAPO when training LLMs for mathematical reasoning. It introduces a model-free approach that constructs token-level contingency signals from rollouts, uses Fisher's exact test and Information Gain to quantify token-correctness associations, and combines these with a direction score to form key-token values that augment the standard rollout advantage. Empirically, integrating KTAE with GRPO or DAPO yields higher accuracy on five math benchmarks and achieves shorter, more efficient responses, sometimes surpassing larger models. The method is lightweight, interpretable, and broadly applicable to reasoning tasks, offering a practical route to improve reasoning performance without additional modeling costs.

Abstract

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.

Figures (10)

• Figure 1: Model performance comparison. KTAE is a plug-and-play method that introduces no additional model. It provides token-level advantage estimation for existing RL algorithms such as GRPO and its variants. "GRPO+KTAE" and "DAPO+KTAE" denote GRPO and DAPO combined with KTAE respectively, both RL-tuned on the Qwen2.5-Math-7B model. Detailed results in Table \ref{['math-all-7B']}.
• Figure 2: Visualization of key-token-values computed by KTAE for a correct rollout whose final result was unparsable and thus received a final reward of 0. Red shading indicates negative token associations with producing a correct rollout, with darker red representing stronger negative influence; Green shading indicates positive associations.
• Figure 3: The outline of KTAE algorithm. It constructs a contingency table based on the correctness of the sampled rollouts, and then calculates the token-level advantage and adds it to the GRPO's rollout-level advantage.
• Figure 4: The metric curves of test accuracy, mean response length, and generation entropy of combining DAPO and GRPO with KTAE.
• Figure 5: Training results after removing different components from KTAE.