기술 요약: 효율적인 추론을 위한 대규모 추론 모델용 동적 사고 토큰 선택 (DYNTS)

1. 문제 정의

DeepSeek-R1 및 Gemini-3-Pro와 같은 대규모 추론 모델(LRM)은 최종 답변을 생성하기 전 명시적인 추론 과정(Chain-of-Thought)을 생성함으로써 문제 해결 능력을 향상시킨다. 이러한 방식은 효과적이지만, 자기회귀적 디코딩(autoregressive decoding) 과정에서 발생하는 키-값(KV) 캐시 상태의 누적으로 인해 상당한 메모리 점유율과 계산 오버헤드를 초래한다.

기존의 KV 캐시 압축 방법들은 LRM에 적용될 때 두 가지 주요 한계점에 직면한다:

1. 시나리오 불일치: 긴 컨텍스트 프리필링(prefilling)을 위해 설계된 방법들은 LRM의 전형적인 특징인 짧은 프리필링 및 긴 디코딩 시나리오에 적합하지 않다.
2. 성능 저하: 긴 디코딩에 특화된 방법들은 전체 KV 캐시 베이스라인의 추론 성능을 따라잡지 못하는 경우가 많다. 이들은 대개 휴리스틱 규칙이나 국소적 어텐션 지표에 의존하는데, 이는 최종 답변에 필수적인 토큰과 불필요한 구문적 구조(syntactic scaffolding)를 구분하지 못한다. 이로 인해 결정적인 의사결정 토큰이 의도치 않게 삭제되어 추론 정확도가 크게 떨어지는 결과가 발생한다.

2. 방법론: DYNTS

저자들은 **동적 사고 토큰 선택(Dynamic Thinking-Token Selection, DYNTS)**을 제안한다. 이 방법은 추론 과정 중 불필요한 토큰을 제거하면서 "의사결정에 결정적인" 사고 토큰만을 식별하고 유지하도록 설계되었다.

2.1. 핵심 통찰: 추론에서의 파레토 법칙

어텐션 맵 분석을 통해 저자들은 추론 과정에서의 희소성(sparsity) 현상을 관찰했다:

• 질문(Question) 토큰은 일관되게 높은 중요도 점수를 보인다.
• 사고(Thinking) 토큰은 매우 희소한 분포를 보인다. 사고 토큰 중 약 20~30%만이 최종 답변 토큰으로부터 유의미한 어텐션을 받으며, 이들이 "중추적 노드(pivotal nodes)" 역할을 수행한다. 나머지 토큰들은 최종 결과에 거의 기여하지 못하며, 주로 구문적 구조를 형성하는 역할에 그친다.
• 이러한 관찰은 파레토 법칙과 일치하며, 소수의 결정적인 사고 토큰이 모델을 정답으로 이끄는 데 핵심적인 역할을 함을 시사한다.

2.2. 중요도 예측기 (Importance Predictor)

사고 토큰의 실제 중요도(최종 답변으로부터의 어텐션 가중치로 정의됨)는 전체 추론 과정이 완료된 후에만 알 수 있으므로, DYNTS는 LRM의 마지막 레이어에 통합된 학습 가능한 경량 중요도 예측기를 도입한다.

• 구조: LRM의 최종 은닉 상태(hidden state)에 부착된 경량 다층 퍼셉트론(MLP) 헤드.
• 학습: 예측기는 MATH 데이터셋을 사용하여 학습되며, 이때 정답 중요도 점수(답변 토큰에서 사고 토큰으로의 어텐션 가중치를 합산하여 계산)를 감독 신호(supervision signal)로 사용한다. 손실 함수는 예측값과 실제 값 사이의 평균 제곱 오차(MSE)를 최소화한다.
• 기능: 추론 중에 예측기는 생성되는 각 사고 토큰의 중요도 점수를 실시간으로 동적으로 추정한다.

2.3. 이중 윈도우 KV 캐시 선택 전략

DYNTS는 컨텍스트 유지와 효율성 사이의 균형을 맞추기 위해 이중 윈도우 메커니즘을 사용하여 메모리를 관리한다:

