PaperlayerView exact revision on arXiv

Source-linked AI summary

TokenSkip: Controllable Chain-of-Thought Compression in LLMs

Heming Xia, Chak Tou Leong, Wenjie Wang, Yongqi Li, Wenjie Li

arXiv:2502.12067v3cs.CLcs.AI

TL;DR

As CoT grows in pursuit of stronger reasoning, whether all tokens contribute equally remains an open question. TokenSkip compresses CoT by skipping less-important tokens, reducing Qwen2.5-14B-Instruct’s GSM8K reasoning tokens by 40% (from 313 to 181) with less than a 0.4% performance drop.

  • Problem

    As longer CoT is explored for stronger reasoning, whether all CoT tokens contribute equally to deriving answers remains an open question.

    Source: §1 · pp. 1–2

  • Method

    TokenSkip uses token importance to prune less-important CoT tokens and fine-tunes models on trimmed trajectories for adjustable compression.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    Across tested models and math benchmarks, TokenSkip compresses CoT while maintaining robust reasoning; on GSM8K, Qwen2.5-14B-Instruct reduces reasoning tokens by 40% (from 313 to 181) with less than a 0.4% performance drop.

    Source: §1 · p. 2

  • Takeaways & Limitations

    TokenSkip offers a way to balance CoT token efficiency and reasoning accuracy at adjustable compression ratios.

    Source: §6 · p. 8

    The token-importance measure was not trained specifically on mathematical data, which may affect compression effectiveness for numerical tokens and mathematical expressions.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Chain-of-Thought (CoT) has been proven effective in enhancing the reasoning capabilities of large language models (LLMs). Recent advancements, such as OpenAI's o1 and DeepSeek-R1, suggest that scaling up the length of CoT sequences during inference could further boost LLM reasoning performance. However, due to the autoregressive nature of LLM decoding, longer CoT outputs lead to a linear increase in inference latency, adversely affecting user experience, particularly when the CoT exceeds 10,000 tokens. To address this limitation, we analyze the semantic importance of tokens within CoT outputs and reveal that their contributions to reasoning vary. Building on this insight, we propose TokenSkip, a simple yet effective approach that enables LLMs to selectively skip less important tokens, allowing for controllable CoT compression. Extensive experiments across various models and tasks demonstrate the effectiveness of TokenSkip in reducing CoT token usage while preserving strong reasoning performance. Notably, when applied to Qwen2.5-14B-Instruct, TokenSkip reduces reasoning tokens by 40% (from 313 to 181) on GSM8K, with less than a 0.4% performance drop. We release our code and checkpoints in https://github.com/hemingkx/TokenSkip.

1 Introduction

TokenSkip addresses the costs of long chain-of-thought outputs by exploiting unequal token importance: it trains models on pruned reasoning trajectories to skip less important tokens while connecting critical reasoning steps.

  • Recent systems including o1 and DeepSeek-R1 demonstrate that scaling CoT from hundreds to thousands of reasoning steps could continuously improve LLM reasoning (OpenAI et al., 2024; DeepSeek-AI et al., 2025).Source: §1 · p. 1
  • Longer CoTs increase inference latency and key-value-cache memory proportionally and attention computation quadratically; step-skipping may conflict with test-time scaling and impair reasoning (Ding et al., 2024; Liu et al., 2024; Jin et al., 2024; Merrill and Sabharwal, 2024; OpenAI, 2024; Snell et al., 2025).Source: §1 · pp. 1–2
  • TokenSkip prunes less important tokens from original CoT trajectories at adjustable compression ratios, then uses supervised fine-tuning to teach target models to trim redundant tokens.Source: §1 · p. 2
  • On MATH-500, LLaMA-3.1-8B-Instruct reduces CoT tokens by 30% with less than a 4% performance decline, achieving a 1.4× inference speedup.Source: §1 · p. 2
  • With LoRA, TokenSkip fine-tunes 0.2% of Qwen2.5-14B-Instruct’s parameters; training takes approximately 2 hours for the 7B model or 2.5 hours for the 14B model on two 3090 GPUs.The compressed CoT training data is no larger than the original training set, with 7,473 GSM8K examples and 7,500 MATH examples.Source: §1 · p. 2
  • TokenSkip reduces Qwen2.5-14B-Instruct GSM8K reasoning tokens by 40% (313 to 181) with less than a 0.4% performance drop.Source: §1 · p. 2

