PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Wei Sun, Wen Yang, Pu Jian, Qianlong Du, Fuwei Cui, Shuo Ren, Jiajun Zhang

arXiv:2505.16826v2cs.AIcs.CL

TL;DR

GRPO and DAPO assign the same rollout-level advantage to every token, despite differences in token importance within reasoning sequences. KTAE estimates token importance from sampled-rollout correctness without adding models, then combines it with rollout-level advantage; KTAE-enhanced methods improve average performance across five mathematical reasoning benchmarks while producing shorter responses.

  • Problem

    GRPO and DAPO assign identical rollout-level advantages to every token, although token contributions can differ and process-level reward models are costly and difficult to scale.

    Source: §1 · p. 2

  • Method

    KTAE statistically estimates token-correctness associations from sampled rollouts and combines token importance with rollout-level advantage, without adding models.

    Source: §1 · p. 2

  • Results

    GRPO+KTAE and DAPO+KTAE improve average performance across five major mathematical reasoning benchmarks and reduce response length without length-penalty rewards.

    Source: §1 · p. 2

  • Takeaways & Limitations

    KTAE offers finer-grained advantage information without extra models, with the paper reporting lower training costs and less susceptibility to reward hacking.

    Source: §1 · pp. 2–3

    Experiments focus on 1.5B- and 7B-parameter models and mathematical reasoning, leaving larger-scale performance and broader-domain applications unverified.

    Source: A Limitations · p. 16

Abstract

from arXiv · show

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.

1 Introduction

GRPO and DAPO use rollout-level advantages that overlook differences in token importance, while process-level reward models can be costly and difficult to scale. KTAE estimates token-level importance from sampled-rollout correctness without adding a model, and its combinations with GRPO or DAPO improve average performance across five mathematical reasoning benchmarks.

1 Introduction

  • GRPO and DAPO assign one rollout-level advantage to every token, although token importance varies and incorrect rollouts may diverge only in their final steps.Process-level reward models can provide finer signals but are costly, difficult to scale, limited in signal accuracy, and prone to reward hacking.Source: §1 · p. 2
  • Across five major mathematical reasoning benchmarks, GRPO+KTAE and DAPO+KTAE improve average performance and reduce response length without length-penalty rewards.Figure 1 compares the methods on Qwen2.5-Math-7B; the text reports the average improvement and shorter responses.Source: §1 · p. 2, Figure 1 · p. 2
  • KTAE uses rollout correctness and token occurrence to estimate token-correctness associations statistically, then combines token importance with rollout-level advantage for token-level estimates.The method uses tools such as Fisher’s exact test and Information Gain to quantify association strength and direction.Source: §1 · p. 2
  • KTAE provides finer-grained advantage information without extra models, lowering training costs and using statistical analysis to distinguish token importance.The paper describes the estimates as interpretable and less susceptible to reward hacking because they retain final-answer correctness and rollout-level advantage.Source: §1 · pp. 2–3
  • For a correct rollout assigned reward 0 because its answer was unparsable, Figure 2 visualizes positive and negative token associations, with darker red indicating stronger negative association.The figure illustrates token-level associations in a case where the rule-based reward misclassified the rollout.Source: §1 · p. 2, Figure 2 · p. 3

2 Preliminary

The preliminary section frames language-model reinforcement learning as token-level sequence generation optimized for reward, then introduces GRPO and its DAPO enhancement. GRPO derives advantages from sampled-group rewards, while DAPO adds mechanisms tailored to mathematical reasoning.

2.1 Reinforcement Learning in LLM

  • LLM reinforcement learning models generation as a token-level MDP and maximizes expected reward with a KL penalty, typically β=0 in recent mathematical reasoning tasks.The penalty discourages large shifts from a reference policy.Source: §2.1 · p. 3

2.2 GRPO

  • GRPO samples rollouts, derives token advantages from their group rewards, and uses a clipped policy objective without a separate value model.The clipping term constrains the updated policy relative to the old policy to enhance training stability.Source: §2.2 · pp. 3–4

2.3 DAPO

  • DAPO adds Clip-Higher, dynamic sampling, token-level policy-gradient loss, and overlong-response penalties to GRPO for mathematical reasoning tasks.The overlong-response shaping is intended to prevent catastrophic repetition loops.Source: §2.3 · p. 4

3 KTAE: A Model-Free Algorithm to Key-Tokens Advantage Estimation

KTAE addresses GRPO’s uniform per-rollout token advantages by estimating token importance from sampled-rollout correctness, without adding models. It combines that importance with GRPO’s rollout-level advantage to produce token-level estimates; it can also be combined with DAPO.

