월~금 오전 7시, 일요일 오전 8시 — AI 뉴스와 용어를 보내드립니다메일로 받아보기

METAL LAB

KV 캐시를 확 줄여도, 8개짜리 '복구용 토큰'을 추가해 원래 성능에 가깝게 되돌리는 방법

arXiv:2608.012472026-08-01

RestoreKV: Recovering Full-Cache Behavior Under Aggressive Query-Agnostic KV Cache Eviction

KV 캐시를 확 줄여도, 8개짜리 '복구용 토큰'을 추가해 원래 성능에 가깝게 되돌리는 방법

긴 문맥을 다루는 LLM은 메모리를 아끼려고 KV 캐시(과거 토큰들의 계산 결과 저장소)를 압축해 버리는데, 압축이 심하면 성능이 크게 떨어진다. RestoreKV는 어떤 정보를 남길지 고르는 기존 방식은 그대로 두고, 문맥 전체를 한 번 훑어 압축으로 사라질 정보를 요약한 '복구 캐시'를 추가로 만들어 같은 메모리 예산 안에서 성능을 되살린다. Qwen3-4B에서 KVzip 방식의 RULER-4K 정확도를 38.2에서 73.2로 끌어올렸고, 추가 시간 비용은 전체 압축 과정의 0.5% 미만이다.

METAL LAB 해설 도표

RestoreKV 처리 흐름

증거 상태측정 결과가 보고됨

  1. ① 문맥 프리필긴 문맥을 모델에 넣어 전체 KV 캐시(과거 정보 저장소)를 생성한다.
  2. ② 복구 토큰 통과8개의 학습된 복구 토큰이 LoRA가 켜진 상태로 전체 KV 캐시를 한 번 훑어 압축 후에도 남길 요약 정보(복구 캐시)를 만든다.
  3. ③ 기존 압축 방식 적용KVzip 등 기존 방식이 정한 중요도 점수에 따라 남길 원본 KV를 고르되, 복구 캐시가 차지할 자리만큼 슬롯을 남긴다.
  4. ④ 예산 맞춤 캐시 결합선택된 원본 KV와 복구 캐시를 합쳐 원래와 동일한 전체 메모리 예산의 최종 캐시를 만든다.
  5. ⑤ 질의응답·생성LoRA는 꺼지고 원래의 고정된 모델이 최종 캐시를 이용해 여러 질문에 답한다.
METAL LAB이 원문을 바탕으로 재구성한 해설 도표이며, 논문 저자의 원문 figure가 아닙니다.

무엇을 했나

  1. 문맥을 한 번 압축해 여러 질문에 재사용하는 'query-agnostic' KV 캐시 방식은 압축을 세게 하면 성능이 급격히 떨어지는 문제가 있다.
  2. RestoreKV는 기존의 '어떤 KV를 남길지 고르는' 방식(KVzip 등)을 그대로 두고, LoRA라는 가벼운 추가 학습 모듈로 8개의 '복구 토큰'이 전체 KV 캐시를 한 번 훑어 압축 후에도 문맥 정보를 유지하는 소형 '복구 캐시'를 만든다.
  3. 이 복구 캐시는 원래 남기려던 KV 자리 중 일부를 차지하므로 전체 메모리 예산은 그대로 유지되며, 복구 캐시를 만들 때만 LoRA가 켜지고 이후 질문 처리·생성 단계는 원래 모델 그대로 동작한다.
  4. 학습은 전체 캐시를 쓰는 '정답 모델(teacher)'의 답변 분포를 압축 캐시를 쓰는 모델이 따라 하도록 하는 자기증류 방식으로, 전체 파라미터의 0.4%만 학습한다.
  5. Qwen3-4B, Qwen3-8B, Qwen3-0.6B, Llama-3.1-8B 등 4개 모델과 RULER-4K, QASPER, QuALITY, LongHealth 등 4개 벤치마크, KVzip·KVzip+·ContrastKV·SnapKV·H2O 5개 압축 방식 조합에서 대부분 성능이 개선됐다.
