Technische Samenvatting: RATEQUANT

Probleemstelling

Grote Taalmodellen (LLM's) vereisen het cacheën van alle eerder berekende key-value (KV) paren tijdens de generatie, wat een geheugenbottleneck creëert die lineair schaalt met sequentielengte, batchgrootte en modeldiepte. Hoewel het kwantiseren van de KV-cache dit footprint verkleint, wijzen bestaande methoden uniform dezelfde bit-breedte toe aan elke attention-head. Deze aanpak negeert de goed gedocumenteerde heterogeniteit in head-belang.

Recente pogingen tot toewijzing van gemengde precisie op het niveau van laag of kanaal vertrouwen op heuristische regels die gekoppeld zijn aan specifieke kwantisators, zonder theoretische garanties of overdraagbaarheid. De auteurs identificeren een dieper, eerder onherkend obstakel: mismatch van het vervormingsmodel. Verschillende kwantisators volgen verschillende vervorming-rates-curves, gemodelleerd als $D(b) = \alpha \beta^{-b}$. De vervormingsrate $\beta$ varieert aanzienlijk tussen ontwerpen (bijvoorbeeld van 3,6 in TURBOQUANT tot 5,3 in QuaRot). Het naïef toepassen van een toewijzingsstrategie voor gemengde precisie, afgeleid van het vervormingsmodel van de ene kwantisator, op een andere, keert de volgorde van marginale winst van attention-heads om, wat vaak resulteert in prestaties die slechter zijn dan uniforme kwantisatie.

Methodologie

RATEQUANT adresseert deze uitdagingen door de toewijzing van bit-breedte voor per-head KV-cache te formaliseren als een rate-distortion optimalisatieprobleem. Het kader bestaat uit vier fasen:

1. Probleemformulering en Optimale Toewijzing

De auteurs modelleren de Mean Squared Error (MSE) van per-head kwantisatie als een exponentiële functie van de bit-breedte: $D(b) = \alpha \beta^{-b}$. Het doel is om de gewogen som van vervormingen over $N$ kwantisatiegroepen (heads) te minimaliseren, onderworpen aan een totaal bit-budget $B$:
$$\min_{b} \sum_{i=1}^{N} w_i \cdot D(b_i) \quad \text{s.t.} \quad \sum b_i = B$$
waarbij $w_i$ de gevoeligheid van head $i$ voorstelt.

Met behulp van Lagrangiaanse stationariteit leiden de auteurs een gesloten oplossing af via reverse waterfilling:
$$b^*_i = \bar{b} + \frac{\ln w_i - \ln \bar{w}}{\ln \beta}$$
waarbij $\bar{b}$ de gemiddelde bit-breedte is en $\bar{w}$ het geometrische gemiddelde van de gevoeligheden. Deze oplossing wijst meer bits toe aan heads met hogere gevoeligheid. Theoretisch is de haalbare vervormingsreductie begrensd door de verhouding van het rekenkundig gemiddelde tot het geometrisch gemiddelde (AM/GM) van de head-gevoeligheden.

2. Gevoeligheidsschatting

Om de gewichten $w_i$ te bepalen, maakt RATEQUANT gebruik van gevoeligheid op basis van gradiënten. Het berekent de gekwadrateerde gradiëntnormen van de KV-projectie-uitgangen met betrekking tot de causale taalmodelleringverlies:
$$w_{l,h} = \mathbb{E}_{x \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=1}^{T} \left\| \frac{\partial \mathcal{L}}{\partial K_{l,h,t}} \right\|^2 \right]$$
De auteurs tonen aan dat op gradiënten gebaseerde proxies superieur zijn aan op activatie gebaseerde metrics (zoals kanaalmagnitude), die versterking van fouten meten in plaats van directe impact op het verlies. Gevoeligheid op basis van gradiënten onthult een U-vormig patroon van belang (hoge gevoeligheid in vroege en late lagen), terwijl op activatie gebaseerde metrics een monotoon toenemend patroon met diepte tonen.

3. Hele Toewijzing en Uitbreidingen

Omdat bit-breedtes gehele getallen moeten zijn, wordt de continue oplossing verfijnd met een gierig algoritme voor marginale winst (Algoritme 1). Dit algoritme wijst iteratief bits toe aan de head die de grootste reductie in gewogen vervorming biedt, totdat het budget uitgeput is.

Twee kritieke uitbreidingen maken RATEQUANT kwantisator-agnostisch:

