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

Thank you for your valuable feedback on our work. We've carefully considered the weaknesses you raised and offer the following clarifications:

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.

Importantly, 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.

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.

We 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.

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.

We 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.

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.

We 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.

Thank you for your valuable feedback on our work. We've addressed the weaknesses you raised below:

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.

Currently, 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.

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.

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.

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:

As 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.

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.

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."

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.

We would be very grateful if our responses address your concerns and you would consider increasing our score!

Thank you for your valuable feedback on our work. We've addressed the weaknesses you raised below:

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.

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.

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.

We 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.

Dear Reviewer GKsR,

We hope this message finds you well.

We are deeply grateful for your thorough review and acknowledgment of our work. We have provided detailed clarifications in response to your valuable feedback.

As the discussion period draws to a close, we would like to hear your thoughts on our response, including whether it adequately addresses your concerns. If you have any updated thoughts, we would be grateful to hear them.

Thank you again for your time and thoughtful engagement.

Best regards,

Authors

Thank you very much for your valuable feedback on our work. We'd like to address the weaknesses you pointed out as follows:

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.

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:

As 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.

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.

We would be very grateful if our responses address your concerns and you would consider increasing our score!