Figure 1: RestoreKV narrows the gap to full-cache behavior under aggressive compression at negligible one-time cost. Applied to KVzip+, RestoreKV reaches 86.4 RULER accuracy at 16× compression on the KVPress Benchmark while adding negligible compression-time overhead.
Figure 1: RestoreKV narrows the gap to full-cache behavior under aggressive compression at negligible one-time cost. Applied to KVzip+, RestoreKV reaches 86.4 RULER accuracy at 16× compression on the KVPress Benchmark while adding negligible compression-time overhead.
Table 1: Generalization across KV cache eviction methods. RestoreKV improves five different base methods on Qwen3-4B under matched KV budgets, with larger gains typically observed at tighter budgets. Colored subscripts indicate absolute changes from the corresponding baseline.
MethodRULER-4KQASPERQuALITYLongHealth
r=0.2r=0.1r=0.05r=0.2r=0.1r=0.05r=0.2r=0.1r=0.05r=0.2r=0.1r=0.05
KVzip91.480.138.238.233.918.066.058.546.371.570.061.0
KVzip + Ours93.5 +2.188.8 +8.773.2 +35.039.0 +0.837.7 +3.828.7 +10.767.0 +1.063.3 +4.852.4 +6.175.8 +4.373.8 +3.867.5 +6.5
KVzip+91.384.851.638.132.423.366.560.051.973.267.863.2
KVzip+ + Ours93.3 +2.089.7 +4.970.7 +19.140.9 +2.838.5 +6.131.6 +8.366.8 +0.362.9 +2.957.9 +6.074.2 +1.073.0 +5.271.5 +8.3
ContrastKV91.680.738.040.131.917.866.262.250.774.269.049.0
ContrastKV + Ours92.3 +0.784.3 +3.640.2 +2.240.9 +0.836.0 +4.121.6 +3.867.6 +1.463.3 +1.151.9 +1.275.2 +1.070.5 +1.554.0 +5.0
SnapKV33.820.612.722.419.014.353.953.042.541.036.836.2
SnapKV + Ours37.7 +3.926.3 +5.714.3 +1.629.9 +7.524.0 +5.020.7 +6.459.4 +5.551.8 −1.249.9 +7.451.8 +10.841.8 +5.038.5 +2.3
H2O8.03.53.222.013.814.459.250.441.560.248.837.0
H2O + Ours17.3 +9.311.6 +8.15.7 +2.530.4 +8.424.7 +10.921.3 +6.960.9 +1.754.9 +4.547.7 +6.264.8 +4.659.0 +10.246.8 +9.8
Figure 2: Overview of RestoreKV. After context prefill, a single LoRA-adapted restore pass processes n restore tokens (n=8 by default) with causal access to the full KV cache, generating n​L​H context-conditioned KV pairs. These restore states are combined with the retained context states under the same total KV budget. LoRA is active only during restore-cache generation; subsequent query processing and decoding use the original frozen backbone. Training (offline): The restore-token embeddings and LoRA adapters (0.4% of the 4B backbone) are optimized by self-distillation from the frozen full-cache teacher, Inference: The trained restore tokens and LoRA adapters generate the restore cache once for each new context.
Figure 2: Overview of RestoreKV. After context prefill, a single LoRA-adapted restore pass processes n restore tokens (n=8 by default) with causal access to the full KV cache, generating n​L​H context-conditioned KV pairs. These restore states are combined with the retained context states under the same total KV budget. LoRA is active only during restore-cache generation; subsequent query processing and decoding use the original frozen backbone. Training (offline): The restore-token embeddings and LoRA adapters (0.4% of the 4B backbone) are optimized by self-distillation from the frozen full-cache teacher, Inference: The trained restore tokens and LoRA adapters generate the restore cache once for each new context.
Table 2: Component and LoRA-target ablations on RULER-4K (Qwen3-4B). The first row is KVzip without a restore cache; the last row is the full RestoreKV configuration. “Fixed” replaces all eight learned restore-token embeddings with the embedding of the line-break token \n.
KV budget ratio r
Restore emb.LoRA target#Params0.20.10.05
91.480.138.2
Learned0.02M (0.0005%)92.782.642.1
Learnedq/k/v4.0M (0.1%)93.488.272.4
Learnedq/k/v/o5.9M (0.15%)93.588.472.0
Fixedq/k/v/o+MLP16.5M (0.4%)93.588.071.9
Learnedq/k/v/o+MLP16.5M (0.4%)93.588.873.2
Figure 3: Performance across four model backbones and four benchmarks. RestoreKV and RestoreKV+ are applied to KVzip and KVzip+, respectively, under the same KV budget. Both methods reduce compression-induced performance loss, with larger gains under tighter cache budgets.
Figure 3: Performance across four model backbones and four benchmarks. RestoreKV and RestoreKV+ are applied to KVzip and KVzip+, respectively, under the same KV budget. Both methods reduce compression-induced performance loss, with larger gains under tighter cache budgets.
Table 3: Effect of full-context conditioning on RULER-4K with Qwen3-4B. Generating the restore cache from the full cache before eviction adds 8.8 points at r=0.05 over generating it from the already evicted cache.
KV budget ratio r
Method0.20.10.05
KVzip (n=0)91.480.138.2
Generated from evicted cache93.787.664.4
Generated from full KV cache93.588.873.2
Figure 4: Comparison with a Synthesized-Cache Baseline on Qwen3-4B. RestoreKV+ consistently improves token-eviction baselines and remains competitive with AM-fast, a strong synthesized-cache method, at ∼13× lower compression time on 4K-token contexts.
Figure 4: Comparison with a Synthesized-Cache Baseline on Qwen3-4B. RestoreKV+ consistently improves token-eviction baselines and remains competitive with AM-fast, a strong synthesized-cache method, at ∼13× lower compression time on 4K-token contexts.
Table 4: Ablation on the number of restore tokens (RULER-4K, Qwen3-4B). All variants share the same training recipe and total KV budget.
KV budget ratio r
Method0.20.10.05
n=0 (KVzip)91.480.138.2
n=193.085.365.3
n=293.487.670.4
n=493.487.068.5
n=8 (default)93.588.873.2
n=1693.387.869.1
Figure 5: Layer-wise context-attention recovery on RULER-4K (Qwen3-4B, r=0.05). KVzip substantially suppresses attention from the final query token to non-sink context KV pairs. RestoreKV recovers part of the lost attention (shaded), increasing the average context-attention mass from 0.56% to 1.30% and reducing KL divergence from the full-cache model by 16% (→0.167).
Figure 5: Layer-wise context-attention recovery on RULER-4K (Qwen3-4B, r=0.05). KVzip substantially suppresses attention from the final query token to non-sink context KV pairs. RestoreKV recovers part of the lost attention (shaded), increasing the average context-attention mass from 0.56% to 1.30% and reducing KL divergence from the full-cache model by 16% (→0.167).
Table A: Ablation on the training ratio range (RULER-4K, Qwen3-4B). RestoreKV is trained by sampling the retention ratio from 𝒰⁡(rmin,rmax). The first row reports the KVzip baseline without restoration.
Training ratio rangeKV budget ratio r
rminrmax0.40.20.10.05
--93.591.480.138.2
0.0250.2594.393.588.873.2
0.0250.5094.193.186.163.5
0.0500.2594.393.286.565.6
0.0500.5094.092.884.551.5
Figure 6: Efficiency analysis (Llama-3.1-8B-Instruct, 32K context). (a) Restore-cache generation adds only ∼0.04 s (0.5%) and 84 MB (0.4%) on top of any base scorer. (b) Attention latency and KV memory are nearly identical to the base method.
Figure 6: Efficiency analysis (Llama-3.1-8B-Instruct, 32K context). (a) Restore-cache generation adds only ∼0.04 s (0.5%) and 84 MB (0.4%) on top of any base scorer. (b) Attention latency and KV memory are nearly identical to the base method.
Table B: Effect of the query-position offset (RULER-4K, Qwen3-4B). Offset-only KVzip applies the same +8 query offset as RestoreKV without adding restore states. Its performance remains close to standard KVzip, suggesting that the positional offset alone is unlikely to account for RestoreKV’s improvement.
VariantRestore statesQuery offsetKV budget ratio r
0.20.10.05
KVzipNo091.480.138.2
Offset-only KVzipNo+890.580.836.9
RestoreKVYes+893.588.873.2
Figure A: RestoreKV composes with Fast KVzip, a learning-based scorer (RULER-4K, Qwen3-8B). At 16×, Fast KVzip drops to 46.7, RestoreKV restores it to 79.3, above plain KVzip (59.0).
Figure A: RestoreKV composes with Fast KVzip, a learning-based scorer (RULER-4K, Qwen3-8B). At 16×, Fast KVzip drops to 46.7, RestoreKV restores it to 79.3, above plain KVzip (59.0).
Table C: Effect of training–inference evictor mismatch (RULER-4K, Qwen3-4B). The inference evictor is fixed to KVzip, while RestoreKV is trained using the evictor in the first column.
Training evictorInference evictorKV budget ratio r
0.200.100.05
-KVzip91.480.138.2
SnapKVKVzip93.581.642.0
KVzipKVzip93.588.873.2
Figure B: SCBench results on long-context tasks (Qwen3-4B).
Figure B: SCBench results on long-context tasks (Qwen3-4B).
Table D: LongBench results over 16 tasks under aggressive KV compression. Categories are single-document QA (S-QA), multi-document QA (M-QA), summarization (Summ), few-shot learning (Fewshot), synthetic tasks (Synth), and code tasks (Code). Rel. denotes the average score relative to the full KV cache.
MethodS-QAM-QASummFewshotSynthCodeAvgRel.
Full KV cache
Llama-3.1-8B-Instruct44.647.629.254.055.148.146.4100.0%
KV budget ratio = 0.125
KVzip42.743.927.761.944.047.744.796.3%
KVzip + Ours43.643.028.563.945.949.345.798.4%
KV budget ratio = 0.0625
KVzip33.332.424.049.120.641.633.572.1%
KVzip + Ours38.235.427.059.322.843.237.781.2%
Table E: RestoreKV training data statistics (context length in tokens).
SourceSamplesCtx (mean)Ctx range
LongAlpaca2,4887,3064,961–15,237
PG-192,2603,0713,071–3,072
Tulu-3 FLAN1,500898512–9,283
Total6,2484,236512–15,237
Table F: Experimental environment. Hardware and software configurations used for all experiments.
Hardware
GPUNVIDIA RTX PRO 6000 Blackwell (96 GB)
CPUIntel Xeon Gold 6530
Memory503 GiB
Software
OSUbuntu 22.04.5 LTS
FrameworkPyTorch 2.8.0 (CUDA 12.8)
Table G: Sensitivity to Training Seeds. RULER-4K accuracy on Qwen3-4B, reported as mean±std over three training seeds.
KV budget ratio
0.400.200.100.05
KVzip93.4691.3880.0838.23
RestoreKV94.25±0.0493.37±0.1188.21±0.5172.81±0.57
Table H: Evaluation metrics used for each benchmark.
BenchmarkEvaluation metric
RULER-4KOfficial string-matching score
QASPERMaximum token-F1 over reference annotations
QuALITYAnswer-letter accuracy
LongHealthAnswer-letter accuracy
LongBenchOfficial per-task metric (F1, ROUGE-L, or accuracy)
SCBenchOfficial task-specific evaluator

