ReFreeKV는 임계값 불민감한 KV 캐시 압축을 목표로 한 연구다

KV 캐시는 긴 컨텍스트에서 GPU 메모리를 급증시키는 주요 요인으로, 기존 프루닝 기법은 입력별 예산 임계값에 민감해 일반화에 취약했다. 논문은 입력별 임계값 탐색을 제거한 'threshold-free' 목표를 제시하고, 이 목표 하에서 실시간으로 예산을 조정하면서 full-cache 성능을 보존하는 방법을 제시했다. 이로 인해 다양한 도메인과 길이를 혼합한 실제 환경에서도 안정적인 추론 성능을 유지하면서 메모리 사용을 줄일 수 있는 가능성이 확인됐다.

Transformer 계열 모델은 autoregressive 생성 과정에서 이전 토큰들의 key와 value 벡터를 KV cache로 저장하고, 매 디코딩 단계에서 이 캐시를 사용해 attention을 계산한다. 이때 KV cache 크기는 모델 크기와 시퀀스 길이에 비례하여 GPU 메모리를 급격히 소모하며, 따라서 불필요한 캐시 위치를 제거하면 메모리와 계산 비용을 절감할 수 있다. 기존 프루닝 방법들은 어떤 위치를 버릴지 결정하려면 입력별 예산 임계값을 사전에 정해야 하며, 이 값은 데이터셋과 태스크에 따라 크게 달라져 실전에서는 튜닝 비용과 성능 불안정성을 초래했다.

ReFreeKV는 두 단계로 작동한다. 첫 번째 단계는 위치 기반 순위화를 수행해 시퀀스 앞부분의 토큰을 우선 보존 후보로 선정하고, 뒤쪽 토큰은 역순으로 고려해 후보 순서를 만든다. 이 순위화는 토큰의 위치 편향을 이용해 계산 오버헤드를 낮추는 한편 이후 단계의 탐색 공간을 단순화한다. 두 번째 단계는 Uni-Metric이라고 명명한 정지 조건을 사용해 순위화된 후보를 왼쪽에서 오른쪽으로 순차적으로 보존한다. Uni-Metric은 원본 attention 행렬의 Frobenius norm과, 특정 위치까지를 남겼을 때의 축소된 attention 행렬(norm) 사이의 비율 감소를 기준으로 삼는다. 식으로는 i_prune = argmin_j (1 - ||Ã_j||_F / ||A||_F < T) 형태로 정지 위치를 결정하며, 논문은 경험적으로 T=1%를 보편 임계값으로 선택했다. 계산량 절감을 위해 전체 n×n 어텐션을 계산하지 않고 마지막 k 행의 평균으로 1×n 벡터 A'를 얻어 근사 norm을 계산하며, 실험에서 k=1이 효율성과 성능 보존에서 우수함이 확인됐다. 구현 측면에서는 PyTorch의 누적 제곱합과 where 연산을 활용해 모든 레이어의 정지 위치를 병렬로 산출해 프루닝 오버헤드를 최소화했고, 초기 두 개의 bottom layer는 안전성을 위해 항상 전체 KV를 유지하는 설계를 채택했다.

주요 실험은 13개 데이터셋과 세 가지 LLM 계열(Llama3-Instruct, Mistral-7B-Instruct, Qwen2.5-Instruct)을 대상으로 수행됐다. ReFreeKV는 Llama3-8B에서 평균 KV 예산을 63.68%로 자동 할당하면서 전체 데이터셋 평균에서 full-cache 성능을 소폭 초과했고, Qwen2.5-7B에서는 평균 76.02% 예산으로 full-cache 대비 +2.63% 향상을 보고했다. Mistral-7B에서는 평균 약 86.75% 예산으로 1.5% 성능 저하라는 결과가 관찰되어 모델별로 보존 수준 차이가 존재함이 확인됐다. Ablation study에서는 attention 행렬 축소를 위해 k=1을 선택할 때 계산 복잡도가 크게 감소하면서도 성능 손실이 거의 없었고, 위치 기반 순위화가 attention-score 기반 순위화보다 전반적인 견고성과 효율성에서 우수한 결과를 냈다. Latency 측정에서는 PyTorch 병렬 연산으로 프루닝 오버헤드가 기존 방법들과 동등하거나 미미한 수준이며, 배치 처리 시 throughput이 10–20% 향상되는 관측이 보고됐다.

전체 아키텍처는 입력 prefilling 후 한 번만 실행되는 프루닝 파이프라인으로 설계됐다. 각 Transformer 레이어·헤드 단위로 KV 위치를 순위화하고, 순차적 보존 판정을 위한 Uni-Metric 계산은 각 샘플에 대해 병렬로 수행돼 루프 오버헤드를 제거했다. Uni-Metric의 핵심 수학적 기반은 attention 행렬 A의 Frobenius norm ||A||_F이며, 특정 위치 i까지 보존했을 때의 축약 행렬 ||Ã_i||_F와의 비율 감소 1 - ||Ã_i||_F / ||A||_F를 사용해 정지 조건을 판단한다. 어텐션 행렬 전체를 계산할 경우 계산 및 메모리 복잡도는 O(n^2)이나, 이를 줄이기 위해 A의 마지막 k 행을 평균해 얻은 1×n 벡터 A'로 근사 norm을 계산함으로써 실전 복잡도를 O(k) 수준으로 낮췄다. 실험 설정에서는 k=1을 채택해 거의 O(1) 수준의 추가 비용으로 프루닝 판단을 수행했고, PyTorch의 cumsum과 where 연산으로 누적 제곱합 기반의 정지 위치를 병렬로 식별했다. 구현상 안정성 확보를 위해 초기 두 레이어의 KV는 항상 보존하도록 설정했으며, 배치 처리를 위해 짧은 샘플은 패딩으로 맞춰 각 샘플별로 독립적인 보존 포인트를 허용했다.

논문에서 명시한 주요 한계는 ReFreeKV가 달성한 압축비와 이론적 최적 예산 사이에 간극이 존재한다는 점이다. 일부 데이터셋에서는 더 작은 예산으로도 full-cache 성능을 유지할 수 있으나 ReFreeKV는 보수적으로 더 많은 예산을 유지하는 사례가 관찰됐다. 또한 ReFreeKV는 경험적 근거로 near-lossless를 보였으나 성능 저하에 대한 형식적 보증을 제공하지 않으며 특정 백본에서 소폭의 성능 하락이 보고된 점이 한계로 명시됐다.