Making Slow Thinking Faster: Compressing LLM Chain-of-Thought via Step Entropy

In plain words

Reasoning models often think out loud at great length, and much of that text adds little. The model is usually sure of what it is writing, so those steps carry almost no new information. The paper measures how uncertain the model was while writing each step, ranks the steps, and cuts the most predictable ones.

A cut step is replaced by a short [SKIP] marker, so the reasoning keeps its outline and the answer can still be reached. Applied to traces that already exist, this keeps accuracy close to the full trace on the maths benchmarks the authors used, while the thinking gets much shorter.

The second part teaches a model to do this itself. It is first fine-tuned on traces that already contain [SKIP] markers, then trained with reinforcement learning. The reward favours correct answers, a reasonable share of skips, and not too many skips or overly long answers. The goal is a shorter thinking phase and so lower inference cost.

TL;DR

Each blank-line-separated step of a model's reasoning gets a score: the mean entropy of its tokens. Replacing the lowest-entropy 80% of steps with [SKIP] keeps accuracy near the full chain on GSM8K, Math500, AIME 2024 and AIME 2025 for DeepSeek-R1-Distill-Qwen 7B and 14B and Qwen3-8B, with 27-45% fewer thinking tokens on Math500 and AIME tab. 1. A model trained with SFT and GRPO to emit [SKIP] itself uses 35-57% fewer tokens on four benchmarks tab. 3. The code is public but has no license file, and several training settings differ from the paper, so use the paper as the spec and the code only as a reference.

Problem

Long chain-of-thought traces raise latency and cost at inference, and much of the trace is redundant §1. Earlier compression work shortens text at the token level or the whole-trace level, or moves reasoning into hidden states, which gives up explicit, checkable reasoning §1 §2. The open question is how to find which reasoning steps are crucial and which are superfluous §1. The method assumes that a model's own per-token uncertainty is a usable signal for redundancy §3.1.

Method

The method has three parts: a per-step score (step entropy), a pruning rule that replaces the lowest-scoring steps, and a two-stage training recipe that teaches a model to produce pruned traces without the pruning step.

1. Generate a full trace. DeepSeek-R1-Distill-Qwen-7B writes a complete chain-of-thought for each problem. The thinking text is split into steps at blank lines, \n\n §3.2, item 1.
2. Score each step. Every token gets the Shannon entropy of the model's distribution over its vocabulary at that position eq. 1. A step's score is its length-normalised entropy, the mean over its tokens eq. 3. The paper first defines the step score as the plain sum eq. 2 and then switches to the mean, to avoid favouring long steps §3.1.
3. Prune. Sort the steps by score, ascending. Replace the κ×N lowest with [SKIP] and keep the rest verbatim, with κ = 0.8 §3.2, item 3 §4.1.
4. Infer with the compressed trace. The compressed thinking is placed after the user query and a closing </think> delimiter, and the model writes only the final answer §3.2, item 4 fig. 2.
5. Stage 1, SFT. Train on (problem, compressed trace, answer) triples with cross-entropy, so the model learns to write [SKIP] markers §3.3. The training set is 130k problems from DeepScaleR and OpenR1-Math, pruned and filtered to 70k samples of at most 4096 tokens §4.2 app. C.
6. Stage 2, GRPO. For each prompt, sample a group of completions and score each with a composite reward §3.3. The paper sums four parts eq. 4:
   • correctness: 2.0 for the right final answer, else 0.0 eq. 5;
   • skip ratio: 1.0 if the share of [SKIP] among the steps is at least 0.8, 0.5 if it is between 0.5 and 0.8, else 0.0 eq. 6 §4.3;
   • skip count: −1.0 when more than τ_skip = 100 [SKIP] tokens appear §3.3 §4.3;
   • response length: −1.0 when the response exceeds τ_length = 3500 tokens §3.3 §4.3. The paper uses 10k of the 70k SFT samples for this stage app. C.

The evidence for the direction of the rule: removing the lowest-entropy steps keeps accuracy, removing the highest-entropy steps hurts it at once, and random removal falls between the two §4.1 fig. 1.

Two equations from the paper, in its notation:

H(t_{i,j} \mid c_{i,j}) = -\sum_{w \in V} p(w \mid c_{i,j}) \log_2 p(w \mid c_{i,j}) eq. 1

H(S_i \mid S_{<i}) = \frac{1}{M_i} \sum_{j=1}^{M_i} H(t_{i,j} \mid c_{i,j}) eq. 3

Pseudocode for scoring and pruning. The paper gives the rule in prose §3.2:

# Steps are the blank-line-separated pieces of the thinking  [§3.2]
# Each step has its text and the entropy of each token  [eq. 1]
# kappa is the pruning ratio, 0.8 in the paper  [§4.1]

