KV 캐시 가지치기

KV 캐시 가지치기(KV Cache Pruning)는 트랜스포머 기반 언어 모델의 추론 단계에서 사용하는 Key-Value 캐시의 크기를 줄이기 위해 일부 K/V 쌍을 제거하거나 압축하는 기법이다.

개념 및 배경

• 트랜스포머 기반의 자기 회귀 모델에서는 이전 토큰의 Key와 Value를 캐시(KV 캐시)해 두고, 이후 토큰 생성 시 재사용한다.
• 이 방식은 중복 계산을 피하게 해주지만, 문맥 길이가 길어질수록 캐시 크기가 선형적으로 증가하여 메모리 부담과 속도 저하를 초래한다.
• 따라서 필요한 K/V만 남기고 덜 중요한 부분을 제거하거나 압축하는 방식이 KV 캐시 가지치기이다.

주요 전략 및 기법

다양한 방법이 제안되고 있으며, 대표적인 접근들을 소개한다:

• 채널 수준 가지치기 (Channel Pruning)
  • KV 캐시의 채널(특성 차원) 단위 중요도를 평가하여 덜 중요한 채널을 제거
  • 예: ThinK — 쿼리 중심 기준으로 채널을 선택적으로 제거함으로써 메모리 절감 효과 달성

• 구조적 가지치기 (Structural Pruning)
  • KVPruner 등은 특정 블록 단위로 K/V 채널을 제거하는 방식으로, 구조적 압축을 수행함
  • 일부 레이어에서만 K/V 투영(projection)을 유지하고 나머지는 상위 레이어의 캐시를 공유하는 방식도 사용됨 (Cross‑Layer Attention 공유 방식)

• 비정형 희소화 (Unstructured Sparsity)
  • Mustafar와 같은 최신 기법은 무작위 희소화 방식으로 캐시 내 원소들을 제거하되, 복원 가능한 sparse 구조로 처리
  • 비정형 구조는 더 유연하지만 하드웨어 최적화가 어렵다는 단점이 있다.

• 동적 제거 / 적응 제거 (Dynamic / Adaptive Pruning)
  • 동적으로 중요하지 않은 토큰 또는 K/V를 제거하는 방식
  • 예: DMS (Dynamic Memory Sparsification) — 추론 시점에 캐시를 희소화하여 메모리 절감과 속도 개선을 동시에 노림
  • 또 다른 예: LazyLLM — 프리필(prefill) 단계에서 필요 없는 토큰의 K/V 계산을 지연 또는 생략

장점

• 메모리 사용량 감소 — 캐시 크기의 절반 이상을 절감하는 사례도 보고됨
• 추론 속도 향상 — 캐시 접근 비용과 메모리 대역폭 부담이 줄어듦
• 긴 문맥 처리 가능성 확대 — 동일한 메모리 예산으로 더 긴 문맥을 지원할 수 있음

한계 및 고려사항

• 중요한 정보 손실 위험 — 잘못된 제거는 성능 저하로 이어질 수 있다
• 복원 또는 재보정 과정 필요 — 일부 방식은 경미한 fine-tuning이 필요
• 하드웨어 친화성 — 구조적 방법은 최적화가 쉬우나 비정형 방법은 구현 복잡성 상승
• 민감도 레이어 간 격차 — 레이어마다 K/V 민감도 차이가 크므로 균형 잡힌 가지치기 기준이 필요

응용 및 사례

• 자연어 생성 모델에서 긴 문맥을 다루는 경우
• 비전-언어 모델의 인퍼런스 경량화
• 실시간 응답 시스템, 모바일/엣지 장치 등 메모리 제약 환경

같이 보기

• 모델 압축
• 신경망 가지치기
• 제로샷 프루닝
• 양자화 (인공지능)