2 Background and Preliminaries

The preliminary studies find that CoT tokens differ in semantic importance and that LLMs can recover reasoning processes from compressed CoTs. They motivate token skipping while identifying autonomous skipping and shortcut discovery as challenges.

2.1 Token Importance

  • Selective Context (Li et al., 2023) measures token importance using LLM semantic confidence, a measure the authors consider applicable to CoT outputs.The motivating question is whether every CoT token contributes equally to deriving the answer.Source: §2.1 · p. 2
  • LLMLingua-2 (Pan et al., 2024) identifies position-dependent confidence and causal attention as limitations, and uses a bidirectional BERT-like model trained on GPT-4 token labels.Token importance is measured using each token’s predicted probability.Source: §2.1 · p. 3
  • Mathematical equations contribute more to answers than connectors such as “so” and “since,” indicating redundancy that could be trimmed from CoT outputs.The study applies LLMLingua-2 to CoT tokens and observes variation in their semantic importance.Source: §2.1 · p. 3

2.2 CoT Recovery

  • LLaMA-3.1-8B-Instruct can comprehend semantic information in compressed CoTs and restore the reasoning process, maintaining interpretability and enabling full CoT recovery when requested.The examples are shown for GSM8K and detailed in Appendix A.Source: §2.2 · p. 3, Figure 3 · p. 3
  • Autonomously skipping redundant CoT tokens and identifying shortcuts between critical reasoning tokens is a non-trivial challenge; the authors describe this work as the first to explore CoT compression through token skipping.The preliminary analysis motivates trimming redundancy to improve CoT efficiency while retaining the ability to restore CoT from compressed outputs.Source: §2.2 · p. 3

3 TokenSkip

TokenSkip selectively trims less important chain-of-thought (CoT) tokens and fine-tunes LLMs on the resulting trajectories to support controllable compression at adjustable ratios.

3 TokenSkip

  • TokenSkip’s training pipeline generates CoT trajectories, compresses them at sampled ratios, and fine-tunes the LLM to support inference at the selected γ values.Figure 4 illustrates mixed-ratio training and controllable inference for γ ∈ {γ0, . . . , γz}.Source: §3 · p. 3, Figure 4 · p. 4

3.1 Token Pruning

  • TokenSkip ranks CoT tokens by semantic importance and retains tokens meeting or exceeding the empirical γ-quantile threshold for the requested compression ratio.The pruning is motivated by the premise that reasoning tokens contribute differently to deriving an answer.Source: §3.1 · p. 4

3.2 Training

  • Training generates CoTs, filters incorrect answers, and fine-tunes on compressed trajectories at sampled ratios while keeping answers unchanged and including some uncompressed CoTs.The ratio is inserted after each question, and the uncompressed examples use γ = 1 to preserve reasoning capabilities.Source: §3.2 · p. 4

3.3 Inference

  • At inference, the prompt supplies the question and desired γ, after which autoregressive decoding generates compressed CoT tokens followed by an answer.The prompt uses the format Q [EOS] γ [EOS].Source: §3.3 · p. 4

4 Experiments

Experiments show TokenSkip reduces CoT token usage while preserving reasoning performance. Analyses indicate that it skips lower-importance tokens and can improve performance under a fixed length budget.

4.1 Experimental Setup

  • Experiments use LLaMA-3.1-8B-Instruct and Qwen2.5-Instruct on GSM8K and MATH-500, with LLMLingua-2 importance scores, LoRA training, and accuracy, token count, latency, and actual compression ratio metrics.The study uses greedy decoding and measures latency on one NVIDIA 3090 GPU at batch size 1; training takes about 2 hours for the 7B model and 2.5 hours for the 14B model.Source: §4.1 · p. 5