3.1 Building Token-Level Contingency Tables

  • For each token, KTAE counts its presence or absence in correct and incorrect sampled rollouts to build a 2 × 2 contingency table.These counts are used to estimate the token’s association with rollout correctness.Source: §3.1 · p. 4, Figure 3 · p. 5

3.2 Quantifying Association Strength via Hypothesis Testing

  • Fisher’s exact test quantifies association strength from the contingency table, with smaller p-values indicating stronger association.The authors choose it over chi-squared and G-tests because it gives exact probabilities for small rollout sets.Source: §3.2 · p. 5

3.3 Quantifying Association Strength via Information Gain

  • Information Gain measures how much knowing whether a token appears reduces uncertainty about rollout correctness.KTAE combines Information Gain with the Fisher score to quantify association strength.Source: §3.3 · pp. 5–6

3.4 Quantifying Association Direction and Final Importance Score

  • KTAE multiplies association strength by a directional score based on proportion and standardized-frequency differences, then sigmoid-normalizes token values and adds them to GRPO’s rollout advantage.The frequency component distinguishes tokens with similar rates across correct and incorrect rollouts from low-frequency tokens with differing rates; concatenating each correctness group reduces the impact of extreme rollout lengths. KTAE is orthogonal to DAPO and can be combined with it.Source: §3.4 · p. 6

4 Experiment

The experiments report stronger benchmark performance for KTAE-augmented methods, along with shorter responses; ablations also indicate that each KTAE component contributes to accuracy. The reported runtime depends on implementation efficiency, with further engineering optimization planned.

Experiment Setting

  • Validation and ablation used Qwen2.5-Math-1.5B, while baseline comparisons used Qwen2.5-Math-7B.The experiments used math12k and its math-level3-5 subset, respectively.Source: §4 · p. 7

Comparison with Baselines

  • DAPO+KTAE-7B achieved the highest average score across five benchmarks, and both KTAE methods outperformed baselines on AMC.GRPO+KTAE improved over GRPO on four of five benchmarks; DAPO+KTAE improved or remained unchanged on four of five.Source: §4 · pp. 7–8, Table 1 · p. 8

Response Length

  • KTAE reduced response length without a length-penalty reward; GRPO+KTAE-7B generated considerably shorter responses than baseline methods.The reported comparison also describes the highest average score across five benchmarks with the least token budget.Source: §4 · p. 8, Table 2 · p. 9

Computational Cost

  • KTAE's computational cost depends primarily on generated-token count, and its overhead is more noticeable for the 1.5B model than the 7B model.The paper describes the cost as acceptable compared with model-based methods, while noting runtime depends on implementation efficiency and plans further engineering optimization.Source: §4 · p. 8, Table 3 · p. 9

Ablation Analysis

  • Removing any KTAE component reduced test accuracy; excluding IG had the largest accuracy impact and produced the shortest sequences.Removing F or tf decreased accuracy while increasing sequence length; removing tf initially increased entropy, while KTAE avoided GRPO's entropy collapse.Source: §4 · pp. 8–9

Visualization Example

  • Some zero-reward rollouts contained an early correct answer followed by redundant tokens that obscured it.The paper reports observing an “aha moment” phenomenon during KTAE training.Source: §4 · p. 9

5 Related Work

The related work situates KTAE among large reasoning models and reinforcement-learning methods, highlighting the challenge of fine-grained reward and advantage signals for reasoning tokens.

Large Reasoning Models

  • Large reasoning models such as DeepSeek R1 use pure RL to induce long CoT and self-reflection, while later work explored RL variants mainly on smaller models.The section frames fine-grained reward exploration in GRPO as challenging.Source: §5 · p. 9

Reinforcement Learning

  • RL methods progressed from REINFORCE and DPO to constrained or clipped updates in TRPO and PPO, then to GRPO and variants such as DAPO and Dr.GRPO.GRPO and its variants use uniform rollout advantage, ignoring token importance in reasoning.Source: §5 · p. 10

6 Conclusion

The conclusion presents KTAE as a model-free way to turn statistical token-correctness associations into finer optimization signals. It reports strong test performance with a minimum token budget and identifies potential application to other reasoning domains.

6 Conclusion

  • KTAE statistically quantifies each token's association with correct rollouts and combines that signal with GRPO's rollout-level advantage.The conclusion describes the result as token-level advantages for more fine-grained optimization.Source: §6 · p. 10
  • KTAE requires no new models, adds minimal computational overhead, and avoids reward hacking.Source: §6 · p. 10
  • KTAE identifies token importance so the model attends more to key tokens, with excellent test performance using the minimum token budget.Source: §6 · p. 10
  • The conclusion states that KTAE's core idea could apply to many other reasoning domains.Source: §6 · p. 10

