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

Background Transformer inference reuses every previous token's key and value projections at each layer, so implementations cache them rather than recompute — a cache that grows with batch size, context length, and depth, and that becomes the binding constraint on memory and throughput before the model weights do once context or batch size grows large. Existing compression methods split into two families: low-rank methods that factor two-dimensional slices of the cache (Palu factors per-head matrices, xKV factors cross-layer feature blocks) and quantization methods that lower the bit-width of every entry (KIVI, TurboQuant).

배경 Transformer 추론은 매 층에서 이전 모든 토큰의 key·value 벡터를 재사용하기 위해 KV 캐시에 저장해 두는데, 이 캐시는 배치 크기·문맥 길이·층 수에 비례해 커지고 긴 문맥이나 큰 배치에서는 모델 가중치보다 먼저 메모리와 처리량의 한계를 결정짓는다. 캐시를 줄이는 기존 방법은 저랭크 분해(Palu는 헤드별 행렬을, xKV는 층 간 특성 블록을 인수분해)와 양자화(KIVI, TurboQuant 등이 모든 항목의 비트폭을 낮춤) 두 갈래로 나뉜다.

Problem Neither family exploits the fact that the cache at a layer is naturally a third-order tensor -- heads, tokens, and features -- whose three axes carry very different amounts of redundancy. A two-dimensional factorization must commit to one axis of redundancy and ignore the others, and a fixed-bit-width quantizer cannot reach the modest, intermediate compression ratios (like 2-3x) where quality is still free, because even its smallest bit setting already overshoots them.

문제 두 갈래 모두 캐시가 실제로는 헤드·토큰·특성이라는 세 축을 가진 3차 텐서이고 그 세 축이 서로 매우 다른 정도의 중복성을 담고 있다는 사실을 활용하지 못한다. 2차원 분해는 한 축의 중복성에만 전념해 나머지를 무시해야 하고, 고정 비트폭 양자화기는 가장 낮은 설정조차 이미 초과해 버리는 중간 정도의 압축률(2-3배 같은) 구간을 표현할 수 없다.

Prior limitations Low-rank methods like Palu and xKV treat the cache as a collection of matrices rather than a single object seen across heads, tokens, and features at once, and provide no residual term to restore the energy a truncation discards. Fixed-bit-width quantizers like KIVI and TurboQuant are strong at their native bit-rates but structurally cannot express compression ratios below their floor -- a 4-bit scheme, for instance, cannot reach below roughly 4x.

기존 한계 Palu·xKV 같은 저랭크 방법은 캐시를 행렬들의 모음으로만 보고 헤드·토큰·특성을 동시에 보는 하나의 대상으로 취급하지 않으며, 절단 후 남는 에너지를 복원할 잔차 항도 제공하지 않는다. KIVI·TurboQuant 같은 고정-비트폭 양자화기는 특정 비트레이트에서는 강력하지만 그보다 낮은 압축 목표는 원천적으로 표현할 수 없다 — 예를 들어 4비트 스킴은 대략 4배 아래로 내려갈 수 없다.

Goal Measure where redundancy in the KV cache actually lives across its three axes, and use that measurement to design a compressor that leaves the head and layer axes intact, compresses only the token and feature axes, and budgets keys and values separately -- reaching a near-lossless result in the 2-3x band.

목표 KV 캐시의 중복성이 실제로 세 축 중 어디에 있는지 측정하고, 그 측정에 근거해 헤드·층 축은 그대로 두고 토큰·특성 축만 압축하면서 키와 값에 별도의 예산을 주는 압축기를 설계해, 2-3배 같은 중간 압축률 구간에서 사실상 무손실인 결과를 얻는 것이다.

Method Taking the singular values of each 1D unfolding of a real KV cache (Mistral-7B, layer 15) shows the head axis is essentially incompressible (singular-value ratio ~1.4-1.5), while the token and feature axes decay substantially -- and within those, keys compress sharply on the feature axis (ratio 27.4) while values compress only mildly (ratio 2.7), an asymmetry driven by projection geometry. Following this measurement, JoLT pins the head and layer axes as identity, truncates only the token and feature axes via a partial Tucker decomposition, restores the discarded truncation energy with a Johnson-Lindenstrauss-rotated low-bit residual (0/2/4/8 bits), and solves a single Lagrangian dual that jointly allocates Tucker ranks and residual bit-width, per layer group and separately for keys and values, under one shared byte budget.