실제로 확인된 결과

  • Qwen3-4B, 예산 비율 5%에서 RestoreKV가 KVzip의 RULER-4K 정확도를 38.2에서 73.2로, KVzip+ 기반은 51.6에서 70.7로 올렸다.
  • KVzip+에 적용했을 때 KVPress 벤치마크에서 16배 압축 상태로 RULER 정확도 86.4를 달성했고, 32K 문맥 기준 추가 시간 비용은 전체 압축 시간의 0.5% 미만, 메모리 증가는 84MB(0.4%)였다.
  • 학습 기반 스코어러인 Fast KVzip에 적용하니 16배 압축에서 RULER-4K 정확도가 46.7에서 79.3으로 올라 일반 KVzip(59.0)보다도 높아졌다.
  • 합성 캐시 방식인 AM-fast와 비교해 r=0.05에서 RULER-4K 70.7 대 52.8로 앞섰고, 4K 문맥 캐시 생성 시간은 0.74초로 AM-fast(9.68초)보다 약 13배 빨랐다.
  • 레이어별 분석에서 문맥에 대한 평균 주의 집중도가 압축 후 0.56%까지 떨어졌던 것을 RestoreKV가 1.30%로 회복시키고, 전체 캐시 모델과의 KL 발산을 0.198에서 0.167로 줄였다.

