🤖 AI

RAP: KV-Cache Compression via RoPE-Aligned Pruning

이 논문은 정렬된 쌍으로 키-값(Key-Value) 캐시 채널을 제거함으로써 회전 위치 임베딩(Rotary Position Embedding)의 의미론적 구조를 보존하는 구조적 프루닝 방법인 RoPE-Aligned Pruning(RAP)을 소개하며, 이를 통해 정확도를 희생하지 않으면서도 롱 컨텍스트 LLM 추론을 위한 상당한 메모리 및 연산 절감 효과를 달성한다.

원저자: Jihao Xin, Tian Lyu, David Keyes, Hatem Ltaief, Marco Canini

게시일 2026-08-04

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Jihao Xin, Tian Lyu, David Keyes, Hatem Ltaief, Marco Canini

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

매우 긴 이야기를 기억하려고 노력하고 있다고 상상해 보세요. 새로운 문장을 들을 때마다, 당신은 이전의 문장들이 어떻게 연결되는지 이해하기 위해 그 모든 문장의 목록을 머릿속에 계속 간직해야 합니다. 인공지능의 세계에서 이 정신적인 목록을 "KV 캐시(KV cache)"라고 부릅니다. 이것은 마치 똑똑한 로봇이 대화의 중요한 세부 사항들을 적어두어, 이야기의 끝에 도달했을 때 시작 부분을 잊어버리지 않도록 하는 거대한 공책와 같습니다. 문제는 이야기가 길어질수록 이 공책이 엄청나게 커진다는 점입니다. 공책을 넘기는 데 너무 많은 메모리가 사용되고 너무 많은 두뇌 에너지가 필요해서, 로봇은 속도가 느려지거나 심지어 공간이 완전히 부족해지기 시작합니다. 과학자들은 로봇이 충돌하지 않고 긴 주제에 대해 대화할 수 있도록, 중요한 부분은 놓치지 않으면서도 이 공책을 줄이는 방법을 항상 찾고 있습니다.

로봇이 이야기의 사건이 어디서 일어나는지 이해하기 위해 사용하는 한 가지 인기 있는 기술은 "회전 위치 임베딩(Rotary Position Embedding, RoPE)"이라고 불립니다. RoPE를 로봇이 기억과 함께 추는 특별한 춤이라고 생각해 보세요. 단순히 "문장 5"라는 숫자를 적는 대신, 로봇은 두 숫자를 한 쌍으로 묶어 훌라후프처럼 빙글빙글 돌립니다. 이 회전은 현재 문장에서 그 문장이 얼마나 떨어져 있는지를 로봇에게 정확히 알려줍니다. 이 춤의 결정적인 규칙은 이 숫자 쌍이 반드시 함께 있어야 한다는 것입니다. 만약 둘을 분리하면 회전이 깨지고, 로봇은 시간 순서에 대해 혼란을 느끼게 됩니다.

이제, RAP(RoPE-Aligned Pruning)이라는 새로운 방법이 등장했습니다. 이 아이디어를 개발한 연구자들은 기존의 로봇 메모리 공책을 줄이려는 시도들이 치명적인 실수를 저지르고 있다는 점을 알아차렸습니다. 당신의 배낭 안에 이 춤추는 쌍들이 가득 차 있다고 상상해 보세요. 기존의 방식들은 배낭에서 무작위로 단일 항목들을 버림으로써 공간을 절약하려고 했습니다. 하지만 항목들이 쌍으로 춤을 추고 있었기 때문에, 단 하나만 버려도 남은 파트너는 손을 잡을 상대 없이 허공에서 홀로 빙글빙글 돌게 됩니다. 춤이 깨져버리고, 로봇의 기억은 쓸모없게 됩니다.