방법 Mistral-7B 15번째 층의 캐시를 세 방식으로 unfold해 특이값을 측정한 결과, 헤드 축은 거의 비압축(특이값 비 ~1.4-1.5)인 반면 토큰·특성 축은 실질적으로 감쇠하고, 그 안에서도 key는 특성 축이 날카롭게(비 27.4) value는 완만하게(비 2.7) 압축 가능하다는 비대칭이 드러난다. 이 측정에 따라 JoLT는 헤드·층 축은 항등으로 고정한 채 토큰·특성 축만 partial Tucker 분해로 절단하고, 절단이 버린 에너지를 Johnson-Lindenstrauss 회전 후 저비트(0/2/4/8비트)로 양자화한 잔차로 복원하며, 하나의 Lagrangian dual이 층 그룹별·K/V별로 Tucker 랭크와 잔차 비트폭을 하나의 바이트 예산 아래 함께 결정한다.

Key idea The central move is putting rank and residual bits on the same error scale so one optimizer can shift budget freely between them, and between keys and values, rather than fixing that split in advance -- like routing water from a single pipe to wherever it is needed instead of filling two separate reservoirs (rank, residual bits) independently. Because values must spread energy across directions to preserve content (a flat spectrum) and are therefore 2-3x harder to compress than keys at every ratio, this freedom to move budget across the asymmetry is what a fixed split could never exploit.

핵심 아이디어 핵심은 랭크를 사는 것과 잔차 비트를 사는 것을 같은 오차 단위로 놓고 하나의 최적화기가 둘 사이, 그리고 key와 value 사이에서 예산을 자유롭게 옮기게 한 것이다 — 마치 두 개의 별도 저수지(랭크, 잔차 비트)에 물을 채우는 대신 하나의 수도관에서 필요한 곳으로 물을 돌리는 것과 같다. Value는 정보 보존을 위해 에너지를 넓게 퍼뜨려야 해서(spectrum이 평평) key보다 2-3배 압축하기 어렵다는 사실이 이 자유로운 배분의 실질적 이득을 만든다 — 고정된 분배였다면 이 비대칭을 활용할 수 없었을 것이다.

Validation Evaluated on Mistral-7B-v0.3 (GQA) and LLaMA-2-13B (MHA), using WikiText-2+C4 perplexity, 8-shot chain-of-thought GSM8K, and RULER single-needle retrieval out to 16K context, with perplexity averaged over seeds {0,1,2} whose standard deviation stays at or below 0.009 PPL in every cell. Ablations removing the residual, swapping the joint dual for a greedy allocator, or using uniform instead of per-group rank each independently confirm that design choice matters.

검증 Mistral-7B-v0.3(GQA)와 LLaMA-2-13B(MHA) 두 모델, WikiText-2+C4 perplexity, GSM8K 8-shot chain-of-thought, RULER 단일-needle 검색(16K 문맥까지)으로 검증하며, perplexity는 시드 {0,1,2} 평균으로 시드 간 표준편차가 모든 셀에서 0.009 이하임을 확인한다. 잔차 제거·greedy 배분·균일 랭크로 바꾸는 소거 실험이 residual과 joint dual 배분이 각각 독립적으로 기여함을 보인다.

Results In the 2-3x band, perplexity, GSM8K accuracy, and RULER retrieval all sit within statistical noise of the uncompressed baseline -- at 2x, relative Frobenius error is 0.009 (K) and 0.006 (V), roughly an order of magnitude below 4-bit quantization (0.08/0.13) or cross-layer SVD (0.08/0.24) at the same budget. Beyond that free zone the two architectures split sharply: the GQA model (Mistral) degrades gradually, while the MHA model (LLaMA) falls off a cliff between 4x and 5x, with GSM8K collapsing from 28.13% baseline to 18.5%, 6.0%, and 1.5% at 4x, 6x, and 8x.

결과 2-3배에서는 perplexity·GSM8K 정확도·RULER 검색 정확도 모두 비압축 baseline과 통계적 잡음 범위 안에 머문다 — 2배에서 상대 Frobenius 오차가 key 0.009, value 0.006로 4비트 양자화(0.08/0.13)나 층 간 SVD(0.08/0.24)보다 한 자릿수 낮다. 이 free zone을 넘어서면 두 아키텍처가 갈라진다: GQA(Mistral)는 완만히 열화되는 반면 MHA(LLaMA)는 4배와 5배 사이에서 GSM8K 정확도가 28.13%에서 18.5%, 6.0%, 1.5%(4/6/8배)로 급격히 무너진다.

