Thin Keys, Full Values: Reducing KV Cache via Low-Dimensional Attention Selection

이 논문은 어텐션 메커니즘에서 쿼리와 키의 역할이 값보다 차원이 낮아도 충분하다는 통찰을 바탕으로, 키 차원을 축소하여 KV 캐시를 대폭 줄이면서도 모델 성능을 거의 유지하는 비대칭 어텐션 기법을 제안하고 이를 통해 7B 파라미터 모델에서 25GB 의 캐시 절감 및 동시 사용자 60% 증가 효과를 입증했습니다.

핵심

1. 문제: AI 의 '기억장'이 너무 무겁다

지금까지의 AI(트랜스포머) 는 문장을 이해할 때 세 가지 역할을 하는 데 똑같은 크기의 공간을 할당했습니다.

• 질문 (Query): "무엇을 찾아야 하지?"
• 열쇠 (Key): "이 정보가 뭐지?" (검색용)
• 내용물 (Value): "정작 중요한 정보는 뭐야?" (실제 데이터)

기존 방식은 이 세 가지가 모두 거대한 책장 (높은 차원) 을 차지합니다. 특히 AI 가 긴 문장을 처리할 때, 과거의 모든 '열쇠 (Key)'를 메모리에 쌓아두어야 하는데, 이 **메모리 (KV Cache)**가 너무 커져서 한 번에 많은 사용자를 받거나 긴 글을 읽을 때 컴퓨터가 터져버리는 문제가 생깁니다.

2. 해결책: "검색용 열쇠는 얇게, 내용물은 그대로"

이 논문은 **"질문과 열쇠는 아주 작게 만들어도 되는데, 내용물은 그대로 두면 된다"**는 사실을 발견했습니다.

🗝️ 비유: 도서관 사서와 책장

• 기존 방식: 사서가 책을 찾을 때, 책 표지 (Key) 에 책의 전체 내용 (Value) 을 다 적어두고, 그걸 찾기 위해 거대한 책장 전체를 훑어봅니다. 책장이 너무 커서 사서가 지칩니다.
• 이 논문의 방식:
  1. 얇은 열쇠 (Thin Keys): 사서가 책을 찾을 때, 책 표지에 책의 제목과 번호만 아주 작게 적습니다. (예: "소설, 3 번 선반"). 이 정보는 아주 작아도 책을 찾는 데 충분합니다.
  2. 꽉 찬 내용물 (Full Values): 하지만 책장 안의 **책 내용 (Value)**은 그대로 두어, 필요한 정보를 뽑아낼 때 풍성하게 가져옵니다.

즉, **"무엇을 찾을지 결정하는 과정 (선택)"**은 아주 간단한 정보만 있으면 되지만, **"실제 정보를 전달하는 과정 (전달)"**은 풍부한 정보가 필요하다는 것입니다.

3. 실험 결과: "작은 열쇠로 큰 성과"

연구진은 이 아이디어를 다양한 AI 모델에 적용해 보았습니다.

• 단순한 게임: "10 번 전의 단어를 찾아라" 같은 게임은 열쇠를 1 개만 써도 100% 성공했습니다. (위치만 기억하면 됨)
• 복잡한 언어: 위키피디아 같은 긴 글을 학습할 때도, 열쇠의 크기를 4 분의 1 로 줄였을 때 AI 의 실력 (언어 이해도) 은 거의 떨어지지 않았습니다. (오직 4% 정도만 감소)
• 기존 모델 적용: 이미 만들어진 거대 AI(GPT-2, Mistral-7B 등) 에도 이 기술을 적용할 수 있었습니다. 열쇠만 압축하고, 아주 적은 데이터로 다시 훈련시키니, 기억 공간은 75% 줄이면서도 성능은 거의 그대로 유지되었습니다.

4. 왜 이것이 중요한가? (실제 효과)

이 기술이 적용되면 어떤 일이 일어날까요?

• 메모리 폭탄 해결: 긴 문장 (예: 128,000 단어) 을 처리할 때, 한 사용자의 AI 가 차지하는 메모리가 25GB나 줄어듭니다.
• 더 많은 사용자: 같은 그래픽 카드 (GPU) 로 약 60% 더 많은 사람이 동시에 AI 를 쓸 수 있게 됩니다.
• 비용 절감: 서버를 더 많이 살 필요가 없어져서 AI 서비스 비용이 크게 내려갑니다.

5. 요약: "선택은 간단하고, 전달은 풍부하게"

이 논문의 핵심 메시지는 **"모든 것을 똑같이 크게 만들 필요는 없다"**는 것입니다.