4.2 Main Results

  • Qwen2.5-14B-Instruct trims 40% of GSM8K tokens with less than 0.4% performance drop; at a 0.5 compression ratio, degradation is 2%.Larger Qwen2.5-Instruct models show less performance degradation at higher compression ratios.Source: Figure 5 · p. 5, §4.2 · p. 5
  • On GSM8K, TokenSkip reaches a 0.53 actual compression ratio with a 10% performance drop and 1.8× latency speedup, while Truncation at 0.5 causes 79% accuracy loss.On MATH-500, TokenSkip reduces CoT tokens by 30% with less than 4% performance drop; Truncation at 0.5 causes a 21% accuracy drop.Source: §4.2 · pp. 5–6, Table 1 · p. 6

4.3 Analysis

  • Adding 0.3 and 0.4 compression ratios weakens ratio adherence, likely because excessive trimming removes critical information and hinders compression training.The More Ratio variant also adheres less well overall than TokenSkip’s default settings.Source: §4.3 · p. 6
  • Skipped GSM8K CoT tokens skew toward lower importance scores, whereas retained tokens predominantly have higher scores.The comparison uses Qwen2.5-14B-Instruct generations at γ = 1.0 and γ = 0.7.Source: Figure 7 · p. 6, §4.3 · pp. 6–7
  • LLMLingua-2 outperforms Selective Context as TokenSkip’s importance metric; GPT-4o results suggest further metric gains, but API costs limit large-scale use.The paper describes LLMLingua-2’s BERT-size model as a cost-effective, efficient alternative.Source: Figure 8 · p. 7, §4.3 · p. 7
  • With the original LLM length budget on MATH-500, TokenSkip at compression ratios 0.7, 0.8, and 0.9 improves performance by 1.3–2.6 absolute points over the original LLM.This setting compares against TokenSkip evaluated with length budgets adjusted by compression ratios.Source: Figure 9 · p. 7, §4.3 · p. 7
  • Case studies show TokenSkip trims redundant tokens within reasoning steps while retaining numeric values and equations, rather than shortening CoT by reducing the number of steps.Examples include skipping semantic connectors and minimally contributive expressions.Source: §4.3 · p. 7

5 Related Work

Related work reduces CoT overhead through discrete or latent reasoning compression and addresses long-prompt costs through prompt compression. TokenSkip instead prunes CoT tokens by semantic importance using a small LM.

Efficient CoT

Prompt Compression

  • Prompt compression uses lightweight LMs to rewrite prompts, continuous tokens to represent long contexts, or filters to remove low-information tokens.Selective Context measures token importance with perplexity, while LLMLingua-2 (Pan et al., 2024) uses a small bidirectional LM trained on GPT-4 compression data, also serving as TokenSkip’s token-importance metric.Source: §5 · p. 8

6 Conclusion

TokenSkip compresses chains of thought by selectively skipping less important tokens while preserving critical ones, enabling adjustable compression ratios. Experiments across various LLMs and tasks validate its effectiveness.

  • TokenSkip is introduced as a simple, effective approach for controllable Chain-of-Thought compression.Source: §6 · p. 8
  • By selectively skipping less important CoT tokens while preserving critical ones, TokenSkip enables compressed outputs with adjustable ratios.This is intended to balance reasoning efficiency and accuracy.Source: §6 · p. 8
  • Extensive experiments across various LLMs and tasks validate TokenSkip’s effectiveness.Source: §6 · p. 8

Limitations

Computational constraints left TokenSkip untested on larger and long-CoT models, while its token-importance measure was not trained specifically on mathematical data and may be less effective on numerical tokens and expressions.

Limitations

  • Computational constraints prevented experiments with Qwen2.5-32B-Instruct, Qwen2.5-72B-Instruct, and long-CoT models such as QwQ-32B-Preview.The authors anticipate TokenSkip’s potential could be further realized on these models and plan to explore them in future work.Source: Limitations · p. 9
  • Token importance was measured using LLMLingua-2 (Pan et al., 2024), which was not trained specifically on mathematical data and may handle numerical tokens and mathematical expressions less effectively.This limitation may affect compression effectiveness.Source: Limitations · p. 9

Ethics Statement · Appendix

The authors report that their publicly released, English-language datasets protect user privacy and contain no personal information. They also state that the research artifacts are permissively licensed, used as intended, and that the work meets ACL ethics.

Ethics Statement

A CoT Recovery · B Experimental Details · B.1 Implementation Details

