기술 요약: xKV - 정렬된 특이 벡터 추출을 통한 교차 레이어 KV-Cache 압축

1. 문제 제기

장문맥 대규모 언어 모델 (LLM) 은 대규모 정보 검색 및 코드베이스 디버깅과 같은 고급 애플리케이션을 가능하게 하지만, 추론 중 Key-Value(KV) 캐시의 메모리 소비라는 치명적인 병목 현상에 직면해 있습니다. 문맥 길이가 수십만 또는 수백만 토큰으로 확장됨에 따라 KV-Cache 메모리 발자국은 감당하기 어려워져, 동시 추론 요청과 처리량을 심각하게 제한합니다.

기존 완화 전략은 일반적으로 두 가지 범주로 나뉩니다:

• 레이어 내부 (Intra-layer) 방법: 양자화, 토큰 제거, 단일 레이어 저랭크 분해와 같은 기술은 각 레이어의 KV-Cache 를 독립적으로 압축합니다. 일정 정도 효과적이기는 하지만, 레이어 간 (across layers) 잠재적 중복성을 활용하지 못합니다.
• 레이어 간 (Cross-layer) 방법: Cross-Layer Attention(CLA) 과 YOCO 와 같은 접근법은 비용이 많이 드는 아키텍처 수정과 사전 학습을 요구합니다. MiniCache 와 같은 다른 방법들은 토큰별 코사인 유사도에 기반한 인접 레이어의 사후 병합에 의존합니다. 그러나 실증 분석에 따르면, 레이어 간의 토큰별 유사성은 종종 상당한 정확도 저하 없이 강력한 압축을 지원하기에 너무 약합니다.

2. 방법론: xKV

저자들은 모델 재학습이나 아키텍처 변경 없이 레이어 간 깊은 구조적 중복성을 활용하는 사후 학습 (post-training), 플러그 앤 플레이 (plug-and-play) 압축 방법인 xKV를 제안합니다.

2.1. 핵심 관찰: 정렬된 특이 벡터

이 논문은 토큰별 코사인 유사도에 기반하여 레이어가 구별된다는 가정에 도전합니다. **Centered Kernel Alignment(CKA)**를 사용하여 저자들은 개별 토큰 임베딩은 다를 수 있지만, 주요 특이 벡터(토큰 공간의 주요 변이를 정의하는 기저) 는 인접 레이어 간에 매우 정렬되어 있음을 입증합니다. 이 높은 CKA 점수는 점별 유사도가 낮더라도 KV-Cache 부분 공간의 기하학적 구조는 공유됨을 나타냅니다.

2.2. 교차 레이어 분해 (CLF)

이 정렬을 활용하여 xKV 는 모델의 $N$개 레이어를 크기 $W$의 연속된 윈도우로 분할합니다. 각 윈도우에 대해 교차 레이어 분해를 수행합니다:

1. 연결 (Concatenation): 윈도우 내 모든 레이어의 KV-Caches 를 수평으로 연결합니다.
2. SVD: 연결된 캐시에 대해 랭크-$r$ 특이값 분해 (SVD) 를 계산합니다.
3. 분해: 결과는 윈도우용 공유 토큰 기저($A_k$) 와 레이어별 재구성 행렬($B_\ell$) 을 산출합니다.
4. 저장: $W$개의 풀랭크 캐시를 저장하는 대신, 시스템은 하나의 공유 기저와 $W$개의 더 작은 재구성 행렬을 저장하여 메모리 복잡도를 $O(W \cdot L \cdot d)$에서 $O(L \cdot r + W \cdot r \cdot d)$로 줄입니다.

2.3. 선택적 재구성 (SR)

디코드 단계에서 모든 토큰을 재구성하는 계산 병목 현상 (시퀀스 길이 $L$에 선형적으로 비례) 을 해결하기 위해 xKV 는 **선택적 재구성 (SR)**을 도입합니다.

• 메커니즘: 밀집 재구성 대신 SR 은 랜드마크 기반 청크 선택 (ShadowKV 와 유사) 을 사용하여 현재 쿼리와 가장 관련 있는 토큰의 희소 부분집합 ($S$) 을 식별합니다.
• 실행: 공유 기저와 레이어별 행렬을 사용하여 이러한 중요한 토큰에 해당하는 행만 재구성합니다.
• 이점: 이는 재구성 비용을 시퀀스 길이에 대해 상수로 줄여 종단 간 속도 향상을 가능하게 합니다.

