기술 요약: CriticalKV

문제 정의

트랜스포머 아키텍처 기반의 대규모 언어 모델 (LLM) 은 자기회귀 생성 과정에서 Key-Value(KV) 캐시와 관련된 2 차 계산 및 저장 비용으로 인해 상당한 배포 과제를 안고 있습니다. 시퀀스 길이가 증가함에 따라 KV 캐시는 급격히 확장되어 높은 GPU 메모리 오버헤드와 증가된 I/O 지연을 초래합니다.

최근의 캐시 제거 방법들 (예: H2O, Scissorhands, SnapKV) 은 비중요한 항목을 제거함으로써 이를 완화하려 시도하지만, 경험적 휴리스틱에 크게 의존합니다. 구체적으로, 이러한 방법들은 더 높은 어텐션 가중치를 가진 항목이 가장 중요하다고 가정합니다. 저자들은 이러한 가정이 공식적인 근거가 부족하다고 주장하며 두 가지 핵심 질문을 제기합니다:

1. KV 캐시 항목의 중요성을 실제로 결정하는 기준은 무엇인가?
2. 중요 항목을 식별하는 데 어텐션 가중치만 의존하는 것으로 충분한가?

이 논문은 현재 방법론들이 값 (value) 상태와 사전 훈련된 파라미터 행렬의 중요한 역할을 간과하고 있어, 비최적의 제거 결정과 상당한 품질 저하 (압축 손실) 로 이어질 수 있다고 주장합니다.

방법론

이 논문은 출력 섭동 분석을 기반으로 중요한 KV 캐시 항목을 식별하는 공식적인 프레임워크인 CriticalKV를 소개합니다.

이론적 프레임워크

저자들은 전체 KV 캐시를 중요한 항목의 부분 집합으로 대체할 때 어텐션 출력의 섭동을 최소화하는 것을 중요한 항목 식별 문제로 정의합니다.

• 섭동 정의: 출력 섭동 $L$은 전체 캐시의 어텐션 출력 ($o$) 과 선택된 부분 집합 ($\hat{o}$) 사이의 $L_1$ 거리로 정의됩니다: $L = \|o - \hat{o}\|_1$.
• 상한 유도: 이론적 분석을 통해 저자들은 이 섭동에 대한 상한 $\theta$를 유도합니다. 중요하게도, 이 상한은 섭동이 어텐션 가중치 ($A$) 뿐만 아니라 투사된 값 상태 ($V W_O$) 에 의해서도 영향을 받음을 보여줍니다. 여기서 $W_O$는 출력 투사 행렬입니다.
• 핵심 통찰: 유도된 상한은 섭동을 최소화하려면 어텐션 가중치와 파라미터 행렬을 통해 투사된 값 상태의 크기를 모두 고려해야 함을 보여줍니다. 어텐션 가중치만 의존하는 것은 이론적으로 비최적인 "완화된" 상한을 최소화하는 것입니다.

알고리즘: 섭동 제약 선택

이 이론을 실용화하기 위해 저자들은 최악의 경우 섭동 유도 상한을 최소화하도록 설계된 2 단계 탐욕 선택 알고리즘 (알고리즘 1) 을 제안합니다:

1. 1 단계 (높은 어텐션 가중치): 캐시 예산의 일부 ($b' = b \times \alpha$, 여기서 $\alpha=0.5$) 를 할당하여 가장 높은 어텐션 가중치를 가진 항목을 선택합니다. 이 단계는 선택된 부분 집합의 누적 어텐션 가중치가 전체의 50% 를 초과하도록 하여 이론적 상한에 대한 핵심 가정을 충족시킵니다.
2. 2 단계 (섭동 제약): 남은 예산 ($b''$) 은 어텐션 가중치와 투사된 값 상태의 $L_1$ 노름의 곱 ($A_i \|V_{i,:}\|_1$) 을 함께 최대화하는 항목을 선택하는 데 사용됩니다. 이 단계는 유도된 섭동 상한의 최소화를 직접적으로 목표로 합니다.

이 알고리즘은 기존 캐시 제거 파이프라인 (예: SnapKV, AdaKV, HeadKV) 에 그들의 근본적인 누적 과정을 변경하지 않고 통합될 수 있는 플러그 앤 플레이 향상 기능으로 설계되었습니다.

주요 기여

1. 공식적 문제 정의: 이 논문은 경험적 휴리스틱을 넘어 출력 섭동의 관점에서 중요한 캐시 항목 선택 문제를 공식화한 최초의 연구입니다.
2. 이론적 통찰: 중요한 항목을 식별하는 데 어텐션 가중치만으로는 부족하며, 파라미터 행렬을 통해 투사된 값 상태가 동등하게 필수적임을 규명했습니다.
3. 새로운 알고리즘: 최악의 경우 출력 섭동을 최적화하는 섭동 제약 선택 알고리즘을 제안했습니다.
4. 보편적 향상: 이 방법은 기존 최첨단 (SOTA) 제거 기술에 대한 보편적 향상으로 기능하며, 무시할 수 있는 계산 오버헤드만 요구합니다.

실험 결과

저자들은 Ruler 및 LongBench 벤치마크의 29 개 데이터셋에서 세 가지 서로 다른 LLM(Llama-3.1-8B, Mistral-7B-v0.3, Qwen2.5-32B) 에 대해 세 가지 최첨단 제거 방법 (SnapKV, AdaKV, HeadKV) 과 CriticalKV 를 통합하여 평가했습니다.

• 성능 향상: 평균적으로 알고리즘은 모든 테스트 구성에서 압축 손실을 절반 이상 감소시켰습니다.
  • Ruler 벤치마크 (캐시 크기 40%) 에서 Qwen2.5-32B 에 HeadKV 를 통합한 경우 평균 손실이 13.7% 에서 3.4% 로 감소했습니다.
  • LongBench 에서 이 방법은 긴 의존성 도메인의 90 개 테스트 사례 중 88 개에서 성능을 향상시켰습니다.
• 견고성: 개선 효과는 다양한 캐시 크기 (2.5% 에서 100% 까지) 와 다른 모델 아키텍처 전반에 걸쳐 일관되었습니다.
• 효율성: 알고리즘은 무시할 수 있는 계산 오버헤드를 도입합니다. 32K 컨텍스트 길이에서 배치 크기 1 일 때 토큰 생성 시작까지의 시간 (TTFT) 은 약 0.06 초만 증가했으며, 디코딩 지연은 기존 제거 방법과 동일하게 유지되었습니다.
• 섭동 분석: 실증적 분석은 이 알고리즘이 Llama 모델의 어텐션 헤드의 92% 이상에서 출력 섭동을 감소시키고, 최종 레이어 은닉 상태의 섭동을 현저히 낮춘다는 것을 확인했습니다.

중요성과 주장

이 논문은 캐시 제거를 위한 새롭고 공식적으로 근거가 있는 관점을 제공한다고 주장합니다. 직관적인 어텐션 가중치 휴리스틱에서 출력 섭동의 이론적 최소화로 초점을 이동시킴으로써, CriticalKV 는 현재 제거 전략의 근본적인 한계를 해결합니다.

저자들은 그들의 작업이 다음과 같음을 강조합니다:

• 캐시 제거에서 값 상태가 중요한 이유에 대한 이론적 근거를 제공합니다.
• 메모리 제약 하에서 생성 품질을 실질적으로 향상시키는 실용적이고 낮은 오버헤드의 솔루션을 제공합니다.
• 출력 섭동에 대한 이론적 분석을 통해 중요한 캐시 선택을 최적화하는 미래 연구에 유망한 길을 엽니다.

코드는 공개적으로 이용 가능하며, 저자들은 후속 연구에서 그들의 접근 방식을 강화한 버전이 저장소에 호스팅된다고 note 합니다.