README.md

---
license: mit
---

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

<div align="center">
<br>
<a>Wei Sun</a>,
<a>Wen Yang</a>, 
<a>Pu Jian</a>,

<a>Qianlong Du</a>,
<a>Fuwei Cui</a>,
<a>Shuo Ren</a>,
<a>Jiajun Zhang</a>
<br> Institute of Automation Chinese Academy of Sciences <br>
    
![Mathematical-Reasoning](https://img.shields.io/badge/Task-Mathematical--Reasoning-red) ![Reinforcement-Learning](https://img.shields.io/badge/Method-Reinforcement--Learning-red) <a href='https://arxiv.org/abs/2505.16826'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a> <a href='https://github.com/xiaolizh1/KTAE'><img src='https://img.shields.io/badge/Project-Github-red'></a>
</div>

## 🔖 Overview
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 ($\textit{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.

<p align="center">
    <img src=https://cdn-uploads.huggingface.co/production/uploads/654a05493ee6a84ff2fd3fc1/-dygcr2iG28gI6jRh0KG1.png width="100%" height="100%">
</p>


In summary, the KTAE algorithm offers several advantages:

+ KTAE provides more fine-grained advantage information without introducing extra models, resulting in lower training costs.

+ KTAE directly computes the importance differences between tokens using statistical analysis methods, offering strong interpretability.

+ KTAE's key-token value is computed based on the correctness of the final answer and retains the original rollout-level advantage, making it less susceptible to reward hacking.

+ KTAE can make the model pay more attention to key tokens and reduce the learning of irrelevant tokens, which can effectively reduce the response length.

## 🔥 Update 
- [21/05/2025]🔥Key-token Advantage Estimation is coming!

## 📃 Contents

- [Models](#Available_Models)
- [Setup](#Setup)
- [Preparation](#Preparation)
- [Train](#Train)
- [Inference](#Inference)
- [Experiments](#Experiments)
- [Citation](#citation)

## 🧠 Available Models

| Model Size | DAPO+KTAE | GRPO+KTAE |
|------------|--------------|--------------|
| 1.5B       | <a href="https://huggingface.co/SunW7777/DAPO_KTAE_1.5B"><img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" alt="HF" width="20"/> DAPO_KTAE_1.5B</a> | <a href="https://huggingface.co/SunW7777/GRPO_KTAE_1.5B"><img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" alt="HF" width="20"/> GRPO_KTAE_1.5B</a> |
| 7B         | <a href="https://huggingface.co/SunW7777/DAPO_KTAE-7B"><img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" alt="HF" width="20"/> DAPO_KTAE-7B</a>   | <a href="https://huggingface.co/SunW7777/GRPO_KTAE-7B"><img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" alt="HF" width="20"/> GRPO_KTAE-7B</a>   |



## 📷 Setup

Please follow the instructions below to install the required packages.


1. Clone this repository

```bash
https://github.com/xiaolizh1/KTAE.git
```

2. Install Package

```bash
conda create -n KTAE python=3.10 -y
conda activate KTAE
cd KTAE
pip install -r requirements.txt
```

## 📈 Train

Our training is mostly performed on [Verl](https://github.com/volcengine/verl) code base and makes some changes. 

##  📌 GRPO+KTAE

```bash
bash examples/grpo_trainer/run_qwen2.5_7b.sh    #train 7b model
bash examples/grpo_trainer/run_qwen2.5_math_1.5b.sh   #train 1.5b model
```

##  📌 DAPO+KTAE

```bash
bash recipe/dapo/run_dapo_qwen2.5_7b.sh    #train 7b model
bash recipe/dapo/run_dapo_qwen2.5_1.5b.sh    #train 1.5b model
```

##  📌 Merge Model

```bash
cd scripts
bash merge_model.sh      #merge checkpoint
```

## ✅ Evaluation

Our evaluate code is base on [Dr.GRPO](https://github.com/sail-sg/understand-r1-zero)

```bash
cd eval
bash run_eval.sh
```

## 👀 Experiments

We provide some results in this section. More detailed results can be found in our paper.
<div align=center>
<img width="90%" src=https://cdn-uploads.huggingface.co/production/uploads/654a05493ee6a84ff2fd3fc1/exOLI0iPBFljL6x2ZOIFe.jpeg>
</div>

### Main Results

+ Method validation result.
<div align=center>
<img width="90%" src=https://cdn-uploads.huggingface.co/production/uploads/654a05493ee6a84ff2fd3fc1/kFJiRcr47hGylp29x9pqx.png>
</div>

+ Comparison with baselines on Accuracy.
<div align=center>
<img width="90%" src=https://cdn-uploads.huggingface.co/production/uploads/654a05493ee6a84ff2fd3fc1/N5tKMS6w12ir0geF1oIgz.jpeg>
</div>

+ Comparison with baselines on Efficiency.
<div align=center>
<img width="90%" src=https://cdn-uploads.huggingface.co/production/uploads/654a05493ee6a84ff2fd3fc1/nfypL2d3jS1GuH7mM_v9y.jpeg>
</div>


### 📊 More Analysis
+ Ablation analysis.
<div align=center>
<img width="80%" src=https://cdn-uploads.huggingface.co/production/uploads/654a05493ee6a84ff2fd3fc1/Qphwk_uwp_uIJOTp4RAMw.png>
</div>

+ Visualization example. 
<div align=center>
<img width="80%" src=https://cdn-uploads.huggingface.co/production/uploads/654a05493ee6a84ff2fd3fc1/5aj7uS9uohvvLDoCyFSr0.png>
</div>


## 🔗 Citation

If you find this repo useful for your research, please consider citing the paper

```
@misc{sun2025ktaemodelfreealgorithmkeytokens,
      title={KTAE: A Model-Free Algorithm to Key-Tokens Advantage Estimation in Mathematical Reasoning}, 
      author={Wei Sun and Wen Yang and Pu Jian and Qianlong Du and Fuwei Cui and Shuo Ren and Jiajun Zhang},
      year={2025},
      eprint={2505.16826},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2505.16826}, 
}
```

## 🌈 Acknowledgement

We would like to thank the following repos for their great work:

+ [Verl](https://github.com/volcengine/verl) for providing the training framework
+ [Vllm](https://github.com/vllm-project/vllm) for the efficient inference engine with high throughput
+ [transformers](https://github.com/huggingface/transformers) for providing the model-base and fune-tuning framework


## 🔎 License

This project is released under the Apache 2.0 license. Parts of this project contain code and models from other sources, which are subject to their respective licenses.