Technische Samenvatting: Een JOLT voor de KV Cache

Probleemstelling

In autoregressieve transformer-inferentie vormt de Key-Value (KV) cache de dominante geheugenkosten, schalend als $O(\text{batch} \times T \times L \times n_h \times d_h)$. Voor lange contexten of grote batches wordt de KV cache, in plaats van de modelgewichten, de primaire geheugenbottleneck en beperkt de doorvoer. Bestaande compressiemethoden vallen uiteen in twee families, die beide falen om de structurele redundantie van de cache volledig te benutten:

1. Low-rank methoden: Deze factoriseren tweedimensionale slices (bijv. per-head matrices of cross-layer blokken), maar committeren zich aan een enkele as van redundantie, waarbij de andere assen worden genegeerd.
2. Quantisatie-methoden: Deze verminderen de bitbreedte van elke entry, maar zijn beperkt door vaste bit-rates (bijv. 4-bit), wat een "compressievloer" creëert die voorkomt dat tussenliggende compressieverhoudingen (2–3×) worden bereikt waar de kwaliteit vaak behouden blijft.

De kernobservatie die dit werk motiveert, is dat de KV cache op een laag een derde-orde tensor is ($n_h \times T \times d_h$) waarbij de assen fundamenteel verschillende hoeveelheden redundantie bevatten. Specifiek zijn de head- en layer-assen grotendeels incompressibel, terwijl de token- en feature-assen de meerderheid van de compressibele structuur bevatten. Bovendien zijn values empirisch gezien 2–3× moeilijker te comprimeren dan keys vanwege de verschillen in hun spectrale eigenschappen.

Methodologie: JoLT

De voorgestelde methode, JoLT (Joint Tucker and JL-residual allocation), pakt deze uitdagingen aan door de cache te behandelen als een derde-orde tensor en een gerichte compressiestrategie toe te passen.

1. Partiële Tucker Decompositie

In tegenstelling tot volledige tensor-factorisaties die alle assen comprimeren, past JoLT een partiële Tucker decompositie toe. Het truncheert alleen de token- en feature-assen (die redundantie bevatten), terwijl de head- en layer-assen intact blijven (die incompressibel zijn). Dit wordt bereikt met Sequentially Truncated Higher-Order SVD (ST-HOSVD).

• Motivatie: Het comprimeren van de head- of layer-assen leidt tot aanzienlijke fouten, terwijl het trunceren van de token- en feature-assen de bulk van de redundantie vangt.

2. JL-Geroteerd Residu

Een zuivere low-rank backbone is onvoldoende voor bijna verliesloze getrouwheid, vooral voor values die een vlak spectrum hebben waarbij trunctie significante energie wegwerpt. JoLT herstelt deze weggeworpen energie met een Johnson-Lindenstrauss (JL) geroteerd low-bit residu:

• Het trunctie-residu $R = X - \tilde{X}$ wordt berekend.
• Een willekeurige JL-rotatie wordt toegepast op het residu.
• Het geroteerde residu wordt uniform gekwantiseerd op $b \in \{0, 2, 4, 8\}$ bits.
• Decodering houdt in dat de rotatie wordt omgekeerd. Dit residu fungeert als een hefboom om een lossy backbone om te zetten in een bijna verliesloze backbone, wat vooral waardevol is voor values.

3. Gezamenlijke Lagrangian Allocatie

De methode maakt gebruik van een enkele Lagrangiaanse duale om middelen gezamenlijk toe te wijzen over de twee stadia (Tucker-rangen en residu-bitbreedtes) en over verschillende targets (Keys versus Values, Layer Groups).

• Doel: Minimaliseer de totale reconstructiefout onder een globaal byte-budget.
• Mechanisme: De duale variabele stelt de allocator in staat om bits te verschuiven tussen rangen en residuen, en tussen keys en values, waarbij de asymmetrie in compressibiliteit dynamisch wordt benut (bijv. het zwaar financieren van het value-residu terwijl de key-residu bits worden verminderd).
• Granulariteit: Allocatie vindt plaats per layer group en apart voor Keys en Values.

4. FlashJoLT (Snelle Variant)

Om de computationele kosten van exacte SVD bij lange contexten aan te pakken, vervangt FlashJoLT de exacte token-mode SVD door een randomized low-rank SVD.