• 기존: 질문, 열쇠, 내용물 = 모두 거대함 (비효율적)
• 새로운 방식: 질문과 열쇠 = 얇고 작음 (선택만 하면 됨), 내용물 = 거대함 (정보 전달 필요)

이처럼 AI 의 '검색용 열쇠'를 얇게 만들면, AI 가 기억해야 할 공간이 획기적으로 줄어들어 더 빠르고, 더 저렴하며, 더 많은 사람이 사용할 수 있는 AI 시대가 열리게 됩니다. 마치 거대한 도서관에서 책장만 줄이고 책 내용은 그대로 두어, 사서가 더 빠르게 책을 찾아주는 것과 같습니다.

기술 요약

1. 문제 정의 (Problem)

현대 트랜스포머 (Transformer) 모델의 자기 어텐션 (Self-Attention) 메커니즘은 쿼리 (Query), 키 (Key), 값 (Value) 에 모두 동일한 차원 ($d_q = d_k = d_v = d_{model}$) 을 사용하는 대칭적인 구조를 따릅니다. 이는 설계상의 관례일 뿐 필수적인 요구사항이 아닙니다.

• 메모리 병목 현상: 대규모 언어 모델 (LLM) 을 긴 컨텍스트 (Long Context) 에서 추론할 때, 가장 큰 메모리 병목은 **KV 캐시 (Key-Value Cache)**입니다. 생성 과정에서 모든 이전 토큰의 키와 값을 저장해야 하며, 컨텍스트 길이가 길어질수록 메모리 사용량이 선형적으로 증가합니다.
• 비효율적인 대칭성: 어텐션 메커니즘은 두 가지 기능적으로 다른 작업을 수행합니다.
  1. 선택 (Selection, $QK^T$): 어떤 토큰이 중요한지 결정하는 스칼라 가중치 생성. 이는 본질적으로 낮은 차원에서도 충분히 수행 가능함.
  2. 값 전달 (Value Transfer, $attn \cdot V$): 선택된 토큰에서 정보를 추출하고 집계하는 작업. 모델의 전체 표현 능력 (의미, 구문, 위치 정보 등) 을 보존해야 하므로 고차원이 필요함.
• 핵심 가설: "선택" 작업은 "값 전달" 작업보다 본질적으로 낮은 차원 ($O(\log N)$) 만으로도 수행 가능하며, 키 (Key) 와 쿼리 (Query) 의 차원을 줄이고 값 (Value) 의 차원은 유지하는 비대칭 구조가 KV 캐시 크기를 획기적으로 줄일 수 있다.

2. 방법론 (Methodology)

저자들은 **비대칭 어텐션 (Asymmetric Attention)**을 제안하며, 이는 기존 모델에 대한 간단한 수정으로 적용 가능합니다.

• 차원 분리:

  • 선택 차원 ($d_{select}$): 쿼리 ($Q$) 와 키 ($K$) 는 $d_{model}$에서 $d_{select}$ ($d_{select} \ll d_{model}$) 로 투영됩니다.
  • 값 차원 ($d_{model}$): 값 ($V$) 은 원래 모델 차원을 유지합니다.
  • 수식: $Q = XW_Q (d_{model} \to d_{select})$, $K = XW_K (d_{model} \to d_{select})$, $V = XW_V (d_{model} \to d_{model})$.
  • 어텐션 계산은 $Softmax(QK^T/\sqrt{d_{select}})V$로 동일하게 수행되며, $QK^T$ 결과물이 스칼라 가중치이므로 차원 불일치 문제가 발생하지 않습니다.

• 기존 모델 적용 (Post-Training Compression via Factored Keys):

  • 사전 훈련된 모델 (예: GPT-2, Mistral-7B) 에는 **SVD(특이값 분해)**를 사용하여 키 행렬 ($W_K$) 을 저랭크 근사합니다 ($W_K \approx AB$).
  • Factored Key: $K' = XA$를 캐시에 저장 (차원$r$).
  • Query 흡수: $B$ 행렬을 쿼리 투영 ($W_Q$) 에 곱하여 ($W_Q^{new} = W_Q B^T$) 어텐션 점수를 수학적으로 동일하게 유지합니다. 쿼리는 매 스텝마다 새로 계산되므로 캐시 부담은 없습니다.
  • 미세 조정 (Fine-tuning): SVD 로 인한 성능 저하를 복구하기 위해 $Q$와 $K$ 투영 레이어만 소량의 데이터로 3 에포크 정도 미세 조정합니다.

3. 주요 기여 (Key Contributions)

