🤖 AI

SelKV: Selective KV Cache Merging with Per-Token Merge-or-Drop and Attention Compensation

SelKV는 가치 벡터 유사도에 기반하여 토큰을 선택적으로 병합하거나 삭제하고 로짓 바이어스(logit bias)를 통해 어텐션 불균형을 보정함으로써, 원래 캐시의 25%만을 유지하면서도 손실이 거의 없는 생성 품질과 상당한 속도 향상을 달성하는 훈련이 필요 없는 프레임워크이다.

원저자: Soumia Bouyahiaoui, Manel Kara laouar, Aicha Boutorh, Mohamed Hadj Ameur

게시일 2026-07-21

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Soumia Bouyahiaoui, Manel Kara laouar, Aicha Boutorh, Mohamed Hadj Ameur

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

당신이 새로운 장을 쓰면서 10만 페이지에 달하는 방대한 이야기를 기억하려고 애쓰고 있다고 상상해 보세요. 줄거리의 한 지점이 떠오를 때마다, 당신은 정확한 세부 사항을 찾기 위해 공책 전체를 앞뒤로 뒤져야 합니다. 페이지를 더 많이 쓸수록 공책은 점점 더 무거워지고, 결국 팔이 너무 아파서 공책을 들 수조차 없게 됩니다. 이것은 거대 언어 모델(LLM)이 긴 텍스트를 처리하려고 할 때 내부에서 일어나는 현상과 정확히 일치합니다. 이 모델들은 매우 똑똑하지만 메모리 문제를 가지고 있습니다. 모델이 글을 읽을 때, 다음에 올 내용을 이해하기 위해 지금까지 본 모든 것을 디지털 파일 시스템인 '키-값 캐시(Key-Value cache)'에 저장하는데, 텍스트가 길어질수록 이 파일 캐비닛은 점점 커지며 결국 컴퓨터의 모든 메모리를 가득 채우고 모든 동작을 느려지게 만듭니다.

이를 해결하기 위해 과학자들은 두 가지 주요 기술을 시도했습니다. 첫 번째는 중요하지 않다고 생각되는 오래된 페이지들을 버리는 '축출(eviction)'입니다. 두 번째는 공간을 절약하기 위해 유사한 페이지들을 하나로 붙이는 '병합(merging)'입니다. 하지만 두 방법 모두 결함이 있습니다. 페이지를 버리는 것은 결정적인 단서를 잃게 할 수 있고, 페이지를 붙이는 것은 모델이 그 묶음에 얼마나 많은 주의를 기울여야 하는지 잊어버리게 만드는 '흐릿한' 버전을 만들어내곤 합니다. 이는 마치 고양이 사진 열 장을 하나로 붙여놓은 것과 같습니다. 모델은 여전히 그 하나의 붙여진 덩어리에 대해 원래 사진 한 장에 주었을 법한 정도의 주의만 기울일 수 있으며, 이는 그 덩어리가 이제 열 장을 대표하고 있음에도 불구하고 말입니다. 이로 인해 모델은 혼란에 빠지고 실수를 저지르게 됩니다.

여기에 디지털 파일 캐비닛을 위한 똑똑한 사서 역할을 하는 새로운 영리한 방법인 SelKV가 등장했습니다. SelKV는 페이지들을 맹목적으로 붙이거나 쓰레기통에 던지는 대신, '소프트 코사인 게이트(soft cosine gate)'를 사용합니다. 이는 두 페이지가 얼마나 닮았는지 확인하여 어떻게 할지 결정하는 클럽의 가드(bouncer)와 같습니다. 두 페이지가 매우 유사하다면 단단하게 결합됩니다. 만약 두 페이지가 완전히 다르다면, 가드는 메모리를 깨끗하게 유지하기 위해 한 페이지를 내보냅니다. 하지만 진짜 마법은 '어텐션 보상(attention compensation)'에 있습니다. 페이지를 붙이는 작업은 보통 모델이 그 페이지들을 무시하게 만들기 때문에, SelKV는 붙여진 페이지들에 약간의 '볼륨 부스트(volume boost)'를 추가하여 모델이 적절한 양의 주의를 기울일 수 있도록 보장합니다.

연구진은 세 가지 AI 모델을 대상으로 여러 문서에 대한 질문 답변부터 코드 작성까지 16가지 서로 다른 작업에서 이 시스템을 테스트했습니다. 그 결과, 원래 메모리 공간의 25%만 유지함으로써 SelKV가 전체 메모리를 가졌을 때와 거의 대등한 성능을 보였으며, 까다로운 다중 문서 퀴즈에서는 오히려 전체 버전보다 더 나은 성적을 거두었다는 것을 발견했습니다. 선택적으로 행동함으로써 AI가 실제로 이야기의 가장 중요한 부분에 집중할 수 있었던 것으로 보입니다. 또한, 텍스트가 거대해졌을 때(100,000 토큰), 이 방법은 AI의 텍스트 디코딩 속도를 3.3배 더 빠르게 만들었습니다. 이 논문은 이 접근 방식이 특정 유형의 어텐션(GQA라고 불리는)을 사용하는 현대적 AI 모델에 특히 유용하며, 강력한 두뇌를 기억력을 잃지 않으면서도 빠르게 구동할 수 있는 방법을 제시한다고 제안합니다.

