Epiphany-Aware KV Cache Eviction Without the Attention Matrix

기술 요약: 어텐션 행렬을 배제한 에피파니 인지형 KV 캐시 제거 (EPIKV)

문제 정의

추론 모델(예: DeepSeek-R1)이 수만 개의 토큰에 달하는 사고의 사슬(Chain of Thought)을 생성함에 따라, Key-Value (KV) 캐시가 선형적으로 증가하여 배포 시 주요한 메모리 병목 현상이 되고 있습니다. 현재의 KV 캐시 제거 전략은 일반적으로 **어텐션 가중치(attention weights)**를 기준으로 토큰의 순위를 매깁니다. 이 접근 방식은 두 가지 결정적인 한계를 가집니다:

1. 노이즈가 섞인 중요도 지표: 긴 추론 과정에서 어텐션 가중치는 토큰의 중요도를 나타내는 신뢰할 수 없는 지표입니다. "어텐션 싱크(attention sinks)"는 내용과 상관없이 가중치를 흡수하며, 필러(filler) 토큰들은 생성 중에 가중치를 끌어당기지만 나중에 참조되는 경우는 드뭅니다.
2. 구조적 불일치: 어텐션 가중치를 읽으려면 $n \times n$ 어텐션 행렬을 실체화(materialize)해야 합니다. 이는 "이거(eager)" 커널 사용을 강제하여, FlashAttention과 같은 퓨즈드(fused) 커널의 사용을 불가능하게 만듭니다. 결과적으로 메모리 사용량이 제곱($O(n^2)$)으로 증가하여, 일반적인 추론 궤적보다 훨씬 짧은 컨텍스트 길이(예: 8192 토큰)에서도 메모리 부족(OOM) 오류를 유발합니다.

이를 완화하려는 기존 방법들(예: ThinKV, RaaS)은 여전히 어텐션 유래 신호에 의존하며, 다른 방법들(예: LagKV)은 일반적이기는 하지만 추론 특화 최적화가 결여되어 있습니다.

방법론: EPIKV

저자들은 어텐션 가중치 대신 모델의 내부 표현(internal representation) 변화를 기반으로 토큰 점수를 매기는 훈련 불필요(training-free) 제거 정책인 EPIKV(Epiphany-Aware KV Cache)를 제안합니다.

1. 에피파니 점수 (The Epiphany Score)

EPIKV는 어텐션 대신, 연속적인 디코드 단계 사이의 잔차 스트림(residual stream, 은닉 상태) 변화를 측정합니다. 레이어 $l$에서의 토큰 $t$에 대한 신호는 다음과 같이 L2 차이로 정의됩니다:
$$g_l(t) = \|h_l(t) - h_l(t-1)\|_2$$
큰 변화는 "에피파니 토큰(epiphany tokens)"을 나타내며, 이는 모델이 탐색적 추론에서 수렴적 추론으로 전환하거나, 단계를 결론짓거나, 통찰을 확정 짓는 전환점을 의미합니다.

2. 2-밴드 레이어 해부 (Two-Band Layer Anatomy)

저자들은 반사실적 차폐 실험(counterfactual occlusion experiments)을 통해 32개 레이어 구조를 가진 추론 모델에서 특정 레이어 구조를 식별했습니다:

• 밴드 A (레이어 7–13): 이곳에서의 높은 은닉 상태 변화는 토론의 중요도와 양(+)의 상관관계를 가집니다. 이 레이어들은 알려진 사실적 검색 및 피처 라우팅(feature routing) 위치와 일치합니다.
• 밴드 B (레이어 18–25): 이곳에서의 높은 은닉 상태 변화는 중요도와 음(-)의 상관관계를 가집니다. 이 상위-중간 레이어들은 유창하고 놀라움이 적은 지속 과정에서도 활성화되므로, 이곳에서의 큰 변화는 예측 가능하고 폐기 가능한 토큰을 의미합니다.

최종 점수는 다음을 결합합니다:
$$s(t) = \bar{g}_{10}(t) - \bar{g}_{21}(t)$$
여기서 $\bar{g}$는 인과적 윈도우(causal window)에 대한 이동 평균을 나타냅니다.

3. 시간적 추세 교정 (Temporal Trend Correction)

원시 신호는 단조로운 위치적 추세(positional trend, 콘텐츠보다는 위치를 추적함)를 보입니다. 이를 교정하기 위해 저자들은 **인과적 이동 z-점수(causal rolling z-score)**를 적용합니다:
$$z_l(t) = \frac{g_l(t) - \mu_l(t)}{\sigma_l(t) + \epsilon}$$
최종 제거 점수는 $z_{10}(t) - z_{21}(t)$입니다. 이 정규화는 위치 편향을 제거하여, 점수가 절대적 크기가 아닌 로컬 콘텐츠 편차를 반영하도록 보장합니다.

