🤖 machine learning

DistillCache: KL-Guided Adaptive KV-Cache Eviction for Memory-Efficient LLM Inference

DistillCache는 내부 모델 신호와 KL-divergence 보상을 기반으로 KV-캐시 토큰을 적응적으로 제거하는 경량화된 정책을 학습하는 강화 학습 프레임워크를 도입하여, 기존의 휴리스틱 및 동시적 방식들과 비교했을 때 롱 컨텍스트 LLM 추론에 있어 우수한 메모리 효율성과 정확도 유지 성능을 달성한다.

원저자: Asaad Althoubi

게시일 2026-08-11

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Asaad Althoubi

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

친구에게 아주 긴 이야기를 들려주고 있는 동안 그 이야기를 기억하려고 애쓰는 모습을 상상해 보세요. 이야기가 길어질수록, 당신의 뇌는 지금까지 말했던 모든 세부 사항들로 가득 차서 다음에 무슨 말을 할지 생각하는 것조차 힘들어집니다. 이것이 바로 현대의 "대규모 언어 모델(Large Language Models, AI 챗봇)"이 직면한 문제입니다. 이 AI 뇌들은 글쓰기, 코딩, 퍼즐 풀기에 매우 뛰어나지만, 지금까지 말한 모든 내용을 "KV 캐시(KV cache)"라고 불리는 특별한 메모리 영역에 기억하며 작동합니다. 문제는 대화가 길어질수록 이 메모리가 점점 더 커져서, 결국 컴퓨터의 메모리를 가득 채우고 속도를 느릿느릿하게 만든다는 점입니다.

이를 해결하기 위해, 과학자들은 엄격한 사서처럼 행동하며 중요해 보이지 않는 이야기의 옛 페이지들을 버리는 방법을 시도해 왔습니다. 어떤 사서들은 단순히 가장 오래된 페이지를 버리고, 다른 사서들은 사람들이 거의 들여다보지 않은 페이지를 버립니다. 하지만 이러한 규칙들은 너무 단순해서, 이야기의 결말을 쥐고 있는 아주 작은 단어 하나를 실수로 버려버려 AI가 줄거리를 잊거나 엉뚱한 소리를 하게 만들 수도 있습니다. 여기서 큰 질문이 생깁니다. 어떻게 하면 AI가 더 똑똑한 사서가 되도록 가르칠 수 있을까요? 즉, 어떤 기억을 남기고 어떤 것을 버려야 할지 정확히 알아서, 메모리 공간이 부족해지지 않고도 긴 이야기를 계속할 수 있게 하는 방법 말입니다.

여기에, 시행착오 학습이라는 영리한 방식을 사용하여 AI가 이러한 결정을 내릴 수 있도록 가르치는 새로운 방법인 DistillCache가 등장했습니다. 고정된 규칙 책을 따르는 대신, DistillCache는 비디오 게임을 배우는 학생처럼 행동합니다. 이 방식은 오래된 기억을 버리는 다양한 전략들을 시도해 보고, 만약 모든 기억을 다 가지고 있었을 때의 예측값과 현재의 예측값이 얼마나 잘 일치하는지에 따라 "점수"를 받습니다. 만약 AI가 중요한 것을 잊어버려 예측이 빗나가면 낮은 점수를 받습니다. 반대로 적절한 기억을 유지하여 흐름을 잘 따라간다면 높은 점수를 받습니다. 시간이 흐르면서 AI는 엄격한 메모리 제한 내에서 어떤 토큰(텍스트의 작은 조각)을 유지하고 어떤 것을 제거할지를 알려주는 하나의 "정책(policy)", 즉 본능을 학습하게 됩니다.

연구진은 이 방법을 인기 있는 AI 모델인 Mistral-7B에 테스트했습니다. 그 결과, AI가 평소 메모리의 **25%**만 유지하도록 강제했을 때도, DistillCache는 전체 메모리를 모두 사용했을 때의 정확도의 **94.2%**를 여전히 달성할 수 있었습니다. 이는 엄청난 성과입니다. 왜냐하면 단순한 규칙을 사용하는 다른 방식들(H2O나 SnapKV 등)은 동일한 압박 조건에서 정확도가 훨씬 더 크게 떨어졌기 때문입니다. 비슷한 시기에 개발된 다른 똑똑한 학습 기반 방식들과 비교해도, DistillCache는 많은 긴 이야기 작업에서 최대 2.7 포인트의 정확도 차이를 보이며 우위를 점했습니다.

DistillCache가 특별한 이유는 그것이 학습하는 방식에 있습니다. 다른 방식들이 과거에 특정 단어가 얼마나 자주 언급되었는지를 살피는 반면, DistillCache는 단어의 "미래 영향력"을 살핍니다. 이 방식은 "지금 이 단어를 버린다면, AI의 다음 문장이 이상해질까?"라고 묻습니다. 또한, AI의 예측이 원래의 전체 메모리 버전과 여전히 비슷하게 들리는지 확인하기 위해 KL-divergence라는 수학적 척도를 사용합니다. 논문은 이러한 접근 방식이 메모리가 꽉 짜여 있는 상황에서도 AI의 논리 구조를 온전하게 유지하는 데 특히 효과적임을 보여줍니다.

