RippleKV: Cross-Layer KV Cache Allocation via Perturbation Propagation

기술 요약: RippleKV

문제 정의

대규모 언어 모델(LLM)의 긴 문맥 추론은 KV 캐시의 메모리 및 대역폭 요구 사항에 의해 제약을 받으며, 이는 문맥 길이와 트랜스포머 레이어 수에 따라 선형적으로 증가합니다. 기존의 KV 캐시 압축 방식들은 일부 토큰을 유지함으로써 메모리 사용량을 줄이지만, 일반적으로 모든 트랜스포머 레이어에 대해 균등하거나 사전 정의된 캐시 예산을 할당합니다.

최 recent의 레이어 인지(layer-aware) 접근 방식들은 레이어 깊이, 어텐션 통계 또는 표현 변화와 같은 프록시(proxies)를 기반으로 비균등한 예산을 할당하려고 시도합니다. 그러나 저자들은 이러한 프록시들이 고정된 글로벌 예산 하에서 서로 다른 레이어의 상대적인 캐시 요구 사항을 직접적으로 정량화하지 못한다고 주장합니다. 구체적으로, 이들은 특정 레이어에서의 섭동(perturbation)이 최종 모델 출력에 어떻게 전파되는지를 측정하지 않습니다. 결과적으로, 이러한 방식들은 "민감한" 레이어에는 자원을 과소 할당하여(심각한 성능 저하 유발) 성능을 떨어뜨리는 반면, "내성이 있는" 레이어에는 자원을 과잉 할당할 위험이 있습니다. 논문의 실증적 분석에 따르면, 압축 손상은 레이어 깊이에 따라 매우 이질적이며 비단조적(non-monotonic)으로 나타나며, 고정된 패턴을 따르기보다 종종 중간 레이어에서 정점을 찍는 것으로 밝혀졌습니다.

방법론: RippleKV

RippleKV는 구조적 사전 지식(structural priors) 대신 **다운스트림 출력 민감도(downstream output sensitivity)**를 기반으로 고정된 글로벌 캐시 예산을 분배하는 교차 레이어 KV 캐시 할당 프레임워크입니다. 이 방법은 오프라인 민감도 프로파일링과 추론 시 예산 할당의 두 단계로 작동합니다.

1. 오프라인 민감도 프로파일링 (Offline Sensitivity Profiling)

레이어의 국소적 통계를 측정하는 대신, RippleKV는 특정 레이어의 Value 캐시가 최종 예측 분포에 얼마나 기여하는지를 추정합니다.

• 섭동 주입 (Perturbation Injection): 소규모 캘리브레이션 세트에 대해, 이 방법은 각 압축 가능한 레이어 $\ell$을 반복 탐색합니다. 해당 특정 레이어의 Value 캐시($V_\ell$)에 노름 적응형(norm-adaptive) 섭동을 주입하는 동시에, 해당 레이어의 Key 캐시($K_\ell$)와 다른 모든 레이어의 캐시는 변경하지 않고 유지합니다.
• 반응 측정 (Response Measurement): 모델은 섭동이 가해진 입력에 대해 출력을 생성합니다. 영향력은 참조 출력 분포(전체 캐시)와 섭동된 출력 분포 사이의 쿨백-라이블러(Kullback-Leibler, KL) 발산(divergence)을 계산하여 측정됩니다.
• 민감도 프로필 (Sensitivity Profile): KL 발산은 캘리브레이션 세트 전체에 걸쳐 평균화되어 각 레이어에 대한 민감도 점수 $s_\ell$을 생성합니다. 이 점수는 레이어의 캐시를 수정하는 것이 최종 예측에 미치는 엔드 투 엔드(end-to-end) 효과를 반영합니다. 저자들은 이 신호가 실제 압축 손상과 강한 상관관계(평균 스피어만 $\rho \approx 0.799$)를 보임을 발견했으며, 이는 레이어 인덱스 상관관계($\rho \approx 0.359$)보다 훨씬 뛰어난 성능을 보였습니다.

2. 민감도 유도 예산 할당 (Sensitivity-Guided Budget Allocation)

민감도 프로필은 척도 불변(scale-invariant) 지수 매핑을 사용하여 레이어별 예산 승수로 변환됩니다:

• 정규화 (Normalization): 민감도 점수는 $[0, 1]$ 범위로 정규화됩니다.
• 지수 매핑 (Exponential Mapping): 비율 파라미터 $r$은 할당 격차를 조절합니다. 정규화된 점수는 승수 $w_\ell = r^{s_\ell}$로 매핑됩니다. $r$ 값이 높을수록 민감한 레이어와 내성이 있는 레이어 간의 예산 격차가 커집니다.
• 클리핑 및 정규화 (Clipping and Normalization): 승수는 제한된 구간 $[1/\gamma, \gamma]$ 내로 클리핑되며, 레이어 예산의 합이 글로벌 예산 $B$와 일치하도록 재조정됩니다.
• 추론 (Inference): 추론 중에 사전 계산된 예산이 적용됩니다. 기본 토큰 선택 전략(예: SnapKV)은 변경되지 않으며, 할당된 예산에 따라 레이어당 유지되는 토큰의 수만 조정됩니다.

주요 기여

1. 실증적 통찰: 본 논문은 레이어별 압축 손상이 깊이에 따라 비단조적이며, 제어된 Value 캐시 섭동에 대한 다운스트림 출력 민감도가 레이어 위치나 국소 통계보다 예산 할당을 위한 더 충실한 신호임을 입증합니다.
2. RippleKV 프레임워크: 섭동 전파를 기반으로 캐시 예산을 할당하는 방법을 소개합니다. Value 캐시 섭동이 최종 출력에 미치는 영향을 통해 레이어 민감도를 추정함으로써, 고정된 전역 제약 조건 하에서 모델별 예산을 도출합니다.
3. 성능 및 효율성: 광범위한 실험을 통해 RippleKV가 여러 모델 제품군과 캐시 예산에 대해 우수한 평균 성능을 달면서도 기존 압축 방식과 유사한 런타임 및 메모리 효율성을 유지함을 보여줍니다.

실험 결과

저자들은 세 가지 모델 제품군(Llama-3.1-8B-Instruct, Qwen2.5-7B-Instruct, Mistral-7B-Instruct-v0.3)에 대해 10%, 20%, 30%의 캐시 예산으로 LongBench에서 RippleKV를 평가했습니다.

• 성능: RippleKV는 평가된 모든 설정에서 가장 높은 평균 성능을 달성했습니다. 10% 캐시 예산의 Llama-3.1-8B-Instruct에서 35.07점을 기록하여, 가장 강력한 베이스라인인 PyramidKV보다 0.83점 앞섰습니다. 또한 예산이 증가함에 따라 Full Cache 베이스라인과의 격차를 지속적으로 좁혔습니다.
• 일반화: 이 방법은 다양한 아키텍처와 작업 카테고리(단일/다중 문서 QA, 요약, 퓨샷, 합성 데이터, 코드)에 걸쳐 잘 일반화되었습니다.
• 효율성: RippleKV는 추론 지연 시간(latency) 및 처리량(throughput) 측면에서 베이스라인과 대등하거나 약간 더 나은 성능을 보였습니다. 피크 메모리 사용량은 다른 압축 방식(예: StreamingLLM, SnapKV)과 유사하여, 용량 재분배가 총 메모리 오버헤드를 증가시키지 않음을 확인했습니다.
• 어블레이션 (Ablation): 레이어별 예산 할당을 제거하거나(균등 예산 사용), 다운스트림 민감도 신호를 레이어 국소 반응으로 대체했을 때 성능 저하가 발생하여, 비균등 할당 전략과 특정 민감도 지표의 유효성을 검증했습니다.

의의 및 주장

본 논문은 다운스트림 출력 민감도가 구조적 사전 지식이나 레이어 국소 통계보다 효과적인 교차 레이어 KV 캐시 할당 근거를 제공한다고 주장합니다. 섭동이 최종 출력으로 어떻게 전파되는지에 집중함으로써, RippleKV는 자원을 너무 민감하거나 너무 내성이 있는 레이어에 잘못 할당할 수 있는 고정된 깊이 기반 스케줄의 함정을 피합니다.

저자들은 RippleKV가 추론 중 추가적인 모델 평가를 필요로 하지 않는 플러그 앤 플레이(plug-and-play) 솔루션임을 강조합니다. 이는 민감도 프로필이 오프라인에서 한 번만 계산되기 때문입니다. 이 방법은 예산 할당을 토큰 스코어링으로부터 분리하여 기존 압축 전략과 통합될 수 있도록 합니다. 이 연구는 캐시 섭동의 전파를 이해하는 것이 메모리 제약 조건 하에서 긴 문맥 LLM 추론을 최적화하는 데 매우 중요하다는 점을 시사합니다.