KVpop: 미래 주의 기반 예측형 온라인 정리로 KV 캐시를 고정 예산으로 압축하는 방법

긴 컨텍스트에서 Transformer의 KV 캐시는 메모리와 실행 시간에서 병목을 유발한다. KVpop은 고정된 머리별 예산을 유지하면서도 미래에 실제로 사용될 토큰을 예측하여 불필요한 KV 항목을 제거하므로 긴 생성에서 메모리 사용과 지연을 동시에 줄인다. 이로 인해 대형 모델의 장문 생성과 장시간 추론이 현실적으로 가능해지는 방향으로 효용성이 증대된다.

Transformer의 autoregressive 디코딩은 이전 토큰들의 key와 value를 KV 캐시에 누적해 향후 쿼리에서 재사용한다. 이 캐시는 시퀀스 길이에 따라 선형으로 커지므로 긴 문서나 수십만 토큰 생성에서는 메모리와 대역폭의 병목을 초래한다. 따라서 실용적으로는 유한한 예산 내에서 어떤 토큰을 남기고 어떤 토큰을 버릴지 결정하는 전략이 필요하다. KVpop의 핵심은 ‘미래에 실제로 얼마나 참조될 것인가’를 직접 학습 타깃으로 삼는 점이다. 구체적으로 보호 윈도우를 벗어난 이후의 쿼리들이 해당 키에 할당하는 누적 attention 확률을 미래 주의 질량으로 정의하고, 이 값을 전치 주의 연산으로 효율적으로 계산해 스코어러의 학습 타깃으로 사용한다. 이렇게 하면 과거의 국소적 신호만으로 점수를 매기는 대신 경계 시점에 누적된 근미래 증거를 반영해 유지 결정을 내릴 수 있다. 학습된 스코어러는 추론 시 단지 상위 k를 선택하는 우선순위만 제공하며, 추론 단계에서는 학생 스코어와 고정된 예산으로 희소 마스크를 구성해 실제 KV를 유지한다. 스코어 계산은 가벼운 헤드별 모듈로 한정되어 백본 가중치와 분리하여 최적화되므로 베이스 모델의 추론 오버헤드를 거의 추가하지 않는다. 또한 머리별로 균일한 예산을 강제하면 실행 효율성이 확보되어 긴 생성 길이에서 더 낮은 지연과 일정한 VRAM 차트를 얻는다.

전체 접근은 세 가지 구성 요소로 나뉜다: 고정된 머리별 예산(sink s, 최근 윈도우 w, 장기 top-k), 스코어러 모듈, 그리고 학습 시 전치 주의를 이용한 미래 주의 타깃 산출이다. 각 attention 헤드는 항상 sink 토큰과 최근 w 토큰을 보존하고, 나머지 후보는 스코어러가 매긴 중요도 순으로 top-k에 들면 유지된다. 이 구조는 추론 시 고정된 메모리 한계를 보장하면서도 학습으로 장기 유틸리티를 반영할 수 있게 한다. 스코어러 학습은 경계 손실(boundary-aware retention loss)을 사용해 이루어진다. 새로 평가 가능한 토큰(t_new)과 교사 컷오프 토큰(t_bnd) 간의 점수 차이에 대한 pairwise logistic(softplus) 손실이 쿼리·헤드별로 평균화되어 계산된다. 이 손실은 캐시 멤버십을 실제로 바꾸는 단일 비교에 집중하기 때문에 샘플당 계산 비용이 상수 시간으로 유지된다. 미래 주의 타깃은 밀집 S×S 확률 행렬을 구성하지 않고 전치 주의 호출로 계산한다. 학생의 희소 LSE를 근사으로 재사용하고, 쿼리·키 역할을 바꾼 attention 호출로 키별 열 합(log-sum-exp)을 얻은 뒤 N_t로 정규화해 m_t를 복원한다. 이 과정은 학습 전용이며 추론 오버헤드를 남기지 않는다. 또한 running top-k는 정적 우선순위 정렬과 Fenwick tree를 사용해 쿼리별 컷오프를 온라인으로 계산하여 매 쿼리마다 전체 재정렬을 피한다.

