🤖 machine learning

A JoLT for the KV Cache: Near-Lossless KV Cache Compression via Joint Tucker and JL-Residual Allocation for LLMs

JoLT는 토큰 및 피처 축에 부분 터커 분해(partial Tucker decomposition)를 적용하고, 존슨-린덴스트라우스 회전 저비트 잔차(Johnson-Lindenstrauss rotated low-bit residual)를 통해 버려진 정보를 복원함으로써, 퍼플렉서티(perplexity) 및 다운스트림 태스크에서의 베이스라인 성능을 유지하기 위해 단일 바이트 예산(unified byte budget) 하에 최적화된 방식으로 대규모 언어 모델의 KV 캐시를 손실이 거의 없는 2-3배 압축을 달성한다.

원저자: Rahul Krishnan, Volker Schulz

게시일 2026-07-15

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Rahul Krishnan, Volker Schulz

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

당신이 거대하고 똑똑한 로봇 두뇌(거대 언어 모델, LLM)를 운영하며 이야기를 들려주고 있다고 상상해 보세요. 이야기를 계속 이어가기 위해, 로봇은 현재 대화에서 자신이 말했던 모든 단어를 기억해야 합니다. 이 로봇은 이 기억을 **KV 캐시(KV Cache)**라고 불리는 특별한 "노트북"에 보관합니다.

문제는 이야기가 길어질수록 이 노트북이 거대해진다는 것입니다. 노트북이 너무 커지면 로봇의 메모리를 전부 잡아먹기 시작하여, 속도를 늦추게 됩니다. 이는 마치 마라톤을 하면서 배낭 안에 도서관 전체를 넣고 뛰려는 것과 같습니다. 결국, 당신은 멈춰 서게 될 것입니다.

과학자들은 이전에 이 노트북을 줄이려고 시도한 적이 있습니다. 어떤 이들은 전체를 아주 작은 상자에 구겨 넣으려 했고(양자화, quantization), 다른 이들은 행이나 열의 패턴을 찾아 페이지를 요약하려 했습니다(저계수 방법, low-rank methods). 하지만 이 논문의 저자인 라훌 크리슈난(Rahul Krishnan)과 폴커 슐츠(Volker Schulz)는 다른 방법들이 놓친 것을 발견했습니다. 바로 이 노트북이 단순히 평평한 종이 뭉치가 아니라는 점입니다. 그것은 헤드(Heads, 다양한 사고 방식), 토큰(Tokens, 단어), **특징(Features, 세부 사항)**이라는 세 가지 면을 가진 3D 블록입니다.

그들은 이 두 면은 지루하고 반복적인 내용으로 가득 차 있어 쉽게 압축할 수 있지만, 나머지 면들은 고유하며 이를 압축하면 로봇의 지적 능력을 잃을 수 있다는 것을 발견했습니다.

위대한 발견: "JoLT" 방식

팀은 JoLT(Joint Tucker and JL-residual allocation)라고 불리는 새로운 기술을 발명했습니다. 이것은 당신의 배낭을 싸주는 매우 똑똑한 짐 싸기 서비스와 같습니다.