1. 이론적 및 실증적 분석: 어텐션 '선택' 작업이 $O(\log N)$ 차원만으로도 충분함을 증명했습니다. 위치 기반 선택 (1 차원/헤드), 콘텐츠 기반 검색 ($\log_2 N$ 차원), 언어 모델링 ($d_{model}/4$) 실험을 통해 이를 입증했습니다.
2. 비대칭 어텐션 제안: $QK$와$V$의 차원을 분리하는 간단한 'Drop-in' 수정을 제안했습니다. WikiText-103 에서 $d_{select} = d_{model}/4$로 설정 시 QK 파라미터는 75% 감소하고 퍼플렉시티 (Perplexity) 는 4.3% 만 증가했습니다.
3. SVD + QK 미세 조정 파이프라인: 기존 모델에 재학습 없이 SVD 를 적용한 후, $QK$ 레이어만 경량 미세 조정을 수행하여 거의 모든 품질 손실을 복구하는 방법을 제시했습니다.
4. 확장성 검증: 10M 파라미터 모델부터 7B (Mistral-7B) 파라미터 모델까지, 그리고 Vanilla Transformer, LLaMA, Mistral(GQA) 등 다양한 아키텍처에서 동일한 성능 저하 비율을 보임으로써 방법론의 일반성을 입증했습니다.

4. 실험 결과 (Results)

• 알고리즘적 작업:
  • 위치 선택 (Copy-back): 헤당 1 차원 ($d_{select}=4$) 만으로도 100% 정확도 달성.
  • 콘텐츠 기반 검색 (Key-Value Retrieval): 헤당 2 차원 ($d_{select}=8$) 에서 100% 정확도 달성 ($\log_2 N$ 이론과 부합).
• 언어 모델링 (WikiText-2/103):
  • $d_{select} = d_{model}/4$ 설정 시, QK 파라미터 75% 감소 대비 퍼플렉시티 증가분은 4.3% 에 불과했습니다.
  • WikiText-2(소규모 데이터) 에서는 과적합으로 인해 성능 저하가 더 작게 나타났으나, WikiText-103(대규모 데이터) 에서 실제 비용이 명확히 드러났습니다.
• SVD 압축 및 미세 조정 (GPT-2 & Mistral-7B):
  • GPT-2 (124M): 키만 SVD 압축 (Rank 192, $d_{model}/4$) 후 $QK$ 미세 조정 시, 초기 27.6% 의 퍼플렉시티 증가가 **1.8%**로 감소했습니다.
  • Mistral-7B (7.2B): Rank 256 (75% 키 캐시 감소) 에서 미세 조정 후 잔류 오차가 **2.0%**에 불과했습니다. 이는 GPT-2 결과와 거의 일치합니다.
• KV 캐시 절감 효과:
  • 7B 모델, 128K 컨텍스트 기준: 사용자당 25GB의 KV 캐시 절감.
  • 동일한 GPU 에서 약 60% 더 많은 동시 사용자 처리 가능.
  • 1M 컨텍스트 기준: 사용자당 196GB 절감 (전체 100 사용자 기준 19.6TB 절감).

5. 의의 및 결론 (Significance & Conclusion)

• KV 캐시 최적화의 새로운 패러다임: 기존 KV 캐시 압축 기법 (양자화, MQA/GQA 등) 과 직교 (Orthogonal) 하는 새로운 접근법입니다.
  • GQA 와의 결합: 헤 수를 줄이는 GQA 와 결합 가능.
  • 양자화와의 결합: 차원을 줄인 후 정밀도를 낮추는 양자화와 결합 시 최대 16 배의 키 캐시 압축 효과 달성 가능.
• 실용적 배포 경로:
  1. 제로 코스트: 기존 모델에 SVD 만 적용 (재학습 없음, 25% 캐시 절감, ~2% 품질 손실).
  2. 경량 미세 조정: SVD + $QK$만 3 에포크 미세 조정 (75% 캐시 절감, <2% 품질 손실).
  3. 처음부터 학습: 새로운 모델 설계 시 $d_{select} = d_{model}/4$로 설정.
• 아키텍처 설계 원칙: $d_q = d_k = d_v$라는 오랜 관례를 깨고, 선택 (Selection) 과 값 전달 (Value Transfer) 의 차원 요구 사항을 분리하는 것이 효율적인 LLM 설계의 핵심임을 시사합니다.

이 논문은 긴 컨텍스트 추론 환경에서 메모리 비용을 획기적으로 줄이면서도 모델의 품질을 거의 유지할 수 있는 실용적이고 이론적으로 타당한 솔루션을 제시합니다.