• Vervormingskalibratie: In plaats van een universeel $\beta$ aan te nemen, past RATEQUANT een per-kwantisator vervormingsmodel ($\alpha, \beta$) aan met behulp van een kleine kalibratieset (16 sequenties). Dit voorkomt de faalmodus "mismatch van het vervormingsmodel".
• Scheiding van K/V: Er wordt erkend dat keys en values vaak verschillende kwantisatieschema's gebruiken (bijvoorbeeld per-kanaal symmetrisch voor keys versus per-token asymmetrisch voor values in KIVI). Het kader behandelt ze als $2N$ onafhankelijke componenten met verschillende vervormingscurves. Dit stelt het algoritme in staat om asymmetrische bit-budgetten te ontdekken (bijvoorbeeld het toewijzen van meer bits aan keys dan aan values).

4. Pipeline

Het proces omvat eenmalige offline kosten:

1. Gevoeligheidskalibratie: Voorwaartse en achterwaartse passes op een kleine dataset (~1,6s voor een 8B-model).
2. Vervormingsmodellering: Aanpassen van $\alpha$ en $\beta$ bij 5 bit-breedtes (<0,1s).
3. Toewijzing: Uitvoeren van het gierige algoritme (<0,01s).
4. Inferentie: Toepassen van de statische bit-breedtes met nul runtime-overhead.

Belangrijkste Resultaten

Experimenten werden uitgevoerd op Qwen3-modellen (4B, 8B, 32B) met behulp van WikiText-2 perplexiteit (PPL) en downstream taken (ARC-Challenge, HellaSwag, enz.).

• Oplossing van Mismatch: Het toepassen van een mismatchend vervormingsmodel (bijvoorbeeld het gebruik van TURBOQUANT's $\beta$ voor KIVI) degradeerde KIVI's PPL van 49,3 naar 87,0 bij 2,5 gemiddelde bits. Gekalibreerd RATEQUANT loste dit op, met een reductie van PPL naar 14,9 (een reductie van 70%).
• Prestatiewinst:
  • KIVI (2,5 bits): RATEQUANT verbeterde PPL van 49,3 naar 14,9.
  • QuaRot (2,5 bits): Verbeterde PPL met 6,6 punten.
  • TURBOQUANT (4,0 bits): Herstelde 66% van de door kwantisatie veroorzaakte degradatie voor Qwen3-8B.
• Scheiding van K/V: Voor KIVI bij 2,5 bits was de door RATEQUANT ontdekte optimale toewijzing 2,85 bits voor keys en 2,15 bits voor values, wat significant beter presteerde dan gezamenlijke toewijzing.
• Efficiëntie: De methode voegt geen runtime-overhead toe. Kalibratie duurt minder dan 2 seconden op een enkele GPU.
• Downstream Nauwkeurigheid: Bij 4,0 bits behaalde RATEQUANT een gemiddelde nauwkeurigheid van 64,2% op downstream taken, wat bijna overeenkomt met de FP16-baseline (64,4%) en 89,8% van de nauwkeurigheidskloof herstelt.

Betekenis en Claims

Het artikel beweert dat RATEQUANT het eerste kader is dat rate-distortion theorie toepast op KV-cache kwantisatie op per-head granulariteit, terwijl het tegelijkertijd de structurele verschillen tussen keys en values adresseert.

De centrale bevinding is dat mismatch van het vervormingsmodel een kritieke faalmodus is in toewijzing van gemengde precisie. Zonder kalibratie van de vervormingscurve ($\beta$) aan de specifieke basis-kwantisator, kunnen strategieën voor gemengde precisie schadelijk zijn. Door de combinatie van:

1. Gevoeligheid op basis van gradiënten (superieur aan op activatie gebaseerde proxies),
2. Empirische vervormingskalibratie (om te matchen met de kwantisator), en
3. Scheiding van K/V (om asymmetrische foutkarakteristieken te hanteren),

transformeert RATEQUANT kwantisatie van gemengde precisie KV naar een algemeen, kwantisator-agnostisch toewijzingslaag. De auteurs stellen dat de AM/GM-verhouding van head-gevoeligheden dient als een praktische voorspeller voor wanneer toewijzing van gemengde precisie aanzienlijke winsten zal opleveren. Het werk suggereert dat het rate-distortion perspectief kan worden gegeneraliseerd naar resource-toewijzing in andere heterogene neurale netwerkcomponenten buiten KV-caches.