Research map / Large language models / Reasoning and chain of thought

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

NeurIPS 2025 · Large language models › Reasoning and chain of thought · approach: RL with verifiable rewards for reasoning

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes KTAE, a model-free token-level advantage estimator that uses contingency-table statistics (Fisher's exact test, information gain, BM25-style term frequency) over sampled rollouts' correctness, added to GRPO/DAPO rollout advantages.

Problem

GRPO and DAPO assign a single rollout-level advantage to every token, which is too coarse to credit or penalize individual reasoning tokens in mathematical reasoning RL, while process reward models are costly and prone to reward hacking.

Approach

Estimate per-token importance without extra models by building a 2x2 contingency table of token occurrence versus rollout correctness across the G sampled rollouts, score association strength with Fisher's exact test and information gain, sign the score with a BM25-style term-frequency direction, and add the sigmoid-normalized result to the GRPO rollout advantage.

Key results

• DAPO+KTAE-7B reaches 52.5 average across five benchmarks vs 51.7 for DAPO-7B and 50.3 for GRPO-7B, with 83.2 on MATH-500.
• GRPO+KTAE-1.5B averages 41.7 vs 40.6 for GRPO-1.5B; DAPO+KTAE-1.5B averages 43.4 vs 41.2 for DAPO-1.5B and exceeds R1-Distill-Qwen-1.5B (41.5).
• GRPO+KTAE-7B reduces average response length to 718.6 tokens from 771.0 for GRPO-7B.

Measured effects reported in the paper

Surprising findings

• Removing the information gain component produced the largest accuracy drop and the shortest sequences in ablation.
• Per-step training time increases substantially with KTAE (e.g., DAPO 363s to 643s on 1.5B), and the authors report GPU utilization below 1% in the current implementation.

Limitations

• Only evaluated on 1.5B and 7B models; scaling to larger models unverified.
• Only validated on mathematical reasoning with binary rewards; generalization to other domains untested.
• Training overhead is non-trivial due to limited parallelism in the current implementation.

Experimental setup

Datasets: MATH12k MATH level 3-5 subset

Benchmarks: AIME24 MATH-500 AMC Minerva OlympiadBench

Metrics: zero-shot greedy pass@1 average accuracy mean response length training time per step

Baselines: GRPO DAPO Simple-RL-Zoo PRIME-Zero OpenReasoner-Zero Oat-Zero R1-Distill-Qwen-1.5B Qwen2.5-Math-Instruct

Builds on: GRPO DAPO PPO BM25 Fisher's exact test

Closest papers by meaning

• Conditional Advantage Estimation for Reinforcement Learning in Large Reasoning Models (ICLR 2026 · similarity 0.89)
• Act Only When It Pays: Efficient Reinforcement Learning for LLM Reasoning via Selective Rollouts (NeurIPS 2025 · similarity 0.88)
• Do Not Let Low-Probability Tokens Over-Dominate in RL for LLMs (ICLR 2026 · similarity 0.88)
• Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM’s Reasoning Capability (ICML 2025 · similarity 0.87)
• STAPO: Stabilizing Reinforcement Learning for LLMs by Silencing Rare Spurious Tokens (NeurIPS 2026 · similarity 0.87)
• STARE: Surprisal-Guided Token-Level Advantage Reweighting for Policy Entropy Stability (NeurIPS 2026 · similarity 0.87)