def step_entropy(step):
    # Length-normalised: mean token entropy over the step  [eq. 3]
    return sum(step.token_entropies) / len(step.token_entropies)

def prune_low_entropy_steps(steps, kappa):
    # Score every step  [§3.2]
    scores = [step_entropy(step) for step in steps]
    # Rank from lowest to highest entropy  [§3.2]
    order = sorted(range(len(steps)), key=lambda i: scores[i])
    # Number of steps to prune: kappa times N, rounded down (the code
    # floors; rounding is not stated in the paper)  [§3.2]
    num_pruned = int(kappa * len(steps))
    pruned = set(order[:num_pruned])
    # Replace pruned steps by [SKIP]; keep the others verbatim  [§3.2]
    return "".join(
        "[SKIP]" if i in pruned else steps[i].text
        for i in range(len(steps))
    )

Pseudocode for the GRPO reward for one completion. Each rule is in §3.3 and the thresholds are in §4.3:

# Composite reward for one sampled completion  [§3.3] [eq. 4]
# completion: thinking text, then the final answer
# answer: ground truth A*

def reward(completion, answer):
    # Correctness: 2.0 for the right final answer, else 0.0  [eq. 5]
    extracted = extract_final_answer(completion)
    total = 2.0 if extracted == answer else 0.0
    # Skip ratio: [SKIP] count over the number of steps in the thinking
    # (the paper does not say how steps are counted)  [eq. 6]
    skips = count_skip_tokens(completion)
    steps = count_steps(completion)
    ratio = skips / max(1, steps)
    # Tiered reward, thresholds tau_high = 0.8 and tau_low = 0.5  [§4.3]
    if ratio >= 0.8:
        total += 1.0
    elif ratio >= 0.5:
        total += 0.5
    # Penalty of -1.0 above tau_skip_num = 100 skip tokens  [§3.3] [§4.3]
    if skips > 100:
        total -= 1.0
    # Penalty of -1.0 above tau_length = 3500 tokens  [§3.3] [§4.3]
    if response_tokens(completion) > 3500:
        total -= 1.0
    return total

Hyperparameters

Values are as the paper states them unless the Notes column says otherwise. "Not stated" means the paper and its appendix are silent.

Results

Author-reported. Setup: full chain-of-thought as baseline, evaluated with Pass@1 on GSM8K, MATH500, AIME 2024 and AIME 2025 §4.2 tab. 1. The values are single numbers per cell; the paper reports no run-to-run variance in the text read for this card.

Table 1: static pruning, 80% of lowest-entropy steps tab. 1. Parenthesised values are the paper's percentage changes.

Table 2: large-scale check, accuracy (%) of DeepSeek-R1-7B on the full sets tab. 2.

Table 3: static pruning vs SFT and SFT+GRPO, DeepSeek-R1-7B tab. 3.

Table 4: comparison with other compression methods, relative to full CoT tab. 4. Values are as printed. The caption calls them percentage changes; for this paper's row, the accuracy figures equal absolute point differences from Table 3 (88.17 to 85.00 is 3.2 points), so read the accuracy column as points.

Table 5: reward ablation for GRPO, DeepSeek-R1-7B tab. 5 app. F. Rc is correctness, Rsr the skip ratio, Rsn the skip count, Rrl the response length.

Table 6: model-aware study, 80% pruning tab. 6 app. G.

Table 7: QwQ-32B at 80% and 90% pruning, with no-thinking baseline tab. 7 app. G.

Table 8: QwQ-32B on MMLU subsets tab. 8 app. G.

Table 9: DeepSeek-R1-7B on MMLU subsets tab. 9 app. G.

Other reported results, in prose:

• Pruning the high-entropy steps instead reduces accuracy at once; past a 40% mask ratio it falls below the no-thinking baseline. Random pruning starts to fall at 40% §4.1 fig. 1.
• Pruning tokens rather than steps drops accuracy after a 20% token mask ratio, while step pruning holds up to 40% of thinking tokens fig. 3 §4.2 discussion.
• Replacing pruned steps: at 90% on DeepSeek-R1-7B the text reports 74% for the full chain, 68% for direct deletion and [MASK], and 70% for [SKIP] and "and then" app. E fig. 4.

Paper vs code

The official repository is staymylove/COT_Compresstion_via_Step_entropy, linked from the abstract, at commit 48cbfcf8ad6cb14086ff5cfde07c978fc51b80ed. It has no LICENSE file, so the code is cited by path and line only, with no quoted snippets. Its README gives the run order README.md.

Entropy and step scores

