기술 요약: PolyKV: KV 캐시 압축을 위한 이질적 유지 및 할당

문제 정의

대규모 언어 모델(LLM)의 장문맥 추론(Long-context inference)은 Key-Value(KV) 캐시의 증가로 인해 심각한 메모리 병목 현상에 직면해 있습니다. 기존의 KV 캐시 제거(eviction) 방법들은 특정 토큰의 부분 집합을 유지함으로써 메모리 사용량을 효과적으로 줄이지만, 일반적으로 모든 트랜스포머 레이어에 대해 균일한 압축 정책과 균일한 캐시 예산을 적용합니다.

이러한 균일한 설계는 **트랜스포머 레이어의 내재된 이질성(heterogeneity)**을 간과합니다. 선행 연구에 따르면 서로 다른 레이어는 어텐션 패턴, 표현된 추상화, 문맥적 역할 및 프루닝(pruning)에 대한 민감도 측면에서 서로 다릅니다. 따라서 단일 글로벌 제거 규칙은 다음과 같은 사실을 반영하지 못하는 거친 추상화에 불과합니다:

1. 서로 다른 레이어는 서로 다른 제거 패턴(예: heavy hitters vs. attention sinks)을 선호할 수 있습니다.
2. 서로 다른 레이어는 KV 캐시 감소에 대한 민감도가 다르며, 이는 각기 다른 캐시 용량을 필요로 함을 의미합니다.

현재의 적응형 방법들은 비균일한 예산을 탐색하기 시작했으나, **제거 패턴(eviction pattern)**과 **캐시 예산(cache budget)**을 레이어별로 결합된 설계 선택지로 취급하는 데에는 여ยัง 미치지 못하고 있습니다.

방법론: PolyKV

저자들은 PolyKV를 제안하며, 이는 KV 캐시 제거를 **레이어별 설계 공간(layer-wise design space)**으로 재정의합니다. PolyKV는 각 레이어와 추론 단계(prefill 및 decoding)에 대해 다음 두 가지 결정을 명시적으로 결합합니다:

1. 무엇을 제거할 것인가: 해당 레이어에 적합한 특정 제거 패턴(예: H2O, TOVA, PyramidKV)을 선택합니다.
2. 얼마나 유지할 것인가: 해당 레이어에 특정 캐시 예산을 할당합니다.

핵심 구성 요소

• 설계 공간 정식화: PolyKV는 각 레이어 $l$이 각 단계 $\phi$에 대해 제거 패턴 $r_{l,\phi}$와 예산 $b_{l,\phi}$를 할당받는 전략 $\Pi$를 정의합니다. 전체 예산은 고정된 글로벌 제한 $B$에 의해 제약됩니다.
• 오프라인 교정(Offline Calibration): 철저한 온라인 탐색을 피하기 위해, PolyKV는 작은 교정 세트를 사용하여 고정된 정적 전략을 생성합니다. 이는 두 가지 범주의 레이어별 신호를 수집합니다:
  • 제거 선호도 신호(Eviction Preference Signals): 어떤 제거 패턴이 해당 레이어의 FullKV 계산을 가장 잘 보존하는지 측정합니다. 지표로는 변형된 활성화 값과 FullKV 활성화 값 사이의 코사인 유사도(cosine similarity), 그리고 특정 패턴에 의한 Perplexity(PPL) 민감도(특정 패턴에 의한 성능 저하 정도)가 포함됩니다.
  • 캐시 예산 민감도 신호(Cache Budget Sensitivity Signals): KV 감소에 의해 레이어가 얼마나 강하게 영향을 받는지 추정합니다. 지표로는 어텐션 엔트로피(분산된 어텐션은 더 큰 예산을 요구함)와 PPL 민감도(집계된 최악의 경우 저하)가 포함됩니다.
• 전략 구축:
  • 패턴 선택: 각 레이어-단계 쌍에 대해, 가장 높은 선호도 점수(예: 가장 높은 코사인 유사도 또는 가장 낮은 PPL 저하)를 가진 제거 패턴을 선택합니다.
  • 예산 할당: 예산은 민감도 신호에 비례하여 할당됩니다. 민감도(또는 엔트로피)가 높은 레이어는 최소 예산 제약($b_{min}$)을 준수하는 범위 내에서 더 많은 예산을 받습니다.