GPT-4o could understand compressed CoT and correctly restore the original reasoning process. The implementation uses LLMLingua-2-generated training data and LoRA fine-tuning, with specified training and inference settings.

A CoT Recovery

  • GPT-4o correctly restored the original CoT process from compressed content in a GSM8K recovery example.Figure 12 presents the recovery result; Figure 11 shows a prompt designed for LLaMA-3.1-8B-Instruct to recover compressed CoT from a GSM8K problem.Source: A CoT Recovery · p. 12, Figure 12 · p. 12, Figure 11 · p. 12

B Experimental Details · B.1 Implementation Details

  • LLMLingua-2 (Pan et al., 2024) supplies the token-importance metric, and each training sample receives a randomly selected compression ratio γ from {0.5, 0.6, 0.7, 0.8, 0.9, 1.0}.Source: §B.1 · p. 12
  • LoRA (Hu et al., 2022) fine-tuning uses rank r=8 and scaling parameter α=16 for 3 epochs on both datasets.Source: §B.1 · p. 12
  • Training uses a peak learning rate of 5e-5 with cosine decay, AdamW, and a warmup ratio of 0.1.Source: §B.1 · p. 12
  • Training is implemented with LLaMA-Factory (Zheng et al., 2024), while inference for TokenSkip and all baselines uses Huggingface transformers.Source: §B.1 · p. 12
  • Inference sets max_len to 512 tokens for GSM8K and 1024 for MATH4.Source: §B.1 · p. 12
  • Experiments use PyTorch 2.1.0, CUDA 12.1, two 24GB NVIDIA GeForce RTX 3090 GPUs, and an Intel Xeon Platinum 8370C CPU with 32 cores.Source: §B.1 · p. 12

B.2 Detailed Results with Qwen … C.1 Out-of-domain Evaluation

With Qwen2.5-Instruct, larger models show less performance degradation at higher compression ratios; TokenSkip also preserves strong accuracy in out-of-domain evaluation after fine-tuning on MATH.

B.2 Detailed Results with Qwen

  • As Qwen2.5-Instruct model scale increases, performance degrades less at higher compression ratios, suggesting larger models better identify shortcuts between critical reasoning tokens.For MATH-500, TokenSkip's length budget is adjusted to max_len×γ because many samples reach the maximum length; GSM8K uses no adjustment.Source: §B.2 · p. 12
  • Table 2 reports accuracy, average CoT token count, and actual compression ratio (ActRatio) for the Qwen2.5-Instruct series.The table covers the series' evaluation on GSM8K.Source: Table 2 · p. 12, §B.2 · p. 12

B.3 Detailed Prompts

  • The main-experiment prompts are detailed in Table 3 for prompt-based baselines.The supplied passage identifies the table's scope but does not provide prompt contents.Source: §B.3 · p. 13, Table 3 · p. 13

C.1 Out-of-domain Evaluation

  • TokenSkip was evaluated on in-domain MATH-500 and out-of-domain GSM8K and MMLU-STEM after LLaMA-3.1-8B-Instruct was fine-tuned on MATH training data.MMLU-STEM covers diverse STEM subjects from the full MMLU dataset.Source: §C.1 · p. 13
  • With 40% token trimming, TokenSkip achieves performance comparable to the original LLM on MMLU-STEM; at compression ratio 0.5, degradation is 0.4% absolute.The results also indicate that the model adheres closely to specified compression ratios while preserving accuracy.Source: §C.1 · p. 13

C.2 Evaluation Beyond Math

On CommonsenseQA, TokenSkip demonstrates generalization beyond mathematical reasoning by reducing CoT length on Qwen2.5-Instruct models without degrading performance.

C.2 Evaluation Beyond Math

  • The CommonsenseQA evaluation used 9,700 training-set samples and measured performance on the validation set.CommonsenseQA is a multiple-choice dataset requiring diverse commonsense knowledge to predict correct answers.Source: §C.2 · p. 13
  • 50% CoT-length reduction on CommonsenseQA with Qwen2.5-Instruct models produced no performance degradation.The reported results are in Table 5.Source: Table 5 · p. 13, §C.2 · p. 13

Related papers

Loading 2502.12067v3…