TokenSkip: Controllable Chain-of-Thought Compression in LLMs

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.

한 줄 요약

TokenSkip는 CoT 내 불필요 토큰을 선택적으로 건너뛰며 추론 효율성을 향상시키는 LLM 압축 기법이다.

핵심 기여도

CoT 내 토큰의 의미적 중요도가 다르다는 통찰을 바탕으로 TokenSkip을 제안.
Qwen2.5-14B-Instruct에서 GSM8K 데이터셋에서 40% 토큰 감소 (313 → 181)와 0.4% 미만 성능 저하 달성.
LoRA를 활용해 전체 파라미터의 0.2%만 fine-tuning하여 훈련 비용 절감.
압축률과 성능 저하 간의 균형을 실험적으로 검증.

핵심 아이디어

기존 CoT는 추론 성능 향상을 위해 길이를 늘리지만, 이는 추론 지연을 유발한다. 본 연구는 CoT 내 토큰이 모두 동일한 중요도를 갖지 않는다는 점을 발견하고, TokenSkip이라는 새로운 접근법을 제안한다. 이는 의미가 적은 토큰을 선택적으로 건너뛰어 CoT를 압축하면서도 핵심적인 추론 토큰은 유지함으로써, 추론 효율성과 정확도를 동시에 달성하는 방식이다. TokenSkip은 압축률을 조절할 수 있는 학습된 모델로, 학습 단계에서는 원본 CoT 트레이젝토리에서 불필요 토큰을 제거한 데이터로 fine-tuning을 수행한다.

기술적 접근법

**TokenSkip**: CoT 내 불필요 토큰을 건너뛰는 학습된 압축 기법.
**LoRA**: 전체 파라미터의 0.2%만 fine-tuning하여 훈련 비용 절감.
**데이터셋**: GSM8K (7,473 예제), MATH-500 (7,500 예제) 사용.
**훈련 환경**: 2개의 NVIDIA 3090 GPU에서 7B 모델은 약 2시간, 14B 모델은 2.5시간 소요.
**평가 지표**: 정확도, CoT 토큰 수, 추론 지연, 압축률.

주요 결과

**Qwen2.5-14B-Instruct**: GSM8K에서 CoT 토큰 40% 감소 (313 → 181), 성능 저하 0.4% 미만.
**LLaMA-3.1-8B-Instruct**: MATH-500에서 CoT 토큰 30% 감소, 성능 저하 4% 미만, 추론 속도 1.4× 향상.
압축률과 성능 간의 일관된 관계를 실험적으로 검증.

의의 및 한계

TokenSkip은 CoT 추론 효율성을 향상시키는 새로운 방식으로, 토큰 수를 줄이면서도 성능 저하를 최소화함으로써 LLM의 실용성을 높인다. 특히, LoRA 기반의 저비용 훈련과 재현 가능한 압축 데이터셋은 실제 배포 환경에서 유용하다. 그러나 압축률이 높아질수록 성능 저하가 발생할 수 있으며, 모든 CoT가 동일한 압축률에 반응하지 않을 가능성도 있다. 또한, TokenSkip은 수학 문제에 초점을 맞춘 실험으로, 다른 유형의 추론 작업에서는 효과가 달라질 수 있다.

실용적 활용

TokenSkip은 대규모 언어 모델을 활용한 실시간 추론 시스템, 특히 수학 문제 해결, 코드 생성, 복잡한 질의 응답 등에서 추론 지연을 줄이면서도 정확도를 유지할 수 있는 기술로 활용 가능하다. 특히, 클라우드 기반 LLM 서비스나 모바일 기기에서의 효율적인 추론에 적합하다.