기술 요약: Still: 단일 순전파를 통한 분할 상환형 KV 캐시 압축 (Amortized KV Cache Compaction in a Single Forward Pass)

1. 문제 정의

KV(Key-Value) 캐시는 수일간 지속되는 코딩 에이전트나 레포지토리 규모의 추론과 같은 장기적 관점(long horizons)에서 대규모 언어 모델(LLM)을 배포할 때 발생하는 주요 메모리 병목 현상입니다. 현재의 배포 전략은 다음과 같은 "전부 아니면 전무(all-or-nothing)" 식의 선택에 직면해 있습니다:

• 무손실 캐싱(Lossless Caching): 캐시가 컨텍스트 길이에 따라 선형적으로 증가하여 결국 메모리 한계에 도달합니다.
• 손실성 대안(Lossy Alternatives): 파인튜닝, 검색 증강 생성(RAG), 또는 문서 요약과 같은 방법들은 모델의 내부 표현 구조와 충실도(fidelity)를 포기합니다.

기존의 KV 캐시 압축 방법은 두 가지 범주로 나뉘며, 각각 중대한 한계를 가집니다:

1. 선택 방법(Selection Methods): 토큰에 점수를 매기고 일부를 제거합니다 (예: H2O, SnapKV). 이들은 경량화되어 있지만, 원래 토큰 서브셋에 구조적으로 묶여 있어 새로운 정보를 합성할 수 없으며, 엄격한 압축 예산 하에서는 성능이 급격히 저하됩니다.
2. 합성 방법(Synthesis Methods): 콘텐츠를 결합하여 새로운 압축 표현을 생성합니다 (예: Attention Matching, Cartridges). 이들은 표현력이 뛰어나지만, 일반적으로 추론 시점에 컨텍스트별 최적화(fitting)를 요구하므로, 압축이 반복적으로 수행되어야 하는 반복적이고 장기적인 궤적(iterative, long-horizon trajectories)에는 계산 비용이 많이 들어 비현실적입니다.

본 논문은 분할 상환형(amortized)(한 번 학습하여 반복 적용 가능하며, 컨텍스트별 최적화가 필요 없음), 합성형(synthetic)(단순 선택을 넘어 정보를 혼합할 수 있음), 그리고 추론 중에 실행할 수 있을 만큼 **경량(lightweight)**인 방법을 위한 간극을 식별했습니다.

2. 방법론: Still

저자들은 동결된 사전 학습 모델(frozen pre-trained models)을 위해 설계된 학습된 프리미티브(learned primitive)인 Still을 소개합니다.

아키텍처

Still은 각 레이어마다 별도로 존재하는 작은 Perceiver 모듈로 구성되며, 이는 동결된 베이스 모델을 대상으로 학습됩니다.

• 입력: 각 트랜스포머 레이어에 대해, 컴팩터는 프리픽스(prefix)의 전체 레이어별 KV 캐시($K, V$)를 받습니다. 키(keys)는 먼저 회전 제거(un-rotated)되어 RoPE(Rotary Position Embeddings)를 제거함으로써 위치 정보가 없는 프레임에 놓입니다.
• 잠재 쿼리(Latent Queries): 컴팩터는 레이어당 $t$개의 학습된 잠재 쿼리($Z$) 뱅크를 유지합니다. 이 쿼리들은 전체 KV 캐시에 크로스 어텐션(cross-attend)합니다.
• 프로세싱: 잠재 쿼리는 다음을 포함하는 pre-norm 블록을 통과합니다:
  1. 크로스 어텐션(Cross-Attention): 잠재 쿼리에서 전체 KV 캐시로의 어텐션.
  2. 셀프 어텐션(Self-Attention): 잠재 쿼리 간의 어텐션.
  3. 피드 포워드(FFN): 선택적 서브레이어.