• Paper defines per-token entropy as the full-vocabulary Shannon entropy in bits eq. 1. Code uses only the sampled token's term, −p·ln p with p from the returned logprob build_complete_thinking_dataset.py build_complete_thinking_dataset.py. It requests top-5 logprobs but does not use them for the score build_complete_thinking_dataset.py.
• Paper's step score is the mean over tokens eq. 3. Code does the same build_complete_thinking_dataset.py build_complete_thinking_dataset.py.
• Step boundaries: paper uses \n\n §3.2. Code splits only where a token equals the literal string ĊĊ build_complete_thinking_dataset.py. Tokens that merge the newline with other text do not split.
• Empty steps: a delimiter with no preceding text is stored as a step with score 0.0 build_complete_thinking_dataset.py. Steps with the same text share a dictionary key and overwrite each other build_complete_thinking_dataset.py.

Pruning

• Paper does not say how κ×N is rounded. Code uses int(), which floors sort_step_entropy.py. Ties are broken by original order sort_step_entropy.py.
• Code's default mask ratio is 0.8 build_masking_thinking_dataset.py, matching §4.1.
• The prompt for the masked-trace inference is built by feeding the chat template's prefix plus the masked thinking and \n</think>\n\n to the completions endpoint build_masking_thinking_dataset.py. The paper describes the same format §3.2, item 4.
• messages=[...], in the masking script ends with a comma, which makes messages a one-element tuple, not a list build_masking_thinking_dataset.py. The chat template may then render a different prompt from the one intended. Not run; check before using that script.

SFT (Stage 1)

• Paper: 3 epochs app. C. Code: num_train_epochs=1 sft.py.
• Paper: filter out samples over 4096 tokens app. C. Code: truncate to 2048 with padding to the maximum length sft.py. Long samples are cut, not dropped.
• Paper: cross-entropy on (problem, compressed trace, answer) §3.3. Code sets labels to the full token sequence, including the prompt and pad tokens (pad is the EOS token) sft.py sft.py. So the loss covers the prompt too.
• Paper: LoRA r = 16, α = 16 app. C. Code matches, plus dropout 0.1 sft.py.
• The SFT data is loaded from the Hugging Face dataset Jianyuan1/cot-data sft.py. The paper's 70k-sample set is not shown to be this dataset, and the code does not repeat the 4096-token filter.
• SFT script does not register [SKIP]. Only the GRPO script adds [SKIP], <think> and </think> when missing and resizes the embeddings grpo_train.py grpo_train.py. Whether the SFT data contains [SKIP] as a single token is not checkable from the repo.

GRPO (Stage 2)

• Paper: G = 14 samples per prompt app. C. Code: num_generations=8 grpo_train.py.
• Paper: KL coefficient 0.04 app. C. Code never passes a KL coefficient to the trainer; the kl_penalty=0.1 field is defined but unused grpo_train.py. Its effective value is whatever the installed TRL version uses by default; trl is unpinned requirement.sh.
• Paper: AdamW app. C. Code does not set an optimiser, so the trainer default applies.
• Paper: learning rate not stated. Code: 1e-5 grpo_train.py. The shell script sets LEARNING_RATE=5e-5 but never passes it run_grpo.sh.
• Paper: GRPO epochs not stated. Code: 1 grpo_train.py. The shell script sets EPOCHS=3, unused run_grpo.sh.
• Paper: 8 GPUs of 80 GB each app. C. Code: CUDA_VISIBLE_DEVICES=0..6 and num_processes: 7 run_grpo.sh accelerate_config.yaml, with a separate vLLM server on GPU 7 run_trl_vllm.sh. The server-mode option is commented out grpo_train.py, so the server is not referenced by the training script.
• QLoRA: the default is on grpo_train.py, but the model is loaded without 4-bit quantisation (the load line is commented out) grpo_train.py. The adapter targets all attention and MLP projections grpo_train.py.
• Skip-count penalty: paper −1.0 §3.3. Code −2.0 grpo_train.py.
• Correctness reward is 2.0 and skip-ratio tiers match grpo_train.py grpo_train.py. Length penalty matches grpo_train.py.
• Reward weights: the paper's text sums the four parts §3.3. Code passes the four functions with no weights grpo_train.py.
• Answer extraction takes the text from the first </think>. If none exists, find returns −1 and the slice keeps only the last character grpo_train.py.
• Skip-ratio denominator: code counts \n\n separators plus [SKIP] tokens grpo_train.py. The paper's N_steps is not defined beyond "steps of the thinking" §3.3.
• Length penalty is applied to the completion's token count after tokenising the full sequence grpo_train.py. The completion is cut at 4096 tokens grpo_train.py. The paper does not state a completion cap.
• The GRPO prompt is a fixed instruction that asks for [SKIP] in redundant steps grpo_train.py. The paper does not give its prompt.
• Training data: the 10k-sample file is a default path that is not in the repo grpo_train.py. How it was built is not shown.

