기술 요약: COMPILERKV

1. 문제 정의

대규모 언어 모델 (LLM) 은 컨텍스트 길이에 비례하여 선형적으로 증가하는 키 - 밸류 (KV) 캐시로 인해 메모리 병목 현상에 직면해 있습니다. 프리필 전용 KV 압축은 배포 친화적인 솔루션으로 등장했는데, 이는 프리필 단계가 끝난 후 토큰의 부분 집합을 선택하여 전체 디코딩 패스 동안 고정하는 방식입니다. 이 접근법은 런타임 메모리 사용량을 고정 비용으로 전환하며, 디코딩 시 개입이 전혀 필요하지 않습니다.

그러나 해당 논문은 기존 프리필 전용 방법들 (예: SnapKV, PyramidKV, HeadKV) 에 근본적인 통계적 불일치가 있음을 지적합니다. 이러한 방법들은 단일하고 노이즈가 많은 프롬프트의 통계를 사용하여, 특히 헤드별 신뢰도(어떤 어텐션 헤드가 검색에 결정적인지) 와 프롬프트 수준 압축 민감도(특정 프롬프트가 얼마나 압축될 수 있는지) 와 같은 보정 신호를 온라인으로 추정하려 합니다. 저자들은 이것이 잘못된 통계 단위라고 주장합니다:

• E1 (단일 프롬프트 노이즈): 프롬프트별 어텐션 점수는 heavy-tailed 분포를 가지며 스케일 편향이 있습니다. 큰 크기가 반드시 의미론적 중요성을 보장하지 않으며, 작은 추정 오차가 되돌릴 수 없는 선택 오류를 초래합니다.
• E2 (헤드 이질성): 어텐션 헤드는 기능적으로 특화되어 있습니다. 일부는 '검색 결정적'인 반면, 다른 일부는 일관되게 노이즈가 많습니다. 단순한 집계 방식은 노이즈가 많은 헤드가 점수를 지배하게 만듭니다.
• E3 (프롬프트 수준 위험 분산): 서로 다른 프롬프트는 서로 다른 압축 한계를 가집니다. 고정된 유지 규칙은 고엔트로피 또는 고퍼플렉시티 프롬프트에서 실패하여 치명적인 정보 손실을 초래합니다.

기존 방법들은 이러한 신호들을 모든 프롬프트에 대해 온라인으로 재추정하며, 이러한 신호들이 단일 프롬프트 내의 안정성이 아니라 교차 프롬프트 규칙성을 보인다는 사실을 활용하지 못합니다.

2. 방법론: COMPILERKV

저자들은 COMPILERKV를 제안합니다. 이는 보정 corpora 에서 오프라인으로 보정 테이블을 컴파일하는 KV 유지 정책으로, 온라인으로 추정하는 대신 이를 컴파일합니다. 핵심 통찰력은 헤드 신뢰도 순위와 위험 특성에서 안전한 임계값으로의 매핑이 한 번 학습하여 재사용할 수 있는 안정적인 아키텍처 사전 지식이라는 점입니다.

이 프레임워크는 프리필 단계 동안 작동하는 세 가지 결합된 구성 요소로 이루어져 있습니다:

A. 안정화된 토큰 유틸리티 (E1 대응)

학습 없이 일시적인 어텐션 스파이크와 헤드 간 스케일 편향을 제거하기 위해, 이 방법은 파라미터가 없는 유틸리티 점수를 계산합니다:

• 윈도우 집계 어텐션: 프롬프트 말단의 관찰 윈도우 $\Omega$에 걸쳐 어텐션을 집계하여 일시적인 스파이크를 부드럽게 만듭니다.
• 헤드 상대 정규화: 헤드의 평균 값 노름으로 토큰 점수를 정규화하여 스케일 편향을 제거합니다.
• 공식: $u^{(l,h)}_t = \alpha_t \cdot \rho^{(l,h)}_t$, 여기서 $\alpha_t$는 윈도우 질량이고 $\rho$는 정규화된 값 노름입니다.

B. 헤드 이질성 테이블 (E2 대응)

모든 헤드를 동등하게 취급하거나 온라인으로 중요도를 재추정하는 대신, COMPILERKV 는 연속적인 신뢰도 가중치의 동결된 테이블 $W_{head}[l, h]$를 컴파일합니다.

• 오프라인 컴파일: 보정 corpora 에서 **보수적 Q-학습 (CQL)**을 통해 학습됩니다. 상태는 $(l, h, B_l)$이며, 행동은 이산화된 신뢰도 가중치입니다.
• 메커니즘: 추론 중 토큰의 유틸리티는 해당 헤드의 신뢰도로 재조정됩니다. 중요한 점은 집계가 가중치 최대 풀링($\hat{u}^{(l)}_t = \max_h (u^{(l,h)}_t \cdot W_{head}[l, h])$) 을 사용한다는 것입니다. 이는 '검색 결정적'인 헤드로 하여금 거부권을 갖게 합니다: 단일 결정적 헤드가 토큰에 강력하게 어텐션한다면, 노이즈가 많은 헤드의 합의와 관계없이 해당 토큰은 유지됩니다.

C. 위험 적응형 임계값 게이트 (E3 대응)

다양한 압축 한계를 가진 프롬프트를 처리하기 위해, 위험 적응형 임계값 게이트는 프롬프트 수준의 위험 특성을 유지 임계값으로 매핑합니다.