메인 벤치마크에서 KVpop은 Qwen3-4B와 Qwen3-8B에 대해 높은 유지율을 보였다. 논문은 AIME와 HMMT에서 CR=75%일 때 Qwen3-4B에서 약 95% 유지, Qwen3-8B에서 95% 유지를 보고했고 CR=88%에서는 Qwen3-4B에서 94% 및 Qwen3-8B에서 99%를 보고했다. 표 1의 상대 성능 비교에서 KVpop은 학습 없는 휴리스틱이나 DMS보다 평균 성능이 높았다. 지연된(stateful) 스코어러 유무에 대한 절단 실험에서 mLSTM 기반 스코어러는 지연 읽기를 허용할 경우 토큰 정확도에서 약 0.2 포인트의 추가 이득을 얻었다. 이 결과는 스코어를 삽입 시점이 아니라 보호 윈도우 경계에서 내릴 때 근미래 문맥이 가치를 더한다는 점을 실험적으로 확인한다. 또한 stateless MLP 스코어러도 경쟁력 있는 결과를 보여 스코어러 선택에서 비용-성능 균형이 가능함을 시사했다. 추론 효율성 평가에서 밀집 attention은 생성 길이에 따라 VRAM이 선형 증가했으나 KVpop과 DMS는 메모리 성장이 억제되었다. 구체적으로 논문은 밀집 모델이 16k에서 18GB, 131k에서 36GB로 상승한 반면 DMS와 KVpop은 약 19%만 증가해 약 19GB 수준에 머물렀다고 보고했다. KVpop은 같은 예산에서 DMS보다 긴 생성 길이에서 더 낮은 지연과 높은 처리량을 유지했다는 결과가 보고되었다.

아키텍처 차원에서는 KVpop이 기본 Transformer 구조를 변경하지 않고 각 KV 헤드에 경량 스코어러를 추가하는 형태이다. 각 헤드는 sink s, 보호 윈도우 w, 장기 top-k로 구성된 예산 B = s + w + k를 유지하며, 추론 시에는 이 보존 규칙만 적용되어 희소 마스크가 구성된다. 스코어러 입력은 key와 value의 연결이며 파라미터는 백본과 분리되어 유지 손실만으로 업데이트된다. 미래 주의 타깃 수학적 정의는 각 그룹 쿼리 헤드(g)에 대해 토큰 t의 후속 쿼리들에서 해당 키로 향하는 확률의 평균 m^{(h,g)}t로 정해진다. 로그 스케일로 변환한 뒤 그룹 축에서 Agg_g( log(eps + m) )를 취해 r^{tgt}{h,t}를 얻는다. 실무적으로 Agg_g는 max를 사용했고, 실험에서는 작은 상수 eps와 per-head 감쇠 γ_h, 감쇠 단계 n을 도입해 오래된 토큰이 무한히 유지되는 것을 방지했다. 전치 주의 구현은 원래의 쿼리-키 내적 ℓ(d,t)를 그대로 재생성하되 쿼리와 키 역할을 바꿔 키별 log-sum-exp를 얻는다. 학생의 희소 LSE(tilde LSE)를 재활용해 근사치를 사용함으로써 밀집 LSE를 계산하는 비용을 회피했다. running top-k는 토큰별 정적 우선순위를 한 번 계산한 뒤 Fenwick tree로 쿼리별 컷오프를 온라인으로 업데이트하여 O(S log S) 시간·O(S) 공간으로 컷오프를 산출한다. 스코어러 설계는 stateless MLP와 stateful mLSTM 버전이 있으며 stateful 경우 보호 윈도우 내에서 메모리가 누적되다가 토큰이 경계에 도달할 때 읽어 점수를 계산한다. 경계 손실은 softplus 기반 pairwise 로지스틱 형태로 새로 평가 가능한 토큰과 컷오프 토큰 간의 점수 차이를 정규화 온도 τ로 나눈 값에 가중치 ω_{q,h}를 곱해 평균화한다. 구현상 scorer 파라미터는 코사인 스케줄로 학습하고 베이스 모델은 KL distillation loss로 고정에 가깝게 유지하였다.

논문은 mLSTM 기반의 stateful 스코어러를 예시로 사용했으나 다른 메모리 기반 모듈들에 대한 광범위한 비교는 수행되지 않았다. KVpop은 밀집 attention을 전제한 사후 레트로핏 방식으로, 처음부터 캐시 친화적 구조로 학습된 모델과의 성능·효율 비교는 남은 과제로 표기되었다. 또한 머리별 균일 예산은 GPU 실행 효율을 높였지만 계층별·헤드별 혼합적 예산 설계가 품질-효율 균형을 더 개선할 가능성이 있다는 점이 언급되었다.