KTAE: A Model-Free Algorithm to Key-Tokens Advantage Estimation in Mathematical Reasoning
Authors: Wei Sun
Wen Yang
Pu Jian
Qianlong Du
Fuwei Cui
Shuo Ren
Jiajun Zhang
Venue: NeurIPS 2025 poster ()
Keywords: Large language model;Reinforcement learning;Mathematical reasoning
Primary area: applications
TL;DR: A novel algorithm that estimates fine-grained, token-level advantages in reinforcement learning without introducing additional models.

Abstract: Recent advances have demonstrated that integrating reinforcement learning with rule-based rewards can significantly enhance the reasoning capabilities of large language models (LLMs), even without supervised fine-tuning (SFT). 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. 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.

КТАE: A Model-Free Algorithm to Key-Tokens Advantage Estimation in Mathematical Reasoning
Sunn Qianlong Du' Fuwei Cui ,Shuo Ren' ,Jiajun Zhang
Wei Sun 1.2 Wen Yangl2, Jian
Institute of Automation, Chinese Academy of Sciences
School of Artificial Intelligence, University of Chinese Academy of Sciences
Wuhan Al Research
{sunwei2023, yan.:e22232. f jaapuu2023 shuo rrn, ren, www..fcuiieiiiei...... Quii..a.cee ac. {qianlong du,
ac.
Project Page: httts://github. http://ggtthub..mo//I.. /TTAE, git
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 rein- forcement learning algorithms such as GRPO and its variants like DAPO, suffer from coarse granularity issue when computing the advantage. Specifically, they compute rollout-level advantages that assign identical values to every token within sequence, failing to capture token-specific contributions and hindering effective learning. To address this limitation, we propose Key-token Advantage Estimation (KTAE)- novel algorithm that estimates fine-grained, token-level advantages without introducing additional models. KTAE leverages the correctness of sampled rollouts and applies statistical analysis quantify the importance of individual to- kens within a sequence to the final outcome. This quantified token-level importance is then combined with the rollout-level advantage obtain more fine-grained token-level advantage estimation. Empirical results show that models trained with GRPO+KTAF and DAPO+KTAE outperform baseline methods across five mathe- matical reasoning benchmarks. Notably, they achieve higher accuracy with shorter responses and even surpass R1-Distill Bestil-Qwen-lle 1.5B using the same base model.
1 Introduction
Notably, large reasoning language models (LRMs) like OpenAl ol [1] and DeepSeck RI [2] have demonstrated the capability to solve complex mathematical reasoning problems that challenge even human experts. This progress marks significant step toward Artificial General Intelligence (AGI) [3] 4 [. These reasoning language models often exhibit behaviors such as reff-reflection and self- verification within the reasoning chain, which are critical for enhancing reasoning accuracy. DeepSeck RI aptly refers to the critical turning points that lead to improved performance as "aha moments". The emergence and cultivation of such moments greatly facilitated by the application of reinforcement learning (RI.) or distilled from more powerful LRMss 6
7 2 For instance, DeepSeek applied RL directly the base language model, using simple rule-based reward function to encourage the model to explore and unlock its reasoning potential [8] through self-exploration.
Corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Performance of 7B Models Across Benchmarks
Qwen? Funus-278S-rONE
80
Oat Zero
SimpleRL, Zoro 78
GRPO-KIAE
Accuracy (%)
60
D
AIME24
AMC
MATH-500
Minerva
Olympiad
Average
Figure 1: Model performance comparison. KTAE is a plug-and-plal method that introduces no additional model. 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 T
As a mainstream RL algorithm, Group Relative Policy Optimization (GRPO) 9 differs from the Proximal Policy Optimization (PPO) [TO] by eliminating the need for a separate critic model. Instead, it estimates the advantage of each token using the rewards obtained from set of generated rollouts. However, due the absence of critic model, GRPO computes a rollout-level advantage assigning the same advantage value to every token within single rollout. This limitation also persists in its improved variant, DAPO [T]. In practice, the importance of each token in a complete Chain-of-Thoughn (CoT) reasoning sequence varies, and we often observe that incorrect rollouts may only diverge from the correct reasoning path in the final steps. Consequently, applying uniform advantage value across all tokens in a rollout lacks granularity and may hinder effective learning. Prior efforts have explored using process-level reward models to provide more fine fine-graing grained signals [12 [3 [1] 15 [6). However, as highlighted by DeepSeek [2], training fine-grained reward models is costly, difficult to scale, has limited capacity to provide accurate signals, and prone to reward hacking [17]
To address these challenges, we propose the Key-token Advantage Estimation (KTAE) algorithm. KTAE introduces no additional models, and instead leverages the correctness of sampled rollouts and the occurrence of each token within them to construct a contingency table. Then, using statistical methods such as Fisher's 's exact test and Information Gain (IG), it quantifies the strength of association between each token and correct rollout. Subsequently, by combining the token's frequency and the reward assigned to its corresponding rollout, KTAEE further quantifies the direction of this association's 's contribution. Finally, these measures are combined (e.g.. through multiplication) to yield 'key-token value" for each token. As shown in Figure 2 when a correct rollout is incorrectly classified as incorrect by the rule, KTAF can still highlight the positively contributing tokens through computing the key-token values. In contrast, GRPO assigns the same negative ad advantages to all tokens in such case. Moreover, KTAE can effectively distinguish between tokens irrelevant to problem solving, such as 'First' and 'denote", and those highly relevant to problem solving, such 'complement" and 'ratio'. Furthermore, KTAE compatible with GRPO and DAPO. The resulting key-token values are then added to the rollout-level advantage computed by GRPO obtain more fine ome-grained token-level advantage estimate. As illustrated in Figure ! improved performances
integrating KTAE with either GRPO or DAPO yields average across five major mathematical reasoning benchmarks. Moreover, KTAE not only improves test accuracy but also very effectively reduces response length without any length penalty reward, resulting in extremely high reasoning efficiency.
In summary, the KTAE algorithm offers several advantages:
KTAE provides more fine- fine-grained advantage information without introducing extra models, resulting in lower training costs.

---

2. KTAE quantifies the importance differences between tokens using statistical analysis meth- ods, offering strong interpretability.
3. KTAE's key-token value computed based on the correctness of the final answer and retains
the original rollout-level advantage, making it less susceptible to reward hacking,
4. KTAE helps the model to focus on key tokens and reduce the learning of irrelevant tokens, which can effectively reduce the response length.
Two complementary angles, A and B, have measures in the ratio of 7 to 23, respectively. ihat the ratio of the measure of the complement angle to the measure the complement of angle B? Express your answer as common fraction.
Question:
1.
Answer:
23/7
Anturann:
Fntt
The
EDV
angitll
eiggrasd
1he
0.2
and
Agpi Aop--A OL orpesii
Thansfone
Right answer just without (boxed []l
0.0
Figure 2: Visualization of key-token- -aaeauus computed by KTAE for correct rollout whose final result was unparsable and thus received a final reward of O. Red shading indicates negative token associations with producing correct rollout, with darker red representing stronger negative influence; Green shading indicates positive associations.
Preliminary
2.1 Reinforcement Learning in LLM
When applying Reinforcement Learning (RL) language models, text generation is modeled as a token-level Markov Decision Process (MDP). At each timestep t, the state consists of the input prompt q and the previously generated tokens policy generates the next token og as action and generation ends upon producing an end- of-sequence token or reaching maximum length T. The full sequence [oI..... orl is then evaluated by reward function R(g. o) = Elal r(st, or). The RL objective aims to maximize an entropy-rege regulaizee expected reward 18
(i-11,,.e.,, i.e., St [q: 01...... . 0t-1]. The
J (m²) (g) [R(g, o)] BE [DLL (m)((
(I)
Here, Tref is reference policy, xx denotes the KI. divergence, and B controls the penalty strength. This KI. term, central RLHF, discourages large shifts from the reference distribution preserve fluency and diversity. In recent mathematical reasoning tasks
is typically set to 0.
GRPO
Group Relative Policy Optimization (GRPO) 9 is simplified based on PPO and eliminates the need for value model. Given an input q, GRPO samples G rollouts {o1...., oG} from the old policy and computes their cumulatíve rewards R (Ri,- Ra). These rewards are then used to estimate
advantages A; e.g., by comparing each R, to baseline derived from R. The optimization objective for GRPO defined as follows:
TeRvo (
E BBwD. (o.))
lod
] (2)
min (O),,)) clip (1), -,, +5)+5.))).) -DD (Tollmmer) -DoaaTTol
3

---

where ri,( mo molor (l..0.c)))
Teas (o (..0 and Ai, is the advantage estimate derived from the group rewards
R, defined as A = R. R) The clipping term with clip ratio [20] aims to constrain the new policy within the trust region of the old policy, enhancing training stability. By eliminating the dependency on the value model Va. GRPO aims substantially reduce training costs while striving
std(R)
to maintain optimization effectiveness comparable to traditional PPO.
2.3 DAPO
Dynamic Sampling Policy Optimization (DAPO) [ is enhancement algorithm of GRPO, specifically tailored for tasks involving mathematical reasoning. To mitigate the phenomenon of entropy collapse, DAPO introduces the "Clip-Higher" method, which raises the upper bound of the clipping function. It incorporates "Dynamic Sampling" to prevent scenarios where all G sampled rollouts exhibit identical preference outcomes (e.g., all positive all negative). A "Token-Level Policy Gradient Loss" is employed to stabilize the training process. Additionally. DAPO introduces "overlong reward shaping" to penalize excessively long responses, thereby preventing the model from falling into catastrophic repetition loops.
(DA = Bg (ypDD,(,,30 (o:) Mo (-19)
G lo.1
min (ris(0)Arr, clip (ip(Fr.(())) 1- Elow, 1 1+ Ehigh Arie)
(3)
1== 0i
1=1
s.t.
0 <<< (or es_eviivaen o;)) <G.
where (0) and A,, are the same as GRPO. Elow and Ehigh represent the upper and lower bounds of Ti. after decoupling.
3 KTAE: A Model-Free Algorithm to Key-Tokens Advantage Estimation
GRPO's 's advantage estimation has a relatively coarse granularity. It assigns the same advantage value to every token within the same rollout. However, tasks requiring complex reasoning steps, such as mathematical reasoning, the importance of different tokens within rollout can vary significantly. To address this, we propose the KTAE (Key-tokens Advantage Estimation) algorithm. Without additional models, KTAE quantifies the importance of different tokens by analyzing the statistical associations within the set of sampled rollouts (correct vs. incorrect). It then integrates this quantified token importance with rollout-level advantage estimates (computed by GRPO) to produce fine-grained, token-level advantage estimations. this section, we will introduce its calculation process in detail.
3.1 Building Token-Level Contingency Tables
For given problem, we sample set of G rollouts {or..... og) each with corresponding rule- based reward (R).... Ra indicating its correctness, following the same approach as GRPO. We divide these rollouts into a correct set 27 and an incorrect set zr. For each token oi in the sampled rollouts og, we examine its occurrence across all rollouts and construct a x 2 contingency table summarizing the counts of correct and incorrect rollouts that contain or do not contain oij. An example contingency table is shown in Figure 3
Here, G is the total number of rollouts. For token Oij. we use the statistics Gogs bogs Cots doys based on its occurrence across the rollouts sets. Gos the count of rollouts in 2r where Oi appears at least once, Gog Σ I(og xT); Coz is the count of rollouts in øT where oij does not appear Co.,
(xT) do., bo. and do., calculated in the same way. The total count: ao,, + bo,, + Co.; + do., = G. We then use the statistics from this contingency table to compute the association between the occurrence of o;y and the policy sampling correct rollout.
3.2 Quantifying Association Strength via Hypothesis Testing
We begin by quantifying the association using hypothesis testing. We set the null hypothesis (Ho) as: the occurrence of oi and the conrectness of its rollout have association. We use Fisher's exact

---

T/F Rollouts
GRPO Advantage
Contingency Tables
KTAE
01
011
K1
A
O2
Policy Model
:
Oij
K₂
A₂
2T
Goy
Doy
K
As
Σ I(ays Φr)
Σ
Ilass E 2xp)
Question
0s
Cay :
day
Σ Iay & ar)
Σ
Ilass & ztp)
KG
Ag
:
2F
Figure 3: The outline of KTAE algorithm. It constructs 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.
test 21 to compute the p-value, which the probability observing the current contingency table more extreme one, assuming Hy is true. The formula for Fisher's exact test p-value is:
(masy boss) Gols Kaniy trais
) (Nors tdhos)
(mang Boeg Menry dang Maory + Co.y Mbmg dor Goss dosslbors Coss Idos !N!
Fisher(og)
(4)
In practice, this calculation is performed in log-space to handle large factorials (see appendix D for details). smaller p-value indicates stronger evidence against the null hypothesis, meaning a stronger association between the occurrence of token Osy and rollout correctness. Since effective p-values often concentrated in a small range, we employ a transformation function to quantify the association strength and amplify the impact of small p-values. We define the association score of Fisher's test as:
F(o.)) = {o - 2-Fisher(ou) if Fisher(Org) #
(5)
if Fisher((os) =
When p = 1 (complete no association), the score is O; when p approaches O (strong association), the score approaches 1. Fisher's exact test chosen over chi-squared G-tests because the latter provide less accurate p-values for small sample sizes N, while Fisher's 's test offers an exact probability 16).
calculation even with small G (e.g.,
3.3 Quantifying Association Strength via Information Gain
To complement the statistical test with an iiformationnthhorede perspective, we compute the Informa- tion Gain (IG) between the occurrence of the token O;y and rollout correctness. Let Y be a random variable representing rollout correctness, and X.,; be variable indicating whether the token oij appears in a rollout. The entropy of rollout correctness HYY is:
H(Y) Gois + Coss loga
(aay Coss
bou + doss loga
bois + dous
(6)
N
N
N
N
The conditional entropy of rollout correctness given whether token os appears, H (Yo)oo,)).
H(Y)Xo.,))
Aosy + bouy
aoiy
loga
aoss
boss loga
bas
N
Ooi
bost
Gon + bos
doss + bos
Coi boss
Cosy + dors N
Cosz
Cois
dous
doi
+
loga
log2
Coi
dois
Con + dois
Coi + dons
Cos + dois
(7)
5

---

The Information Gain (IG) is defined as IG(o,)) - HYY (
A higher IG value indicates that knowing whether the token oij appears reduces the uncertainty about rollout correctness more significantly, suggesting a stronger association with correctness. Otherwise, means the association weaker. Through Fisher's exact test and Information Gain, we have quantified the strength of association between the occurrence of token oi and rollout correctness (e.g., via linear combination f1 F(oi;) + l12 (G(oo))
3.4 Quantifying Association Direction and Final Importance Score
However, both F(oz)) and [G()))) can only quantify the strength of association between the occur- rence of otj and the correct rollout, they cannot quantify the direction of this association (i.e., positive or negative association). For the detailed proof, see Appendix [ To determine the direction of the association and further quantify token importance, adapt the Term Frequency calculation idea from BM25 [2 to compute standardized frequency scores for the token Oij within the set correct and incorrect rollouts.
Specifically, we concatenate all correct rollouts into a single long sequence and all incorrect rollouts into another. We compute the term frequency (tf) of the token Oi in these two concatenated sequences, denoted as (or)) and (o.)) Based on tfs we compute standardized frequency scores:
(k1 + 1) (fr))(o)))
T (TT)(
(8)
An (1 b b
leste
) + (fTT)(o)))
Tettere
Here, T/F refers to the correct or incorrect rollouts. lent and lenp are the average lengths of the concatenated correct and incorrect rollouts, respectively, and lenavg is the average rollout length across all G rollouts. iı and b are adjustable parameters controlling the influence of term frequency and length normalization (can be set empirically or tuned). Treating all correct/incorreet rollouts as single sequences reduces the impact of individual rollouts with extreme lengths.
The token directional score (oor)) combines a measure of effect size based on proportion differences and a measure based on standardized frequency score differences. We use Cohen's h effect size (aresin / aresin vy) metric to measure the difference in the proportion of correct rollouts
boy
versus incorrect rollouts (Bory does where oG appears. Simultancously, we consider the ratio difference of the standardized frequency scores. The final formula
Gosy Coys
Gos aresin
bosi
(TF((s)) + ha (Fr( TF
TFF(0.)
D(ois) = arrssin
(9)
Goo;
+
Cos3
boss
+ doe;
This combination aims to capture different aspects of importance: when the token o.y 's frequency is similar in correct and incorrect rHigh frequency generic tokens), its importance might be better reflected by the probability difference in where it appears, hence the dominance of the arcsin square root proportion difference term ( Cohen's 's h effect size); when the token oiy" 's frequency differs significantly (especially for low-frequency but critical tokens), the frequency ratio better reflects its discriminative power, increasing the importance of the frequency score ratio term.
Theoretically, both the Fisher score F(o:s) and the Information Gain IG((o;)) are strictly greater than zero, while the directional score D(o;)) spans the full real range (-oo, +oo). To derive the final token-level relevance score, we multiply the magnitude of correlation (e.g.. F(ors) ) or IG(o,;)) )) by
the directionality score D(o,;), ). which reflects whether the token is positively negatively associated with correct rollouts. Finally get key-token--e token-value of o;₃ is (h, F(o,,) + h2 IG(o,))) Do.. Positive key-token-values represent positive association direction.
To stabilize training and constrain the output range, we apply a sigmoid normalization to the resulting -ey-tokkn-values. These normalized values are then added to the rollout-level advantage computed by GRPO, thereby producing the final token-level advantage:
ÁKTAE ACRPO + a((1 F(o.;) + h22 IG(o;)) Do,)) 0.5
(10)
KTAE is an algorithm for estimating the advantage of tokens, which computes the key-token-vall value through the rollouts obtained from sampling. The complete implementation process is shown in
Algorithm|] T It is orthogonal to the improvement strategy of DAPO, and can be combined with DAPO in addition. An schematic diagram is shown in Appendix F

