기술 요약: 어텐션 매칭을 통한 빠른 KV 압축

1. 문제 제기

장문맥을 위한 대규모 언어 모델 (LLM) 의 확장성은 현재 키 - 값 (KV) 캐시의 메모리 요구 사항에 의해 병목 현상을 겪고 있습니다. 자기회귀 (autoregressive) 트랜스포머 아키텍처에서 KV 캐시는 모든 이전 토큰에 대한 키와 값을 유지해야 하므로, 장시간 설정 (예: 멀티 세션 대화, 에이전트 코딩) 에서는 요청당 기가바이트 단위의 메모리에 도달하는 경우가 많습니다.

KV 캐시 크기를 줄이기 위한 기존 접근법은 일반적으로 두 가지 범주로 나뉘며, 둘 다 상당한 한계를 가지고 있습니다:

• 토큰 공간 방법: 토큰 제거 (H2O, SnapKV), 토큰 병합 (KVMerger), 그리고 요약과 같은 기술들은 높은 압축 비율에서 급격히 성능이 저하됩니다. 특히 요약은 손실이 매우 크며, 정밀한 정보 추출이 필요한 작업의 하류 성능을 해칩니다.
• 잠재 공간 방법: Cartridges(Eyuboglu 등, 2025) 와 같은 최근 연구는 완전한 문맥 성능을 맞추기 위해 잠재 공간에서 컴팩트한 KV 캐시를 학습할 수 있음을 보여주었습니다. 그러나 이 접근법은 모든 문맥에 대해 엔드 - 투 - 엔드 경사 기반 최적화 (프refix-튜닝) 에 의존하므로, 문맥당 여러 GPU 시간을 요구하는 등 prohibitively 느리고 비용이 많이 듭니다.

열려 있는 문제는 문맥당 경사 최적화의 계산 비용을 들이지 않고도 하류 모델 동작을 보존하는 잠재 공간에서의 빠르고 고품질의 문맥 압축을 달성하는 것입니다.

2. 방법론: 어텐션 매칭

저자들은 참조 쿼리 집합에 대해 원본 캐시 $(K, V)$의 어텐션 출력과 어텐션 질량을 재현하는 컴팩트한 KV 캐시 $(C_k, C_v)$를 구성하는 어텐션 매칭 기반의 방법론 군을 제안합니다.

핵심 목표

목표는 임의의 쿼리 $q$와 임의의 추가된 고정/미래 토큰 $(K_{fixed}, V_{fixed})$에 대해 어텐션 출력이 대략적으로 보존되도록 하는, 크기 $t \ll T$인 컴팩트한 캐시 $(C_k, C_v)$를 찾는 것입니다. 저자들은 연결된 블록에 대한 어텐션의 분해를 활용하여, 참조 쿼리 집합 $Q_{ref}$에 대해 다음 두 가지 양을 매칭하는 것으로 충분함을 보여줍니다:

1. 로컬 어텐션 출력: 값들의 가중 합.
2. 어텐션 질량: 정규화되지 않은 어텐션 로짓의 합 (소프트맥스의 분모).

컴팩트한 캐시는 토큰 수가 적어 ($t < T$) 총 어텐션 질량을 자연스럽게 매칭할 수 없는 경우 (예: $q=0$일 때, 질량은 $T$ 대 $t$) 의 불일치를 처리하기 위해 이 방법은 토큰별 스칼라 편향 (bias) $\beta$를 도입합니다. 최적화 목표는 다음과 같습니다:
$$\frac{\exp(qK^\top)V}{\sum \exp(qK^\top)} \approx \frac{\exp(qC_k^\top + \beta)C_v}{\sum \exp(qC_k^\top + \beta)}$$
$$\sum \exp(qK^\top) \approx \sum \exp(qC_k^\top + \beta)$$

알고리즘 파이프라인

이 프레임워크는 경사 하강법을 압축 시점에 사용하지 않고, 문제를 효율적인 폐쇄형 (closed-form) 하위 문제로 분해합니다:

1. 참조 쿼리 생성 ($Q_{ref}$):

   • 반복 프리필 (Repeat-prefill): 문맥 자체 (또는 반복된 버전) 에 프리필 패스를 실행하여 쿼리 활성화 값을 추출합니다.
   • 자기 학습 (Self-study): 문맥에 조건부인 합성 상호작용을 생성하여 쿼리 분포를 확장합니다.
   • 온-폴리시 (On-policy): 레이어를 순차적으로 압축하고 부분적으로 압축된 모델에서 쿼리를 추출하여 분포 변화를 줄입니다.

2. 키 선택 ($C_k$):

   • 최고 어텐션 키: 참조 쿼리 전반에 걸쳐 가장 높은 어텐션 점수를 받은 원본 키를 선택합니다 (RMS 를 통해 집계).
   • 직교 매칭 추적 (OMP): 어텐션 질량 매칭 목표를 가장 잘 충족하도록 키를 탐욕적으로 선택한 후, 가중치를 적합시키기 위해 비음수 최소제곱법 (NNLS) 을 적용합니다.

3. 편향 및 값 적합 ($\beta, C_v$):

   • $\beta$ 적합: 원본 키의 어텐션 질량을 매칭하기 위해 비음수 최소제곱법 (NNLS) 을 통해 해결합니다.
   • $C_v$ 적합: 선택된 키와 적합된 편향을 기반으로 어텐션 출력을 매칭하기 위해 일반 최소제곱법 (OLS) 을 통해 해결합니다.

비균일 압축

