컴파일러KV: 위험 적응형 KV 압축 프레임워크

초록

컴파일러KV는 사전 단계에서 오프라인 경험을 테이블 형태로 컴파일해, 프리필 전용 KV 캐시 압축을 수행한다. 헤드 이질성을 고려한 신뢰 가중치와 프롬프트 위험을 정량화한 엔트로피·퍼플렉시티 기반 임계값 게이팅을 결합해, 512 토큰 예산 하에서 전체 KV 성능의 97.7%를 회복하고 기존 최고 방법보다 평균 5.2점 이상 향상한다.

상세 분석

본 논문은 장기 컨텍스트를 필요로 하는 LLM에서 KV 캐시 메모리의 선형 증가 문제를 해결하기 위해, “프리필‑only” 압축을 하나의 일회성 의사결정 문제로 재정의한다. 기존 방법들은 전역적인 압축 비율을 적용하거나, 순간적인 어텐션 스코어에만 의존해 헤드별 기능 차이를 무시함으로써, 메모리 예산이 촉박할 때 고엔트로피 프롬프트에서 정보 손실이 급격히 발생하는 ‘테일 실패’를 초래한다.

컴파일러KV는 두 가지 핵심 모듈을 도입한다. 첫째, Head Heterogeneity Table은 오프라인 컨텍스추얼 밴딧 학습을 통해 각 레이어·헤드(l,h)에 대한 신뢰 가중치 wₗ,ₕ를 학습한다. 이 가중치는 헤드별 유틸리티 점수 uₗ,ₕ(t) 에 곱해져, 기능적으로 중요한 검색 헤드가 압축 과정에서 과소평가되지 않도록 보장한다. 논문은 이를 정량적으로 설명하기 위해, 헤드 가중치를 적용한 어텐션 근사 오차에 대한 상한(Theorem 4.1)을 제시하고, 이 상한이 헤드 이질성을 반영할수록 압축 후 손실이 지수적으로 감소함을 증명한다.

둘째, Risk‑Adaptive Threshold Gating은 프롬프트 수준의 위험을 두 가지 통계량—어텐션 엔트로피 H와 로컬 퍼플렉시티 P—의 결합으로 추정한다. H는 토큰 간 어텐션 분포의 확산 정도를, P는 현재 토큰 예측의 불확실성을 나타내며, 두 값이 동시에 높을수록 압축 위험이 크다고 판단한다. 위험 점수 R = λ₁·H + λ₂·P를 정규화한 뒤, 사전에 정의된 임계값 테이블에 매핑해 토큰 보존 비율 τ(R) 을 결정한다. 이 과정은 오프라인에서 미리 학습된 LUT 형태로 저장되므로, 실제 추론 시 추가 연산 비용이 거의 없으며, 프리필 단계에서 바로 적용할 수 있다.

유틸리티 점수 자체는 글로벌 평균 어텐션 (\bar A_{j,t})와 상대값 크기 ρₗ,ₕ(t) 을 곱해 정의한다. (\bar A_{j,t})는 모든 레이어·헤드에 걸친 어텐션 평균으로, 단일 헤드의 스파이크 노이즈를 억제한다. ρₗ,ₕ(t)는 토큰의 value 벡터 L2 노름을 시퀀스 평균으로 정규화해 스케일 편향을 제거한다. 이렇게 얻은 기본 유틸리티 uₗ,ₕ(t) 에 헤드 가중치 wₗ,ₕ를 곱해 최종 점수 sₗ,ₕ(t)=wₗ,ₕ·uₗ,ₕ(t) 를 산출하고, 위험 기반 임계값 τ(R) 보다 높은 토큰만을 선택한다.

실험에서는 LongBench 16개 데이터셋과 4가지 LLM(7B13B) 백본을 대상으로 512 토큰 예산을 적용했다. 결과는 전체 KV 대비 97.7% 복구율을 보였으며, 특히 복잡한 요약·다중 홉 추론 작업에서 기존 최고 방법(SnapKV, PyramidKV 등)보다 평균 5.2점, 최대 7.8점의 BLEU/ROUGE 향상을 기록했다. Ablation study는 헤드 가중치와 위험 게이팅 각각이 독립적으로 성능을 23% 끌어올리지만, 결합될 때 시너지 효과가 나타난다는 것을 확인한다. 또한, 오프라인 컴파일 단계는 약 50K 프롬프트(≈100M 토큰)로 충분히 수렴하며, 추론 시 O(1) 테이블 조회만으로 실시간 오버헤드가 미미함을 입증한다.

요약하면, 컴파일러KV는 (1) 헤드 이질성을 정량화한 신뢰 가중치, (2) 프롬프트 위험을 정밀하게 측정하는 엔트로피·퍼플렉시티 기반 임계값, (3) 오프라인 RL 기반 테이블 컴파일이라는 세 축을 결합해, 메모리 제약이 심한 장기 컨텍스트 상황에서도 안정적이고 효율적인 KV 압축을 구현한다.