• 구현: PolyKV는 훅(hook) 기반 프레임워크로 구현되었습니다. 이는 기본 모델 코드를 수정하지 않고도 다양한 제거 방법을 모듈식으로 통합할 수 있도록 생성 파이프라인의 공통 지점(pre-kernel, post-kernel, cache update)에 로직을 주입합니다.

주요 기여

1. KV 제거의 재정의: 본 논문은 KV 캐시 제거를 공유된 메모리 제약 하에서 제거 패턴의 선택과 캐시 용량 할당을 결가한 레이어별 설계 공간으로 재구성했습니다.
2. PolyKV 프레임워크: 이 프레ка임워크는 제거 패턴과 예산을 레이어 단위로 결합하여 이 설계 공간을 실행 가능한 시스템으로 만듭니다. 이는 통합된 인터페이스를 통해 여러 제거 프리미티브를 통합하고, 전략 구축을 가이드하기 위해 오프라인 교정 신호를 사용합니다.
3. 실증적 검증: 저자들은 이질적인 레이어별 전략이 고정된 메모리 예산 하에서 균일한 단일 정책 베이스라인보다 우수한 성능을 보임을 입증했습니다.

실험 결과

실험은 LLaMA-3.1-8B 및 Qwen3-8B를 대상으로 LongBench 및 RULER 벤치마크에서 수행되었습니다.

• LongBench 성능 (평균 512-토큰 예산 기준):
  • LLaMA-3.1-8B: PolyKV는 가장 강력한 단일 정책 베이스라인과 FullKV 사이의 성능 격차 중 **54.5%**를 회복했습니다.
  • Qwen3-8B: PolyKV는 격차의 **25.7%**를 회복했습니다.
  • PolyKV는 예산 변화(128–1024 토큰) 전반에 걸쳐 가장 강력한 베이스라인보다 일관되게 우수한 성능을 보였으며, 개선 폭은 1.7%에서 6.4% 사이였습니다.
• RULER 성능:
  • 결과는 엇갈렸습니다. PolyKV는 LLaMA에서는 베이스라인을 약간 상회했으나, Qwen에서는 (3.98 포인트 차이로) 베이스라인보다 낮은 성능을 보였습니다. 저자들은 이를 RULER 태스크의 특성, 즉 이질적 정책에 의해 폐기될 수 있는 희소한 증거 토큰(sparse evidence tokens)에 의존하는 특성 때문이라고 설명합니다.
• 절제 연구(Ablation Studies):
  • 제거 패턴 vs. 예산: 주요 성능 향상은 특히 prefill 단계에서 서로 다른 제거 패턴을 레이어에 할당하는 데서 발생했습니다.
  • 예산 할당: 비균일한 예산 할당은 추가적인 이득을 제공하지만, 모델 의존성이 높고 패턴 선택보다 견고함이 떨어집니다. 이는 전체 예산이 민감한 레이어를 고립시키지 않을 만큼 충분히 클 때 가장 효과적입니다.
• 태스크 분석: 성능 향상은 광범위한 문맥 이해를 요구하는 태스크(예: QA, few-shot learning)에서 가장 두드러졌습니다. 희소한 토큰의 정확한 검색(예: RULER의 NIAH)이나 엄격한 로컬 형식 준수(예: 코드 생성)를 요구하는 태스크에서는 성능이 저하되었습니다.

의의 및 주장

본 논문은 KV 캐시 제거를 단일 글로벌 유지 정책이 아닌 구성적이고 레이어 인지적인 최적화 문제로 다루는 것이 더 우수한 결과를 낳는다고 주장합니다.

• 이질성이 핵심이다: 본 연구는 트랜스포머 레이어가 동질적이지 않음을 입증하며, "원 사이즈(one-size-fits-all)" 방식의 제거 전략은 최적이 아님을 보여줍니다.
• 정적 효율성: 오프라인 교정을 사용함으로써 PolyKV는 온라인 정책 탐색에 따른 런타임 오버헤드 없이 이러한 이질적 이점을 달성하여 실제 배포가 가능하게 합니다.
• 한계점: 저자들은 현재의 교정 신호가 정확한 검색 태스크에 필요한 "희소하지만 결정적인" 증거 토큰을 항상 보존하지는 못한다고 솔직하게 언급하며, 이는 신호 설계의 향후 개선 방향을 제시합니다.

요약하자면, PolyKV는 제거 전략과 캐시 용량을 각 트랜스포머 레이어의 특정 특성에 맞춤으로써, 장문맥 LLM 추론에서 메모리와 품질 사이의 트레이드오프를 유의미하게 개선할 수 있음을 보여줍니다.