👉 LightThinker 👈
LightThinker: Thinking Step-by-Step Compression
[](https://github.com/zjunlp/LightThinker)
[](https://opensource.org/licenses/MIT)
Github •
📄arXiv •
𝕏 Blog
## Table of Contents
- 👀[Overview](#overview)
- 🔧[Installation](#installation)
- 🏃[Quick Start](#quick-start)
- 🎁[Acknowledgement](#acknowledgement)
- 🚩[Citation](#citation)
## 👀Overview
LLMs have shown remarkable performance in complex reasoning tasks, but their efficiency is hindered by the substantial memory and computational costs associated with generating lengthy tokens. In this paper, we propose LightThinker, a novel method that enables LLMs to dynamically compress intermediate thoughts during reasoning. Inspired by human cognitive processes, LightThinker compresses verbose thought steps into compact representations and discards the original reasoning chains, thereby significantly reducing the number of tokens stored in the context window. This is achieved by training the model on when and how to perform compression through data construction, mapping hidden states to condensed gist tokens, and creating specialized attention masks.
## 🔧Installation
```bash
git clone https://github.com/zjunlp/LightThinker
cd LightThinker
conda create -n lightthinker python=3.9 -y
conda activate lightthinker
pip install -r requirements.txt
cd data && unzip data.zip && cd ..
```
## 🏃Quick Start
> First, we train the model to learn how to compress (step 1). Then, we perform inference on the test set to obtain output results (step 2). Finally, we evaluate the output results (step 3).
### Step 1. Training
To execute the training, run the following command:
```bash
bash train.sh
```
Currently, the script's parameters are set to run on a machine with 4 A800 GPUs. If you encounter OOM (Out Of Memory) issues, please reduce the `micro_batch_size` and `max_length`. For other parameters in the script, please refer to the [documentation](./ARGS.md).
### Step 2. Inference
To execute the inference, run the following command:
```bash
bash inference.sh
```
Here, you need to modify the script file's `model_tag`, `model_short_tag`, `ckpt`, `output_tag`, and `split_size`. For details regarding the script's parameters, please refer to the [documentation](./ARGS.md).
### Step 3. Evaluation
> [!NOTE]
> If this is your **first time** conducting an evaluation, please execute the following code first:
```bash
python evaluation/init.py
```
To execute the evaluation, run the following command:
```bash
method=""
tokenizer_path=""
comp_config=""
model_type=""
dataset=""
bos_token=""
eos_token=""
cache_size=1024
file1=""
file2=""
file3=""
file4=""
python evaluation/eval_file.py \
--method $method \
--tokenizer_path $tokenizer_path \
--comp_config $comp_config \
--model_type $model_type \
--dataset $dataset \
--files $file1 $file2 $file3 $file4 \
--cache_size $cache_size \
--bos_token $bos_token \
--eos_token $eos_token \
--interaction
```
Please note that if you set `split_size>1` in the second step, the number of file i here should match the value of `split_size`. It should be noted that manual evaluation was conducted during the assessment. Use the `--interaction` flag to enable manual evaluation. The `cache_size` parameter is used for `H2O` and `SepLLM`, but not for `LightThinker` or `AnLLM`.