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

Longer CoT can improve reasoning but increases inference costs, motivating a way to balance efficiency and accuracy. TokenSkip measures token importance, skips less important tokens through compressed-CoT fine-tuning, and achieves substantial token reductions with limited performance loss across models and tasks. The strongest reported result is a 40% reduction on GSM8K for Qwen2.5-14B-Instruct with less than a 0.4% performance drop.

  • Problem

    Longer CoT improves reasoning but increases latency, memory, and attention costs, while aggressive skipping can impair reasoning performance.

  • Method

    TokenSkip measures semantic token importance, prunes less important tokens from CoT trajectories, and fine-tunes LLMs on compressed examples to skip redundant tokens at adjustable ratios.

  • Results

    40% fewer reasoning tokens, from 313 to 181, produces less than a 0.4% performance drop for Qwen2.5-14B-Instruct on GSM8K.

  • Takeaways & Limitations

    TokenSkip enables controllable CoT compression while preserving strong reasoning performance across the evaluated models and math benchmarks.

  • Takeaways & Limitations

    Experiments excluded larger and long-CoT models, and the LLMLingua-2 importance metric was not specifically trained for mathematical data.

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 efficiency–accuracy trade-off in long Chain-of-Thought reasoning by selectively skipping less important tokens and learning shortcuts between critical tokens. Experiments show substantial token reductions with limited reasoning-performance loss across models and benchmarks.

  • Motivation: Longer CoT sequences improve reasoning but increase inference latency, memory usage, and attention computation, especially for thousands of reasoning steps.Prior skipping methods may impair reasoning performance, leaving efficiency and accuracy difficult to balance.
  • Approach: TokenSkip analyzes token-level semantic importance and trains LLMs on compressed CoT trajectories to skip redundant tokens at adjustable ratios.The method prunes unimportant tokens from original trajectories, then applies supervised fine-tuning to target models.
  • Results: 40% fewer reasoning tokens, from 313 to 181, yields less than a 0.4% performance drop for Qwen2.5-14B-Instruct on GSM8K.This is the headline result reported for TokenSkip.
  • Results: 30% fewer CoT tokens on MATH-500 with less than a 4% performance decline produces a 1.4× inference speedup for LLaMA-3.1-8B-Instruct.The experiments cover LLaMA-3.1-8B-Instruct, the Qwen2.5-Instruct series, GSM8K, and MATH-500.
  • Practical considerations: TokenSkip uses low-cost fine-tuning and compressed training data no larger than the original set, supporting efficient and cost-effective deployment.For Qwen2.5-14B-Instruct, LoRA fine-tunes 0.2% of model parameters.

2 Background and Preliminaries

This section examines whether CoT tokens contribute equally to reasoning and reviews token-importance measurements for identifying redundancy. It finds varying token contributions, limitations in unidirectional confidence measures, and evidence that compressed CoTs can be restored.

  • CoT efficiency: Token redundancy in CoT outputs motivates removing less useful tokens to improve CoT efficiency.The section frames whether every token contributes equally to deriving the answer as a central research question.
  • Token importance: Selective Context measures token importance using an LLM’s semantic confidence and can be applied to tokens generated in CoTs.Here, x_i denotes a token and M_L denotes the LLM used to compute token confidence.
  • Token importance: LLMLingua-2 addresses position dependence and incomplete context in causal-LM importance scores by using a bidirectional BERT-like language model.Its method uses GPT-4 token labels and a token-classification objective to estimate importance probabilities.
  • Token importance: Mathematical equations tend to contribute more to final answers, whereas semantic connectors such as “so” and “since” generally contribute less.This observation supports treating semantic importance as uneven within CoT outputs.
  • CoT recovery: LLMs can restore CoT processes from compressed outputs, preserving interpretability and allowing complete reasoning to be presented when needed.Examples using LLaMA-3.1-8B-Instruct demonstrate recovery from compressed CoTs.
  • CoT recovery: The analysis motivates token skipping as a non-trivial approach for trimming redundant tokens and learning shortcuts between critical reasoning tokens.The authors describe this work as the first to explore CoT compression through token skipping.

3 TokenSkip

TokenSkip compresses chain-of-thought trajectories by pruning less important tokens and fine-tuning the target LLM on compressed trajectories with adjustable ratios. At inference, users specify a ratio to obtain compressed CoT outputs autoregressively.

  • Training: Mixed-ratio fine-tuning enables controllable CoT inference at desired compression ratios.The training process uses compressed CoTs with multiple ratios sampled from the ratio set.
  • Token pruning: TokenSkip prunes redundant CoT tokens according to their semantic importance, retaining tokens that meet a compression-dependent threshold.Importance scores are ranked, and the empirical γ-quantile determines which tokens remain.
  • Training: Training generates target-LLM CoTs, filters incorrect answers, and prunes the remaining trajectories using ratios sampled from a predefined ratio set.Each compressed sample includes the question, compressed CoT, answer, and its compression ratio.
  • Training: Answer tokens remain unchanged during compression, while some original CoTs are retained with γ set to 1 to preserve reasoning capabilities.Compression is applied solely to CoT sequences.
  • Inference: At inference, TokenSkip accepts a question followed by a desired γ and autoregressively predicts an output containing CoT tokens and an answer.The prompt format is Q [EOS] γ [EOS].

4 Experiments