4. FlashAttention 호환성

결정적으로, EPIKV는 다음만을 사용하여 점수를 계산합니다:

1. 이미 고대역폭 메모리(HBM)에 캐싱된 키(keys)와 값(values).
2. 표준 포워드 패스(forward pass)에 의해 노출되는 레이어별 은닉 상태(hidden states).
   EPIKV는 output_attentions=True를 요청하지 않으며 어텐션 행렬을 절대 실체화하지 않습니다. 따라서 EPIKV는 어텐션 기반의 "이거(eager)" 커널으로의 폴백(fallback) 없이 표준 FlashAttention 추론 스택 내에서 실행될 수 있습니다.

주요 기여

1. 2-밴드 해부학 식별: 은닉 상태 변화가 중간 레이어(7–13)에서는 중요도와 양의 상관관계를 갖는 반면, 상위-중간 레이어(18–25)에서는 음의 상관관계를 갖는다는 것을 발견하였으며, 이는 어텐션 기반 신호보다 우수한 성능을 보였습니다.
2. 위치 불가지론적 점수 산정: 인과적 이동 z-점수를 통해 원시 은닉 상태 분산에 내재된 위치적 추세를 제거함으로써, 고품질의 제거 성능을 회복함을 입증했습니다.
3. 어텐션 없는 성능 동등성: EPIKV는 어텐션 행렬을 실체화하지 않고도 MATH-500 및 AIME-2024에서 가장 강력한 어텐션 기반 베이스라인(ThinKV, H2O)과 대등하거나 이를 능가하는 성능을 보여줍니다.
4. 엔지니어링 효율성: 동일한 예산에서 EPIKV는 어텐션 기반 제거보다 최대 2.8배 빠르며, 어텐션 기반 점수 산정이 초래하는 메모리 벽(memory wall) 문제를 피합니다.

결과

실험은 DeepSeek-R1-Distill-LLaMA-8B를 사용하여 MATH-500 및 AIME-2024 데이터셋에서 수행되었습니다.

• 정확도:
  • MATH-500에서 4096-토큰 캐시 사용 시, EPIKV는 **72%**의 정확도를 달성하여 ThinKV(71%)와 대등하며 H2O(67%)를 능가했습니다.
  • AIME-2024에서 8192-토큰 캐시 사용 시, 지연 정규화된(lag-normalized) KV 변형 EPIKV는 **37%**에 도달하여 가장 우수한 어텐션 기반 베이스라인(33%)을 앞질렀습니다.
• 속도:
  • AIME-2024 (8192 토큰)에서, EPIKV는 문제를 처리하는 데 440초가 소요되었으며, 이는 가장 빠른 어텐션 베이스라인인 ThinKV(721초)와 RaaS(1239초)에 비해 1.6배에서 2.8배 빠른 속도를 나타냅니다.
• 메모리 확장성:
  • 어텐션 기반 점수 산정 방식은 $O(n^2)$ 어텐션 행렬으로 인해 약 8192 토근에서 80GB A100 GPU를 고갈시키지만, EPIKV(FlashAttention 활용)는 48GB 메모리 사용량만으로 65,536 토큰까지 확장 가능합니다 (실행 가능한 컨텍스트 길이가 16배 증가).
• 베이스라인의 실패 모드: H2O와 같은 어텐션 기반 방법은 엄격한 예산(예: 1024 토큰) 하에서 점진적으로 성능이 저하되는 대신, 빈 출력을 생성하거나 구조화되지 않은 텍스트를 생성하며 붕괴됩니다.

의의 및 주장

본 논문은 EPIKV가 표준 FlashAttention 서빙 스택에서 KV 캐시 제거를 배포 가능하게 만든다고 주장합니다. 어텐션 행렬으로부터 중요도 점수 산정을 분리함으로써, 이 방법은 다음과 같은 효과를 제공합니다:

• 어텐션 가중치를 실체화하는 구조적 비용을 제거합니다.
• 어텐션 가중치가 노이즈가 많은 추론 궤적에서 더 견고한 신호를 제공합니다.
• 기존 하드웨어에서 훨씬 더 긴 컨텍스트 윈도우와 높은 배치 크기를 가능하게 합니다.

저자들은 2-밴드 해부학이 경시대회 수학 문제에 특화되어 있으며 문제 난이도에 따라 변화한다(예: GSM8K의 경우 더 이른 레이어로 이동)는 점을 명시하며, 이는 신호가 고정된 레이어 인덱스가 아니라 추론의 공고화(consolidation)를 추적함을 시사한다고 밝혔습니다. 또한 향후 연구를 위한 검증 리소스로서 반사실적 중요도 라벨을 공개합니다.