KV 캐시를 테스트타임 학습으로 압축해 미지의 질문 대응력을 지킨다
Latent Compaction with Test-Time Training
무엇인가
긴 문맥을 다루는 LLM 추론에서 키-값(KV) 캐시는 배치 크기를 제한하고 디코딩 스텝마다 메모리 대역폭을 잡아먹는 주된 병목이다. 이 논문이 푸는 문제는 명확하다. 문서를 미리 압축해 두되, 그 뒤에 들어올 질문이 무엇인지 모르는 상태에서도 원래 정보를 잃지 않아야 한다는 것이다. 기존 접근은 크게 두 갈래다. H2O, SnapKV, PyramidKV, KVZip+처럼 원본 KV 항목 중 일부를 골라 남기는 축출(eviction) 방식과, Cartridges나 Attention Matching처럼 압축된 표현 자체를 학습하는 방식이다. 저자들은 후자의 연장선에서 "압축된 KV 캐시를 하나의 신경망처럼 취급하고 테스트타임 학습으로 직접 최적화하면 어떨까"라는 질문을 던진다.
어떻게 동작하나
제안 방법인 Latent Compaction은 모델 가중치를 전부 동결한 채, 문서 하나에 전용으로 붙는 학습 가능한 KV 슬롯 묶음을 만든다. 각 슬롯은 키·값 벡터와 함께 스칼라 어텐션 바이어스를 하나씩 갖는데, 이 바이어스는 압축으로 달라진 어텐션 질량을 보정하는 역할을 한다. 압축 캐시의 점수는 ℓ(q)=sK_cmp·q+b_cmp로 계산되고 출력은 V_cmp의 소프트맥스 가중합으로 나온다. 질문 토큰과 생성 토큰은 예산 밖에 일반 KV 쌍으로 그대로 덧붙여지고, 최종 출력은 압축 캐시의 정규화 상수와 새 토큰의 정규화 상수를 로그 공간에서 가중 평균해 하나의 소프트맥스로 합쳐진다. 즉 문서 부분만 압축하고 이후 토큰은 손대지 않는 구조다.
무엇과 다른가
초기화는 세 종류의 항목을 섞는다. 프리저버(preserver)는 KVZip+ 재구성 점수를 모든 레이어·KV 헤드에 걸쳐 통합 랭킹해 상위 항목으로 채우며, 레이어·헤드마다 개수가 달라질 수 있다. 앵커(anchor)는 교사 어텐션의 정규화된 최대값이 높은 원본 항목과 문서 첫 위치(어텐션 싱크)로 구성된다. 다이버시파이어(diversifier)는 키 공간에서 가장 먼 점을 반복 선택하는 최원거리 우선 탐색으로 뽑아 키 공간 커버리지를 넓힌다. 예산 배분은 레이어·헤드당 평균 B개 중 B/4를 앵커·다이버시파이어에 고정 할당하고, 나머지 LH(B−M)개 슬롯을 전역 KVZip+ 랭킹으로 채운다. 앵커·다이버시파이어 예산의 25%가 앵커다.
어떻게 쓰나
학습은 두 단계다. 1단계는 레이어별로 독립적으로 수행하는 어텐션 피팅으로, 교사 쿼리들에 대해 압축 캐시의 어텐션 출력과 로그 정규화 값을 전체 캐시의 것에 MSE로 맞춘다. 이때 프리저버의 키·값·바이어스는 고정되고 앵커·다이버시파이어만 갱신된다. 2단계는 전체 캐시 항목을 함께 미세조정하는 종단간 증류로, 교사가 생성한 상위 20개 토큰과 나머지를 하나의 꼬리 범주로 묶은 그룹화 KL 손실에 교사의 그리디 토큰에 대한 언어모델링 손실을 더해 최소화한다. 감독 신호는 문서에 대한 일반적인 프롬프트 480개(문서 요청 24종 × 응답 지시 20종)로 만든 교사 롤아웃에서 나오며, 하위 평가 질문과 정답은 일절 쓰이지 않는다. 학습이 끝난 캐시는 여러 질문에 재사용된다.
전제와 한계
실험은 Qwen3-1.7B를 BF16 연산·FP32 KV 캐시로 고정하고 진행했다. 1단계는 학습률 3×10⁻³로 1,000회, 2단계는 10⁻⁴로 1,000회 업데이트한다. RULER에서는 32배 압축에서 96.7% 정확도를 기록해 KVZip+의 40.0%를 크게 앞섰다. LongBench에서는 Qwen3-0.6B·1.7B·4B 모두 4배~32배 구간에서 텍스트 요약 베이스를 앞질렀고, 16배·32배에서는 KVZip+도 앞섰다. 4배 압축에서는 KVZip+와 비슷하지만 압축이 강해질수록 격차가 벌어졌다. Phonebook Lookup에서는 16배 압축에서 특히 큰 향상을 보였고, 같은 계열의 0.6B·4B·8B와 다른 계열인 MiniCPM5로도 절차가 이식됐다. 다만 이때는 각 모델이 자기 자신의 교사 출력으로 캐시를 최적화하는 방식이며, 최적화된 캐시 자체를 모델 간에 옮기는 실험은 아니다.
절제 실험은 각 구성 요소의 기여를 수치로 보여준다. 1단계를 생략하면 모든 문맥 길이에서 정확도가 급락하고, 1단계에서 프리저버의 키·값을 바꾸도록 허용해도 정확도가 크게 떨어진다. 로그 정규화 손실을 빼도 비슷한 폭으로 나빠진다. 2단계는 4K·8K·16K에서 각각 6.75, 15.50, 7.75%포인트를 올렸고, 여기에 언어모델링 손실을 더하면 4.75, 0.25, 5.25%포인트가 추가로 개선됐다. 상위 20개+꼬리 목표를 전체 어휘 목표로 바꾸면 4K는 그대로, 8K는 0.50%포인트 오르고 16K는 0.50%포인트 내려 일관된 이득이 없어, 저장 비용이 적은 그룹화 목표가 채택됐다. 학습량은 1단계를 1,000회에서 2,000회로 늘려도 61.00%에서 61.25%로 거의 변하지 않았고, 2단계는 1,000회에서 정점을 찍고 이후 과적합 기미를 보였다. 어텐션 바이어스를 제거하면 4K·8K·16K에서 0.75, 3.00, 0.25%포인트가 떨어지고, 추론 시 정규화 가중치를 단순 평균으로 바꾸면 정확도가 0.25% 이하로 무너진다.
실무 관점에서 이 논문의 쓸모는 "문서를 미리 넣어두고 여러 질문으로 재사용하는" 워크로드에 있다. 사내 문서 QA, 코드베이스 인덱스, 장기 대화 메모리처럼 같은 컨텍스트를 반복 서빙하면서 압축률을 높이고 싶은 경우, 문서당 한 번의 준비 비용을 지불하고 캐시를 고정해 재사용하는 이 방식이 맞는다. 반대로 질문마다 컨텍스트가 달라지는 일회성 요청에는 준비 비용이 그대로 손실이다. 도입 전에 확인할 것은 세 가지다. 첫째, 논문의 압축률은 논리적 KV 쌍 예산 기준이며 실제 메모리 절감이나 추론 속도 향상을 측정한 것이 아니라는 점. 둘째, 바이어스·정밀도·패딩 때문에 실제 메모리 사용량은 별도로 재야 한다는 점. 셋째, 문서당 480개 프롤프트 롤아웃과 2,000회 업데이트라는 준비 비용이 서비스 지연 예산에 맞는지다.
저자들이 밝힌 한계도 분명하다. 문서별 피팅 자체가 여전히 비싸고, 압축이 극단으로 갈수록 정확도가 무너진다. 실제로 가장 공격적인 압축률에서는 텍스트 요약이 RULER에서 Latent Compaction을 앞섰고, 64배 이상에서는 모든 압축 방법의 Phonebook Lookup 정확도가 0에 수렴했다. LongBench 점수도 압축이 강해질수록 계속 떨어져 비압축 모델과의 격차가 남는다. 향후 과제로 준비 비용 절감과 더 넓은 과제로의 확장을 제시한다. 코드는 공개돼 있다.