스마트한 압축 (부분 턱커, Partial Tucker): 전체 3D 블록을 통째로 구겨 넣는 대신, JoLT는 데이터를 살펴보고 이렇게 말합니다. "'헤드'와 '레이어'는 고유하고 소중하니 그대로 두자. 하지만 '토큰'과 '특징'은 군더더기로 가득 차 있네." 그래서 이 두 특정 면만을 압축합니다. 이는 거대하고 폭신한 베개를 가져와서 가장자리의 탄탄한 부분은 그대로 둔 채 가운데의 공기만 빼내는 것과 같습니다.
안전망 (JL-residual): 베개를 누르면 공기가 빠져나갑니다. 만약 그냥 내버려 둔다면 베개는 납작해지고 쓸모없어질 것입니다. JoLT는 이 "빠져나간 공기"(손실된 정보)를 포착하여 JL-residual이라고 불리는 작고 매우 효율적인 안전망에 저장합니다. 이 그물은 패킹 능력이 매우 뛰어나서 단 몇 비트만으로도 누락된 세부 사항을 담아낼 수 있습니다.
완벽한 균형 (라그랑주 쌍대, Lagrangian Dual): 여기에 마법 같은 부분이 있습니다. 로봇에게는 사용할 수 있는 공간에 대한 엄격한 예산(예: 1바이트)이 있습니다. JoLT는 수학적인 "스마트 할당기"를 사용하여 베개를 얼마나 압축할지, 그리고 안전망에 얼마나 많은 공간을 할당할지를 결정합니다. 이 도구는 어떤 메모리 부분(Key)은 압축하기 쉽고, 다른 부분(Value)은 고집스러워 더 많은 안전망 공간이 필요하다는 것을 깨닫습니다. 그리고 최선의 결과를 얻기 위해 예산을 동적으로 이동시킵니다.

그들이 증명한 것 (그리고 증명하지 못한 것)

저자들은 이 기술을 두 가지 유명한 로봇 두뇌인 Mistral-7B(Grouped-Query 스타일 사용)와 LLaMA-2-13B(Multi-Head 스타일 사용)에서 테스트했습니다.

"자유 구역(Free Zone)": 그들은 메모리를 2~3배(2–3×) 줄여도 로봇의 성능이 전혀 떨어지지 않는 스위트 스팟(sweet spot)을 발견했습니다. 이는 "손실이 거의 없는(near-lossless)" 상태였습니다.
- GSM8K(수학 문제) 및 RULER(텍스트 속에서 바늘 찾기)와 같은 테스트에서, 압축된 로봇은 압축되지 않은 로봇과 통계적 오차 범위 내에서 정확히 동일한 점수를 기록했습니다.
- 메모리 재구성 과정에서의 오류는 매우 작았습니다. Key의 경우 약 0.009, Value의 경우 약 0.006이었습니다. 이는 4비트 양자화나 교차 레이어 SVD와 같은 이전 방법들보다 약 10배(한 자릿수) 더 나은 결과입니다.
"절벽(The Cliff)": 그들은 또한 한계점도 발견했습니다. 만약 메모리를 너무 세게 압축하려고 하면(3배 이상) 상황이 엉망이 됩니다.
- Mistral 로봇은 압축을 진행함에 따라 성능이 완만하게 저하되며 우아하게 무너졌습니다.
- 하지만 LLaMA 로봇은 4배에서 5배 사이의 압축 구간에서 "절벽"을 만났습니다. 성능이 급격히 추락하여, 퍼플렉시티(perplexity)가 5.39에서 9.07로 크게 뛰어올랐습니다(이는 단어 예측 능력이 훨씬 나빠졌음을 의미합니다).

빠른 버전: FlashJoLT

완벽한 압축을 계산하는 데는 시간이 걸립니다. 이를 해결하기 위해 그들은 FlashJoLT를 만들었습니다. 매번 무거운 수학 계산을 완벽하게 수행하는 대신, 주요 패턴을 빠르게 추측하는 "무작위화된" 지름길을 사용합니다.

결과: 메모리 압축 속도가 5~13배 더 빠르지만, 품질은 느린 완벽한 버전과 정확히 동일하게 유지됩니다.

그들이 제외한 것

이 논문은 이 특정 문제에 대해 무엇이 잘 작동하지 않는지 명확히 밝히고 있습니다:

모든 것을 압축하기: 세 가지 면(헤드, 토큰, 특징)을 모두 압축하려고 시도하는 것은 나쁜 생각입니다. "헤드"와 "레이어"는 너무 고유해서, 이를 압축하면 로봇의 두뇌에 해를 끼칩니다.
고정 비트 양자화: 모든 숫자의 비트 수를 단순히 낮추는 것(예: 모든 것을 4비트로 강제하는 것)은 2~3배 압축의 "스위트 스팟"에 도달할 수 없습니다. 압축이 충분히 되지 않거나 품질을 너무 많이 잃게 됩니다.
일률적인 적용: "Key"와 "Value"를 동일하게 취급할 수 없습니다. "Value"는 훨씬 압축하기 어렵기 때문에(2~3배 더 어려움), 더 많은 공간 예산이 필요합니다.

결론

저자들은 실제 하드웨어(A100 GPU)에서 이를 측정하였으며, JoLT가 로봇의 지능을 해치지 않으면서 메모리를 2~3배 줄일 수 있는 손실이 거의 없는(near-lossless) 방법을 제공한다는 것을 확인했습니다.

하지만 저자들은 이것이 모든 것에 대한 마법의 해결책은 아니라고 신중하게 언급합니다.

이것은 "자유 구역"(2~3배)에서는 잘 작동하지만, 특정 로봇 유형(LLaMA 등)에서 더 강하게 밀어붙이면 품질이 급격히 떨어집니다.
메모리 저장 공간은 작아졌지만, 로봇은 말을 할 때마다 메모리를 "압축 해제"하는 수학적 계산을 여전히 수행해야 합니다. 저자들은 이것이 실제 환경에서 진정으로 실용적이 되려면, 압축된 메모리를 먼저 풀지 않고도 직접 읽을 수 있는 특수한 컴퓨터 칩(Fused Kernels)을 설계해야 한다고 제안합니다.

요약하자면, JoLT는 긴 대화 중에 엄청난 공간을 절약해 주는 매우 정교하고 수학적인 패킹 기술이지만, 한계가 있으며, 제 속도를 내기 위해서는 미래의 하드웨어로부터 약간의 도움이 필요합니다.

기술 요약: KV 캐시를 위한 JOLT

문제 정의

자기회귀적 트랜스포머 추론(autoregressive transformer inference)에서 키-값(Key-Value, KV) 캐시는 $O(\text{batch} \times T \times L \times n_h \times d_h)$ 로 스케일링되는 지배적인 메모리 비용을 나타냅니다. 긴 컨텍스트나 대규모 배치 환경에서, 모델 가중치보다 KV 캐시가 주요 메모리 병목 현상이 되어 처리량을 제한합니다. 기존의 압축 방식은 두 가지 계열로 나뉘지만, 둘 다 캐시의 구조적 중복성을 완전히 활용하지 못합니다:

저계수(Low-rank) 방식: 이들은 2차원 슬라이스(예: 헤드별 행렬 또는 레이어 간 블록)를 인수분해하지만, 하나의 축에 대해서만 중복성을 확정함으로써 다른 축들을 무시합니다.
양자화(Quantization) 방식: 모든 엔트리의 비트 폭을 줄이지만, 고정된 비트율(예: 4비트)에 제한되어, 품질이 유지될 수 있는 중간 단계의 압축 비율(2–3×)을 달 achieves하는 데 한계가 있습니다.

본 연구의 동기가 되는 핵심 관찰은 특정 레이어의 KV 캐시가 3차 텐서( $n_h \times T \times d_h$ )이며, 각 축이 서로 매우 다른 양의 중복성을 보유하고 있다는 점입니다. 구체적으로, 헤드(head)와 레이어(layer) 축은 거의 압축 불가능한 반면, 토큰(token)과 피처(feature) 축은 대부분의 압축 가능한 구조를 보유하고 있습니다. 또한, 값(values)은 그 스펙트럼 특성의 차이로 인해 키(keys)보다 압축하기가 2–3배 더 어렵다는 것이 경험적으로 확인되었습니다.

방법론: JoLT

제안된 방법인 JoLT(Joint Tucker and JL-residual allocation)는 캐시를 3차 텐서로 취급하고 타겟팅된 압축 전략을 적용하여 이러한 과제들을 해결합니다.