어디에 쓸 수 있나

  • 긴 문서를 미리 압축해 캐시로 저장해 두고 여러 질문에 반복 재사용하는 문서 QA·요약 서비스에서 메모리 절감과 정답 품질 사이 균형을 개선하는 용도
  • 기존에 쓰고 있는 KVzip, KVzip+, SnapKV, H2O 등 다양한 압축 방식 위에 추가 모듈로 얹어 재학습 없이 압축 손실을 줄이는 용도
  • 학습형 경량 스코어러(Fast KVzip 등)와 결합해 스코어링 속도는 유지하면서 압축 성능 저하를 완화하는 용도

한계와 남은 검증

  • 학습 시 사용한 압축기(evictor)와 추론 시 사용하는 압축기가 다르면 성능 회복 폭이 크게 줄어드는 것으로 나타나, 압축 방식이 바뀔 때마다 재학습이 필요할 수 있다.
  • 실험은 Qwen3 계열(0.6B~8B)과 Llama-3.1-8B, 문맥 길이 최대 약 15K~32K 범위, 그리고 SCBench의 약 104K 토큰 사례에 한정되어 더 큰 모델이나 다른 아키텍처에서의 검증은 아직 없다.
  • 훈련 데이터는 LongAlpaca, PG-19, Tulu-3 Flan에서 만든 약 6.2천 개 문맥-질문 쌍에 한정되어 있어 다른 도메인 문서에 대한 일반화는 별도 검증이 필요하다.
  • 복구 토큰 개수(n=8)나 LoRA 랭크 등 설정은 특정 값으로 고정해 실험했으며, 다른 설정에서의 최적성은 추가 탐색이 필요하다.