• 출력: 최종 잠재 상태는 선형 헤드를 통해 크기가 $t \times d$인 압축된 키($C_k$)와 값($C_v$)으로 투영됩니다 (여기서 $t \ll T$, $T$는 원래 시퀀스 길이).
• 위치 처리(Position Handling): 생성된 압축 키는 캐시에 다시 쓰여지기 전에 선택된 출력 위치에서 다시 회전(re-rotated)됩니다. 값(Values)은 회전되지 않은 상태를 유지합니다.
• 바이어스 채널(Bias Channel): 선택적 바이어스 헤드($\beta$)가 지원되지만, 배포를 단순화하기 위해 모든 보고된 결과에서는 비활성화($\beta=0$)되었습니다.

학습

• 목적 함수(Objective): 컴팩터는 "교사(teacher)"(전체 압축되지 않은 캐시를 가진 동결된 베이스 모델)와 "학생(student)"(압축된 캐시를 사용하는 동일한 모델) 사이의 forward KL divergence를 사용하여 학습됩니다. 손실(loss)은 정답 측 토큰(answer-side tokens)에 대해서만 마스킹됩니다.
• 데이터: 학습에는 4개 도메인(금융, 프로젝트 구텐베르크, 법률, 코드)의 추출형 다중 선택 QA 데이터셋을 사용합니다.
• 효율성: 베이스 모델의 가중치는 동결된 상태로 유지되며, 컴팩터 파라미터만 학습됩니다. 컴팩터는 베이스 모델 파라미터 수의 약 1%를 추가합니다 (예: Qwen3-4B의 경우 약 50M 파라미터).

운영 모드

1. 단일 패스 압축(Single-Pass Compaction): 모델이 전체 프리필(prefill)을 실행하여 전체 캐시를 생성한 후, Still에 의해 압축됩니다. 이후 생성은 압축된 캐시를 사용하여 계속됩니다.
2. 반복적 청크 압축(Iterative Chunked Compaction): 시간이 지남에 따라 도착하는 장기 컨텍스트를 위해, Still은 순환 스케줄(recurrent schedule)을 지원합니다. 모델은 토큰의 청크를 처리하며, 각 청크가 끝날 때마다 누적된 캐시를 고정된 크기로 압축합니다. 이를 통해 시스템은 컴팩터가 충분히 긴 반복적 호라이즌에 대해 학습되었다면, 물리적 용량을 훨씬 초과하는 컨텍스트를 처리할 수 있습니다.

3. 핵심 기여

• 분할 상환형 합성(Amortized Synthesis): Still은 분할 상환형 합성(한 번의 인코더 학습)이 분할 상환형 선택(한 번의 스코어러 학습)보다 장기 컨텍스트 압축에 우수함을 입증합니다. 이는 선택 방식의 "서브셋 제한(subset-bound)" 한계를 극복하고, 컨텍스트별 최적화 방식의 "컨텍스트별 최적화 비용" 문제를 피합니다.
• 단일 순전파(Single Forward Pass): 매번 새로운 컨텍스트에 대해 최적화 단계가 필요한 컨텍스트별 합성 방법과 달리, Still은 학습된 모듈을 통한 단일 순전파를 통해 압축을 수행합니다.
• 반복적 장기 호라이즌 능력(Iterative Long-Horizon Capability): 압축이 순전파 과정이기 때문에, Still은 새로운 컨텍스트가 도착함에 따라 반복적으로 적용될 수 있으며, 이는 컨텍스트별 피팅 방식으로는 불가능한 장기적 관점의 운용을 가능하게 합니다.
• 일반화(Generalization): 이 방법은 모델 규모(4B에서 32B), 아키텍처(dense 및 MoE), 그리고 어텐션 메커니즘(global 및 sliding-window) 전반에 걸쳐 전이됩니다.

4. 결과

논문은 8배에서 200배의 압축 비율과 8k에서 128k의 컨텍스트 길이에 대해 Still을 평가합니다.

