오늘은 NeurIPS 2025에 Oral로 선정된 KVzip이라는 논문에 대해서 Review해 볼 것이다.
https://github.com/snu-mllab/KVzip
GitHub - snu-mllab/KVzip: [NeurIPS'25 Oral] Query-agnostic KV cache eviction: 3–4× reduction in memory and 2× decrease in la
[NeurIPS'25 Oral] Query-agnostic KV cache eviction: 3–4× reduction in memory and 2× decrease in latency (Qwen3/2.5, Gemma3, LLaMA3) - snu-mllab/KVzip
github.com
Background
요즘 딥러닝에서는 Scaling을 하면 성능이 좋아진다는 것이 거의 정설로 받아들여지고 있다. 기존의 Scaling이라고 하면, 딥러닝 훈련시에 데이터의 다양성을 늘리는 것, 모델의 파라미터 개수를 늘리는 것을 뜻했다. 하지만 요즘에는 Scaling의 개념이 확장되었는데, 특히 LLM 같은 경우는 inference시에 input text의 길이를 늘리는 것도 Scaling이라고 한다.
input text의 길이를 늘리면 inference 성능이 좋아지기 때문에, 긴 input text를 효율적으로 처리하도록 모델을 design하는 연구들이 계속해서 나오고 있었다. 긴 input text를 잘 처리할 수 있는 모델이란, 같은 길이의 input text를 처리할 때 기존의 모델들에 비해서 Performance degradation이 거의 없으면서 메모리 사용량과 FLOP을 줄이는 것을 말한다. 보통은 sparse attention을 한다고 생각하면 된다. 가장 간단한 예시로는 decoding stage에서 attention을 할 때, 최근 K개의 토큰의 key만 저장해두고, 최근 K개의 key랑만 attention을 하는 window-based cache방법이 있다.
하지만 이러한 간단한 방법은 너무 큰 performance degradation을 동반하기 때문에, 조금 더 발전된 방법이 등장하게 된다. 이 방법은 pre-fill은 원래처럼 진행하고, pre-fill 후에 중요한 kv cache만 남기고 나머지 kv cache들은 eviction을 한 후, 중요한 kv cache들 만을 attention해서 decoding을 하자는 아이디어에서 출발한다. 보통 LLM의 input은 [CTX:Q]로 이루어지는데(CTX는 content를 담고있는 글이고, Q는 그 content안의 내용에 대해서 질문하는 것이다), 발전된 방법은 [CTX:Q] pre-fill이 끝나고 나서, query와 attention score가 높았던 CTX안의 토큰들의 kv와 query의 kv만 남겨두면, decoding stage에서도 효과적으로 query와 관련있는 token들만 보고 decoding을 할 수 있다는 아이디어이다.
하지만 이 방법은 한계가 존재하는데, 어떤 한계가 존재하는지 알아보기 위해 concrete example을 생각해보자. 보통의 LLM inference의 경우에는 하나의 CTX에 대해서 여러 Query가 주어지는 경우가 많다. 만약 [CTX:Q1]과 [CTX:Q2] 가 연속적인 input으로 LLM에 들어온다고 상상해보자. 그러면 우리는 [CTX:Q1]의 decoding stage를 효과적으로 하기 위해서, Q1과 잘 반응하는 CTX안의 token들만 남겨놓은 상태로 decoding stage를 진행했을 것이다. 이제 [CTX:Q2]에 대한 inference를 진행해야 하는데, 우리에게는 두 가지 선택지가 있다.
첫 선택지는 현재 메모리에 저장된 CTX의 토큰들은 Q1과 잘 반응하는 토큰들만 남아있는데, Q2와 잘 반응하는 CTX의 토큰들은 Q1과 잘 반응하는 CTX의 토큰들과 다를 수 있기 때문에, [CTX:Q2] 전체 prefill을 처음부터 다시 진행하고 CTX에 Q2와 잘 반응하는 CTX의 토큰들만을 남긴 후 [CTX:Q2]에 대한 decoding stage를 진행하는 것이다. 이 경우 prefill을 다시 해야하므로, 추가적인 overhead가 발생하는 한계가 존재한다.
두 번째 선택지는 추가적인 계산 overhead를 막기 위해서 그냥 Q1과 잘 반응했던 CTX의 토큰들을 그대로 [CTX:Q2]의 decoding stage에 이용하는 것이다. 이 경우에는 Q1과 Q2가 CTX에서 참조하는 부분이 다를 경우 많은 Performance degradation이 수반되는 한계가 존재한다.
그러면 이쯤에서 우리는 자연스럽게 생각할 수 있다. 어 그러면 그냥 Query에 상관없이 CTX에서 중요한 kv pair만 남겨둘 수 있는 방법은 없을까? 그것이 이 논문에서 풀고자 하는 문제이다.
Background - optional
long context를 KV cache에 저장하는 것이 얼마나 많은 memory를 잡아먹는지 한번 구체적인 숫자로 느껴보도록 하자. FP16 Qwen2.5-14B에 토큰 길이가 120,000인 CTX을 prefill 했을 때 생기는 KV cache의 memory consumption을 계산해보자. 전체 memory consumption := {2 x (# of layers) x (# of tokens) x (# of kv heads) x (kv head dim) x (2 bytes)} = {2 x 48 x 120,000 x 8 x 128 x 2 bytes} = 23,592,960,000 Bytes(23GiB)정도 이다. 이렇게 구체적인 숫자로 보면, decoding stage를 위해서 cache eviction을 하는 것이 decoding stage에서 FLOP과 memory consumption에 얼마나 큰 도움을 줄 수 있는지 감이 올 것이다.
Methods
자 그럼 이제 본격적으로 문제를 풀어보자. 우리가 하고싶었던 걸 다시 remind해보면, CTX를 pre-fill하면서 정말 필요한 k,v만 남기는 것이다. 여기서 저자들은 이런 생각을 하게 된다. 결국 우리가 하고싶은게, CTX에서 가장 필요한 정보를 k,v cache로 encode 해놓고, 뒤에 decoding stage에서 이 encode된 정보를 잘 이용하는것이라면 이게 예전에 딥러닝에서 많이 쓰던 방법인 auto-encoder로 미리 general한 정보를 encoding 해놓고, down-stream task에서 Linear Probing해서 쓰는 느낌이라는 생각을 하게 된다. 따라서 auto-encoder처럼 eviction한 kv cache로 부터 다시 full input context를 만들 수 있으면, 충분히 좋은 encoding이라고 생각하게 되는 것이다.
조금 더 구체적으로 말하면, eviction이 끝난 evicted KV cache를 pre-fill해둔 LLM 모델에 "Repeat the previous context"라는 query를 줬을 때 output으로 이전 input context를 전부 복원할 수 있으면 evicted KV cache는 원래의 context정보를 잘 encoding하고 있는것이라고 생각할 수 있다.
그런데 auto-encoding 방법들하고는 다른 점이 존재한다. 왜냐하면 auto-encoding에서는 그냥 간단히 dim을 줄이면서 MLP를 통과시키는 것이 encoding을 하는 방법이었고, 그냥 decoding 역시 다시 MLP를 통과시켜서 re-generation loss를 걸어주면서 from scratch부터 훈련시키면 됐는데, 지금은 LLM은 이미 훈련이 잘 되어있고, 당연히 re-generation을 잘 하기 위해서 context마다 LLM 자체를 fine-tuning시키는 것은 너무 infisible하니까, re-generation을 잘 하기 위한 eviction strategy만을 learning base로 하든 heuristic으로 하든 해야할 것이다. 저자들은 heuristic strategy를 택했다.
방법은 다음과 같다. context 전체의 KV cache를 가지고 repeat problem을 푼 후에, repeat problem을 풀면서 attention이 강하게 됐던 CTX 내의 kv cache들만 남기는 것이다. 이 kv cache들만 남기면 반대로 repeat problem을 문제없이 잘 풀 수 있을것이라고 생각한 것이다. (굉장히 타당한 논리이다.)
이제 구현 detail에 대해서 설명해보겠다. CTX를 evicted_kv cache로 encoding하기 위해서 첫번째로 해야 할 것은, CTX를 pre-fill하는 것이다. 이것은 그냥 원래 LLM prefill 하듯이 병렬로 잘 하면 된다. 그 다음으로 할 것은 이제 잘 prefill되어있는 CTX kv cache를 바탕으로, "Repeat the previous context"라는 query를 넣어서 output으로 원래 context를 내보내는 문제를 풀면서 ctx내의 attention이 많이 되는 k,v를 찾는 작업이다. 이때 그냥 auto-regressive하게 문제를 풀게되면 당연히 원래의 context가 output으로 나오지 않을 것이다. 따라서 다음토큰을 예측할 때 이전토큰 까지는 gt를 사용하는 teacher-forced decoding방식을 채택해서 문제를 풀게 된다. 이러면 두 마리 토끼를 한번에 잡을 수 있는데, 이유는 output을 만드는 과정이 병렬로 한번에 처리되기 때문이다. 구체적인 수식은 논문을 참고하면 되는데, 개인적으로는 논문의 수식이 조금 헷갈리게 적혀있어서(사실 output으로 나오게 되는 원래의 context가 teacher-forced decoding을 쓰는 바람에 input이라고 명명이 되어있다.), 그냥 위에 Figure 4를 보는게 더 좋은 것 같다.
방법론을 전부 읽고 나니 떠오르는 생각을 정리해보면, 지금 input이 context이고 output으로 input을 re-generation하는 pre-text task를 풀어서 그때 중요했던 kv cache를 남긴 후에 여러 downstream task(queries)에 적용한다는 생각이 든다. 그렇다면, 꼭 CTX에 대한 pre-text task가 re-generation이어야만 하는 이유가 있을까? 더 좋은 pre-text task가 있지 않을까? 라는 생각이 막 든다.
Methods : technical-details
위의 방식대로 re-generate를 위한 LLM inference를 하게 되면, 기존의 context만큼의 kv cache와 추가로 repeat prompt + 기존의 context만큼의 kv cache가 만들어지는 매우 huge한 inference가 되게 된다. 따라서 re-generate task를 chunk단위로 하게 된다. chunk단위로 할 때에는 repeat prompt를 "Repeat the previous context starting with" + "이전 chunk의 마지막 8개 token"을 주고 re-generation task를 수행하도록 한다. (pseudo code를 보면 for문을 돌면서 cnunk단위로 inference 하기 때문에 전체 병렬은 하지 않고, 한 chunk내에서만 병렬이 가능하다)
논문에서는 원래 pre-fill 기반의 method들 보다 kv cache eviction을 위한 process가 더 오래걸린다는 것은 인정하면서도, 여러 query에 대해서 inference를 할 때에는 재계산 필요없음 + 높은 압축률을 바탕으로 자신들의 방법이 훨씬 효율적이라는 것을 강조하고 있다. 밑의 그림에서 점선은 이전 pre-fill기반 방법들의 y축 값을 의미한다.
Experiments : detailed explanation of why this method works
다음으로는 이 논문의 방법론이 실제로 working 하는지에 대해 분석한 결과들이다. 첫 번째 실험은 CTX의 전체 kv cache를 전부 준비하고, 여러 query에 관한 decoding을 할 때 서로 다른 query들이 CTX의 어떤 kv를 많이 참조하는지 비교한 결과이다.
그림을 보면, QA에서 중요하게 다뤄졌던 kv pair는 모두 Repeat 에서도 중요하게 다뤄졌다. 또, Summarization에서 중요하게 다뤄졌던 kv pair는 대체적으로 Repeat에서도 중요하게 다뤄졌다. 또 Reasoning에서 중요하게 다뤄졌던 kv pair는 Repeat에서도 중요하게 다뤄졌다. 반면에 QA-1과 QA-2는 서로 중요하게 생각하는 kv가 다른걸 알 수 있다. 따라서 다른 task들에게서 중요하게 생각하는 kv pair는 대체로 Repeat task에서도 중요하게 다뤄졌음을 알 수 있다. (꼭 이게 Repeat task여야만 하는가에 대한 생각이 계속 든다 더 좋은게 뭐가 있을진 아직 잘 모르겠다)
다음 실험으로는 얼마나 이 방법이 KV cache를 잘 압축할 수 있는지에 대해 설명하는 실험이다. 만약 re-generation task에서 원래 context의 모든 kv cache가 중요하다고 해버리면, 사실 의미가 없다. 따라서 얼마나 필요로 하는 kv pair가 적은지 또한 매우 중요한 factor이다. 밑의 그림을 보면, pre-fill에서의 query-aware attention score를 계산하는 방법들보다 훨씬 더 좋은 압축률을 보요주고 있는 것을 볼 수 있다.
Experiments - main
기존 baseline들에 비해서 kv zip의 성능이 좋다는 것을 실험 결과를 통해 증명하고 있다. (Qwen2.5-7B-1M)
밑의 그림을 통해서 여러 모델 architecture에서도 자신들의 방법이 좋다는 것을 보여주고 있다.
Idea of Mine
마지막에 Analysis 부분을 보면, pre-text task가 Recon task말고 다른 task일 때 SQuAD 데이터셋에 대한 Accuracy를 측정한게 있는데, 굉장히 simple한 pre-text task만을 담고 있다. 훨씬 더 잘 설계된 pre-text task가 있다면, recon task보다 좋을 수 있지 않을까?
Overall Review
전체적으로 잘 쓰여진 논문이다. query-agnostic하게 evicted kv cache를 만드는 선구자격 논문이 아닐까 싶다. 왜 자신들의 방법론이 잘 통하는지에 대해서 그저 수치뿐만이 아니라 여러 task들에 대해서 CTX 안의 kv pair의 attention score correlation을 보여준 것도 굉장히 좋았던 것 같다. (요즘은 사실 이런 추가적인 분석이 기본이 된 느낌이긴 하다)
리뷰 끝!