TokenSkip: Controllable Chain-of-Thought Compression in LLMs

https://arxiv.org/abs/2502.12067

TokenSkip: Controllable Chain-of-Thought Compression in LLMs

 

2025 emnlp main에 붙은 논문이네요 

 

여기선 사후에 CoT를 잘라내는 것이 아닌 중요도가 낮은 토큰을 없앤 데이터로 llm을 fine-tuning하여 압축된 cot를 직접 생성하게 만들려고 합니다. 

동일하게 긴 reasoning trace를 지적합니다. 

본인 trace에서 pruining을 진행하고, 학습하여 효율성을 높인다! 입니다. 

 

데이터 생성 후 틀린 Trajectory는 제거하고, 정답을 맞춘 데이터만 compressor에 넘어가게 됩니다.

여기서 LLMLingua-2의 bidirectional token classifier(GPT-4로 학습된 모델)을 통해 각 토큰에 대해 중요도를 계산하고, 지정된 비율에 맞춰 토큰을 제거합니다.

학습할 땐 비율을 0.5 ~ 1 까지 다 진행하여 원본도 학습해 reasoning 능력이 손상되지 않도록 진행함. 

학습은 LoRA로 진행하고, 입력에는 압축률또한 들어가서 압축률에 맞는 추론을 진행하도록 함. 

스킵율이 커질 수록 성능은 약해지지만 토큰도 감소하게 된다. 

0.7일 때 약간 성능이 증가하는 모습에서 중요 토큰 사이의 short cut을 잘 학습했다고 저자는 말한다. 

0.3, 0.4에선 제대로 비율대로 압축하지 못하는 것을 보이기도 한다. 

누가 토큰을 고르냐도 중요하긴 하네요 

더 좋은 모델을 사용했을 때 정답을 유지하면서 토큰 감소량이 큽니다. 

여기선 실제 CoT 압축 사례를 보여준다. 

 

