추론 인식 압축

추론 인식 압축 (Reasoning-Aware Compression)

추론(reasoning) 모델을 표준 캘리브레이션 세트로 프루닝하면 일반 LLM을 프루닝할 때보다 손상이 더 크고, 모델을 오히려 더 느리게 만들 수 있어요. Reasoning-Aware Compression (RAC)는 프루닝 솔버가 캘리브레이션하는 대상을 바꿔 이 문제를 해결합니다. SGLang으로 생성한 모델 자신의 chain-of-thought를 캘리브레이션하는 것이죠.

출처: 문서

본문

Reasoning Models Can be Accurately Pruned Via Chain-of-Thought Reconstruction(ICLR 2026)에서 온 내용입니다.

문제 (The problem)

SparseGPT와 Wanda 같은 원샷(one-shot) 프루닝 방법은 Calibration activation matrix X에 대한 레이어별 재구성 오류를 최소화해 제거할 가중치를 고릅니다:

min_{W'} || W X - W' X ||_F^2    s.t.  ||W'||_0 <= S

X는 관례적으로 프롬프트 토큰으로 만듭니다 — C4의 일부 또는 태스크 프롬프트 집합. 일반적인 서빙 워크로드에서는 프롬프트가 토큰 수를 지배하므로 이는 공정한 근사입니다.

추론 모델은 그 비율을 뒤집습니다. 쿼리 하나당 수천 개의 chain-of-thought 토큰을 생성하므로, 프루닝된 모델이 실행할 거의 모든 forward pass는 모델이 스스로 생성한 토큰에 대한 것입니다. 프롬프트에만 캘리브레이션하면 모델이 거의 방문하지 않는 분포에 맞춰 프루닝된 가중치를 최적화하는 셈이에요.

결과는 우아한 정확도 감소가 아닙니다. 프루닝된 모델은 두서없이 말하기(ramble) 시작합니다 — 더 긴 chain of thought를 생성하면서 그리고 덜 정확하게 답하므로, 프루닝이 지연을 줄이는 대신 늘립니다. MATH-500에서 50% 희소성으로, C4-캘리브레이션된 DeepSeek-R1-Distill-Qwen-7B는 원래 가속하려던 dense 모델보다 평가에 거의 여섯 배 오래 걸립니다.

해결책 (The fix)

RAC는 캘리브레이션 중에 dense 모델의 자체 on-policy rollout을 샘플링하고 프롬프트·디코드 activation을 함께 재구성합니다:

X_RAC = [ X_prompt , X_decode ]

솔버는 건드리지 않으므로, 기존 SparseGPT 또는 Wanda 워크플로우에 그대로(drop-in) 적용할 수 있는 변경입니다.

DeepSeek-R1-Distill-Qwen-7B, MATH-500, 50% 희소성의 SparseGPT, 1M 캘리브레이션 토큰:

Calibration set acc@1 Eval wall clock
Dense (no pruning) 0.936 23.3 min
C4 0.744 135.0 min
Task prompts only 0.812 115.6 min
RAC (prompts + on-policy CoT) 0.900 35.3 min

DeepSeek-R1-Distill-Qwen (1.5B–32B)과 Qwen3 (1.7B–14B) 전반에서, 논문은 RAC가 50% 희소성에서 dense 정확도의 최대 95%를 유지하며 prompt-only 캘리브레이션보다 최대 17포인트 개선한다고 보고합니다.

사용하기 (Using it)

SGLang은 이 레시피를 실행 가능한 예제로 examples/usage/reasoning_aware_compression에 세 단계로 제공합니다:

Phase Script What it does
I rac_collect_traces.py sgl.Engine이 on-policy CoT 트레이스를 캘리브레이션 세트로 샘플링
II rac_prune.py llm-compressor가 그 activations에 대해 SparseGPT/Wanda 실행
III rac_serve_and_eval.py SGLang이 sparse 체크포인트를 서빙하고 MATH-500 채점

Phase I이 비용이 큰 단계입니다 — 논문 예산은 1M on-policy CoT 토큰 — 그리고 SGLang의 배치 생성이 여기서 일을 합니다. Phase II는 프루닝 솔버를 llm-compressor에 위임하는데, 이것은 SGLang 의존성이 아닙니다. pip install "llmcompressor>=0.12.0"로 별도 설치하세요.

cd examples/usage/reasoning_aware_compression

python rac_collect_traces.py \
    --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
    --dataset open-r1/OpenR1-Math-220k --prompt-column problem \
    --target-tokens 1000000 --output-dir ./rac_traces_math

python rac_prune.py \
    --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
    --calibration ./rac_traces_math/traces.jsonl \
    --sparsity 0.5 --output-dir ./rac_pruned_50

python -m sglang.launch_server --model-path ./rac_pruned_50

예제 README는 같은 프롬프트로 논문의 prompt-only 베이스라인을 만들어 캘리브레이션 전략을 정면으로 비교할 수 있게 안내합니다.

프루닝된 추론 모델 평가 (Evaluating a pruned reasoning model)

정확도만으로는 위에서 설명한 실패 모드를 숨깁니다. 프루닝된 추론 체크포인트를 비교할 때는 항상 **평균 완성 길이(mean completion length)**와 wall clock을 정확도와 함께 보고하세요 — 2점 덜 맞으면서 chain of thought를 세 배로 뱉는 모델은 좋은 트레이드가 아닙니다. rac_serve_and_eval.py는 세 가지를 모두 보고합니다.

  • Quantization — 서빙 시점에 적용되는 모델 압축의 다른 축.

더 알아보기 (Learn more)