Comparison At matched compression budget, reconstruction error is an order of magnitude below both 4-bit quantization and cross-layer SVD (xKV), and a direct multilinear-format comparison shows Tucker beats CP, t-SVD, and tensor-train at every ratio on both architectures (the best alternative format still reaches only ~0.28 value error at 2x, versus JoLT's 0.006). Unlike Palu or xKV, JoLT deliberately leaves the head and layer axes uncompressed -- a choice validated by giving a full four-mode allocator complete freedom and observing it drives head and layer ranks back to full size on its own.

비교 동일 압축 예산에서 4비트 양자화·층 간 SVD(xKV) 대비 재구성 오차가 한 자릿수 낮고, CP·t-SVD·tensor-train 같은 다른 다중선형 분해 방식 중에서도 Tucker가 모든 압축률에서 최선임을 직접 비교로 확인한다(최선의 대안조차 2배에서 value 오차가 약 0.28로 JoLT의 0.006보다 훨씬 나쁨). Palu·xKV 같은 저랭크 방법과 달리 헤드·층 축은 건드리지 않고 압축 대상에서 제외하는데, 4축 전체를 자유롭게 배분하게 해도 최적화기가 스스로 헤드·층 랭크를 원래 크기로 되돌린다는 사실로 이 선택을 정당화한다.

Significance In a regime where the KV cache, not the model weights, sets the throughput ceiling at long context or large batch size, this spectral measurement of uneven axis-wise compressibility reframes the first question any cache-compression design should ask: which axis actually holds the redundancy. It also surfaces an important distinction for judging real deployment value -- a storage saving here does not yet automatically translate into a decoding-latency saving.

의의 긴 문맥·큰 배치에서 서빙 처리량의 상한을 정하는 것이 모델 가중치가 아니라 KV 캐시인 상황에서, 캐시의 세 축이 서로 다른 정도로 압축 가능하다는 이 스펙트럴 측정은 캐시 압축을 설계할 때 첫 질문이 되어야 할 것을 보여준다: 어느 축이 실제로 중복되어 있는가. 다만 이 저장 공간 절감이 곧바로 디코딩 지연 절감으로 이어지지는 않는다는 점도 함께 드러나는데, 이는 압축의 실제 배포 가치를 가늠할 때 중요한 구분이다.

Limitations The authors state the most important limitation themselves: decoding must reconstruct the cache from its Tucker factors at every step, and "today's attention kernels are not built to absorb" that cost -- so the storage win is not yet a deployable inference-latency path. Validation is confined to a single A100 GPU, two models (7B/13B scale), and contexts up to 16K, and the MHA architecture falls off sharply past 4x, so no single compression ratio is safe across architectures.

한계 가장 중요한 한계는 저자들이 스스로 명시한 것으로, 디코딩 각 단계마다 Tucker 인자로부터 캐시를 다시 복원해야 하는데 "오늘날의 attention 커널은 이를 흡수하도록 만들어지지 않았다" — 즉 저장 공간 절감이 아직 실제 지연시간 절감으로 이어지는 배포 가능한 경로가 아니다. 검증은 단일 A100 GPU, 두 모델(7B·13B급), 최대 16K 문맥으로 한정되며, MHA는 4배를 넘어서면 급격히 무너져 아키텍처에 무관한 하나의 안전한 압축률을 제시하지 못한다.

Future work The authors explicitly name a fused kernel -- one that computes attention directly from the encoded factors without materializing the reconstructed cache -- as the key remaining systems direction, framing it as the gap between today's storage win and a deployable latency win. Calibration robustness across text domains is stated as "in progress," and the long-context cap policy is validated only to 8192 tokens.

향후 과제 저자들은 인코딩된 인자로부터 직접 attention을 계산해 캐시를 명시적으로 복원할 필요가 없는 융합 커널을 다음 시스템 과제로 명시적으로 지목하며, 이것이 현재 저장 공간 이득과 실제 배포 가능한 지연시간 이득 사이의 간극이라고 밝힌다. 여러 텍스트 도메인에 걸친 캘리브레이션 견고성 검증은 "진행 중"이라 명시하고, 장문맥 cap 정책은 현재 8192 토큰까지만 검증되었다고 밝힌다.

Released resources No public code or checkpoint release was found -- no project page or repository beyond the arXiv listing itself was located.

자원 공개 논문에 공개된 코드나 체크포인트 링크는 확인되지 않았다 — arXiv 원문 외 별도 프로젝트 페이지나 저장소는 발견되지 않았다.