하지만 이 논문은 이것이 모든 것을 해결하는 마법 지팡이는 아니라는 점을 주의 깊게 명시하고 있습니다. 이 똑똑한 사서를 훈련시키는 데는 시간과 추가적인 컴퓨터 자원(강력한 그래픽 카드로 약 130시간)이 필요합니다. 왜냐하면 AI가 차이점을 배우기 위해 모든 메모리를 가진 상태와 메모리를 줄인 상태를 각각 한 번씩, 즉 매 단계마다 두 번씩 연습해야 하기 때문입니다. 또한, DistillCache가 일반적인 긴 이야기에는 훌륭하지만, RLKV라는 다른 방식이 중간 정도의 메모리 수준에서는 특정 논리 퍼즐에서 약간 더 나은 성능을 보이기도 했습니다. 다만, 메모리가 극도로 제한될 때는 DistillCache가 그 격차를 따라잡았습니다.

결론적으로, DistillCache는 AI의 메모리를 관리하는 최선의 방법은 정적인 규칙이 아니라, "내가 여전히 말이 되고 있는가?"를 끊임없이 확인하는 학습된 직관임을 시사합니다. 이를 통해 AI 모델이 방대한 양의 컴퓨터 메모리 없이도 훨씬 더 긴 대화와 문서를 처리할 수 있게 하며, 잠재적으로 긴 문맥을 다루는 AI를 더 빠르고 접근 가능하게 만듭니다. 연구진은 또한 이 "똑똑한 사서"가 한 종류의 AI 모델에서 훈련되었더라도 추가 훈련 없이 다른 모델에도 사용될 수 있다는 점을 발견했는데, 이는 이러한 메모리 본능이 보편적일 수 있음을 시사합니다.

기술 요약: DistillCache

문제 정의

Transformer 기반 대규모 언어 모델(LLM)은 자기회귀적 디코딩(autoregressive decoding) 과정에서 토큰 표현을 재계산하는 것을 피하기 위해 키-값(KV) 캐시에 의존한다. 그러나 KV 캐시의 크기는 시퀀스 길이에 따라 선형적으로 증가하며, 이는 긴 문맥(long-context) 추론의 컨텍스트 윈도우와 처리량(throughput)을 제한하는 심각한 메모리 병목 현상을 초래한다.

기존 솔루션들은 주로 정적인 어텐션 통계나 위치 신호에 기반하여 토큰을 유지하는 휴리스틱 제거 전략(예: H2O, SnapKV)에 의존한다. 저자는 이러한 휴리스틱이 특히 장거리 의존성이나 구조적 추론이 포함된 시나리오에서, 토큰의 중요성이 과거의 어텐션 패턴이 아닌 다운스트림 생성 결과에 따라 결정됨에도 불구하고, 토큰의 미래 예측 영향력을 포착하는 데 실패하는 경우가 많다고 주장한다.

방법론: DistillCache

본 논문은 DistillCache를 제안한다. 이는 KV 캐시 제거를 엄격한 메모리 예산 하에서의 순차적 의사결정 문제로 정식화한 강화 학습(RL) 프레임워크이다. 휴리스틱 방식과 달리, DistillCache는 전체 캐시 모델의 출력 분포를 보존하는 토큰의 부분 집합을 유지하도록 정책(policy)을 학습한다.

핵심 구성 요소

상태 공간 (특징 표현):
캐시 내 각 토큰에 대해, 정책은 디코딩 중에 사용 가능한 모델 내부 신호로부터 유도된 5차원 특징 벡터를 관찰한다:
- $\alpha_i$ : 모든 헤드와 레이어에 걸쳐 해당 토큰이 받은 평균 어텐션 질량(attention mass).
- $\|v_i\|_2$ : 토큰 가치 벡터(value vector)의 $\ell_2$ 노름.
- $\sigma_i$ : 헤드 간에 수신된 어텐션 질량의 분산.
- $H_i$ : 해당 토큰에 어텐션을 주는 최근 쿼리들의 어텐션 가중 평균 엔트로피 (문맥 의존적 중요성 포착).
- $pos_i$ : 정규화된 시퀀스 위치.
행동 공간:
경량 멀티 레이어 퍼셉트론(MLP) 정책 네트워크는 캐시된 각 토큰에 스칼라 점수를 할당한다. 훈련 중에는 이 점수에 검멜(Gumbel) 노이즈를 더하여 확률적 랭킹을 유도함으로써, 플래켓-루스(Plackett–Luce) 분포로부터 예산( $B$ )만큼의 상위- $B$ 토큰 부분 집합을 샘플링할 수 있게 한다. 추론 시에는 상위- $B$ 토큰이 결정론적으로 선택된다.
훈련 아키텍처 (교사-학생 증류 설정):
정답(ground-truth) 레이블 없이 정책을 훈련하기 위해, 저자는 교사-학생 증류 방식을 채택한다:
- 교사 캐시(Teacher Cache): "골드(gold)" 다음 토큰 분포( $P_t$ )를 계산하기 위해 압축되지 않은 전체 토큰 이력을 유지한다.
- 학생 캐시(Student Cache): 정책의 결정에 따라 압축된 이력을 유지하며 압축된 분포( $\hat{P}_t$ )를 계산한다.
- 보상 신호: 각 단계에서의 보상은 교사와 학생 분포 사이의 음의 쿨백-라이블러(KL) 발산으로 정의된다: $R_t = -D_{KL}(P_t \parallel \hat{P}_t)$ . 이 목적 함수는 어텐션 점수와 같은 대리 신호에 의존하는 대신, 매 디코딩 단계마다 분포적 충실도(distributional fidelity)를 직접 최적화한다.
최적화:
정책은 엔트로피 정규화와 분산 감소 베이스라인(Leave-One-Out)을 사용하는 REINFORCE 알고리즘을 통해 최적화된다. 훈련 과정은 오프라인으로 수행되며, 일시적으로 메모리 사용량이 두 배가 되지만, 추론 시에는 특징 추출 및 정책 점수 산출을 위한 미미한 오버헤드만 발생하며 단 한 번의 순방향 패스(forward pass)만 필요하다.

