KV 캐시를 활용한 빠른 샘플링과 효율적 추론

초록

본 논문은 대규모 언어 모델(LMM)에서 자동 회귀 디코딩 시 생성되는 키‑밸류(KV) 캐시를 별도의 임베딩으로 재활용함으로써, 추가 연산·메모리 비용 없이 두 가지 작업—Chain‑of‑Embedding과 Fast/Slow Thinking 전환—을 구현한다. 실험 결과, Llama‑3.1‑8B‑Instruct와 Qwen2‑7B‑Instruct에서 경쟁력 있는 정확도를 유지하면서 토큰 생성량을 최대 5.7배 절감한다.

상세 분석

이 연구는 KV 캐시가 단순히 과거 토큰의 키·밸류 텐서를 저장해 어텐션 연산을 가속화한다는 기존 인식을 넘어, “경량 표현”으로서의 활용 가능성을 체계적으로 탐색한다. 핵심 아이디어는 KV 텐서를 층·헤드·위치 차원을 적절히 풀링하여 토큰‑레벨 임베딩을 만든 뒤, 이를 기존 Hidden‑State 기반 방법과 동일한 수학적 프레임워크에 대입하는 것이다.

첫 번째 응용인 Chain‑of‑Embedding(CoE)에서는 기존 연구가 레이어별 히든 상태 평균을 사용해 ‘경로 임베딩’과 변위·각도 메트릭(Δr, Δθ)을 계산한다. 저자들은 이를 토큰 차원으로 전환해 KV‑CoE를 정의한다. 구체적으로 각 레이어 l의 키·밸류 K(l,t), V(l,t)를 flatten하고, 모든 레이어에 걸쳐 평균을 취해 e_t를 얻는다. 이후 e_t 간의 유클리드 거리와 코사인 각을 구해 Δr_t, Δθ_t를 산출하고, 기존 CoE‑R·CoE‑C와 동일한 가중합 혹은 복소수 평균으로 최종 점수를 만든다. 이 과정은 추가적인 활성값 저장이나 재연산이 필요 없으며, 메모리 사용량이 거의 0에 가깝다.

두 번째 응용인 Fast/Slow Thinking Switch는 토큰 생성 과정 중 KV‑CoE 점수를 실시간으로 모니터링해 ‘난이도’를 추정한다. 점수가 사전에 정의된 임계값을 초과하면 ‘Slow’ 모드(더 많은 토큰·깊은 추론)로 전환하고, 그렇지 않으면 ‘Fast’ 모드(간단한 토큰 예측)로 유지한다. 이 전환은 별도의 리워드 모델이나 재디코딩 없이 제어 토큰만 삽입해 구현된다. 실험에서는 Qwen3‑8B와 DeepSeek‑R1‑Distil‑Qwen‑14B에 적용해 평균 토큰 수를 3.25.7배 줄였으며, 정확도 저하가 12% 수준에 머물렀다.

KV‑Cache 기반 임베딩이 전통적인 학습된 임베딩보다 전반적인 의미론적 일관성에서는 열위에 있지만, (i) 후보 집합이 제한된 상황, (ii) 상대적 순위 판단이 핵심인 작업에서는 충분히 강력함을 보였다. 특히, KV‑Cache는 d_head ≪ d_model이라는 차원 축소 특성에도 불구하고, 레이어·헤드 정보를 모두 포함하므로 로컬 패턴(예: 토큰 간 변화량) 탐지에 유리하다.

또한, 메모리 효율성 측면에서 Figure 1이 보여주듯, 동일 컨텍스트 길이에서 “Model + KV Cache” 구성은 “Model + KV Cache + Hidden States” 대비 최대 1.86배 적은 VRAM을 사용한다. 이는 대규모 서비스 환경에서 비용 절감과 응답 지연 감소에 직접적인 이점을 제공한다.

한계점으로는 KV‑Cache가 전역적인 의미론적 비교가 필요한 검색·재구성 작업에는 부적합하고, anisotropic 특성으로 인해 정규화·스케일링이 필요하다는 점을 인정한다. 향후 연구에서는 KV‑Cache에 경량 정규화 레이어를 추가하거나, 사전 학습 단계에서 KV‑Cache 친화적 목표 함수를 도입해 표현력을 강화할 여지가 있다.

요약하면, 이 논문은 “이미 존재하는 KV‑Cache를 어떻게 재활용하느냐”라는 실용적 질문에 대한 구체적 솔루션을 제시하고, 두 가지 실제 응용에서 비용‑효율성·성능‑트레이드오프를 정량적으로 입증함으로써 LLM 추론 파이프라인에 새로운 설계 패러다임을 제시한다.