Evaluation

• Paper: Pass@1, sample count not stated §4.2. Code: one sample per problem build_complete_thinking_dataset.py.
• Code drops samples with no </think> token, because tokens.index raises and the exception handler returns no result build_complete_thinking_dataset.py build_complete_thinking_dataset.py. The paper does not say how such samples were handled.
• Token counts: token_num_calculate.py counts subword tokens of the stored text with the base DeepSeek tokenizer token_num_calculate.py token_num_calculate.py. Unless the tokenizer knows [SKIP], each marker is several subwords, not one. The paper's note "contains the Unicode characters" is not explained.

Other paper-internal points

• Math500 SFT+GRPO token reduction: tables and text say 35% tab. 3 §4.3. From the table values, 2439.42 vs 3703.83 is 34.1%, which the appendix also reports app. F.
• GSM8K baseline tokens: 298.33 in tab. 1 and 296.09 in tab. 5.
• Appendix G text gives a 0.6% minimum DeepSeek-R1 token reduction; the smallest in tab. 6 is 1.3%.
• The paper gives one equation for the step sum eq. 2 and a normalised one eq. 3 and says the normalised one is used §3.1.
• Some accuracies are simple fractions of small sets (e.g. 35.71 = 5/14, 33.33 = 1/3, 30.00 = 3/10), and Math500 88.17 is not k/500 for any whole k. The per-benchmark sample count is therefore not stated and may differ from the full benchmark.

Pitfalls

• Entropy definition. The paper's entropy is over the full vocabulary, but the released code scores only the sampled token. The two can rank steps differently. If you reproduce the paper, implement eq. 1 and eq. 3. If you want the code's numbers, implement its surrogate and say so.
• Theory is a bound. The lemma and theorem bound conditional mutual information by step entropy §3.1 app. B. They say low-entropy steps carry little information only in the sense of an upper bound. They do not prove redundancy, and the empirical pruning results carry the claim.
• Segmentation changes the result. Steps are blank-line blocks. Other boundaries (sentences, clauses) could change entropies and results app. H. Empty or repeated steps are scored 0.0 or merge by text, so they get pruned first or are lost build_complete_thinking_dataset.py build_complete_thinking_dataset.py.
• Rounding and ties. κ×N is floored in code and ties keep original order sort_step_entropy.py sort_step_entropy.py. Short traces can prune fewer steps than the ratio suggests.
• Fixed κ = 0.8 is a fixed choice. The paper tunes it on 50 samples from one model and reports it works across three others. It may not transfer to other tasks or model families §5 app. H.
• Noise in the reported numbers. Single samples per problem, unstated set sizes and small per-cell differences mean a few points are within noise (e.g. the 0.5-point gap between 57.14 and 56.67 on AIME 2024 in tab. 3).
• Hidden costs. Each training trace is generated in full first, with top-5 logprobs, up to 32768 tokens build_complete_thinking_dataset.py. GRPO samples 8 or 14 completions per prompt, up to 4096 tokens each. Pruning itself is cheap; data generation and RL dominate.
• Unpinned dependencies. requirement.sh pins vLLM but not TRL or DeepSpeed requirement.sh requirement.sh. GRPO argument names and defaults may change with TRL versions. Pin them before reproducing.
• Special-token handling. The paper treats [SKIP] as one token app. E. The repo registers it only in the GRPO script grpo_train.py, and the token-count script uses the base tokenizer. Check how [SKIP] tokenises in each stage before comparing token counts or training runs.
• Qwen3 and QwQ runs. The README mentions Qwen3 and QwQ options, but only the DeepSeek-R1-Distill-Qwen-7B serving config is in the repo run_vllm_LM.sh README.md. Reproducing the other models means writing their configs.

Lineage

• Builds on: GRPO, a group-relative policy-gradient method §3.3 (arXiv:2402.03300); LoRA for the adapters app. C; DeepSeek-R1 and Qwen3 as the base models §4.2 (arXiv:2501.12948, arXiv:2505.09388). No wiki card for these yet.
• Compared against: token-efficient prompts (Lee et al. 2025), LC-Prompt and TokenSkip (Xia et al. 2025), CoT-Valve (Ma et al. 2025), R1-Compress (Wang et al. 2025) tab. 4 app. D. No wiki cards for these yet.
• Related latent-reasoning work: Coconut and iCoT, which move reasoning out of the text §2. No wiki cards for these yet.
• Superseded by: none known from the wiki's content.