1. 질문 윈도우 ($W_q$): 모든 프롬프트 토큰의 KV 캐시를 저장한다. 이들은 무한한 중요도 점수를 부여받으며 절대 삭제되지 않는다.
2. 로컬 윈도우 ($W_l$): 국소적인 언어적 일관성을 보장하기 위해 가장 최근에 생성된 토큰들을 유지한다.
3. 선택 윈도우 ($W_s$): 사고 토큰의 이력을 저장한다. 전체 KV 캐시 예산에 도달하면, 시스템은 이 윈도우 내에서 예측된 중요도 점수가 가장 높은 상위 $k$개의 토큰을 유지하고 나머지는 삭제한다.

3. 주요 기여

• 토큰 희소성의 발견: 본 논문은 LRM에서 오직 소수의 사고 토큰(약 20~30%)만이 최종 답변에 유의미한 영향을 미치며, 대다수는 불필요하다는 것을 경험적으로 입증하였다.
• DYNTS 프레임워크: 정답에 대한 사전 지식 없이도 결정적인 추론 단계를 유지할 수 있도록, 학습된 중요도 예측기와 동적 KV 캐시 선택 전략을 결합한 새로운 방법을 제시하였다.
• 이론적 분석: 저자들은 삭제되는 토큰의 양이 모델의 은닉 차원 및 레이어 수와 비교하여 특정 임계값을 충족할 경우, 중요도 예측기가 도입되어 발생하는 오버헤드보다 불필요한 토큰을 삭제함으로써 얻는 계산적 이득이 더 크다는 것을 증명하는 "손익분기 조건(Break-Even Condition)"을 도출하였다.

4. 실험 결과

이 방법은 두 가지 LRM(DeepSeek-R1-Distill-Llama-8B 및 DeepSeek-R1-Distill-Qwen-7B)을 대상으로 6개의 벤치마크(AIME24/25, AMC23, GPQA-D, GK23EN, MATH500)에서 평가되었다.

• 추론 정확도: DYNTS는 최신 KV 캐시 압축 방법들(예: R-KV, SnapKV, H2O)보다 우수한 성능을 보인다. 동일한 메모리 예산 하에서, 가장 우수한 베이스라인 대비 Pass@1을 2.6% 향상시켰다. 결정적으로, DYNTS의 성능은 전체 KV 캐시 베이스라인과 대등한 수준(예: R1-Llama에서 61.9% vs. 61.6%)을 보여주며, 공격적인 압축이 추론 능력을 저해하지 않음을 입증하였다.
• 효율성 이득:
  • 지연 시간(Latency): 바닐라 트랜스포머(vanilla Transformers) 대비 추론 지연 시간을 1.84–2.62배 감소시켰다.
  • 메모리: 피크(peak) KV 캐시 메모리 점유율을 3.32–5.73배 줄였다.
  • 처리량(Throughput): 시퀀스 길이가 길어짐에 따라 최대 4.51배의 속도 향상을 달ach하며, 주기적인 삭제를 통해 자원 소비를 효과적으로 제한한다.
• 절제 연구(Ablation Studies): 실험을 통해 (예측기를 통한) 결정적 사고 토큰만을 유지하는 것이 무작위 또는 bottom-k 선택보다 우수함을 확인하였다. 로컬 윈도우나 질문 토큰을 제거할 경우 성능이 크게 저하되었으며, 이는 이중 윈도우 설계의 필요성을 입증한다.

5. 의의 및 주장

본 논문은 DYNTS가 자원이 제한된 환경에서 대규모 추론 모델의 핵심적인 배포 병목 현상을 해결한다고 주장한다. 사고 과정의 고유한 희소성을 활용함으로써, 이 방법은 다음과 같은 기능을 가능하게 한다:

• 확장성: 추론 성능을 희생하지 않으면서 KV 캐시 점유율을 획기적으로 줄임으로써, 제한된 메모리를 가진 하드웨어에서도 복잡한 추론 작업을 실행할 수 있게 한다.
• 효율성: 동적, 중요도 기반의 토큰 선택이 추론 작업에서 정적 또는 휴리스틱 기반의 압축보다 우월함을 이론적, 경험적으로 입증하였다.
• 일반화 가능성: 수학적 추론을 바탕으로 학습되었음에도 불구하고, 과학적 QA 및 코드 추론(HumanEval)을 포함한 다양한 데이터셋에서 견고한 성능을 보여줌으로써 토큰 희소성의 원리가 전이 가능하다는 것을 시사한다.

저자들은 본 연구가 추론 성능을 선형적인 메모리 증가와 분리함으로써, 강력한 LRM을 엣지 디바이스에 배치할 수 있는 실질적인 경로를 제공한다고 결론짓는다.