• 속도-품질 프런티어(Speed-Quality Frontier): Qwen 및 Gemma 모델에서 Still은 유리한 측면의 속도-품질 프런티어를 점유합니다. 선택 기반 방법들이 타이트한 예산 하에서 성능이 저하되는 반면, Still은 전체 행렬에 걸쳐 정확하고 빠르게 작동하며, 컨텍스트별 합성 방법들은 너무 느립니다.
• RULER 벤치마크: 매칭된 학습 설정에서 RULR 벤치마크를 수행했을 때, Still은 18개 셀 중 16개에서 가장 강력한 분할 상환형 선택 베이스라인인 KV-Distill보다 8~22 포인트 더 높은 정확도를 보였습니다.
• 전이 성능(Transfer Performance): 혼합 도메인 MCQ로 학습하고 RULER로 테스트하는 제로샷 전이 설정에서, Still은 학습-테스트 불일치에도 불구하고 KV-Distill에 비해 상당한 우위를 유지합니다.
• 요약(Summarization): Still은 자유 형식 요약을 지원합니다. HELMET (multi_lexsum)에서 Still은 8k–64k 컨텍스트에 대해 전체 컨텍스트 이득의 74–95%를 보존하며, Attention Matching 및 KV-Distill보다 뛰어난 성능을 보입니다. LongBench v1에서 KV-Distill과의 쌍별 비교 결과 60%의 승률을 기록했습니다.
• 반복적 압축(Iterative Compaction): Still은 반복적 청크 압축을 가능하게 합니다. 만약 학습 호라이즌이 배포 호라이즌에 비해 너무 짧으면 성능이 저하되지만 (예: 8k로 학습된 컴팩터가 128k에서 붕괴됨), 더 긴 반복적 호라이즌(16k, 32k)에서 학습하면 모델이 물리적 용량보다 훨씬 많은 컨텍스트를 흡수할 수 있습니다.

5. 의의 및 주장

본 논문은 분할 상환(amortization)이 장기 컨텍스트 캐시 압축을 실행 가능하게 만들며, 합성(synthesis)이 압축된 상태를 극한의 압축에서도 유용하게 만든다고 주장합니다.

• 실행 가능성(Tractability): 최적화 도구(컴팩터)를 한 번 학습함으로써, 시스템은 컨텍스트별 방식의 누적되는 피팅 비용을 피할 수 있으며, 이는 반복적인 장기적 배포를 가능하게 합니다.
• 표현력(Expressivity): 합성은 압축된 캐시가 임의의 많은 위치로부터 정보를 혼합할 수 있게 하여, 토큰 선택의 정보 밀도 한계를 극복합니다.
• 실용성(Practicality): Still은 베이스 모델의 가중치를 수정하지 않고도 정답 선택 및 자유 형식 생성을 지원할 만큼 가볍고, 추론 중에 반복적으로 호출될 수 있습니다.

저자가 인정한 한계점

• 무손실이 아님: 가장 어려운 128k 설정에서는 여전히 전체 컨텍스트 추론이 더 우수합니다.
• 학습 호라이즌 제약: 반복적 압축은 자유로운 외삽(extrapolation)이 아닙니다. 컴팩터의 성능는 학습 호라이즌에 의해 제한됩니다. 1M 토큰을 지원하려면 그에 비례하여 긴 호라이즌에서 학습해야 하며, 이는 비용이 많이 듭니다.
• 정확한 검색(Exact Retrieval): 현재의 순환 방식은 광범적 MCQ 작업에 대한 의미론적 증거를 보존하더라도, 정확한 검색 및 "바늘 찾기(needle-in-a-haystack)" 작업에는 취약합니다.
• 메모리 vs 지연 시간(Latency): 단일 패스 모드에서 Still은 유지되는 캐시 크기는 줄이지만, 초기 전체 컨텍스트 프리필 지연 시간은 제거하지 못합니다.
• 배포 요구사항: 각 베이스 체크포인트는 고유하게 학습된 컴팩터를 필요로 하며, 이는 파라미터 및 학습 데이터 오버헤드를 추가합니다.