Experiments evaluate TokenSkip across models, benchmarks, baselines, compression settings, and analyses of token importance, ratio adherence, length budgets, and qualitative behavior.

  • Experimental Setup: TokenSkip is evaluated on LLaMA-3.1-8B-Instruct and Qwen2.5-Instruct using GSM8K and MATH-500.Training uses the respective dataset training sets; MATH evaluation uses MATH-500.
  • Experimental Setup: TokenSkip uses LLMLingua-2 for token importance, randomly samples compression ratios from 0.5 to 1.0, and trains with LoRA.Training takes approximately 2 hours for the 7B model and 2.5 hours for the 14B model on 3090 GPUs.
  • Main Results: Qwen2.5-14B-Instruct loses less than 0.4% performance with 40% token trimming on GSM8K.At compression ratio 0.5, performance degradation is 2%.
  • Main Results: 79% GSM8K and 21% MATH-500 accuracy drops occur for Truncation at compression ratio 0.5, whereas TokenSkip preserves reasoning performance.Prompting baselines fail to achieve desired compression ratios; TokenSkip adheres to varied target ratios.
  • Analysis: Skipped tokens have lower importance while retained tokens have higher importance, supporting TokenSkip’s selective token-discarding behavior.At lower ratios of 0.3 and 0.4, ratio adherence degrades, likely because excessive trimming loses critical information.
  • Analysis: TokenSkip using LLMLingua-2 outperforms the Selective Context variant, while GPT-4o indicates that token-importance metrics could improve further.GPT-4o’s API cost makes it impractical for large-scale datasets, whereas LLMLingua-2 offers a cost-effective alternative.
  • Analysis: With the original maximum length budget, TokenSkip outperforms the original LLM by 1.3 to 2.6 absolute points at compression ratios 0.7, 0.8, and 0.9 on MATH-500.The comparison uses the same length budget rather than a reduced budget based on compression ratio.

5 Related Work

Related work reduces CoT or prompt redundancy through simplification, skipping, parallel generation, latent representations, concise prompting, and token-importance filtering.

  • Efficient CoT: Efficient-CoT research simplifies or skips reasoning steps, generates them in parallel, or compresses CoTs into latent representations.These approaches aim to reduce computational overhead while retaining CoT benefits.
  • Prompt Compression: Prompt compression methods generate concise prompts, use implicit continuous tokens, or filter low-information tokens from natural-language inputs.Selective Context uses LLM perplexity, while LLMLingua-2 uses a small bidirectional language model for token-importance measurement.
  • Prompt Compression: TokenSkip differs by applying semantic token-importance filtering to CoT generation, enabling adjustable compression ratios for reasoning outputs.The supplied related-work passages establish the neighboring prompt-compression approaches but do not provide comparative results here.

6 Conclusion

The paper introduces TokenSkip for controllable CoT compression by selectively skipping less important tokens while preserving critical ones.

  • Conclusion: TokenSkip enables LLMs to generate compressed CoTs with adjustable ratios by selectively skipping less important tokens and preserving critical ones.The stated goal is to balance reasoning efficiency and accuracy.
  • Conclusion: Extensive experiments across LLMs and tasks validate TokenSkip’s effectiveness.

Limitations

The study did not evaluate TokenSkip on larger Qwen2.5 models because of computational constraints. Its token-importance measure was derived from LLMLingua-2 and was not specifically trained on mathematical data, which may affect compression effectiveness.

  • Larger Qwen2.5-32B-Instruct and Qwen2.5-72B-Instruct models were not evaluated because of computational constraints.
  • The authors believe TokenSkip could achieve a more favorable reasoning-performance and CoT-token-usage trade-off on those larger models.
  • The token-importance measure, derived from LLMLingua-2, was not specifically trained on mathematical data and may affect compression effectiveness.

Ethics Statement

The experiments use publicly released English datasets labeled through human interaction, with privacy protections and permissively licensed scientific artifacts.

  • The datasets are publicly released and labeled through interaction with humans in English.
  • The datasets contain no personal information, and user privacy is protected.
  • The scientific artifacts are available for research with permissive licenses and are used consistently with their intended purpose.

A CoT Recovery

The paper evaluates recovery of compressed CoT using prompts for LLaMA-3.1-8B-Instruct and GPT-4o, with GPT-4o correctly restoring the original reasoning process.

  • GPT-4o correctly restores the original CoT process from compressed CoT content.
  • The recovery experiment asks a model to recover compressed mathematical CoT output into its original full content.
  • Figure 11 presents the recovery prompt for LLaMA-3.1-8B-Instruct on a GSM8K math problem.
  • Figure 12 presents GPT-4o recovery of compressed CoT for a GSM8K math word problem.
  • LLMLingua-2 is used as the token-importance metric for generating compressed CoT training data.

B.2 Detailed Results with Qwen

Results on Qwen2.5-Instruct and out-of-domain benchmarks show that TokenSkip can compress CoT while preserving reasoning performance, with larger models degrading less at higher compression.

  • Larger Qwen2.5-Instruct models show less performance degradation at higher compression ratios.This indicates stronger ability to identify shortcuts between critical reasoning tokens.
  • Table 2 reports accuracy, average CoT token count, and actual compression ratio for Qwen2.5-Instruct comparisons.
  • TokenSkip maintains strong generalizability out of domain while adhering closely to specified compression ratios and preserving accuracy.
  • 40% token trimming on MMLU-STEM yields performance comparable to the original LLM.
  • At a compression ratio of 0.5, MMLU-STEM performance declines by only 0.4% absolute.

C.2 Evaluation Beyond Math

TokenSkip generalizes beyond mathematical reasoning: on CommonsenseQA, Qwen2.5-Instruct models reduce CoT length by 50% without performance degradation.

  • CommonsenseQA requires diverse commonsense knowledge to predict correct answers.
  • The CommonsenseQA evaluation uses 9,700 training samples and evaluates on the validation set.
  • TokenSkip reduces CoT length by 50% on CommonsenseQA without performance degradation.
  • Table 5 reports accuracy, average CoT token count, and actual compression ratio for Qwen2.5-Instruct models on CommonsenseQA.
Loading 2502.12067v3…