1. 부분 턱커 분해 (Partial Tucker Decomposition)

모든 축을 압축하는 전체 텐서 인수분해와 달리, JoLT는 부분 턱커 분해를 적용합니다. 이는 중복성을 가진 토큰 및 피처 축만을 절단(truncate)하고, 압축 불가능한 헤드 및 레이어 축은 온전하게 유지합니다. 이는 순차적 절단 고차 SVD(Sequentially Truncated Higher-Order SVD, ST-HOSVD)를 사용하여 수행됩니다.

동기: 헤드나 레이어 축을 압축하면 상당한 오차가 발생하지만, 토큰과 피처 축을 절단하면 중복성의 대부분을 포착할 수 있습니다.

2. JL 회전 잔차 (JL-Rotated Residual)

순수한 저계수 백본(low-rank backbone)만으로는 특히 절단 시 상당한 에너지를 버리게 되는 평탄한 스펙트럼을 가진 '값(values)'에 대해 근사 무손실 충실도를 확보하기에 불충분합니다. JoLT는 Johnson-Lindenstrauss(JL) 회전 저비트 잔차를 사용하여 이 버려진 에너지를 복구합니다:

절단 잔차 $R = X - \tilde{X}$ 를 계산합니다.
잔차에 무작위 JL 회전을 적용합니다.
회전된 잔차를 $b \in \{0, 2, 4, 8\}$ 비트로 균등 양자화합니다.
디코딩 시 회전을 역으로 적용합니다. 이 잔차는 손실이 있는 백본을 근사 무손실 백본으로 전환하는 레버 역할을 하며, 특히 '값(values)'에 유용합니다.

3. 공동 라그랑주 할당 (Joint Lagrangian Allocation)

이 방법은 두 단계(Tucker 랭크와 잔차 비트 폭) 및 두 가지 대상(Keys vs. Values, 레이어 그룹) 전반에 걸쳐 자원을 공동 할당하기 위해 단일 **라그랑주 쌍대(Lagrangian dual)**를 사용합니다.

목표: 전체 바이트 예산 제약 조건 하에서 총 재구성 오차를 최소화합니다.
메커니즘: 쌍대 변수는 랭크와 잔차 사이, 그리고 키와 값 사이에서 비트를 동적으로 이동시켜, 압축성의 비대칭성(예: 키 잔차 비트는 줄이고 값 잔차에는 비트를 집중 투자)을 활용할 수 있게 합니다.
세밀도: 할당은 레이어 그룹별로 수행되며, 키와 값에 대해 별도로 이루어집니다.

4. FlashJoLT (빠른 변형)

긴 컨텍스트에서 정확한 SVD의 계산 비용 문제를 해결하기 위해, FlashJoLT는 정확한 토큰 모드 SVD를 **무작위 저계수 SVD(randomized low-rank SVD)**로 대체합니다.

상위 $q_{cap}$ 개의 토큰 모드 방향만을 계산합니다.
**테일 질량 계수(tail-mass accounting)**를 채택하여, 누락된 스펙트럼 테일을 합성 잔차 특이값으로 추가함으로써 할당기가 주어진 랭크가 포착한 에너지를 과소평가하는 것을 방지합니다.
이 변형은 품질 저하 없이 압축 시간을 5–13배 단축합니다.

주요 기여

스펙트럼 분석: 헤드와 레이어 축은 압축 불가능하며, 토큰과 피처 축은 중복성을 보유하고 있고, 값이 키보다 훨씬 압축하기 어렵다는 것을 식별한 경험적 연구.
JoLT 알고리즘: 부분 턱커 분해와 JL 회전 잔차를 결합하고, 랭크와 비트의 공동 할당을 위해 단일 라그랑주 쌍대를 사용하는 새로운 압축 알고리즘.
FlashJoLT: 품질을 보존하면서 압축 시간을 획기적으로 줄이는 무작위 SVD 변형.
경험적 검증: GQA (Mistral-7B) 및 MHA (LLaMA-2-13B) 모델에 대한 종합적인 연구를 통해 "프리 존(free zone)"의 근사 무손실 압축을 입증.