해결하려는 문제	긴 Chain-of-Thought(CoT)는 추론 성능을 높이지만, autoregressive decoding으로 인해 생성 토큰 수, latency, KV-cache 메모리, attention 계산량이 증가한다. 기존의 간결화 프롬프트는 길이를 정확히 제어하지 못하고, 단순 truncation은 핵심 계산이나 정답까지 잘라 성능을 크게 떨어뜨린다.
핵심 연구 질문	CoT의 모든 토큰이 최종 정답 도출에 동일하게 중요한가? 저자들은 숫자·수식·핵심 개념은 상대적으로 중요하지만, 접속사·관사·반복 설명 등은 중요도가 낮아 제거 가능하다고 본다.
핵심 아이디어	중요도가 낮은 CoT 토큰을 제거한 압축 trajectory로 LLM을 fine-tuning하여, 추론 시 모델이 불필요한 토큰을 직접 건너뛰고 중요한 reasoning token 사이의 shortcut을 생성하도록 학습한다. 이는 reasoning step 자체를 삭제하기보다 각 step 내부의 언어적 중복을 압축하는 방식이다.
Token importance 측정	기본적으로 LLMLingua-2의 bidirectional token classifier를 사용한다. 전체 문맥을 양방향으로 보고 각 토큰의 중요 확률을 계산하기 때문에, causal LM perplexity 기반 Selective Context보다 위치 편향과 단방향 문맥 한계를 줄인다. TokenSkip이 매 trajectory를 GPT-4로 직접 라벨링하는 것은 아니며, GPT-4o는 별도 비교 실험에서 상한선에 가까운 compressor로 사용된다.
Step 1: CoT 생성 및 필터링	대상 LLM이 각 학습 문제에 대해 원본 CoT와 답을 생성한다. 이후 정답이 틀린 trajectory는 제거하여 올바른 reasoning 결과만 압축 학습에 사용한다.
Step 2: Token pruning	각 CoT 토큰 c_i의 중요도 I(c_i)를 계산하고, 지정된 비율 γ에 따라 threshold를 정한다. 중요도가 threshold 이상인 토큰만 유지하여 압축 CoT c~를 만든다. 논문에서 γ는 compression ratio라고 부르지만 실질적으로는 원본 토큰의 유지 비율에 가깝다.
압축률 설정	학습 시 γ∈{0.5, 0.6, 0.7, 0.8, 0.9, 1.0} 중 하나를 sample마다 무작위로 선택한다. γ=1.0은 원본 CoT, γ=0.5는 대략 절반 정도의 토큰을 유지하는 설정이다.
Step 3: 학습 형식	입력은 Question [EOS] γ [EOS], target은 Compressed CoT + Original Answer로 구성한다. CoT만 압축하고 최종 답은 변경하지 않는다. 원본 CoT가 포함된 γ=1.0 데이터도 함께 학습하여 기존 reasoning 능력 손상을 완화한다.
학습 방법과 비용	LoRA 기반 SFT를 사용한다. 주요 설정은 LoRA rank 8, alpha 16, 3 epochs, learning rate 5×10^{−5}이다. Qwen2.5-14B에서는 약 0.2%의 파라미터만 학습하며, RTX 3090 두 장에서 7B는 약 2시간, 14B는 약 2.5시간이 소요된다.
추론 방식	추론 시 질문과 원하는 γ만 입력하면 모델이 처음부터 압축된 CoT를 생성한다. 전체 CoT를 먼저 생성한 뒤 후처리하는 방식이 아니므로, inference 단계에서는 LLMLingua-2나 별도 compressor가 필요 없다.
평가 모델·데이터	LLaMA-3.1-8B-Instruct, Qwen2.5-3B/7B/14B-Instruct를 사용한다. 주요 벤치마크는 GSM8K와 MATH-500이며, 추가적으로 MMLU-STEM과 CommonsenseQA에서 일반화 성능을 평가한다.
비교 방법	Be concise, Only use numbers or equations, 단어 축약 프롬프트, 고정 비율 길이 제어 프롬프트, 단순 output truncation과 비교한다.
대표 결과: Qwen2.5-14B / GSM8K	원본은 정확도 93.1%, 평균 CoT 313.11 tokens이다. TokenSkip γ=0.6은 180.68 tokens로 약 40% 이상 줄이면서 정확도 92.7%, 즉 0.4%p 하락만 보였다. γ=0.7에서는 218.62 tokens로 줄이면서 정확도가 오히려 93.4%로 나타났다.
대표 결과: LLaMA-3.1-8B / GSM8K	원본은 86.2%, 213.17 tokens, 5.96초이다. γ=0.7에서는 82.5%, 150.12 tokens, 4.36초로 약 1.4× speedup을 얻었다. γ=0.5에서는 113.05 tokens와 1.8× speedup을 달성하지만 정확도는 78.2%로 하락한다.
대표 결과: LLaMA-3.1-8B / MATH-500	원본은 48.6%, 502.60 tokens, 16.37초이다. γ=0.7에서는 46.7%, 349.13 tokens, 11.55초로 토큰 약 30% 감소, 정확도 1.9%p 하락, 약 1.4× speedup을 보였다. 복잡한 수학 문제에서는 GSM8K보다 압축에 따른 성능 저하가 크다.
기존 방법 대비 장점	간결화·길이 제어 프롬프트는 목표 비율을 제대로 따르지 못했다. MATH-500에서 일부 프롬프트의 실제 유지 비율은 0.94∼0.97에 머물렀다. 반면 truncation은 목표 길이는 맞추지만 GSM8K의 0.5 설정에서 정확도가 86.2%에서 7.0%로 급락했다. TokenSkip은 압축률 제어와 정답 성능 보존을 동시에 달성한다.
중요 분석 1: 실제 중요 토큰을 남기는가?	원본 CoT에만 나타난 skipped token은 낮은 importance score에 집중되고, 압축 CoT에 유지된 token은 높은 score에 집중되었다. 이는 모델이 단순히 짧게 생성하는 것이 아니라, 학습된 중요도 패턴에 따라 불필요한 토큰을 생략한다는 근거다.
중요 분석 2: 모델 크기 효과	압축이 강할수록 대형 모델이 더 안정적이었다. 약 50% 유지 시 정확도 하락은 Qwen2.5-3B −9.3%p, 7B −5.4%p, 14B −1.7%p였다. 큰 모델일수록 중요한 reasoning state 사이의 shortcut을 더 잘 학습하는 것으로 해석된다.
중요 분석 3: 동일 token budget에서의 효과	MATH-500에서 TokenSkip에 원본 모델과 동일한 최대 생성 길이를 제공하면 γ=0.7, 0.8, 0.9에서 원본보다 1.3∼2.6%p 높은 정확도를 보였다. 이는 압축 표현이 고정된 token budget을 더 효율적으로 사용할 가능성을 보여준다.
도메인 일반화	MMLU-STEM에서는 약 40% 토큰을 줄여도 원본과 비슷하거나 더 높은 정확도를 보였고, CommonsenseQA의 Qwen2.5-14B에서는 약 50% 토큰을 줄인 γ=0.5에서도 정확도 82.1%가 유지되었다.
CoT 복원 실험	압축 CoT를 LLaMA-3.1-8B 또는 GPT-4o에 입력하면 자연어 형태의 상세 reasoning으로 복원할 수 있음을 정성적으로 보였다. 다만 복원된 설명이 실제 내부 reasoning과 동일한지, 모든 논리 정보가 보존되는지는 정량적으로 검증하지 않았다.
주요 기여	① CoT를 reasoning step 단위가 아닌 token 단위로 압축했다. ② 하나의 모델에서 여러 γ를 지원해 정확도–효율 trade-off를 조절할 수 있다. ③ inference 시 추가 compressor 없이 압축 CoT를 직접 생성한다. ④ 소규모 데이터와 LoRA만으로 재현 가능한 비교적 저비용 방법을 제시한다.
주요 한계	실제 실험의 CoT 길이는 주로 수백 토큰이며, 논문 동기에서 언급한 수천∼수만 토큰의 long-CoT 모델에서는 검증하지 않았다. Qwen2.5-32B/72B, QwQ-32B 같은 대형 reasoning 모델도 제외되었다. 또한 LLMLingua-2는 수학 전용 importance model이 아니며, 그 중요도는 실제 causal reasoning importance와 동일하지 않다.
핵심 해석	TokenSkip은 모델이 “덜 생각하게” 만드는 방식이라기보다, 동일한 논리적 계산을 더 적은 자연어 토큰으로 표현하도록 학습하는 방식이다. 논문의 핵심 성과는 특히 큰 모델에서 CoT의 언어적 중복을 크게 제거해도 reasoning 성능을 상당 부분 유지할 수 있음을 보인 것이다.