왜 중요한가

긴 문서를 매번 다시 읽지 않고 캐시로 재사용하는 서비스에서는 메모리를 아끼려 캐시를 심하게 줄일수록 답변 품질이 떨어지는 문제가 있는데, 이 방법은 기존 압축 파이프라인을 바꾸지 않고 작은 모듈만 얹어 그 손실을 줄인다. 압축 방식 자체를 새로 설계할 필요 없이 기존 시스템에 끼워 넣을 수 있다는 점에서 실무 적용 부담이 적다.

이 논문의 용어

  • KV 캐시 · 트랜스포머 모델이 이전 토큰들을 다시 계산하지 않도록 저장해 두는 key-value 값들의 모음
  • query-agnostic 압축 · 미래에 어떤 질문이 올지 모른 채 문맥을 한 번만 압축해 여러 질문에 재사용하는 방식
  • LoRA · 원래 모델 가중치는 그대로 두고 작은 저랭크 행렬만 추가로 학습해 적은 비용으로 모델을 조정하는 기법
  • 자기증류(self-distillation) · 압축 없이 전체 정보를 쓰는 같은 모델(teacher)의 출력을 압축된 모델(student)이 따라 하도록 학습시키는 방법
  • RULER-4K · 4K 토큰 길이의 문맥에서 검색·추적·집계 능력을 시험하는 13개 과제로 구성된 합성 벤치마크

저자 · Changwoo Baek

arXiv에서 원문 보기

최신 논문

논문 전체 보기 →

METAL LAB 최신 기사

그림 출처: Changwoo Baek et al., arXiv:2608.01247, arxiv-nonexclusive