A Limitations

The paper limits its validation to smaller models and mathematical reasoning, while identifying broader reward settings and token-level attribution as open directions. It also describes broader-impact concerns and a logarithmic-space computation for factorials.

A Limitations

  • KTAE is validated mainly on 1.5B- and 7B-parameter models and mathematical reasoning, leaving larger-scale performance and broader-domain applications for further study.The authors state that KTAE is theoretically applicable beyond rule-based-reward tasks, but this paper provides in-depth analysis and experiments only for mathematical reasoning.Source: A Limitations · p. 16

B Broader Impacts

  • The authors note that reasoning-specialized models may sacrifice general abilities such as summarization, abstraction, and translation.Source: B Broader Impacts · p. 16

C Future Work

  • Experiments focus on binary discrete rewards, with multiclass discrete and continuous rewards left for future exploration.Source: C Future Work · p. 16
  • A single token may carry too little information within a full reasoning path to determine whether it decisively affected correctness or whether its absence would cause failure.Source: C Future Work · p. 16

D Calculating factorials using logarithmic space

  • The factorial computation is converted into additions and subtractions of ln Γ terms, which can be computed in parallel with torch.lgamma() for greater efficiency and numerical precision.Source: D Calculating factorials using… · p. 17

E Why quantify association direction?

Association strength alone does not reveal whether a token is associated with correct or incorrect rollouts. KTAE therefore quantifies association direction and adds token-specific importance to GRPO’s rollout-level advantages.

E Why quantify association direction?

F A schematic diagram of KTAE

G.1 Dataset and Benchmark

The experiments first validate KTAE on Qwen2.5-Math-1.5B, then use a harder MATH subset and a broader benchmark set for 7B evaluation. Comparisons include foundational RL methods and representative alternatives.

G.1 Dataset and Benchmark

  • Initial validation trains Qwen2.5-Math-1.5B Base on MATH12k and tests on MATH500; the later 7B setup trains on MATH Levels 3–5.Source: §G.1 · p. 18
  • The KTAE-7B evaluation covers AIME24, MATH-500, AMC, Minerva, and OlympiadBench.Source: §G.1 · p. 18

G.2 Baselines

  • Baselines include GRPO and DAPO, plus Simple-RL-Zoo, PRIME-Zero, OpenReasonerZero, and Oat-Zero, spanning rule-based, process-reward, and task-tailored approaches.Source: §G.2 · p. 19

G.3 Implementation Details and Hyperparameters

The paper specifies the KTAE-7B training setup and illustrates token attribution in case studies. Its background sections situate the method among reasoning models and RL algorithms, while the prompt and “Aha Moment” examples describe evaluation practice and an observed behavior.

G.3 Implementation Details and Hyperparameters

  • KTAE-7B uses DAPO with KTAE on Qwen2.5-Math-7B, trained in VerL with a 4096-token context, 3072-token generation limit, and 1024-token prompt limit.Source: §G.3 · p. 19

H Case Study

  • In case studies, KTAE assigns positive contribution to correct-answer segments and negative contribution to incorrect or distracting segments, even when later output obscures a correct answer.Source: H Case Study · p. 19

I.1 Large Reasoning Language Models

  • The related-work discussion describes open-source RL replication while identifying fine-grained reward signals in GRPO as an open challenge.Source: §I.1 · p. 21

I.2 Reinforcement Learning

  • GRPO removes the critic model but retains uniform rollout-level advantages, and the paper presents KTAE as a statistical token-level alternative.Source: §I.2 · p. 21

J Prompt

  • The same Qwen-Math prompt template is used for RL training and validation.Source: J Prompt · p. 21

K Aha Moment

  • In the Figure 10 example, KTAE produces a successful answer after reconsidering another solution, while GRPO answers incorrectly; the authors describe this as a degree of self-reflection and error correction.Source: K Aha Moment · p. 21

L Dataset License

Training data comes from MATH, and evaluation covers five mathematical reasoning benchmarks; the authors state that they follow each dataset’s license.

L Dataset License

  • MATH supplies the training data, and evaluation uses AIME24, MATH-500, AMC, Minerva, and OlympiadBench under their associated licenses.The authors state that they strictly adhere to each dataset’s license.Source: L Dataset License · p. 21
  • Figure 10 gives an example of an “Aha moment.”Source: Figure 10 · p. 22
Loading 2505.16826v2…