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

왜 중요한가

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

핵심 기여

경계 기반 예측 타깃과 손실로 고정 예산 유지

KVpop은 토큰이 보호 윈도우를 벗어난 시점에서 해당 토큰의 미래 주의 질량을 타깃으로 삼아 유지 여부를 학습한다. 경계 손실(pairwise logistic boundary loss)은 새로 평가 가능한 토큰과 교사 컷오프 토큰 간의 비교만 최적화하여 학습 비용을 상수 시간으로 유지한다. 이 방식은 유지 결정이 토큰 삽입 시점이 아니라 증거가 누적된 경계 시점에서 이뤄지도록 보장한다.

전치 주의(transposed-attention)를 통한 효율적 타깃 계산

전치 주의는 쿼리와 키의 역할을 바꿔 열 단위 합산을 계산함으로써 S×S 밀집 attention 행렬을 만들지 않고도 각 키의 미래 주의 질량을 복원한다. 학습 중에만 호출되는 이 추가 패스에서는 학생의 희소 LSE 값을 재활용하여 계산 비용을 낮춘다. 그 결과 타깃 산출이 정확도를 유지하면서도 훈련 오버헤드를 최소화한다.

지연된(stateful) 스코어러로 근미래 문맥 통합

KVpop은 mLSTM 기반의 상태ful 스코어러를 옵션으로 지원하여 토큰이 보호 윈도우를 벗어나기 전까지 메모리에 문맥을 축적할 수 있게 한다. 지연된 읽기(delayed readout)는 스코어러가 근미래 증거를 반영한 뒤 유지 결정을 내리게 하여 즉시 점수만 사용하는 stateless 접근보다 정확도를 향상시킨다. 실험에서는 지연 읽기가 토큰 정확도에서 유의미한 개선을 가져왔다.

추론 시 일정한 머리별 예산으로 실행 효율성 확보

모든 KV 헤드에 대해 균일한 예산 B = s + w + k를 강제함으로써 실행 시 퍼헤드 캐시가 편차 없이 유지된다. 이 균일성은 GPU에서의 컴파일 및 실행 효율을 높여 DMS와 같은 동적 게이트 기반 방법보다 긴 생성에서 더 낮은 지연을 가능하게 한다. 실험에서 KVpop은 동일한 예산에서 DMS보다 더 높은 장문 처리 처리량과 안정적 메모리 사용을 보였다.

핵심 아이디어 이해하기

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 실행 효율을 높였지만 계층별·헤드별 혼합적 예산 설계가 품질-효율 균형을 더 개선할 가능성이 있다는 점이 언급되었다.

실무 활용

KVpop은 사후 훈련(retrofit) 방식으로 기존 Transformer 모델에 적용할 수 있어 추론 인프라의 메모리·지연 문제를 완화한다. 장문 생성이 빈번한 애플리케이션에서 KVpop은 메모리 소요를 고정 예산으로 제한하면서도 모델 성능을 상당 부분 유지한다. 제공된 결과는 수학적 추론과 코드 생성 계열에서 전이 가능성을 보였으므로 제품 환경에서 긴 대화나 대규모 시퀀스 생성에 실용적 가치가 있다.

대화형 에이전트에서 긴 대화 히스토리를 유지하면서 GPU 메모리를 일정 수준으로 제한해야 하는 서비스 배포
수학·추론 중심의 장문 자동생성 파이프라인에서 KV 캐시를 압축해 비용과 지연을 줄이는 배치 추론
코드 생성·STEM 추론처럼 중간에 오래 참조되는 토큰이 존재하는 작업에서 성능 저하를 최소화하면서 메모리 예산을 절감

코드 공개 여부: 미확인

키워드

KV cache(KV 캐시)future-attention(미래 주의 질량)transposed-attention(전치 주의)mLSTM(mLSTM)boundary-aware loss(경계 인식 손실)

용어 해설

키-값 캐시(KV Cache): — Transformer의 autoregressive 디코딩에서 이전 토큰의 key와 value를 저장해 재계산을 피하는 구조이다. 토큰 수에 비례해 메모리가 선형 증가하므로 긴 컨텍스트에서는 메모리 병목이 된다. 이 논문에서는 해당 캐시를 고정 예산으로 압축하는 문제를 다룬다.
미래 주의 질량(Future-Attention): — 어떤 과거 토큰이 보호 윈도우를 벗어난 이후 향후 쿼리들로부터 받는 누적 attention 확률의 합이다. 이 값이 크면 해당 토큰이 장기적 유용성을 가진 것으로 간주된다. KVpop은 이 값을 학습 타깃으로 사용해 어떤 토큰을 유지할지 결정한다.
전치 주의 연산(Transposed-Attention): — 쿼리와 키의 역할을 바꿔 열 단위의 합산을 계산하는 추가적인 attention 호출이다. 이 방식을 통해 전체 S×S 밀집 행렬을 만들지 않고도 각 키에 대한 미래 주의 질량을 병렬로 복구한다. 학습 단계에서만 호출되어 추론 오버헤드를 남기지 않는다.
펜윅 트리(Fenwick Tree): — 정적 우선순위에 따라 토큰 순위를 유지하고 쿼리별로 k번째 컷오프 순위를 온라인으로 계산하는 자료구조이다. 단일 토큰이 참여할 때마다 순위를 삽입·조회하여 쿼리마다 재정렬 비용을 피한다. KVpop은 이를 이용해 running top-k를 효율적으로 구현한다.

코드 예제

text

1:  student model f_theta, dense teacher f_theta_bar
2:  sequence x_{1:S}, budget (s,w,k)
3:  H^0 <- Embed(x_{1:S})
4:  L_score <- 0
5:  for l = 1,...,L do
6:    (Q^l,K^l,V^l) <- Proj^l(H^{l-1})
7:    rhat^l <- Scorer^l(K^l,V^l)
8:    r^l(q) <- ApplyDecay(rhat^l,q)
9:    M^l <- TopKMask(r^l,s,w,k)
10:   (H^l, LSE_tilde^l) <- SparseAttn(Q^l,K^l,V^l,M^l)
11:   r_tgt^l <- FutureTarget(Q^l,K^l,LSE_tilde^l,w)
12:   L_score <- L_score + BndLoss(r^l, r_tgt^l, s, w, k)
13: end for
14: L_score <- L_score / L
15: L <- KL(f_theta_bar(x_{1:S}) || LMHead(H^L)) + L_score
16: Update theta using grad_theta L
17: return L

이 코드는 논문의 Algorithm 1에 해당하는 훈련 스텝 의사코드로, KVpop 스코어러 학습 루프와 경계 손실 계산, 전치 주의로 타깃을 계산하는 흐름을 보여준다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

코드 예제

관련 토론

댓글

관련 기사

KV cache를 검색 가능한 벡터 공간으로 라우팅

Grouped Query Attention의 작동 원리와 KV cache 최적화

DeepSeek-V2의 핵심 기술인 Multihead Latent Attention(MLA)의 원리와 구현.

KV Cache 압축 평가를 부풀리는 법

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

KV cache를 검색 가능한 벡터 공간으로 라우팅

Grouped Query Attention의 작동 원리와 KV cache 최적화

DeepSeek-V2의 핵심 기술인 Multihead Latent Attention(MLA)의 원리와 구현.

KV Cache 압축 평가를 부풀리는 법