이 논문의 저자들은 더 똑똑한 방식으로 가방을 싸는 법을 제안합니다. 무작위로 단일 항목을 던져버리는 대신, RAP는 춤추는 쌍들을 살펴보고 한 번에 전체 쌍을 통째로 버리기로 결정합니다. 만약 한 쌍이 별로 중요하지 않다면, 그 듀오 전체가 떠납니다. 만약 한 쌍이 중요하다면, 그 듀오 전체가 남습니다. 이렇게 하면 춤이 온전하게 유지됩니다. 연구진은 이 방법을 30억 개에서 140억 개의 "뇌 세포"를 가진 여러 인기 있는 AI 모델들에 테스트했으며, 이 방법이 매우 잘 작동한다는 것을 발견했습니다. 그들은 원래 크기의 70%만을 유지하면서 메모리 공책을 30% 줄일 수 있었으며, 높은 정확도를 유지했습니다. 비록 원래의 압축되지 않은 모델에 비해 약간의 성능 저하는 있었지만 말입니다.

이것이 왜 더 멋진 일인지 설명하자면, 메모리를 줄이려고 노력하는 다른 방법들은 로봇이 말할 때마다 누락된 조각들을 재구축하기 위해 추가적인 작업을 수행해야 하며, 이는 속도를 늦춥니다. 하지만 RAP는 로봇이 말을 시작하기 전에 공책을 편집하는 것과 같습니다. 그것은 여분의 페이지들을 영구적으로 제거하므로, 로봇은 빈틈을 채우기 위해 추가적인 수학 계산을 할 필요가 없습니다. 이는 더 빠르고, 에너지를 덜 사용하며, 여전히 이야기를 아주 잘 기억합니다. 논문은 다른 기술들이 공간을 절약할 수는 있지만, 종종 로봇의 시간 감각을 망가뜨리거나 속도를 늦춘다는 것을 보여줍니다. RAP는 메모리를 줄이고, 로봇의 속도를 높이면서도, 동시에 이야기의 정확성을 유지할 수 있는 최상위 방법 중 하나이며, 강력한 저계수(low-rank) 경쟁 모델들과 매우 유사한 성능을 보이면서도 독특한 효율성 이점을 제공합니다. 이것은 마치 이야기를 하기 위해 백과사전 전체를 들고 다닐 필요는 없다는 것을 깨닫는 것과 같습니다. 당신은 그저 적절한 장들을 가지고 있으면 되고, 페이지를 반으로 찢지 않으면 됩니다.

기술 요약: RAP (RoPE-정렬 프루닝)를 통한 KV-캐시 압축

1. 문제 정의

대규모 언어 모델(LLM)의 긴 컨텍스트 추론은 키-값(Key-Value, KV) 캐시의 메모리 및 연산 요구 사항으로 인해 점점 더 병목 현상을 겪고 있습니다. 컨텍스트 길이가 확장됨에 따라, KV 캐시는 GPU 메모리 사용량과 어텐션 연산 비용을 지배하게 됩니다. 구조적 프루닝(Structured pruning)은 $W_k$ 및 $W_v$ 프로젝션 가중치에서 덜 유용한 채널을 제거함으로써 캐시를 압축할 수 있는 직접적인 경로를 제공하지만, 기존 방식들은 **회전 위치 임베딩(Rotary Position Embedding, RoPE)**과의 치명적인 불일치 문제에 직면해 있습니다.

현대적인 LLM들은 Query/Key 프로젝션 이후에 RoPE를 적용하여 특징 차원(feature dimensions)을 쌍(pair) 단위로 회전시킵니다. 표준적인 개별 채널 프루닝 방식(예: KVPruner)은 이러한 쌍의 구조를 고려하지 않고 개별 채널을 제거하는 경우가 많습니다. 이러한 "고아(orphaning)"가 된 절반의 쌍은 RoPE의 회전 의미론(rotational semantics)을 깨뜨려, 모델 정확도의 침묵하면서도 파괴적인 붕괴를 초引합니다. 또한, 가중치를 압축하려고 시도하는 기존의 저차원 분해(low-rank decomposition) 방식들(예: SVD-LLM, PaLU)은 RoPE와 행렬 곱의 비가환성(non-commutativity) 때문에 그 확장 행렬을 Query 프로젝션( $W_q$ )과 융합할 수 없어, 디코딩 과정 중 매 단계마다 전체 차원의 Key 벡터를 재구성해야 하는 비용이 발생합니다.