저자들은 서로 다른 어텐션 헤드가 압축에 대해 다른 민감도를 보인다고 관찰합니다. 이에 따라 서로 다른 헤드와 레이어에 서로 다른 유지 비율을 할당하는 비균일 압축 스케줄을 제안합니다. 이 스케줄은 헤드별 민감도 곡선에 기반한 탐욕적 교환 알고리즘을 사용하여 모델당 한 번 사전 계산되므로, 다양한 문맥에 걸쳐 견고하게 적용할 수 있습니다.

청크별 압축

매우 긴 문맥 (예: 60k 개 이상의 토큰) 을 지원하기 위해, 이 방법은 입력의 연속된 청크에 대해 독립적으로 압축을 적용합니다. 전역적 위치 의미와 청크 간 상호작용을 보존하기 위해 기본적으로 KV 기반 청킹 접근법 (전체 문맥 프리필, KV 상태 슬라이싱, 압축, 병합) 을 사용합니다.

3. 주요 기여

• 어텐션 매칭 공식화: 출력 가능도 대신 어텐션 동작 보존을 직접 최적화하는 잠재 공간 압축을 위한 새로운 목적 함수로, 효율적인 선형 대수 문제로 자연스럽게 분해됩니다.
• 폐쇄형 솔루션: 압축 시 엔드 - 투 - 엔드 경사 기반 최적화의 필요성을 제거하고, 최소제곱법 및 NNLS 솔버로 대체합니다.
• 속도 - 품질 파레토 프론티어: 이 접근법은 단일 GPU 에서 초 단위로 최대 50 배의 압축 비율을 달성하며, Cartridges보다 수백 배 빠르면서도 동등하거나 우수한 하류 정확도를 유지합니다.
• 비균일 스케줄링: 압축에 가장 민감한 어텐션 헤드를 우선시하는 모델별 사전 계산된 예산 할당 전략입니다.
• 편향 메커니즘: 어텐션 질량 불일치를 보정하기 위한 토큰별 스칼라 편향을 도입하여, 줄어든 키 집합으로 정확한 질량 매칭을 가능하게 합니다.

4. 결과

이 방법론들은 QuALITY(장문 문서 QA), LongHealth(밀집 임상 기록), QASPER(생성형 QA), 그리고 LongBench v2에서 평가되었습니다.

• 성능 대 속도: QuALITY 에서 50 배 압축 비율로 어텐션 매칭 방법들은 파레토 프론티어를 따라가며, 토큰 제거 기반선 (H2O+, KVzip) 과 요약을 능가합니다. Cartridges의 정확도와 일치하거나 초과하면서도 두 자릿수 (orders of magnitude) 만큼 더 빠릅니다(수시간 대 수분).
• 고압축: 극단적인 비율 (최대 100 배) 에서도 이 방법들은 상당한 하류 정확도를 유지합니다 (예: LongHealth 에서 약 60% 대 요약 기반선의 약 30%).
• 일반화: 이 접근법은 생성형 작업 (QASPER) 과 검색 집약적 벤치마크 (RULER) 로 일반화되며, 가장 빠른 변형 (AM-HighestAttnKeys-fast) 은 RULER 에서 최신 KVzip 을 능가했습니다.
• 효율성: Gemma-3-12B 에서 60k 토큰 문맥의 압축은 키 선택 전략 (OMP 대 최고 어텐션) 에 따라 약 100~600 초가 소요되며, 이는 쿼리 생성 및 키 선택 단계에 의해 지배됩니다.
• 온라인 압축: AIME 에 대한 예비 결과는 중간 궤적에서의 반복적인 압축이 추론 성능을 보존하여 추론 깊이를 물리적 메모리 한계와 분리할 수 있음을 보여줍니다.

5. 중요성 및 주장

이 논문은 어텐션 매칭이 빠르고 손실이 큰 토큰 공간 방법과 느리고 고품질인 잠재 공간 방법 간의 간극을 메우는 실용적이고 효율적인 KV 캐시 압축 원시 (primitive) 를 제공한다고 주장합니다.

• 실용성: 문맥별 경사 하강법을 피함으로써, 이 방법은 지연 시간과 계산 비용이 제약인 실제 배포에서 고충실도 압축을 가능하게 합니다.
• 시스템 통합: 이 접근법은 기존 추론 엔진 (예: FlashAttention, 가변 길이 KV 패킹) 과 호환되며, 임계 경로 밖에서 수행될 수 있습니다 (분리된 압축).
• 한계: 저자들은 이 방법이 Cartridges보다 훨씬 빠르지만, 단순한 휴리스틱보다는 여전히 느리다고 인정합니다. 또한, 극단적인 압축 비율 (예: LongHealth 에서 100 배) 에서는 Cartridges와 같은 경사 기반 방법이 원본 키로 제한되지 않는 더 넓은 표현 공간을 탐색할 수 있어 여전히 더 우월할 수 있습니다.
• 향후 방향: 이 논문은 쿼리 생성 최적화, 부분집합 선택을 넘어 컴팩트한 키의 직접 최적화로 나아가기, 그리고 온라인 장시간 에이전트 설정을 위해 추론 엔진에 압축을 통합하는 등의 향후 작업을 제안합니다.

요약하자면, 이 연구는 어텐션 동작을 직접 매칭함으로써 빠르고 고품질의 문맥 압축이 달성 가능함을 입증하여, 장문맥 LLM 의 메모리 병목 현상에 대한 확장 가능한 솔루션을 제시합니다.