결과

"프리 존" (2–3× 압축)

본 논문은 모델 품질이 압축되지 않은 베이스라인의 통계적 노이즈 범위 내에 머무는 2–3× 압축 비율을 **"프리 존(free zone)"**으로 정의합니다:

Perplexity (당혹도): Mistral-7B (GQA)와 LLaMA-2-13B (MHA) 모두에서, 컨텍스트 길이가 1024인 경우 2–3× 압축 시 퍼플렉시티 드리프트가 무시할 만한 수준입니다.
재구성 오차: 2× 압축 시, JoLT는 키(Keys)에 대해 0.009, 값(Values)에 대해 0.006의 상대 프로베니우스 오차(relative Frobenius error)를 달성합니다. 이는 4비트 양자화 및 교차 레이어 SVD 베이스라인보다 약 10배 낮은 수치입니다.
다운스트림 태스크:
- GSM8K: 압축된 모델(2–3×)은 베이스라인의 95% 신뢰 구간 내에 머뭅니다.
- RULER (Needle-in-a-Haystack): Mistral은 2–3× 압축 시 100% 검색 정확도를 유지하며, LLaMA는 베이스라인의 천장에 부합합니다.

아키텍처별 차이

2–3× 프리 존을 벗어나면 두 아키텍처는 서로 다른 양상을 보입니다:

Mistral (GQA): 완만하게 저하되며, 4× 이상의 정수 배 단계마다 약 4%의 퍼플렉시티 손실이 발생합니다.
LLaMA (MHA): 4×와 5× 사이에서 급격히 저하됩니다. 4×에서 퍼플렉시티가 약 27% 증가하고, 5×에서는 약 68% 증가합니다. LLaMA의 추론 작업(GSM8K)은 높은 압축 비율에서 크게 무너집니다.

효율성

메모리: 지속적인 KV 풋프린트는 압축 비율에 따라 선형적으로 줄어듭니다 (예: 컨텍스-2048 기준 Mistral의 경우 268.4 MB에서 2× 압축 시 134.6 MB로 감소). 피크 프로세스 메모리는 가중치/활성화 함수에 의해 제한됩니다.
속도: FlashJoLT는 정확한 백본 대비 압축 시간을 5–13배 단축합니다.

의의 및 주장

본 논문은 Jo-LT가 KV 캐시의 근사 무손실 2–3× 압축을 제공하여, 모델 품질을 희생하지 않고도 긴 컨텍스트 추론을 위한 메모리 병목 현상을 효과적으로 제거한다고 주장합니다. 캐시를 텐서로 취급하고 키와 값 사이의 특정 스펙트럼 비대칭성을 활용함으로써, JoLT는 기존의 저계수 및 양자화 방식보다 중간 압축 영역에서 우수한 성능을 보입니다.

저자들은 이 방법을 긴 컨텍스트 서빙 비용을 줄이기 위한 직접적인 레버로 제시합니다. 두 가지 주요 경계를 언급합니다:

아키텍처 한계: 이 방법은 GQA 모델에는 견고하지만, MHA 모델의 경우 4× 이상에서 급격한 저하를 보이므로, 2–3× 대역이 안전하고 모델에 구애받지 않는 권장 범위임을 시사합니다.
배포 비용: 저장 용량은 줄어들지만, 현재 디코딩 시 매 단계마다 인수를 통해 캐시를 재구성해야 합니다. 저자들은 추론 성능을 동일하게 유지하고 이 방법의 잠재력을 완전히 실현하기 위한 핵심 다음 단계로 재구성과 어텐션을 결합한 퓨즈드 커널(fused reconstruction-and-attention kernel) 개발을 꼽았습니다.