• 위험 신호: 관찰 윈도우에 걸쳐 계산된 어텐션 엔트로피(구조적 위험) 와 국소 퍼플렉시티(의미론적 위험) 를 사용합니다.
• 오프라인 컴파일: 검색 테이블 (LUT) 이 이산화된 위험 구간 $(b_{ent}, b_{ppl})$을 유지 임계값 $\tau^{(l)}$로 매핑합니다. 이 또한 CQL 을 통해 컴파일됩니다.
• 보수적 편향: CQL 목적 함수는 드물고 고위험 상태 (고엔트로피/고퍼플렉시티) 에서 공격적인 임계값을 패널티로 부과하여, 일반적으로 온라인 휴리스틱이 실패하는 영역에서 더 안전한 유지를 지향하도록 테이블을 편향시킵니다.
• 탄력적 선택: 최종 선택은 임계값 $\tau^{(l)}$ 이상의 유틸리티를 가진 모든 토큰을 유지합니다. 후보 집합이 예산 $B_l$을 초과하면 상위 $B_l$개 토큰으로 제한됩니다. 미달하는 경우 모든 후보가 유지됩니다 (탄력적 과소 유지).

D. 이론적 근거

이 논문은 근사 오차가 버려진 꼬리 질량에 의해 제어된다는 것을 보여주는 공식 분석을 제공합니다 (명제 4.1). 이는 단일 프롬프트 내의 어텐션 행들 간의 독립성을 가정하는 대신, 프롬프트 수준(독립적인 프롬프트 간) 에서 유한 테이블 보정 오차를 제어함으로써 분산을 줄인다고 주장합니다.

3. 주요 기여

1. 문제 재정의: 프리필 전용 압축에서 보정을 위한 통계 단위는 단일 프롬프트가 아닌 교차 프롬프트 규모여야 함을 규명합니다.
2. 오프라인 컴파일: 헤드 신뢰도와 위험 적응형 임계값을 오프라인으로 컴파일하여 휴대 가능한 아키텍처 사전 지식으로 변환하는 방법을 도입합니다.
3. 결합된 정책: 헤드 신뢰도 가중치와 위험 임계값이 동일한 하류 보상에 대해 보정되는 공동 최적화를 제안하며, 이는 독립적인 모듈이 아닙니다.
4. 휴대성: 컴파일된 테이블이 분리된 corpora 간뿐만 아니라 다른 모델 패밀리 (예: LLaMA 에서 Mistral 로) 간에도 최소한의 성능 저하로 전이됨을 입증합니다.

4. 실험 결과

이 방법은 LongBench, RULER, 그리고 Needle-in-a-Haystack (NIAH) 벤치마크를 통해 4 개의 7B–8B 백본 (LLaMA-3, Mistral, Qwen2, InternLM) 에서 평가되었습니다.

• LongBench (512 토큰 예산): COMPILERKV 는 네 가지 백본 모두에서 압축 SOTA를 달성하여, 가장 강력한 베이스라인 (예: ChunkKV, HeadKV) 보다 평균 +1.67 점 (95% CI [+1.08, +2.37]) 더 높은 성능을 보였습니다. 일부 백본에서는 FullKV 성능을 약간 초과했는데, 저자들은 이를 정보 획득이 아닌 노이즈 제거에 기인한다고 설명합니다.
• 압력 환경:
  • RULER (128k 컨텍스트): 고정된 512/32k 캐시 비율 하에서 COMPILERKV 는 약 73 점의 점수를 유지하여 SnapKV(38) 와 KVZip(66.5) 을 크게 능가하는 반면, FullKV(~79) 아래에 머뭅니다.
  • NIAH (32k): 0.89 정확도를 달성하여 SnapKV 의 0.42 와 대조됩니다.
  • 배칭: 32k 입력에 배치 크기 16 에서 COMPILERKV 는 실행 가능하지만, FullKV 는 메모리 부족 (OOM) 으로 실패합니다.
• 전이성:
  • 교차 코퍼스: 헤드 신뢰도 순위는 분리된 보정 corpora 간에 높은 스피어만 상관관계 ($\bar{\rho} \approx 0.90$) 를 보입니다.
  • 교차 모델: 모델 간 직접 테이블 전이 (예: LLaMA-3 에서 Mistral 로) 는 목표 특화 컴파일에 비해 불과 0.4–0.8 LongBench 점의 손실만 초래합니다.
• 효율성: 온라인 오버헤드는 최소화되었습니다: 레이어당 $O(1)$ 테이블 조회와 후보 풀에 대한 단일 정렬입니다. 32k 컨텍스트에서 FullKV 대비 최대 2.45 배의 디코딩 처리량을 달성합니다.

5. 중요성 및 주장

이 논문은 COMPILERKV 가 유지 정책이 계산되는 방식의 통계적 불일치를 수정함으로써 고정된 예산에서 FullKV 와의 격차를 해소한다고 주장합니다. 보정 신호의 계산을 온라인, 노이즈가 많고, 단일 프롬프트 추정으로부터 오프라인, 컴파일된, 교차 프롬프트 사전 지식으로 전환함으로써, 이 방법은 노이즈와 이질성에 대한 선택 과정을 안정화합니다.

저자들은 혁신성이 단순히 '헤드 인식'(이전 방법인 HeadKV 가 가지고 있는 기능) 이 아니라, 연속적인 헤드 신뢰도와 프롬프트 위험 임계값의 공동 오프라인 컴파일에 있다고 강조합니다. 이는 헤드 신뢰도를 튜닝 트릭에서 재사용 가능한 과학적 객체로 변환하여, 한 번 컴파일하고 $O(1)$의 온라인 오버헤드로 모델과 컨텍스트 전반에 배포할 수 있게 합니다. 이 연구는 프리필 전용 압축의 경우, '올바른' 신호는 입력 프롬프트의 일시적 특성이 아니라 모델의 아키텍처 속성임을 시사합니다.