기술 요약: SelKV - 토큰별 병합 또는 탈락 결정 및 어텐션 보상을 통한 선택적 KV 캐시 병합

1. 문제 정의

대규모 언어 모델(LLM)은 자기회귀 추론(autoregressive inference) 과정에서 핵심적인 병목 현상에 직면합니다: 바로 키-값(Key-Value, KV) 캐시입니다. 컨텍스트 길이( $n$ )가 증가함에 따라 KV 캐시의 메모리 점유율은 선형적으로( $O(L \cdot H \cdot d \cdot n)$ ) 증가하며, 이는 장기 컨텍스트 애플리케이션의 배치 크기, 처리량 및 배포 비용을 심각하게 제한합니다.

기존의 압축 방법들은 이를 완화하려고 시도하지만, 다음과 같은 두 가지 주요 한계점을 가집니다:

균일한 병합 또는 탈락 결정(Uniform Merge-or-Drop Decisions): 기존 방식들은 대개 이진 전략을 적용합니다. 즉, 모든 제거된 토큰을 병합하거나(예: KVMerger, WeightedKV) 아니면 완전히 제거합니다(예: H2O, SnapKV). 이러한 세밀함의 부족은 제거된 토큰의 값 벡터(value vector)가 병합 대상과 이질적일 때, 가중 평균이 유지된 의미적 표현을 손상시킨다는 문제를 낳습니다. 반대로, 무분별한 제거는 가치 있는 정보를 손실시킵니다.
어텐션 저하(Attention Sag): 여러 토큰이 하나의 캐시 엔트리로 병합되면, 해당 엔트리는 단일 미병합 토큰과 동일한 소프트맥스 어텐션 질량(softmax attention mass)을 받게 됩니다. 이는 병합된 위치가 병합된 토큰 수( $\sim m$ )에 비례하여 체계적으로 과소 어텐션되는 "어텐션 저하" 현상을 초래하며, 결과적으로 모델이 보존된 정보를 충분히 활용하지 못하게 만듭니다.

2. 방법론: SelKV

저자들은 이러한 한계를 해결하기 위해 설계된, 학습이 필요 없는(training-free) 듀얼 컴포넌트 프레임워크인 SelKV(Selective KV Cache Merging)를 제안합니다. SelKV는 프리필(prefill) 단계 이후 기존 병합 파이프라인에 통합되는 플러그형 모듈로 작동합니다.

핵심 구성 요소

1. 소프트 코사인 게이트(Soft Cosine Gate, 균일한 결정을 해결하기 위함)
이진 병합-또는-탈락 결정 대신, SelKV는 값 벡터 유사성에 기반한 연속적인 토큰별 게이팅 메커니즘을 도입합니다.

메커니즘: 병합 대상으로 라우팅된 각 제거된 토큰에 대해, 시스템은 이들의 값 벡터 간 코사인 유사도를 계산합니다.
게이팅 함수: 소프트 게이트 $g = \max(\cos\_sim, 0)$ $g = max (cos_s im, 0)$ 가 병합 강도를 조절합니다:
- $g \approx 1$ : 높은 유사도 $\rightarrow$ 전체 병합.
- $g = 0$ : 이질적(직교)인 벡터 $\rightarrow$ 토큰 탈락.
- $0 < g < 1$ : 중간 정도의 유사도 $\rightarrow$ 부분 병합.
구현: 이는 학습된 파라미터나 임계값 튜닝을 필요로 하지 않습니다. 유지된 토큰 $i$ 의 병합된 값은 원래의 값과 제거된 토큰들의 기여도를 각각의 게이트에 따라 스케일링하여, 게이트가 적용된 어텐션 가중 평균으로 계산됩니다.

2. 어텐션 보상(Attention Compensation, 어텐션 저하를 해결하기 위き 위해)
병합된 위치의 체계적인 어텐션 부족을 교정하기 위해, SelKV는 디코딩 시점의 로짓 바이어스(logit bias)를 적용합니다.

메커니즘: 원시 병합 횟수(높은 압축률에서 불안정할 수 있음)를 사용하는 기존 방식과 달리, SelKV는 프리필 통계로부터 어텐션 비율( $R_{h,i}$ )을 도출합니다. 이 비율은 병합된 위치가 원래 상태에 비해 받아야 할 총 어텐션 질량을 추정합니다.
교정: 디코딩 중에 어텐션 로짓에 바이어스 항 $\alpha \cdot \log(R_i)$ 를 더합니다. 이는 병합된 위치가 이제 나타내는 정보에 비례하여 적절한 어텐션 질량을 받을 수 있도록 소프트맥스 분포를 재조정합니다.

파이프라인 개요

SelKV 파이프라인은 프리필 이후 6단계로 진행됩니다:

중요도 점수 산정: 어텐션 가중치만 사용하는 대신, VATP를 따라 어텐션 가중도와 값 벡터 크기의 곱을 기준으로 토큰 점수를 매깁니다.
토큰 선택: 레이어당(MHA의 경우 헤드별로, 또는 GQA의 경우 유니온을 통해) 상위 $m_s$ 개의 토큰을 최근 윈도우와 함께 선택하여 유지합니다.
병합 대상 라우팅: 제거된 토큰은 프리필 어텐션 패턴에 따라 위치 버킷 내의 유지된 토큰으로 라우팅됩니다.
선택적 병합: 소프트 코사인 게이트가 각 제거된 토큰의 병합 강도를 결정합니다.
어텐션 보상: 로짓 바이어스를 계산하고 적용합니다.
RoPE 재배치: 키(Key)를 연속적인 위치로 재임베딩합니다(위치 왜곡을 피하기 위해 장기 컨텍스트 평가 시에는 비활성화됨).

3. 주요 기여

토큰별 세밀함(Per-Token Granularity): 의미적 유사성에 따라 병합 강도를 연속적으로 조절할 수 있는 소프트 코사인 게이트를 도입하여, 기존 방식의 "전부 아니면 전무(all-or-nothing)" 식의 결함을 제거했습니다.
안정적인 어텐션 보상: 원시 카운트 기반 방식보다 더 견고하게 어텐션 저하를 교정하는 새로운 어텐션 비율 기반 로짓 바이어스를 제안합니다.
학습이 필요 없는 통합(Training-Free Integration): 이 프레임워크는 미세 조정(fine-tuning)을 요구하지 않으며, 기존의 원샷 압축 파이프라인에 통합될 수 있습니다.
아키텍처 적응성: MHA 및 GQA 아키텍처를 모두 지원하며, GQA 모델을 위한 토큰 선택 유니온 처리를 포함합니다.

4. 실험 결과

프레임워크는 세 가지 모델(LongChat-7B (MHA), LLaMA-3.1-8B (GQA), Gemma-2-9B (GQA))에 대해 LongBench(16개 영어 데이터셋)에서 평가되었으며, KV 캐시의 **25%**만을 유지했습니다.

베이스라인 대비 성능: SelKV는 대표적인 원샷 베이스라인들(SnapKV, LOOK-M, PyramidKV)을 일관되게 능가했습니다.
- GQA 모델에서 SelKV는 풀 캐시 베이스라인에 가장 근접한 성능을 보였으며, 평균 격차는 $-0.67$ (Gemma-2) 및 $-0.73$ (LLaMA-3.1)였습니다.
- MHA 모델(LongChat)의 경우, 특정 31.5k 토큰 설정에서는 PyramidKV가 약간 더 우수한 성능을 보였으나, SelKV 역시 경쟁력 있는 수치를 기록했습니다.
다중 문서 QA: 특히, SelKV는 여러 다중 문서 QA 작업(예: LLaMA-3.1에서의 HotpotQA 및 MuSiQue)에서 풀 캐시 베이스라인을 초과했습니다. 저자들은 이를 선택적 병합이 핵심 증거는 보존하면서 방해되는 컨텍스트를 억제하는 암시적 어텐션 필터 역할을 한다는 증거로 해석합니다.
효율성: 100k 토큰의 컨텍스트 길이에서, 압축된 캐시는 풀 캐시 베이스라인 대비 3.3배의 디코딩 속도 향상을 제공했습니다.
강건성: 방법론은 다양한 압축 비율(10%–90% 유지)과 요약, 퓨샷 학습, 코드 완성 등의 작업 범주 전반에서 강건함을 입증했습니다.

5. 의의 및 주장

본 논문은 SelKV를 성능 저하 없이 장기 컨텍스트 LLM 추론을 가능하게 하는 중요한 진전으로 자리매김합니다. 저자들은 다음과 같이 주장합니다:

선택적 병합은 균일한 전략보다 우월하다: 표현 유사성에 따라 병합 강도를 조절함으로써, SelKV는 무분별한 병합으로 인한 의미적 손상을 방지합니다.
어텐션 보상은 필수적이다: 어텐션 저하를 교정하지 않으면 병합된 토큰이 과소 활용됩니다. SelKV의 비율 기반 교정은 보존된 정보가 효과적으로 접근되도록 보장합니다.
암시적 필터링: 복잡한 다중 문서 QA에서 SelKV가 풀 캐시 베이스라인보다 뛰어난 성능을 보이는 것은, 압축 과정 자체가 노이즈를 필터링하여 모델 성능을 향상시킬 수 있음을 시사하며, 이는 압축이 항상 능력을 저하시킨다는 가정에 도전하는 발견입니다.

저자들은 결론적으로 SelKV가 메모리가 제한된 환경에서 생성 품질을 거의 손실 없이 유지하면서 LLM을 배포할 수 있는 실용적이고 파라미터가 필요 없는 솔루션을 제공하며, 특히 GQA 기반 아키텍처에서 유용하다고 밝힙니다.