주요 기여

캐시 제거를 위한 RL 프레임워크: 정적인 휴리스틱 점수를 넘어, 메모리 제약 조건 하에서 다운스트림 분포 충실도를 위해 토큰 유지를 최적화하는 범용 RL 프레임워크인 DistillCache를 도입하였다.
분포 보존 목적 함수: 매 단계별 KL 발산을 최적화하는 것이 단순한 어텐션 점수와 같은 대리 신호에 의하는 것보다 장기적 트레이드오프를 명시적으로 학습함으로써, 실질적인 성능 향상을 가져온다는 것을 입증하였다.
포괄적 평가: Mistral-7B-Instruct-v0.3에 대한 실험과 LongBench, 추론 벤치마크(GSM8K, BBH)에 대한 Llama-3-8B-Instruct로의 제로샷 전이 성능을 통해 광범위하게 평가하였다.

실험 결과

Mistral-7B-Instruct-v0.3에 대해 25% KV 캐시 예산을 적용하여 평가한 결과:

LongBench 성능: DistillCache는 전체 캐시 정확도의 **94.2%**를 유지하며, 강력한 휴리스틱 베이스라인(H2O, SnapKV) 대비 최대 2.7%p, 동시기의 RL 기반 방법들(ForesightKV, RLKV) 대비 최대 1.4%p 높은 성능을 보였다.
추론 작업: 추론 벤치마크에서 DistillCache는 가장 우수한 동시기 방법들과 경쟁할 만한 수준이며, 공격적인 압축(12.5% 예산) 상황에서는 이들을 능가하였다.
모델 간 전이: Mistral-7B에서 훈련된 정책은 Llama-3-8B-Instruct로 제로샷 전이되어, 매번 모델별 재학습이 필요한 기존 RL 베이스라인들과 달리 휴리스틱 베이스라인들과 대등한 경쟁력을 유지하였다.
효율성: 이 방법은 최대 2.1배의 전체 캐시 처리량을 제공한다. 특징 추출 및 정책 평가를 위해 약간의 오버헤드(~1.4 ms/token)가 발생하지만, 다른 학습된 방법들과 유사한 수준이며 전체 캐시 추론보다 훨씬 빠르다.
오류 분석: DistinkCache는 휴리스틱 및 타 RL 방법들에 비해 반복적인 오류와 오답 발생률이 가장 낮게 나타났으며, 이는 시퀀스의 일관성과 사실적 정확도 보존 능력이 우수함을 시사한다.

의의 및 주장

본 논문은 DistillCache가 자원 제한적인 하드웨어에서 더 긴 문맥 추서(long-context inference)를 가능하게 하기 위해, **학습된 분포 인식 정책(learned, distribution-aware policies)**의 효과성을 입증한다고 주장한다. 과거의 어텐션 통계가 아닌 미래의 분포적 충실도를 대상으로 토큰 유지 결정을 최적화함으로써, 이 프레임워크는 예산 제약이 있는 정책을 통해 다운스트림 생성의 결과를 고려하는 법을 학습한다.

저자는 DistillCache가 단순히 휴리스틱 지표를 학습으로 모사하는 것이 아니라, 모델의 출력 분포 보존을 직접적으로 최적화하는 차별화된 접근 방식임을 강조한다. 훈련 비용(약 130 GPU-시간)과 중간 정도의 예산에서 추론 작업에 대한 RLKV와의 미세한 성능 격차를 인정하면서도, 본 논문은 분포 보존 적응형 제거(distribution-preserving adaptive eviction)가 유망한 방향임을 결론짓는다. 또한, 토큰 수준의 분포 충실도와 헤드 수준의 게이팅 전략을 결리하는 것이 미래 최적화를 위한 보완적인 경로가 될 수 있음을 시사한다.