2. 방법론: RoPE-정렬 프루닝 (RAP)

저자들은 RoPE에 의해 부과된 쌍 단위 구조를 특별히 준수하도록 설계된 구조적 프루닝 프레임워크인 RAP를 제안합니다.

핵심 원리

RAP는 개별 채널을 프루닝하는 대신, 프루닝의 입도를 **RoPE-정렬된 쌍(RoPE-aligned pairs)**으로 제한합니다. 차원의 전체 쌍을 제거함으로써, 이 방법은 회전 연산이 수학적으로 유효하게 유지되고 위치 의미론이 보존되도록 보장합니다.

RAP 파이프라인

RoPE-쌍 점수 산출 (RoPE-Pair Scoring):
- Key 프로젝션 가중치( $W_k$ )에 대해, 열(columns)은 모델의 특정 페어링 전략(인접 또는 반대편 분리)에 따라 RoPE 쌍으로 그룹화됩니다.
- 중요도는 소규모 캘리브레이션 데이터셋에서 계산된 피셔 정보(Fisher information) 점수를 사용하여 측정됩니다. 쌍 $p = (j, j')$ 에 대해, 점수 $\sigma_p$ 는 두 채널의 집계된 루트-피셔 크기(root-Fisher magnitude)입니다.
- RoPE 페어링이 없는 Value 프로젝션( $W_v$ )의 경우, 개별 채널 피셔 점수가 계산됩니다.
예산 할당 (Budget Allocation):
- 글로벌 유지 비율 $\rho$ 가 레이어와 헤드에 걸쳐 배분됩니다.
- 레이어 적응형 (Layer-Adaptive): 레이어는 민감도(집계된 피셔 점수)에 따라 예산을 할당받으며, 더 민감한 레이어에 더 큰 예산을 부여합니다.
- 헤드 균등형 (Head-Uniform): 레이어 내에서 모든 헤드는 배치 GEMM 연산을 용이하게 하기 위해 동일한 수의 유지되는 쌍을 받습니다.
- K/V 균형 (Balanced K/V): Key 차원이 압착되어 긴 컨텍스트 성능을 해치는 것을 방기하기 위해, RAP는 $K$ 와 $V$ 모두에 동일한 헤드당 채널 예산을 할당합니다. 구체적으로, $K$ 가 $m$ 개의 쌍을 유지한다면, $V$ 는 $2m$ 개의 채널을 유지합니다.
구축 및 융합 (Construction and Fusion):
- 유지된 쌍을 원래 인덱스로 매핑하기 위한 이진(binary)이며 쌍을 보존하는 선택 행렬 $E_k$ 가 구축됩니다.
- 결정적으로, $E_k$ 는 쌍 구조를 보존하기 때문에 RoPE 연산과 교환(commute) 가능합니다: $RoPE(X \tilde{W}_k E_k) = RoPE(X \tilde{W}_k) E_k$ .
- 이 교환 가능성 덕분에 프루닝 마스크를 오프라인에서 Query 프로젝션에 융합할 수 있습니다: $\tilde{W}_q = W_q E_k^\top$ .
- 결과: 모델은 프루닝된 차원에서 실행되며, KV 캐시는 저차원의 잠재 키(latent keys)만을 저장합니다. 온라인 재구성이 필요하지 않습니다.
정확도 회복 (Accuracy Recovery):
- 프루닝 중에 손실된 정확도를 회복하기 위해, RAP는 경량 LoRA (Low-Rank Adaptation) 업데이트를 **지식 증류 (Knowledge Distillation, KD)**와 결와 결합하여 사용합니다.
- 프루닝된 모델에 대해 작은 랭크( $r=8$ )의 LoRA 어댑터를 훈련시켜 밀집(dense) 티처 모델과 출력값을 정렬시킨 후, 배포 전에 병합합니다.

3. 주요 기여

RoPE 불일치 식별: 본 논문은 표준 채널 프루닝이 RoPE 쌍을 깨뜨려 정확도 붕괴를 일으키며, 이러한 실패가 런타임 에러 없이 성능 저하로만 나타나는 '침묵의 실패'임을 강조합니다.
RoPE-정렬 프루닝 전략: RAP는 RoPE 쌍을 엄격히 준수하는 프루닝 입도를 도입하여, 복잡한 재학습이나 재구성 없이도 회전 연산이 유효하게 유지되도록 합니다.
재구성이 필요 없는 디코딩: 매 디코딩 단계마다 전체 차원의 키를 재구성해야 하는 저차원 분해 방식과 달리, RAP는 프루닝 마스크를 오프라인에서 $W_q$ 에 융합하여 "재구성 세금(reconstruction tax)"을 제거하고 메모리 및 연산 오버헤드를 줄입니다.
포괄적인 평가: 다양한 모델 제품군(Llama, Mistral, Qwen) 및 스케일(3B~14B)에 대해 검증되어 견고함을 입증했습니다.

4. 실험 결과

공정한 품질 및 효율성 비교를 위해 유지되는 KV 채널 비율 $\rho$ 를 일치시키는 iso-KV 설정 하에 평가가 수행되었습니다.

정확도 vs. 압축 (Accuracy vs. Compression):
- 30% 압축 비율 ( $\rho=0.7$ )에서, RAP는 사용 가능한 수준의 정확도를 유지합니다 (예: Llama-3-8B의 경우 WikiText-2 PPL ~9.12, 베이스라인 8.28 대비). 이는 심각한 정확도 붕괴(PPL ~397)를 겪는 RoPE 무관 채널 프루닝(KVPruner)보다 크게 우수합니다.
- RAP는 낮은 어텐션 비용으로 작동하면서도 강력한 저차원 베이스라인(PaLU, SVD-LLM)과 대등하거나 더 나은 성능을 달-성합니다.
효율성 이득 (Efficiency Gains):
- 파라미터 및 FLOPs: RAP는 유지 비율 $\rho$ 에 따라 어텐션 파라미터와 FLOPs를 선형적으로 감소시킵니다 (예: $\rho=0.7$ 에서 약 70%). 반면, 저차원 방식은 온라인 확장 계수 때문에 더 높은 파라미터 수를 유지하는 경우가 많습니다.
- 지연 시간 (Latency): RAP는 재구성 단계를 피하기 때문에 프리필(Prefill) 및 디코딩 지연 시간을 줄입니다 (예: $\rho=0.7$ 에서 밀집 베이스라인의 83%/77%). 저차원 방식은 특히 디코딩 중에 지연 시간이 증가하는 경향이 있습니다.
견고성 (Robustness):
- RAP는 다양한 모델 크기와 아키텍처 전반에서 성능을 유지합니다.
  로 4비트 가중치 양자화와 결합 가능하며 상호 보완적입니다.
- 커스텀 Triton 커널을 통해 비연속적인 RoPE 인덱싱을 효율적으로 처리하여 표준 PyTorch 인덱싱의 오버헤드를 방지합니다.

5. 의의 및 주장

본 논문은 RAP가 다음을 동시에 달성하는 유일한 방법이라고 주장합니다:

KV 캐시 크기 감소.
모델 파라미터 및 프로젝션 FLOPs 감소.
높은 압축 비율에서도 높은 정확도 유지.
단계별 재구성 오버헤드 없이 작동.

저자들은 RAP가 Ro-based LLM을 위한 "드롭인(drop-in)" 솔루션임을 강조합니다. RoPE의 구조적 제약을 존중함으로써, 기존 프루닝 방식의 "침묵하는 실패" 문제를 해결하고 긴 컨텍스트 추론을 위한 더 효율적인 대안을 제시합니다. 이 연구는 모델 고유의 메커니즘(RoPE와 같은)과 구조적으로 정렬되는 것이 효과적인 모델 압축에 필수적임을 시사합니다.