3. 주요 기여

1. 레이어 간 중복성 발견: 이 논문은 KV-Caches 가 낮은 토큰별 유사성으로 가려진 형태로 레이어 간 매우 정렬된 주요 특이 벡터를 공유한다는 사실을 밝혀냈습니다.
2. xKV 알고리즘: 교차 레이어 분해 (CLF) 를 통해 공유 기저를 추출하여 $W$개 레이어 ($W > 2$) 간 효과적인 정보 공유를 가능하게 하는 학습이 필요 없는 방법입니다.
3. 선택적 재구성: 쿼리와 관련된 토큰만 재구성하는 디코딩 시간 최적화로, 밀집 재구성의 선형 계산 오버헤드를 제거합니다.
4. 플러그 앤 플레이 호환성: 이 방법은 미세 조정이나 아키텍처 수정이 필요 없으므로 기존 사전 학습 모델 (예: Llama-3.1, Qwen2.5, DeepSeek-V2) 에 적용 가능합니다.

4. 실험 결과

저자들은 RULER, LongBench, Multi-turn NIAH 벤치마크를 사용하여 Llama-3.1-8B-Instruct, Qwen2.5-7B/14B-Instruct, DeepSeek-Coder-V2-Lite 에서 xKV 를 평가했습니다.

• 압축 대 정확도:
  • xKV 는 풀 어텐션 베이스라인 대비 최소한의 정확도 손실 (평균 ≤3%) 로 최대 8 배 KV-Cache 압축을 달성합니다.
  • Llama-3.1-8B 에서 xKV 는 8 배 압축 시 RULER 에서 평균 정확도 **88.50%**를 달성하여, MiniCache(1.3 배에서 45.04%) 를 크게 능가하고 SnapKV 및 PyramidKV 와 같은 토큰 제거 방법과 맞먹거나 초과합니다.
  • 멀티 턴 NIAH 작업에서 xKV 는 대화 턴 전반에 걸쳐 안정적인 성능을 유지하는 반면, 토큰 제거 방법은 첫 턴 이후 급격한 감소가 발생합니다.
• 효율성 및 처리량:
  • 메모리: xKV 는 동일한 GPU 메모리 예산 하에서 더 큰 배치 크기를 가능하게 합니다.
  • 지연 시간: 밀집 재구성은 어텐션 지연 시간을 최대 3.6 배 줄입니다.
  • 처리량: 선택적 재구성 (xKV-SR) 과 결합 시, NVIDIA A100 GPU 에서 풀 어텐션 베이스라인 대비 최대 4.23 배의 종단 간 속도 향상을 달성합니다.
  • 비교: xKV-SR 은 유사한 정확도 수준에서 ShadowKV 보다 처리량이 30% 더 높으며, Keys 와 Values 를 모두 GPU 에 유지하여 PCIe 전송 병목 현상을 피합니다.
• 강건성: 이 방법은 DeepSeek-V2 의 Grouped-Query Attention(GQA) 및 Multi-head Latent Attention(MLA) 을 포함한 다양한 모델 아키텍처에서도 효과적입니다.

5. 중요성 및 주장

이 논문은 xKV 가 장문맥 LLM 추론을 위해 메모리 발자국과 지연 시간을 동시에 줄이는 실용적인 플러그 앤 플레이 접근법을 제공한다고 주장합니다. 토큰별 유사성에서 특이 벡터 정렬로 초점을 전환함으로써, xKV 는 이전 방법들이 놓친 미활용 레이어 간 중복성을 해제합니다.

저자들은 그들의 접근법이 다음과 같음을 강조합니다:

• 아키텍처 변경에 필요한 고비용의 사전 학습을 피합니다.
• 유사도 기반 병합 (MiniCache) 의 정확도 한계를 극복합니다.
• 선택적 재구성을 통해 밀집 저랭크 재구성의 계산 병목 현상을 해결합니다.
• ShadowKV 와 같은 오프로딩 기반 방법을 제한하는 PCIe 대역폭 한계를 우회하여 압축된 캐시에 대한 완전한 온디바이스 (GPU) 추론을 가능하게 합니다.

이 작업은 효율적인 장문맥 추론을 위한 기초적인 단계로 제시되며, 채택을 촉진하기 위해 코드가 공개되었습니다.