• Het berekent alleen de top $q_{cap}$ token-mode richtingen.
• Het maakt gebruik van tail-mass accounting, waarbij de weggelaten spectrale staart wordt toegevoegd als een synthetisch residu singular waarde om te voorkomen dat de allocator de energie die door een bepaalde rang wordt gevangen onderschat.
• Deze variant bereikt een 5–13× versnelling in compressietijd met verwaarloosbaar kwaliteitsverlies in de "free zone".

Belangrijkste Bijdragen

1. Spectrale Analyse: Een empirische studie die identificeert dat de head- en layer-assen incompressibel zijn, terwijl de token- en feature-assen redundantie bevatten, waarbij values aanzienlijk moeilijker te comprimeren zijn dan keys.
2. JoLT Algoritme: Een nieuwe compressor die partiële Tucker decompositie combineert met een JL-geroteerd residu, gekoppeld via een enkele Lagrangiaanse dual voor de gezamenlijke allocatie van rangen en bits.
3. FlashJoLT: Een randomized-SVD variant die de compressietijd drastisch vermindert terwijl de kwaliteit behouden blijft.
4. Empirische Validatie: Een uitgebreide studie op GQA (Mistral-7B) en MHA (LLaMA-2-13B) modellen die een "free zone" van bijna verliesloze compressie vaststelt.

Resultaten

De "Free Zone" (2–3× Compressie)

Het paper identificeert een 2–3× compressieverhouding als een "free zone" waar de modelkwaliteit behouden blijft binnen de statistische ruis van de ongecomprimeerde baseline:

• Perplexity: Op zowel Mistral-7B (GQA) als LLaMA-2-13B (MHA) is de perplexity drift verwaarloosbaar bij 2–3× voor contextlengtes tot 1024.
• Reconstructiefout: Bij 2× compressie bereikt JoLT een relatieve Frobenius-fout van 0.009 (Keys) en 0.006 (Values). Dit is ongeveer een orde van grootte lager dan 4-bit quantisatie en cross-layer SVD baselines.
• Downstream Taken:
  • GSM8K: Gecomprimeerde modellen (2–3×) blijven binnen het 95% betrouwbaarheidsinterval van de baseline.
  • RULER (Needle-in-a-Haystack): Mistral behoudt 100% retrieval accuratesse bij 2–3×; LLaMA evenaart zijn baseline plafond.

Architecturale Divergentie

Voorbij de 2–3× free zone divergeren de twee architecturen:

• Mistral (GQA): Degradeert gracieus, met een verlies van ongeveer 4% perplexity per gehele ratio stap voorbij 4×.
• LLaMA (MHA): Degradeert scherp tussen 4× en 5×, waarbij de perplexity toeneemt met ~27% bij 4× en ~68% bij 5×. Redeneringstaken (GSM8K) voor LLaMA storten significant in bij hogere ratio's.

Efficiëntie

• Geheugen: De persistente KV footprint krimpt lineair met de compressieverhouding (bijv. van 268.4 MB naar 134.6 MB bij 2× voor Mistral bij ctx-2048). De piek process memory blijft begrensd door gewichten/activaties.
• Snelheid: FlashJoLT biedt een 5–13× versnelling in compressietijd vergeleken met de exacte backbone.

Betekenis en Claims

Het paper claimt dat JoLT een bijna verliesloze 2–3× compressie van de KV cache biedt, wat effectief de geheugenbottleneck voor lange-context inferentie wegneemt zonder de modelkwaliteit op te offeren. Door de cache als een tensor te behandteren en de specifieke spectrale asymmetrie tussen keys en values uit te buiten, presteert JoLT beter dan bestaande low-rank en quantisatie methoden in het tussenliggende compressieregime.

De auteurs positioneren de methode als een directe hefboom om de kosten van het serveren van lange contexten te verlagen. Ze merken twee primaire grenzen op:

1. Architecturale Limieten: De methode is robuust voor GQA-modellen, maar ervaart scherpe degradatie voor MHA-modellen voorbij 4×, wat suggereert dat de 2–3× band de veilige, model-agnostische aanbeveling is.
2. Implementatiekosten: Hoewel de opslag wordt verminderd, vereist decodering momenteel het reconstrueren van de cache vanuit factoren bij elke stap. De auteurs identificeren de ontwikkeling van een fused reconstruction-and-attention kernel als de cruciale volgende stap om inferentie-pariteit te bereiken en het volledige potentieel van de methode te realiseren.