---

Algorithm 1 Key-token Advantage (Ktima)ion(KTAB)
Input: Set G rollouts {o1..... sampled from policy model, rule-based reward (R),... Ro
weighting parameter , hg
1: Calculate the rollout-level advantage of GRPO AGRPO
2: Summarize all tokens in rollouts into set O
3: Divide the G rollouts into FT and pp sets according to the reward R obtained by each rollout
4: for in O do
5: a a= a=Σ Doo E Exr). Σ I[ E mp).c= -= d Len (zp)
6: Calculate F()) according Eq. 4 and calculate TG(o) according Eq. 6 and Eq.
7: Weighted add F(o) and (G) to get quantized association strength /1 {() + h2 TG(o)
8: Calculate the frequency of in the correct and incorrect rollouts according to Eq.
9: Calculate the quantized association direction according to Eq.
Multiply the association direction and association strength to get the key-token-value of each token, and then add it to AGRPO get ÅKTAE
11: end for
Output: ÅKTAE
4 Experiment
Experiment Setting. Our validation and ablation experiments were conducted on the Qwen2.5- Math-1.5B (23 base model and the comparison experiment with baseline methods is based on Qwen2.5-Math- base model, using math 12k [12 and its subset math-level3-5 respectively. See Appendix G.I for specific details of dataset and benchmark, Appendix G.3 for implementation details and hyperparameters, and Appendix J for prompt details.
G.3
Method Validation Result. Experiments revealed several key performance trends (Fig, 4. KTAE consistently enhances MATH500 test accuracy when integrated with GRPO and DAPO, respectively, Regarding mean response length, the addition of KTAE significantly reduced the response length for both algorithms compared to their original versions. We believe that achieving improved model performance while simultaneously reducing generation cost is more meaningful. In terms of genera- tion entropy, GRPO+K BROK+KTHX KTAE showed accelerated entropy decrease early on but stabilized at a higher level later, beneficial for mitigating entropy collapse [I]. For DAPO+KTAL its entropy value was considerably higher than all other configurations and exhibited a continuous upward trend. While such high entropy contributes increased sampling diversity and avoids entropy collapse, it may also introduce a potential risk of reduced training stability.
Mocnn Responze Langtth (Smaarthing) (o)
DAPCHATAE
0.297
MATH500 Accursc(Smmorting
0.736
DAPO
0.264
884
GR1D ETAE
Generction Entrrop
0.231
0.704
0.188
748
680
0132
632
GRPOWETAE
0.099
68FO
544
10
20
30
50
60
70
80
40 Steps
50
60
70
80
40
Steps
Steps
Figure 4: The metric curves of test accuracy, mean response length, and generation entropy of combining DAPO and GRPO with KTAE.
Comparison with Baselines. In Table [ the DAPO+KTAE-TE model achieved the highest average score across the 5 benchmarks, significantly outperforming others on MATH500. Both GRPO+KTAE and DAPO+KTAE achieved better performance than all baseline methods on AMC (See Appendix G.2 for more details about baselines). GRPO+KTAF showed performance improvements compared to the original GRPO on four out of five benchmarks, with only slight decrease on AIME24 (Considering AIME24 has just 30 problems, this equates to only one fewer correct answer). Likewise,
7

---

compared to the original DAPO, DAPO+KTAE's performance improved or remained unchanged on four out of five benchmarks, experiencing slight decrease only on OlympiadBench. This demonstrates the effectiveness of the KTAE algorithm. This performance was consistent with the 1.5B model, where our model even surpassed R R1-Distill-ween-1. with the same base model.
Table The zero-shot greedy pass@1 performance of the 1.5B and 7B models across five mathemat- ical reasoning benchmarks. All the results above of our reproduction. refers to OlympiadBench: denotes the results from [24]. $ denotes the results from [25). @8k refers the max response length.
1,5B Models
AIME24
MATH-500
AMC
Minerva
Olympiad
Avg
Qwen2.5-Mathe -aath-1-BtB Instruct
10.0
74.2
48.2
26.5
40.2
39.8
Owen2.5-Math-11.5.
16.7
61.8
43.4
15.1
28.4
33.1
RI-Disiil--Wenn-11BBBB..
20.0
77.4
49.4
25.0
35.8
41.5
Oat-Zero-1.5B [25
20.0
74.4
50.6
23.9
37.0
41.2
GRPO-1.5B
16.7
76.0
51.8
22.1
36.3
40.6
GRPO+KTAE-1.5B
26.7
75.4
41.0
27.2
38.2
41.7
DAPO-1.5B
16.7
77.6
47.0
25.7
39.0
41.2
DAPO+KTAE-L.5)
20.0
77.6
50.6
29.0
40.0
43.4
7B Models
AIME24
MATH-500
AMC
Minerva
Olympiad
Avg
1rren.55-Mtt-IIsttee 226
13.3
79.8
50.6
34.6
40.7
43.8
Qwen2.5-Mathh
13.3
57.6
45.0
14.7
23.7
30.9
Murus-2-7B-PRIME 27
23.3
82.2
57.8
36.0
39.9
47.8
Oat-Zero-7B
25
36.7
79.8
61.4
30.5
41.8
50.0
OpenReasoner- Zero 7B
19
13.3
78.4
45.8
27.9
45.0
41.7
SimpleRL-Zero-7B 28
26.7
78.6
59.0
33.8
43.4
48.3
GRPO-7B
36.7
81.0
57.8
32.7
43.2
50.3
GRPO+KTAE-7B
33.3
82.4
65.1
33.8
43.7
51.7
DAPO-7B
36.7
81.8
60.2
34.5
45.3
51.7
DAPO+KTAE-7B
36.7
83.2
63.9
35.3
43.7
52.5
Table Solle2dene demonstrates that model also significantly reduces the length of the response without any length penalty reward. This effect is particularly pronounced for the 7B parameter model, where the GRPO+KTAE model exhibits considerably shorter generation lengths compared the baseline methods. This indicates that the KTAE algorithm enables the model to concentrate more effectively on key tokens that crucial for problem resolution, thereby curtailing the generation of redundant or non-essential tokens. That is to say, KTAE achieved the highest average score across the benchmarks while using the least token budget, demonstrating the highest reasoning efficiency.
As Table 3 KTAE is a model-free algorithm whose computational cost is largely independent of model size, depending primarily on the number of generated tokens. For larger models such the 7B variant, each training step inherently slower, making the relative efficiency loss introduced by KTAE less noticeable. In contrast, for smaller models like the 1.5B variant, the shorter training steps make KTAE'' overhead more apparent. Nevertheless, the computational cost of KTAE remains acceptable compared methods that rely model- computation.
In practice, KTAE's runtime mainly depends on implementation efficiency. CPU-only serial implementation would require several hours, the current version employs only limited tensor parallelism across N rollouts per sampling. resulting in low data parallelism and GPU utilization below 1 %. Further optimization
concatenating computations into larger tensors to better exploit GPU capabilities could substantially improve efficiency. We plan to pursue these engineering optimizations in future work.
Ablation Analysis. Figure the impact of each KTAE component. Removing any component consistently reduced test accuracy. Excluding IG had the largest negative effect accuracy and

---

Table The response length of the 1.5B and 7B models across five mathematical reasoning bench-
marks. All the results above are of our reproduction. refers to OlympiadBench.
1.5B Models
AIME24
MATH-500
AMC
Minerva
Olympiad
Avg
Oat-Zero-1.5B
1198
878
652
692
938
871.6
GRPO-1.5B
1299
635
908
731
958
906.2
GRPO+KTAE-L.BB
1187
884
617
663
890
848.2
DAPO-1.5B
1218
617
950
712
937
886.8
DAPO+KTAE-1.5B
1110
983
582
666
861
840.4
7B Models
AIME24
MATH-500
AMC
Minerva
Olympiad
Avg
Eurus- Eurus-2-7B-PRE 2-7B-PRIME
1498
685
1099
777
1077
1027.2
Oat-Zero-7B
977
658
903
677
892
821.4
OpenReasooner-Zero-77
2300
1193
1901
1269
1871
1706.8
SimpleRL-Zero-7B
1074
634
832
584
881
801
GRPO-7B
989
606
806
641
813
771.0
GRPO+KTAE-7B
941
563
741
577
771
718.6
DAPO-7B
1155
676
969
700
986
897.2
DAPO+KTAE-7B
1013
604
864
607
798
777.2
Table 3: Comparison of the training time of one step of the KTAE algorithm and baseline algorithm under the same hardware conditions (8 NVIDIA H100 80G)
Algorithms
Qwen2.5-7B-MATH Qwen2.5-1.5B-MAN.
GRPO
559.36s
277.69s
GRPO+KTAE
641.98s
371.79s
DAPO
1006.30s
363.36s
DAPO+KTAE
1159.00s
642.74s
produced the shortest sequences. contrast, removing F or tff decreased accuracy while increasing sequence length, though lengths remained shorter than those of the GRPO baseline. For entropy, removing If initially caused notable increase. Importantly, while GRPO suffered from entropy collapse, KTAE avoided this. Overall, IG is key for accuracy and brevity, 1.f supports diversity and stability, and F contributes to overall performance. All components are essential in accuracy.
Visualization Example. Beyond the example illustrated in Figure 2 also observed several zero-reward rollouts. These rollouts are characterized by including the correct answer in their early stages, but subsequently generating large number of redundant tokens, ultimately leading the correct answer being obscured by the subsequent sequence. Appendix [H provides a such example, where clear boundary can be distinctly observed, effectively separating the correct answer from the redundant tokens. This further validates KTAE's accuracy in identifying key tokens. As shown in Appendix K also observed the 'aha moment' phenomenon during the KTAE training process.
5 Related Work
Large Reasoning Models. enable new era of test-time scaling [3 32 and Himann-lkee stepwise reasoning. DeepSeek R1 2 used pure RL induce long Chain-of-Thought (CoT) and self-refiection. Following R1, subsequent explored RL training variants, mainly smaller models. While
Breakthroughs
29 30 in Large Language Models (LLMs)
work [1] 19 28 27
25
33
27 25), exploring more fine-grained reward GRPO challenging. This work proposes token-leve advantage estimation for GRPO and its variants.
RI's paradigm was replicated

---

0.741
Moom Resporte Length (Snggh (Sm2o0Pimm))))
GEPONETAE w/s Fidee COPO RTAE
0242
SIFONTAAE w/s Faten
T14000 Acuroey(rrachrir
BRFONTAAE wie
0.722
748
GAPO KTA
DAPONCTAE 1/
GAPONKTAE
Gonrafion Entrery
0.88
GAPON<TAE
GAPO
6APO
0.664
0.54
0.640
GPPONT WE w/S Fishe
GRPCHTTAL w/S 11
0.67
ORTONTTAE w/ 20
575
GRRCHETAE
0.600
1800
Steps
Figure 5: Training results after removing different components from KTAE.
Reinforcement Learning. RL is key for sequential decision-making, using policy gradient methods. Early methods (e.g., REINFORCE (34, DPO (35)) had high variance. TRPO (0 and PPO [TO] improved stability with constrained/clipped updates, though PPO is costly. GRPO 9 removed
the Critic using group statistics. GRPO variants, like DAPO [T] and Dr.GRPO
25 built on this.
However, GRPO and variants use uniform rollout advantage, ignoring token importance in reasoning.
We propose Key-token Advantage Estimation, linking tokens to correctness statistically, for finer granularity. More related work in Appendix[]
6 Conclusion
This paper introduces KTAE, an algorithm uses statistical analysis to quantify each token's 's association to correct rollouts. By combining this with GRPO's rollout-level advantage, KTAEE computes token-level advantages, thereby providing more fine-grained optimization signals and significantly improving training effectiveness. It requires new models, adds minimal computational overhead, and avoids reward hacking. KTAE can effectively identify the importance of different tokens in the rollout, making the model pay more attention key tokens in the training process, showing excellent test performance utilizing the minimum token budget. Theoretically, the core idea of the KTAE can be applied to many other reasoning domains. Therefore, KTAE still holds significant potential.
Acknowledgements
We thank our colleagues Jianghao Chen, Chong Li, Tengxiao Xi, Tianyu Peng, Xingquan Zhang, Boyu Guan, Jiawei Guo for their insightful and constructive feedback. We thank Qian Li and thank all reviewers for their valuable
Zhenggang Piao for their special assistance. In addition, comments and recognition of our work. This work is supported by National Key R&D Program of China 9022ZD0160610 and the Strategic Priority Research Program of Chinese Academy of Sciences under Grant XDA04080400.
References
[1] Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Camey, al. Openai ol system card. arXiv preprint arXiv:2412. 2024.
[2] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint 2rXiv:2501..1944 2025.
[3] Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, et al. Kimi kl. 5: Scaling reinforcement learning with llms. arXiv preprint
2025.
[4] Qwen Team. Qwq-32b: Embracing the power of reinforcement learning, March 2025.
[5] Trieu H Trinh, Yuhuai Wu, Quoe Le, He He, and Thang Luong. Solving olympiad geometry without human demonstrations, Nature, 677(7995) 2024.
10

---

[6] Yixin Ye, Zhen Huang. Yang Xiao, Ethan Chemn Shijie Xia, and Pengfei Liu. Limo: Less is more for reasoning. arXiv preprint arXiv:2502.03387, 2025.
[7] Niklas Muennighoff, Zitong Yang. Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer. Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling. arXiv preprint urXiv:2501.19393, 2025.
[8] Rosie Zhao, Alexandru Meterez, Sham Kakade, Cengiz Pehlevan, Samy Jelassi, and Eran Malach. Echo chamber: RI post-training amplifies behaviors learned in pretraining. arXiv preprint mTXiv:250..07912, 2025.
[9] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song. Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.033000 2024.
[10] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
[11] Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, al. Dapo: An open-source Ilm reinforcement learning system at scale. arXiv preprini arXiv:2503. /4476, 2025.
[12] Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference Learning Representations, 2023.
[13] Guoxin Chen, Minpeng Liao, Chengxi Li, and Kai Fan. Step-level value preference optimization for mathematical reasoning, 2024.
[14] Wei Sun, Qianlong Du, Fuwei Cui, and Jiajun Zhang. An efficient and precise training data construction framework for process-supervised reward model in mathematical reasoning, 2025.
[15] Peiyi Wang, Lei Li, Zhihong Shao, R. X. Xu, Damai Dai, Yifei Li, Deli Chen, Y. Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce llms step-by-step without human annotations, 2024.
[16] Lunjun Zhang, Arian Hosseini, Hritik Bansal, Mehran Kazemi, Aviral Kumar, and Rishabh Agarwal. Generative verifiers: Reward modeling next-token prediction, 2025.
[17] Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization, 2022.
[18] John Schulman, Xi Chen, and Pieter Abbeel. Equivalence between policy gradients and soft q-leaming, 2018.
[19] Jingcheng Hu, Yinmin Zhang, Qi Han, Daxin Jiang, Xiangyu Zhang, and Heung- Yeung Shum. Open-reasoner zero: An open source approach to scaling up reinforcement learning on the base model. arXiv preprint arXiv:2503.24290 2025.
[20] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pages 1889-1897. PMLR, 2015.
[21] Alan Agresti. A survey of exact inference for contingency tables. Statistical science, 7(1):111- 153, 1992.
[22] Stephen E Robertson, Steve Walker, Susan Jones, Micheline M Hancock-Beaulieu Mike Gatford, et al. Okapi at trec-3. Nist Special Publication Sp, 109:109, 1995.
[23] An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, Keming Lu, Mingfeng Xue, Runji Lin, Tianyu Liu, Xingzhang Ren, and Zhenru Zhang. Qwen2.5-math technical report: Toward mathematical expert model via self self-improvemmnnt arXiv preprint (tt.:24409.1222222. 2024.
11

---

[24] Quy-Anh Dang and Chris Ngo. Reinforcement learning for reasoning in small Ilms: What works and what doesn't. arXiv preprint arXiv:2503. /6219, 2025.
[25] Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin. Understanding rl rl-zero-like training: A critical perspective. arXiv preprint arXiv:2503.20783, 2025.
[26] An Yang. Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, et Qwen2. 5-math technical report: Toward mathematical expert model via self-imprrvement. arXiv preprint arXiv:2409. 1XX1:2209..11122, 2024.
[27] Ganqu Cui, Lifan Yuan, Zefan Wang, Hanbin Wang, Wendi Li, Bingxiang He, Yuchen Fan, Tianyu Yu, Qixin Xu, Weize Chen, al. Process reinforcement through implicit rewards. arXiv preprint arXiv:2502.014555 2025.
[28] Weihao Zeng, Yuzhen Huang, Qian Liu, Wei Liu, Keqing He, Zejun Ma, and Junxian He. Simplerl-zoo: Investigating and taming zero reinforcement learning for open base models in the wild, 2025.
[29] Marah Abdin, Sahaj Agarwal, Ahmed Awadallah, Vidhisha Balachandran, Harkirat Behl, Lingjiao Chen, Gustavo de Rosa, Suriya Gunasekar, Mojan Javaheripi, Neel Joshi, et al. Phi-4- reasoning technical report. arXiv preprint MKKiv:2504.233// 2025.
[30] Zhong-Zu Zhing- Li, Duzhen Zhang. Ming-Liang Zhang, Jiaxin Zhang, Zengyan Liu, Yuxuan Yao, Haotian Xu, Junhao Zheng, Pei-Jie Wang, Xiuyi Chen, et al. From system to system 2: A survey of reasoning large language models. arXiv preprini arXiv:2502. 17419, 2025.
[31] Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoe V. Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling, 2024.
[32] Hritik Bansal, Arian Hosseini, Rishabh Agarwal, Vinh Q. Tran, and Mehran Kazemi. Smaller, weaker, yet better: Training Ilm reasoners via compute-optimal sampling, 2024.
[33] Liang Wen, Yunke Cai, Fenrui Xiao, Xin He, Qi An, Zhenyu Duan, Yimin Du, Junchen Liu, Lifu Tang, Xiaowei Lv, et al. Light-rl: Curriculum sft, dpo and rl for long cot from scratch and beyond. arXiv preprint arXiv:2303.
2025.
[34] Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8:229 2266, 1992.
[35] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 66:53728-537411 2023.
[36] Jia LI, Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Costa Huang, Kashif Rasul, Longhui Yu, Albert Jiang, Ziju Shen, Zihan Qin, Bin Dong, Li Zhou, Yann Fleureau, Guillaume Lample. and Stanislas Polu. Numinamath. [https: htttps://huggingaaee htt.g/aaa.........aa////aaaaaaaa (htts::///ottttt:://// com/
project nucinaa-uumiaaaa/e http://www...........eee//eee/ea dataaattdttttttaaaaaaa pdf)
2024.
[37] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.0383 2021.
[38] Aitor Lewkowyez, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. Solving quantitative reasoning problems with language models. Advances in Neural Information Processing Systems, 35:3843-388 2022.
[39] Zhen Huang, Zengzhi Wang, Shijie Xia, Xuefeng Li, Haoyang Zou, Ruijie Xu, Run-Ze Fan, Lyumanshan Ye, Ethan Chern, Yixin Ye, et al. Olympicarena: Benchmarking multi disti-dis cognitive reasoning for superintelligent ai. Advances in Neural Information Processing Systems, 37:19209-19253, 2024.
12

---

[40] Guangming Sheng. Chi Zhang, Zilingfeng Ye. Xibin Wu, Wang Zhang. Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. Hybridflow: flexible and efficient rlhf framework. arXiv preprint arXiv: 2409. 19256, 2024.
[41] Jianghao Chen, Pu Jian, Tengxiao Xi, Dongyi Yi. Qianlong Du, Chenglin Ding, Guibo Zhu, Chengqing Zong, Jinqiao Wang, and Jiajun Zhang. Chinesewebtext: Large-scale high-quality chinese web text extracted with effective evaluation model. arXiv preprint arXiv:231 2..0::3111.111 1.01149, 2023.
[42] Pu Jian, Donglei Yu, and Jiajun Zhang. Large language models know what is key visual entity: An llm-assisted multimodal retrieval for vqa. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 10939-10956, 2024.
[43] Pu Jian. Junhong Wu, Wei Sun, Chen Wang, Shuo Ren, and Jiajun Zhang. Look again, think slowly: Enhancing visual reflection in vasion-language models, 2025.
[44] Pu Jian, Donglei Yu, Wen Yang, Shuo Ren, and Jiajun Zhang. Teaching vision- -annuage models to ask: Resolving ambiguity in visual questions. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguisties (Volume 1: Long Papers), pages 3619-3638,
2025.
[45] Yuchen Yan, Yongliang Shen, Yang Liu, Jin Jiang, Mengdi Zhang, Jian Shao, and Yueting Zhuang. Inftythink: Breaking the length limits of long-context reasoning in large language models. arXiv preprint ar wwxiv:2503.0669) 2025.
[46] Xuezhi Wang. Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXivv2203.11117, 2022.
[47] Yiwei Li, Peiwen Yuan, Shaoxiong Feng, Boyuan Pan, Xinglin Wang, Bin Sun, Heda Wang, and Kan Li. Escape sky-high cost: Early-stopping self-consistency for multi-step reasoning.
[48] Shuo Ren, Pu Jian, Zhenjiang Ren, Chunlin Leng, Can Xie, and Jiajun Zhang. Towards scientific intelligence: A survey of llm-based scientific agents. arXiv preprint arXiv:2503.24044 2025.
[49] Yuhuai Wu, Elman Mansimov, Roger B Grosse, Shun Liao, and Jimmy Ba. Scalable trust-region method for deep reinforcement learning using kronecker-factrred approximation. Advances in neural information processing systems, 30, 2017.
[50] Xingyu Chen, Jiahao Xu, Tian Liang, Zhiwei He, Jianhui Pang, Dian Yu, Linfeng Song. Qiuzhi Liu, Mengfei Zhou, Zhuosheng Zhang, et al. Do not think that much for 2+ 3=? the overthinking of ol ol-like llms. arXiv preprint arXiv:2412.21187, 2024.
[51] Silei Xu, Wenhao Xie, Lingxiao Zhao, and Pengcheng He. Chain of draft: Thinking faster by writing less. arXiv preprint 2rXiv:2502.11660 2025.
[52] Yang Sui, Yu-Neng Chuang, Guanchu Wang. Jiamu Zhang, Tianyi Zhang, Jiayi Yuan, Hongyi Liu, Andrew Wen, Shaochen Zhong, Hanjie Chen, et al. Stop overthinking: A survey on efficient reasoning for large language models. arXiv preprint wrXiv:2503.16411 2025.
[53] Tengxiao Liu, Qipeng Guo, Xiangkun Hu, Cheng Jiayang, Yue Zhang, Xipeng Qiu, and Zheng Zhang. Can language models learn to skip steps? arXiv preprint arXiv:2411.01855, 2024.
[54] Tingxu Han, Zhenting Wang, Chunrong Fang, Shiyu Zhao, Shiqing Ma, and Zhenyu Chen. baken-bbuget-aware I1m reasoning. arXiv preprint arXiv:2412.18547 2024.
[55] Abhinav Kumar, Jaechul Roh, Ali Nasch, Marzena Karpinska, Mohit lyyer, Amir Houmansadr, and Eugene Bagdasarian. Overthinking: Slowdown attacks reasoning llms. arXiv preprint arXiv:2502.02542 2025.
[56] Bairu Hou, Yang Zhang, Jiabao Ji, Yujian Liu, Kaizhi Qian, Jacob Andreas, and Shiyu Chang. Thinkprune: Pruning long chain-of-thought of llms via reinforcement learning. arXiv preprint arXiv:2504.01296 2025.
13

---

[57] Pranjal Aggarwal and Sean Welleck. L1: Controlling how long reasoning model thinks with reinforcement learning. arXiv preprint arXiv:2503.04697, 2025.
[58] Haotian Luo, Li Shen, Haiying He, Yibo Wang. Shiwei Liu, Wei Li, Naiqiang Tan, Xiaochun Cao, and Dacheng Tao. Ol Length-harmonizz fine-tuning for ol-like reasoning pruning. arXiv preprint arXiv:2501. 12570, 2025.
[59] Yuxiao Qu Matthew YR Yang. Amrith Setlur, Lewis Tunstall, Edward Emanuel Beeching, Ruslan Salakhutdinov, and Aviral Kumar. Optimizing test-time compute via meta reinforcement fine-tuning, arXiv preprint arXiv:2503.07572. 2025.
[60] Yige Xu, Xu Guo, Zhiwei Zeng, and Chunyan Miao. Softcot: Soft chain-of-thought for efficient reasoning with Ilms, arXiv preprint arXiv:2502.12334 2025.
[61] Yuyang Wu, Yifei Wang. Tianqi Du, Stefanie Jegelka, and Yisen Wang. When more is less: Understanding chain-of-thought length in llms. arXiv preprint arXiv:2502.07266 2025.
[62] Jianghao Chen, Wei Sun, Qixiang Yin, Lingxing Kong, Zhixing Tan, and Jiajun Zhang. Ace-rl: Adaptive constraint-enhancee reward for long-foor generation reinforcement learning, 2025.
[63] Jianghao Chen, Zhenlin Wei, Zhenjiang Ren, Ziyong Li, and Jiajun Zhang. Lr 2 bench: Evaluating long-chain reflective reasoning capabilities of large language models via constraint satisfaction problems. arXiv preprinl arXiv:2502. 17848, 2025.
[64] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforce- ment learning. In International conference on machine learning, pages 1928-1933 1937. PmLR, 2016.
[65] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5600 2013.
[66] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference machine learning, pages 1587-1596. PMLR, 2018
[67] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. Pmlr, 2018.
[68] Timothy P Lillicrap. Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv: 1509002971, 2015.
[69] Ziniu Li, Tian Xu, Yushun Zhang, Zhihang Lin, Yang Yu, Ruoyu Sun, and Zhi- Zhi-Quan Luo. Remax: A simple, effective, and efficient reinforcement learning method for aligning large language models. arXiv preprint 2xXx:2310.1150: 2023.
14

---

Appendix
Limitations
16
B Broader Impacts
16
Future Work
16
D Calculating factorials using logarithmic space
16
Why quantify association direction?
17
A schematic diagram of KTAE
18
Implementation Detials
18
G
Dataset and Benchmark
18
G.2 Baselines
19
G.3
Implementation Details and Hyperparameters
19
H
Case Study
19
More Related Work
19
Large Reasoning Language Models
19
1.2 Reinforcement Learning
21
J Prompt
21
Aha Moment
21
Dataset License
15

---

A Limitations
Building upon GRPO, the KTAE algorithm introduces more fine-grained token-level advantage. which guides the model to focus more on key tokens, thereby demonstrating significant performance improvement without introducing additional models. However, our experimental validation primarily focused models with 1.5B and 7B parameters, and the performance of KTAE on larger- scale models has not yet been fully verified. Furthermore, while theoretically KTAE is applicable to any task beyond based on rule-based rewards, this paper only provides in-depth analysis and experimental validation on mathematical reasoning tasks, and potential for application broader domains requires further exploration.
B Broader Impacts
Mathematics represents a pinnacle of human wisdom and serves as the foundation of many scientific disciplines. Our approach aims to empower large language models tackle complex mathematical reasoning problems, bringing their capabilities closer to human expert-level intelligence. By doing so, we seek to advance the development of large language models across scientific fields and support human efforts in driving scientific progress. A common limitation among current large reasoning language models is that while their reasoning capabilities are enhanced, they tend to sacrifice some general abilities like summarization, abstracting, and translation. Consequently, these models may become specialized models primarily focused on the reasoning domain.
C Future Work
The experiments presented this paper primarily focus on tasks based on binary discrete rewards, such as in the fields of mathematics and code generation, where the reward is simply classified as either correct or incorrect. Future work will explore how extend the core idea of KTAF scenarios with multiclass discrete rewards, and even to those with continuous reward values.
Furthermore, although KTAE quantifies the importance of each token provide fine-grained opti- mization signals, the amount of information carried by single token is often insufficient within complete reasoning path. This implies that, even for human experts, it difficult to determine whether the presence of a single token has a decisive impact on the final outcome (correct or incorrect), or whether its absence would necessarily lead failure. Future work will aim address or optimize this problem, for instance.
D Calculating factorials using logarithmic space
The Gamma function is defined by the following integral:
r(2) = r'eee ((t)=-((tee (Re(2)> > 0).
When is a positive integer, this integral reduces directly to the factorial expression. Starting from the integral definition of the Gamma function:
[(2 + 1) = 1 t' dt.
Let = t², du = dtt then du = zt² dt, and v= = -t Applying integration by parts:
T(2 + 1) [- + 1or 1-- 1=- e-le-dd dt.
Ast + oo, -t - 0; as 0, - + 0 (since > 0). Therefore, the boundary term vanishes, and we obtain:
I( + 1) 22).
This recurrence relation is consistent with the factorial identity n! n (n 1)!.
16

---

Fisher's exact test:
(Bog tboss)
(Corcors
+duy)
Boss
Fisher(ois)
Jass tross
Of which:
(ooss boss
(aors + Doyg) (In P(ani +Boyy +1] P(ensy +1)- Fibory +1))
aou
Oos, !bo.,
The above process converts factorial operations into addition and subtraction of the In I function, which can be efficiently computed in parallel using forch.lgamma(", This significantly improves both computational efficiency and numerical precision.
E
Why quantify association direction?
011
coorect
iincrreett
12
Xcorrect
Xincorrect
Appeared Times
Appeared Times
a
b
b
a
Not
Not Appeared
Appeared
C
d
c
Times
Times
Figure 6: Example contingency table after changing position.
In the two contingency tables presented above in Figure 6 [a, bj and [c, d] have been interchanged. It evident that, for these two tokens p. the direction of association with 'obtaining correct rollout" becomes completely opposite association metrics calculated by Equations follows. This indicates that methods such Fisher's 's exact test and Information Gain (IG) can only quantify the strength of association between the token and "obtaining a correct rollout', but fail to reveal its direction. Therefore, we propose the use of Equation 9 to accurately quantify the directionality of this association.
result of this interchange. However, it is noteworthy that the and remain identical. The proof process as
(ath)) ( (ctd)
(a b)((c d)t(a + c)!(b dd)!
Fisher(on))
(ate)
abbleld!
(+a) (dte)
d)!(a (b (b+a)l(d++ d) +
Fisher(o12)
(0+a)
blaldlc!N!
FFiher((on)
a+cc
b + d
b d N b
IG(on))
logz
loga
N
b
log2
log2
N
a +
a a+1
a +
a a+1
d
logz
log2
N
c c
c
c c
c d
+ d
b d
log
loga
N
N
N
a
b
d
d
log
loga
loga
loga
+
b
N
d
N
+ d

---

b+ d
b d
+
IG(012)
log2
loga
N
N
N
N
b
b
b
log2
log2
N
b+
b+ a
b +
b a
d
d
N
d
log2
d
c
d+c
loga
ddc
b+
d
+
d
a+c
log2
log2
N
N
N
N
b
d
d
loga
log2
loga
loga
N
b+0
b+a
c
N
d+c c
IG(on1)
F A schematic diagram of KTAE
01
01
0.08
0.00
9.38
R:
0, SG
O2
B. S4
SW
Ai
O2
0.18
0.28
0.23
0.26
0. BB
Question
Question
0g
-1.23
-1.35
-1
-L
-1
-1.33
-2.21
-1.90
-L
-1.1G
-4,
04
8.
0.56
0.30
0.56
3:50
0.58
04
0.28
0.78
0.62
0.
0.6L
0.21
Anu
Ana
Aos
Aou
Aou
Aor
Aou
Aoa
Aos
Aoi
CRPO
R4
mean (R)
KTAB
A.
GRRO
Key fokentoken_vull value
(R.
GRPO
GRPO-KTAE
Figure 7: An example of comparing KTAE with vanilla GRPO.
As shown Figure / policy model sampled 4 rollouts, among which only 03 final answer was incorrect, while the remaining 3 rollouts all obtained correct results. Each square in the figure represents a token. The left side of the figure displays the advantage value calculated by GRPO for each token. It can be observed that not only are all correct rollouts assigned the same advantage value, but within the same correct rollout, every token is also assigned exactly the same advantage value. This reflects GRPO's characteristic of performing evaluation at the rollout level. Building upon this, when we superimpose the key-token value calculated by KTAE (as shown on the right side of the figure), each token quantified with different importance score, thus significantly differentiating the contribution levels of various tokens within the rollout. This provides finer grained optimization information compared GRPO.
G Implementation Detials
G.1 Dataset and Benchmark
For the initial validation phase of our method, we first utilized the widely-used MATH dataset, specifically the MATHI2K [2] subset as the training set and its corresponding MATH500 as the test set. We conducted experiments on the Qwen2.5-Math-L1. verified the effectiveness of our proposed KTAE method. Subsequently, to enable a more fair and comprehensive comparison with existing baseline methods, and following our initial validation of the method's 's effectiveness, we decided use a more challenging subset of the MATH dataset (specifically problems from Levels 3-5) as the training set, while still using MATH500 as the test set. Under this setup, we trained our model on the larger Qwen2.5-Math-77 Base model 23 To comprehensively evaluate the mathematical reasoning capabilities of the KTAF-7B model, we selected five prominent and widely recognized benchmarks in the field of mathematical reasoning for testing: AIME24 (36), MATH-500 37 12). AMC Minerva 38 and OlympiadBench(39)
Base model[23), which successfully
18

---

G.2 Baselines
In the method validation phase, we aim to comprehensively evaluate the performance of our KTAE. To this end, we first compare them against the foundational GRPO[ and [APO[[I] algorithms to quantify the performance gains introduced by our KTAE mechanism. Furthermore, to evaluate the model trained with our KTAF against existing reinforcement learning training techniques, we also selected the following representative approaches for comparison: 1. Maapll-RR-Zooo222 A bascline method trained on the Owen2.5-Matth- base model using the math-leve13-5 dataset, employing An online process reward model (PRM) update method, characterized by its ability to enable online PRM updates using only policy rollouts and outcome labels through implicit process rewards. 3. OpenReasoonerZero| A zero-RL method based on the Qwen2.5-7B base model, which centrally applies the vanilla PPO algorithm. 4. Oat-Zero [25]: Trained starting from the Swen2.5Math-7B model and utilizing a rule- based reward. It employs improved Dr.GRPO algorithm, which removes the standard deviation in GRPO advantage computation and token-level normalization in policy loss computation. These comparison methods encompass applications of basic RI algorithms, methods based on process rewards, and improved algorithms tailored for specific tasks (such as mathematical reasoning), aiming to evaluate the effectiveness and advancement of our methods from multiple perspectives.
the basic GRPO algorithm and a rule-based reward. 2. PRIME-
G.3 Implementation Details and Hyperparameters
Our KTAE-7B model was trained based on the Qwen2.5-Maie -rten.55-M base model, employing combined approach of DAPO and KTAE. The training utilized the VerL[44] reinforcement learning framework for optimization. During the training process, the model inherited the maximum context length of 4096 from the base model. The specific training hyperparameters were configured as follows: The maximum generation length was set to 3072, and the maximum prompt length was set to 1024. The sum of these two values aligns with the model's maximum context length. The learning rate was fixed at le-6. The training batch size 1024 questions. The number of rollouts sampled per question (G) was set to 16. The sampling temperature was 1.0. For the DAPO method, the clip low redio and clip high redio hyperparameters were set to (0.2, 0.28) in 1.5B models and (0.2, 0.24) in 7B models. respectively. The overlong buffer length was set to 512, and the length penalty coefficient was 1.0. The three hyperparameters h1, h2, and hy for the KTAE method were all set to 1.0, 2.0, 1.0, k and b in Eq 8 is set to 2.0, 0.5. To ensure reproducibility of the experimental results, all random seeds conducted used were set to O. Furthermore, for method validation preliminary experiments, additional training on the 1wen2.5-MMth--1550 base model. The hyperparameters for this training were largely consistent with the 7B model configuration. All our experiments were performed on NVIDIA H100 80G GPUs
H
Case Study
Similar scenarios are observed in the other two examples in Figure 8 In these rollouts, the model also generated portions containing the correct answer, but subsequently produced additional or incorrect content, ultimately obscuring the correct result and leading to a zero reward. In such cases. the KTAE algorithm accurately identifies segments within the rollout that contain the correct answer and evaluates them as having a positive contribution, while assessing incorrect or distracting segments as having a negative contribution.
I More Related Work
I.1 Large Reasoning Language Models
in Large Language Models (LLMs) and Vision-Language Models [42 43 44 have ushered new era of test-time scaling 31 32 46 47 48), enabling models to simulate human-like like, stepwise reasoning processes. OpenAl's O1 I introduces profound paradigm shift, demonstrating that extending the length of each chain can significantly enhance model reasoning performance. DeepSeek RI 2 employed pure RL with rase- reward, guiding LLMMs toward the spontaneous emergence of long Chain-of-Thought (CoT) and self-reflection behaviors. This work established an RL training paradigm starting from base model and open-sources both its
Recent breakthroughs
19

---

Question:
The perimeter a particular square circumference of particular circle are equal What is the ratio of the area the square to the area of the circle?
Express your answer as common fraction in terms of
Answer:
90
30
Faw
murdy
Figure An example visualization of KTAE calculation results.
20

---

training algorithm (GRPO) and model weights, Following the success of RL training demonstrated have explored various RL training algorithms, predominantly focusing on the smaller Qwen2.5 series models. At the same time, lot of work has
by R1, subsequent efforts
27
25
49
been done to alleviate the overthinking 50
51
52
53
54
55
problem of L.R and the problem While this line of work successfully
of too long generation length [56, 57 replicated the RL training paradigm in open-source models, the exploration
58 59
62
63
more fine-grained reward signals in GRPO remains an open challenge. This work introduces novel perspective token-level advantage estimation to enable the seamless integration of GRPO and variants.
1.2
Reinforcement Learning
RL is key for sequential decision-making, using policy gradient methods (On-policy 20 64 49) Off-policy 665
(.)). Early methods (e.g., REINFORCE [34], DPO 135 ReMax 69 had
high variance. Subsequent work focused on stabilization, leading to TRPO 20 and PPO [00, which use constrained updates and clipped objectives, respectively. PPO's reliance on separate Critic model increases computational cost. GRPO [9] addressed this by removing the Critic and using group-level statistics, GRPO variants, such as DAPO [Т] (addressing scaling and entropy collapse) and Dr. GRPO [25] (simplifying GRPO), have built upon this. However, GRPO and its variants assign uniform rollout-level advantage, overlooking token-specific importance in reasoning. To improve this granularity, propose Key-token Advantage Estimation, utilizing statistical methods to quantify the association between individual tokens and rollout correctness.
Prompt
As shown in Figure 9 and validation.
we use the same prompt template (Qwen-Math template) for both RL training
Prompt templates of RL training and Evalutaion
<|im start/ssystem/nPlesss reason step by step, and put your final answer
within \\boxed()<<|im end |m start |>user\n(questionkk im end |>\n<| im start >assistantln
Figure 9: Prompt template in RL training and validation.
K
Aha Moment
As shown in Figure 10 the model's output "we need to find the other solution." and solve the question successfully.It demonstrates a phenomenon similar to the 'Aha moment" mentioned in the Deepseek-RI paper. However, GRPO did not show the "aha moment' phenomenon and answered the question incorrectly. indicating that KTAE has developed a certain degree of self-reflection and error correction capabilities.
L Dataset License
Our training data based on the MATH dataset [37] and we evaluate our model on AIME24 [6), MATH-500 37 12,, AMC (36]. Minerva (38. and OlympiadBench (39]. We strictly adhere to the licenses associated with each dataset.
21

---

Question:: An eculstoul Hangle insorteed cenboiia
For such that one of the varions Hangle concious with vertce the pars saka lungth of thi coulatrr trangle.
Answer 16y3
Girro solution:
GRAOOKTAE solution:
poolen, need find eide of equiaaer stargaa insobbed partoolo with one of oetioi vorto the paraods, wichh a (5,00
soks proble
langt of equlsaee Hongga insorted pariolll
wdl werica vitto cfite partials which in (0,0)
Latk other ve ciw the aquikten sirrge (aa and (e% si>
Lat's der
vertows of aquisiennl tangie as (auvau) (rau
*j fay. and x] - tay.
o| fat and by..
The cittace betaoe (3, go one (ouu inl the sano desteno betwern (8.0) and (audull
The estance (8 g) and you au sares as to deker botween (6, and (Rault))
Let's beeck the atdo bengh the truliiirn tiaegle by Then hove
Let's dende the aide kergh of erpaution thue by than home
of
Bay
by go and
and
Sinsihing the
This
appiibeent dikka belnwwe the por muat veuul.
myp (112)
The decafecce
B B
Suberlluing
Brnce of and of tyy wettttt and 2o
Squaiing (it 2i/? Lifatt
V/e +od"+/1 pu3.
affe] (a) diles =2f -af7.
Expanding ripioandd get Aauuis
Bay bu/buent 2//2
Cancel ng but aj and ola)) fro bolidde wereet 1200yy Mlag 2039j
Recnoniing was ger
Aha moment
Since the sanggi aquitaarnt we
Thay hame
(900y- (alyy 2200 38M 11000 P0 the
+Bi-3 def (4) 1280
This rquation surthee we noad food otherr andutton we our enaeee
to atand and veriows didrict
butel tulri et the scaution get
1 66m) 68/4 6.
dei
=] Gtr] 250e]
Reananging berna get 189a
+N1e] 188) -n
16/3.
Figure 10: An example of 'Aha moment'.
22

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justifiicton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We summarize the main contents of the paper in the abstract and introduction, and introduce our proposed methods and contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
23

---

Justification: We mentioned the current limitations of our method in Appendix A For example, due limited computing resources, we have not tested on models with parameter size greater than 7B; we have not verified it in fields other than mathematics. Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiscless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide a full set of assumptions and complete (and correct) proofs for each theoretical result through a large number of comparative experiments and ablation experiments.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
24

---

Our paper presents our method and implementation details, disclosing enough information reproduce the main experimental results of the paper.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We upload the full code in the supplemental material and will open source our model weights later.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
mublic/guides http://ddee/oeeuubiisseeseeeeees
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
25

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes)
Justification: We present the implementation details and parameter settings our experi- ments in detail in Appendix G.3
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: We did not perform error analysis, but our experiments used a fixed random seed and tested temperature 0 to minimize the experimental error caused by randomness.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
26

---

Justification:
We describe the computational resources used in the experiments in Appendix
G33
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:///errrr...cc/b//b////eeeee
Answer: [Yes]
Justification: The research conducted in this paper complies all respects with the NeurIPS Code of Ethics
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Our paper discuss potential positive impacts in Appendix B
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
27

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper does not release the data or models that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The creators original owners of assets used in the paper (e.g., code, data, models) are appropriately acknowledged, and the licenses and terms use are clearly mentioned and appropriately respected in Appendix for details.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
The code used the article is well documented and uploaded with the code in the supplementary material.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodel as part of their
submissions via structured templates. This includes details about training, license, limitations, etc.
28

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMSs if it is important, original, or non-standard component of the core methods in this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development our paper does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer LLM policy for what should or should not be described.
29


Reviews:
[{"review_id": "Kr44ZraDv5", "reviewer": "Reviewer_3Vhm", "summary": "This paper proposes KTAE: A Model-Free Algorithm to Key-Tokens Advantage Estimation in Mathematical Reasoning | OpenReview, 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. Empirical results show that model trained with GRPO+KTAE and DAPO+KTAE outperform baseline methods across five mathematical reasoning benchmarks.", "questions": "- Can this approach generalize to tasks beyond mathematical reasoning?", "limitations": "YES", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Pros:\n- This paper is well-written and easy-to-understand. The idea is intuitive, reasonable, and effective.\n- The importance of individual tokens within a sequence to the final outcome is a very important problem in reinforcement learning.\n- Experiments are conducted based on GRPO and DAPO, two recent popular RL approaches for LLMs, proving the effectiveness of KTAE.\n- Solid motivations and illustrations in Section 3.\n- Clear visualization examples to show the important tokens found by KTAE.\n\nCons:\n- Experiments are conducted on 1.54B/7B models. It is not clear whether this method also have the scaling law.\n- It would be better if the paper can also analyze the training efficiency (comparing with baselines).", "was_revised": true, "final_justification": "The authors have well addressed my concerns. I would like to keep my positive rating.", "rebuttal": "Thank you very much for your valuable feedback on our work. We'd like to address the weaknesses you pointed out as follows:\n\n**W1:**  Currently, mathematical reasoning tasks commonly employ foundation models with a certain level of mathematical capability for reinforcement learning (RL) training. As a result, related work generally uses Qwen2.5 Math series models as base models, which are only available in 1.5B, 7B, and 72B sizes. Since RL training requires frequent sampling, it demands significant computational resources. Due to our current computational limitations, it's challenging for us to conduct experiments on larger models within a limited timeframe. In future work, we will continue to explore the effectiveness of KTAE on larger-scale models and models with different capabilities.\n\n**W2:** Regarding the training efficiency issue, we conducted supplementary experiments. The average time per step using 8 H100 GPUs for training on the 7B and 1.5B models is as follows:\n\n|                            | Qwen2.5-7B-MATH/s | Qwen2.5-1.5B-MATH/s |\n|----------------------------|-----------------|-------------------|\n| GRPO                       | 559.36          | 277.69            |\n| GRPO+KTAE                  | 641.98          | 371.79            |\n| DAPO                       | 1006.30         | 363.36            |\n| DAPO+KTAE                  | 1159.00         | 642.74            |\n\nAs you can see, KTAE is a model-free algorithm, and its computational time is almost independent of the model scale, depending only on the number of tokens generated. For the 7B model, due to its larger size, each training step takes longer, so the computational efficiency loss introduced by KTAE is not as noticeable. Conversely, for the 1.5B model, since the training steps are shorter, the training efficiency loss from KTAE is more pronounced. Nevertheless, KTAE's computational overhead is still acceptable compared to methods that utilize the model itself. In reality, the computation time for KTAE largely depends on code optimization. If only CPU serial computation were used, it would take several hours. The current code only performs small tensor parallel computations at the level of N rollouts in a single sampling, leading to a very low degree of data parallelism. We haven't yet thoroughly optimized the computational efficiency of the KTAE algorithm from an engineering perspective; GPU utilization when computing KTAE is less than 1%. We can definitely concatenate these into larger tensors, which would further leverage GPU capabilities and significantly reduce computation time. We will attempt further engineering optimization in the future.\n\n\n**Q1:** Regarding your question, we believe that KTAE, in principle, is task-agnostic. It is entirely probability-based, and can be applied to any domain that uses binary rewards for GRPO reinforcement learning training, such as code generation or question answering. It just so happens that in the field of mathematical reasoning, rule-based binary reward methods are mainstream algorithms, and related work has been experimented with mathematical reasoning tasks. For ease of comparison, we also chose the same task for KTAE. In fact, KTAE can be used to calculate key tokens for any task in any domain, as long as a reasonable binary reward objective can be designed.\n\nWe would be very grateful if our responses address your concerns and you would consider increasing our score!"}, {"review_id": "FBTvya5DnA", "reviewer": "Reviewer_8J8v", "summary": "This paper proposes KtAE, a model-free algorithm to estimate token-level advantage of model rollouts to provide finer-grained supervision signals in LLM reinforcement learning. \n\nSpecifically, KtAE constructs contingency tables based on the correctness of sampled rollouts and uses statistical methods (Fisher's exact test and Information Gain) to quantify each token's association with correct outcomes. \n\nThe proposed method can be empirically deployed on existing algorithms GRPO and DAPO to further improve model performance. Interestingly, the proposed method can not only improve the accuracy, but also speed up the inference by generating shorter responses.", "questions": "* How could KtAE be compared and combined with fine-grained reinforcement learning methods? For example, what if we have multiple reward models for different features and different granularities?\n* The relationship between BM25-inspired frequency calculation and token importance is unclear. Could you elaborate further on this?\n* What if some tokens appear rarely across rollouts? How often would such corner cases occur and impact the final performance?\n* Could you provide timing comparisons between KtAE and corresponding baseline to discuss its scalability to larger models and batch sizes?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n* The idea of using statistical association measures to identify key tokens is innovative and economic. The empirical performance is also surprisingly appealing. The case study of \"aha moment\" as a result of KtAE provides the insight that such token-level advantage helps to elicit the self-reflection behaviors of LLMs.\n* Given the improvement in both performance and efficiency (reduced response length), the proposed method demonstrates strong empirical results with reasonable cost, showing potential to be applied at scale.\n* The detailed mathematical formulations and implementation details show technical maturity.\n\n### Weaknesses\n* While the empirical results are strong, KtAE is mainly a combination of statistical measures. Given the sparsity of token-level supervision signals (especially when rollout budget is limited), it is hard to quantify the quality and stability of the token-level advantages. The authors could either (i) provide detailed convergence analysis or theoretical guarantees about the quality of the token-level advantages; or (ii) conduct further ablation study to probe the sensitivity of KtAE to the statistical hyperparameters such as the number of rollouts and length of responses.\n* The token-level advantages are mainly determined by the binary rewards and statistical estimation. This makes KtAE a bit limited to mathematical reasoning (which has determined final answers). \n* KtAE introduces some additional hyperparameters (h1, h2, h3, k1, b) which may need to be adjusted on different datasets. It would be better to include some sensitivity analysis to help understand the generalizability and stability of the method.\n* While the token-level advantage calculation does not require additional model calls, it would still be helpful to discuss the computational overhead from this process to understand the complexity of the method.", "was_revised": true, "final_justification": "The authors have well addressed my concerns. I would like to keep my positive rating.", "rebuttal": "Thank you for your valuable feedback on our work. We've addressed the weaknesses you raised below:\n\n**W1:**  The core idea behind the **KTAE algorithm** is based on the assumption that \"if a token frequently appears in all correct rollouts and rarely appears in all incorrect rollouts, then that token is very important for obtaining correct rollouts.\" We believe this assumption aligns with common understanding, and KTAE models it using statistical methods.\n\nCurrently, mainstream training methods typically limit the number of sampled rollouts per problem to 8 or 16, and our experiments follow this setup. In practice, a smaller 'n' value often leads to situations where all sampled rollouts for a given problem yield identical rewards, meaning no single rollout offers a significant advantage over others. In such cases, KTAE, like GRPO, calculates an advantage value of zero for all tokens. In a sense, the evaluation quality of KTAE is positively correlated with that of GRPO. Theoretically, a larger number of rollouts or a more uniform distribution of binary rewards among the 'n' rollouts, which is analogous to a larger statistical sample, would improve KTAE's evaluation quality. Your point is certainly worth investigating, but due to time and computational constraints, we may not be able to conduct a more fine-grained experimental analysis of KTAE in the short term. In future work, we will focus on analyzing KTAE's sensitivity to statistical hyperparameters.\n\n**W2:**  We believe that **KTAE** is fundamentally task-agnostic. It's entirely probability-based and can be applied to any domain that uses **binary rewards** for **GRPO reinforcement learning (RL) training**, such as code generation or question answering. The reason we've focused on mathematical reasoning tasks is simply that rule-based binary reward methods are prevalent in this area, and related work primarily conducts experiments here. This choice allows for easier comparison. Essentially, KTAE can be used to calculate key tokens for any task in any domain, provided a suitable binary reward optimization objective can be designed.\n\n\n**W3:**  While KTAE does introduce several hyperparameters, we didn't deliberately tune them in our practical use. We simply chose relatively average weight values for each component as hyperparameters, without specifically emphasizing any particular component. Since the principles behind each component are task-agnostic, we believe that these average weights theoretically don't need to be adjusted for different tasks. As you pointed out, KTAE has many hyperparameters, making a detailed sensitivity analysis for each parameter difficult in the short term. We performed a coarse-grained ablation study on each component in our paper, demonstrating their effectiveness. In the future, we will focus on analyzing KTAE's sensitivity to statistical hyperparameters.\n\n**W4&Q4:**  Regarding training efficiency, we conducted supplementary experiments. The average time per step using 8 H100 GPUs for training on the 7B and 1.5B models is as follows:\n\n|                            | Qwen2.5-7B-MATH/s | Qwen2.5-1.5B-MATH/s |\n|----------------------------|-----------------|-------------------|\n| GRPO                       | 559.36          | 277.69            |\n| GRPO+KTAE                  | 641.98          | 371.79            |\n| DAPO                       | 1006.30         | 363.36            |\n| DAPO+KTAE                  | 1159.00         | 642.74            |\n\nAs you can see, **KTAE** is a model-free algorithm, and its computation time is almost independent of the model scale, depending only on the number of tokens generated. For the 7B model, due to its larger size, each training step takes longer, so the computational efficiency loss introduced by KTAE is not as noticeable. Conversely, for the 1.5B model, since the training steps are shorter, the training efficiency loss from KTAE is more pronounced. Nevertheless, KTAE's computational overhead is still acceptable compared to methods that utilize the model itself. In reality, the computation time for KTAE largely depends on code optimization. If only CPU serial computation were used, it would take several hours. The current code only performs small tensor parallel computations at the level of N rollouts in a single sampling, leading to a very low degree of data parallelism. We haven't yet thoroughly optimized the computational efficiency of the KTAE algorithm from an engineering perspective; GPU utilization when computing KTAE is less than 1%. We can definitely concatenate these into larger tensors, which would further leverage GPU capabilities and significantly reduce computation time. We will attempt further engineering optimization in the future.\n\n**Q1:**  If there were multiple reward models for different features and granularities, each rollout would receive multi-dimensional and diverse rewards. Theoretically, KTAE currently only applies to binary rewards. However, if the multiple rewards from different reward models can be combined and compressed into a **binary reward** based on task characteristics, then KTAE would still be applicable.\n\n**Q2:**  We believe that a token appearing in multiple correct rollouts doesn't fully represent its relevance to those correct rollouts. Its importance is truly indicated only when it frequently appears in correct rollouts and rarely appears in incorrect ones. Statistical values calculated solely from contingency tables only consider \"whether it appeared,\" which can magnify the influence of low-frequency words. Therefore, we introduced **term frequency** to balance \"whether it appeared\" and \"whether it appeared frequently.\"\n\n**Q3:**  If certain tokens appear with extremely low frequency across multiple rollouts, KTAE would still consider the impact of their presence or absence on obtaining a correct rollout. Even if they appear very rarely, if they appear in every correct rollout and in no incorrect rollouts, they would still receive a high advantage value. This applies to words in \"aha moments,\" for example. While such situations are uncommon, KTAE accounts for their existence.\n\nWe would be very grateful if our responses address your concerns and you would consider increasing our score!"}, {"review_id": "Ci0cM6nNVL", "reviewer": "Reviewer_GKsR", "summary": "This work studies the problem of doing reinforcement learning finetuning for large language models. It has been observed that many existing methods (e.g., GRPO and DAPO) typically assigns a uniform advantage function for the entire rollout, lacking more fine-grained granularity. This work proposes the KATE framework, which can be directly leveraged within GRPO and DAPO to provide fine-grained, token-level advantages without introducing additional models. In particular, KATE builds a contingency table and uses hypothesis testing for estimating the association between tokens and the final correctness, which is further complemented with the information-theory-based method using conditional entropy. Also, to determine the direction of the association, the idea of term frequency is also incorporated. With all these approaches combined, a final importance score is computed, which is used as the token-level advantage. The effectiveness of KATE is demonstrated via experiments through its combination with GRPO and DAPO on training models with different sizes.", "questions": "I would appreciate the authors' thoughts on the concerns that I raised in the weakness section, in particular, regarding whether the current method simplifies the token-correctness association too much (i.e., ignoring some potentially important factors), how the final score is connected to the advantage, and how the proposed method can be positioned beyond GRPO and DAPO.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "I have a very mixed feeling about the proposed method, which are listed in the following.\n\n\nStrength:\n\n- It is indeed an interesting and important problem of how to assign more fine-grained training signals for the tokens in the rollout during RL finetuning of LLM. I appreciate the authors' valuable efforts in this direction.\n\n- Many different ideas are incorporated and nicely combined. The final score handles both the strength of the association and the direction of the association.\n\n- There is no additional model training required, and the proposed method can be easily incorporated with GRPO, DAPO, and other similar methods.\n\nWeakness:\n-  From one perspective, as mentioned above, it is carefully crafted with many different ideas incorporated. On the other hand, it is slightly against my intuition that \"directly counting\" the token occurrence and the conditional entropy can be a strong indication of the association of the token with the response correctiveness. In particular, it seems that the method ignores the order of the tokens and the connection between tokens, which may be important.\n\n- Also, the constructed method is largely heuristic without many theoretical insights. In particular, the score will be served as an advantage function while its connection with the advantage function is unclear to me, which seems to be more of an artifact created based on intuitions.\n\n- Furthermore, the proposed method seems to be mostly suitable for algorithms like GRPO and DAPO, which performs multiple rollouts from one prompt. It does not really apply for PPO and others (which uses value networks to estimate advantages). This may limit the usefulness of the proposed method. In particular, from a higher-perspective, to obtain a more fine-grained advantage estimation, it essentially (based on my understanding) viewing the problem from the single-step one (bandits) to a multiple-step one (MDP), and algorithms like PPO are well-desighed for the later case (although the current LLM training are mostly treated in a single-step manner)", "was_revised": true, "final_justification": "The authors have elaborated their perspectives on the connection between the heuristics in this work and the advantage function, which facilities my understanding and leads to my raised score. However, I personally maintain a suspicious opinion, as the proposed methods are still largely heuristics.", "rebuttal": "Thank you for your valuable feedback on our work. We've addressed the weaknesses you raised below:\n\n**W1:**  The core idea behind the **KTAE algorithm** is based on the assumption that \"if a token frequently appears in all correct rollouts and rarely appears in all incorrect rollouts, then that token is very important for obtaining correct rollouts.\" We believe this assumption aligns with common understanding, and KTAE models it using statistical methods. While KTAE indeed doesn't account for the **order of tokens** or **dependencies between tokens**, which are crucial factors, it's essentially an improvement upon **GRPO**. The original GRPO assigns the same advantage to every token in a rollout, completely ignoring the differences between tokens, let alone their order or dependencies. KTAE takes a step towards **token-level advantage** based on GRPO, so it's challenging to cover every aspect. Furthermore, quantifying the order of tokens and their dependencies is incredibly difficult. Even trained models struggle to accurately assess the impact of token order and dependencies in mathematical reasoning steps on achieving correct rollouts. Of course, considering these factors would undoubtedly lead to more accurate token-level advantages, and we believe this requires continuous effort and experimentation in future work.\n\n**W2:**  The score computed by KTAE reflects a token’s estimated contribution to producing correct rollouts. We acknowledge that KTAE adopts a heuristic approach, but we would emphasize that it is a pragmatic solution to the pressing limitation of GRPO. Within the GRPO framework, we adopt the view that tokens with greater importance should receive higher advantage values, since they warrant stronger learning signals. This formulation provides a practical way to incorporate token-level discriminability into model-free RL training, complementing GRPO’s rollout-level estimation.\n\n**W3:**  Currently, KTAE is fundamentally an improvement on GRPO, refining GRPO's rollout-level advantage into a **token-level advantage**. Therefore, its applicability is limited to reinforcement learning training based on GRPO. The **PPO algorithm** doesn't require multiple samplings; it can obtain token-level advantage by treating each token as a step through a value model. It would be exciting if a **model-free algorithm** could calculate token-level advantage in PPO, and we will strive to achieve such an algorithm. Furthermore, KTAE is compatible with PPO-like methods. PPO typically uses a value model to compute per-token advantages, but these estimates are often coarse or noisy in reasoning tasks where rewards are sparse and delayed. KTAE can serve as a complementary module: its statistically grounded signals could be used to augment or denoise the token-level advantage estimates from a PPO critic. We leave a more thorough exploration of such integration to future work, as mentioned in the conclusion.\n\nWe sincerely hope that our responses have addressed your concerns, and we would be grateful if you would consider updating your score and potentially increasing it."}, {"review_id": "xs4mC9bqOg", "reviewer": "Reviewer_BQ9i", "summary": "Key-Token Advantage Estimation (KTAE) tackles the coarse-granularity problem of existing RLHF variants (e.g., GRPO, DAPO) by computing fine-grained, token-level advantages instead of assigning a single rollout-level value to every token. The method statistically infers each token’s contribution to the final correctness signal and blends this with the rollout-level advantage, remaining model-free and easy to plug into standard policy-gradient pipelines. Empirically, GRPO+KTAE and DAPO+KTAE outperform their baselines on five mathematical-reasoning benchmarks, achieving higher accuracy with shorter answers and even surpassing R1-Distill-Qwen-1.5B when using the same backbone.", "questions": "Please consider the weaknesses outlined above; if subsequent responses offer convincing clarifications, I will be happy to revisit and potentially raise my score.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- **Fine-grained, token-level credit assignment:** Existing RLHF-style methods propagate a single rollout-level advantage to every token, so they cannot tell which parts of the answer actually “earned” the reward. KTAE attacks this long-standing credit-assignment bottleneck by statistically estimating each token’s marginal contribution to the final correctness signal, giving the optimizer far more informative gradients. \n- **Model-free and compute-efficient:** Unlike token-reward approaches that bolt on an auxiliary critic or a learned reward network, KTAE is purely analytical; it re-weights the existing rollout advantage with a closed-form importance estimate. That means no extra parameters, no new forward passes, and no critic-training instability, making it easy to drop into any policy-gradient pipeline. This lightweight property will appeal to practitioners who already find RLHF training expensive.\n- **Consistent empirical gains on math-reasoning suites:** Across five benchmark datasets for mathematical reasoning, GRPO+KTAE and DAPO+KTAE deliver higher exact-answer accuracy while also shortening the generated solutions. \n\n## Weaknesses\n- **Heuristic, non-learned composition of statistics:** KTAE sidesteps a learned critic by multiplying four hand-selected statistics—Fisher’s test score, information gain, token frequency, and the rollout-level advantage—to obtain a key-token value. While this model-free choice keeps computation cheap, the paper offers no theoretical rationale for using a product, nor does it explore alternative compositions or automatic ways to learn the combination. The ablation study shows each component matters in isolation, but I am left unsure whether such composition structure is robust when ported to new domains.\n- **Surrounding contexts are ignored when scoring tokens:** A token’s value can depend on where it appears (e.g., a numeral “2” inside a denominator vs. as a final answer). Because KTAE treats each token independently, it may mis-attribute credit when the semantic payoff arises only from multi-token spans or long-range dependencies. That limitation becomes more glaring in text-rich domains (dialogue, code) where meaning is rarely confined to a single token.\n- **Sparse diagnosis of performance regressions:** Although KTAE generally boosts accuracy, it drops on two of the five benchmarks (AIME24 for GRPO+KTAE and OlympiadBench for DAPO+KTAE). The paper merely remarks that this is “one fewer correct answer” in a small test set, offering no deeper probe into why the estimator hurt — for instance, whether token-level weighting mis-fires on very short proofs, whether the length-penalty interacts with Olympiad-style solutions, or whether sampling variance alone explains the loss. A concise error taxonomy (e.g., mis-tagged key tokens vs. over-shortened outputs) plus a couple of annotated examples would clarify whether the regression stems from KTAE itself or dataset. \n- **Domain scope is narrow:** All experiments revolve around deterministic, auto-gradable math problems. It is unknown whether KTAE remains effective in open-ended generation where rewards are noisy, delayed, or reference-based (e.g., ChatGPT-style helpfulness RLHF). A broader evaluation—or at least a discussion of potential adaptations—would strengthen claims of generality.", "was_revised": true, "final_justification": "The authors provided reasonable clarifications for my main concerns, including the rationale for the statistic composition and the causes of benchmark regressions. While these do not fully resolve all limitations, I found the responses satisfactory, so I maintain my borderline accept recommendation.", "rebuttal": "Thank you for your valuable feedback on our work. We've carefully considered the weaknesses you raised and offer the following clarifications:\n\n**W1:** We detail the theoretical basis for using **multiplicative combination** in Appendix E. Contingency tables inherently reflect the *likelihood* of an association but don't specify *which* particular metric is associated. Therefore, to determine importance, we must calculate the **direction of association** and multiply it by the **degree of association**. While we acknowledge that learning to automatically determine combination methods could offer benefits, it would introduce additional training overhead. We plan to explore other combination approaches in future work.\n\nImportantly, our goal is to develop a model-free, low-cost, and interpretable mechanism for estimating token-level advantage, which avoids the instability and overhead of training additional reward or value models. KTAE's calculation method is purely **statistical** and entirely **task-domain independent**. As long as a binary reward system is used, KTAE remains applicable when transferred to other domains.\n\n\n**W2:** The core of the KTAE algorithm rests on the assumption that \"if a token frequently appears in all **correct rollouts** and rarely in all **incorrect rollouts**, then that token is highly important for achieving correct rollouts.\" We believe this assumption aligns with common intuition, and KTAE effectively models it using statistical methods.\n\nWe fully acknowledge that KTAE currently doesn't account for **token order** or **inter-token dependencies**. However, KTAE is fundamentally an advancement built upon GRPO. The original GRPO assigns uniform advantage to every token in a rollout, completely overlooking individual token differences, let alone their order or dependencies. KTAE represents a significant step towards **token-level advantage** within the GRPO framework, making it challenging to address all complexities simultaneously. Quantifying token order and dependencies is inherently difficult; even highly trained models struggle to precisely evaluate their impact on correct mathematical reasoning steps. We agree that incorporating these aspects would undoubtedly lead to more accurate token-level advantage, and this is a promising direction worth studying for future research.\n\n**W3:** Our experiments were exclusively trained on the **MATH dataset**, meaning performance on other benchmarks relies entirely on **domain generalization capabilities**. Reinforcement learning training involves extensive sampling, introducing inherent randomness. Consequently, achieving consistent improvements across all benchmarks is challenging. We consider **AIME24** and **OlympiadBench** to be among the most difficult benchmarks, also requiring the longest CoT (Chain of Thought) lengths. DAPO (Decision-aware Policy Optimization) intrinsically includes a length penalty. This led to DAPO+KTAE's length decreasing too rapidly on OlympiadBench, causing the model to prioritize brevity over some accuracy. For GRPO+KTAE on AIME24, the performance decrease was marginal, corresponding to just one fewer correct answer. We attribute such minor fluctuations to the **stochastic nature of reinforcement learning training**.\n\nWe appreciate the suggestion for deeper analysis. We are preparing a brief taxonomy of error modes (e.g., missed key tokens, excessive length reduction) and will include annotated examples to clarify where KTAE succeeds or fails in the final version.\n\n\n**W4:** From a principled standpoint, KTAE is **task-domain agnostic**. It's entirely probability-based, making it applicable to any domain that uses **binary rewards** for **GRPO** reinforcement learning, such as code generation or question answering (QA). In fact, for any task in any domain where a reasonable binary reward objective can be designed, KTAE can be used to identify critical tokens. Even with continuous numerical rewards, KTAE can be applied by binarizing them through thresholding. We will discuss these potential adaptations more explicitly in the final version, and we consider generalizing KTAE beyond math reasoning a central goal of our future work.\n\n\nWe hope these clarifications adequately address your concerns. We would be sincerely grateful if, upon reviewing our response, you would consider reconsidering your score and potentially increasing it."}]

Author rebuttal:

Meta-review: Existing reinforcement learning algorithms like GRPO and DAPO compute rollout-level advantages that assign the same value to all tokens in a sequence, overlooking token-specific contributions. To overcome this, the paper introduces Key-token Advantage Estimation (KTAE), a method that provides fine-grained, token-level advantage estimates without requiring additional models. KTAE works by analyzing the correctness of sampled rollouts and statistically quantifying each token’s contribution to the final outcome. This paper shows that algorithms combining GRPO and DAPO with KTAE outperform baseline methods across multiple benchmark tests. The idea is novel and the contribution is solid. Authors well address reviewers' concern during the rebuttal phase. We therefore recommend acceptance.
Decision: Accept (poster)