Technische Samenvatting: HeadQ: Model-Zichtbare Distorsie en Score-Ruimte Correctie voor KV-Cache Quantisatie

1. Probleemstelling

Huidige methoden voor KV-cache quantisatie optimaliseren voornamelijk voor herstel van opslagruimte, typisch door de Gemiddelde Kwantiseringsfout (MSE) tussen originele en gekwantiseerde sleutels en waarden te minimaliseren. De auteurs betogen echter dat dit doel niet overeenkomt met de manier waarop attention-mechanismen cache-gegevens werkelijk consumeren.

• Sleutels worden gelezen via logits (query-sleutel scores), wat betekent dat alleen specifieke lineaire projecties van de sleutels de attention-verdeling beïnvloeden.
• Waarden worden gelezen via attention-gewogen uitlezing ($\sum p_j v_j$).

Het artikel stelt dat aanhoudende cache-fouten gemeten moeten worden in model-zichtbare coördinaten in plaats van ruwe Euclidische opslagcoördinaten. Een verstoring kan een grote opslag-MSE hebben maar nul impact op attention (als deze ligt in een "softmax-null" richting), of een kleine opslag-MSE maar catastrofale attention-schade (als deze de "score-zichtbare" richting beïnvloedt). Dit verschil leidt tot suboptimale quantisatiestrategieën die bits verspillen aan irrelevante componenten terwijl ze falen in het beschermen van gedragskritieke componenten.

2. Theoretische Grondslag: Model-Zichtbare Geometrie

Het artikel vestigt een geometrisch raamwerk dat onderscheid maakt tussen opslagfout en gedragsfout:

2.1 Sleutel Score Ruimte (Fisher Geometrie)

Voor een vaste query $q$ induceert een sleutelverstoring $\delta k_j$ een scoreverstoring $\delta s_j = q^\top \delta k_j / \sqrt{d}$.

• Softmax-Null Richtingen: Een token-gedeelde translatie (het toevoegen van een constant vector $\mu$ aan alle sleutels) voegt een constante toe aan alle scores. Omdat softmax invariant is voor constante verschuivingen, is deze verstoring softmax-null en veroorzaakt het nul attention-schade, ondanks potentieel grote opslag-MSE.
• Score-Zichtbare Ruimte: Attention hangt alleen af van de gecentreerde actieve coëfficiënten van de sleutelverstoring ten opzichte van de query.
• Fisher Informatie: De lokale tweede-orde geometrie van de softmax-verdeling wordt vastgelegd door de Fisher Informatie Matrix. Het artikel leidt af dat de KL-divergentie tussen de originele en verstoord attention-verdeling lokaal evenredig is met een Fisher-gewogen variantie van scorefouten (Eq. 6), en niet met ruwe sleutel-MSE.

2.2 Waarde Uitlees Geometrie

Waarden worden verbruikt via een gewogen som. Voor vaste attention-gewichten $p$ is de uitleesfout $\Delta o = \sum p_j \Delta v_j$.

• Het artikel stelt een $A^2$-gewogen token-distorsie voor als de geschikte surrogaat voor waardenfout. In tegenstelling tot ruwe waarde-MSE, weegt deze metriek de distorsie van elk token met het kwadraat van zijn attention-kans ($p_j^2$), wat weerspiegelt dat fouten in sterk geattentioneerde tokens meer bijdragen aan de outputfout.

3. Methodologie: HeadQ

HeadQ is een quantisatiemethode aan de sleutelzijde die ontworpen is om de "score-zichtbare" residual te behouden in plaats van de ruwe sleutelvector.

Algorithmisch Overzicht

1. Kalibratie: Voor elke laag en kop wordt een laag-rang basis $U_h$ (rang $r$) geleerd via PCA vanuit kalibratie-queries (post-RoPE).
2. Cache Schrijven:
   • Een standaard low-bit quantizer $B_b$ genereert een basissleutel $\hat{k}^{base}_j$.
   • De residual $e_j = k_j - \hat{k}^{base}_j$ wordt berekend.
   • De residual wordt geprojecteerd op de geleerde basis: $z_j = Q_q(U_h^\top e_j)$, waarbij $Q_q$ een low-bit zijkant quantizer is.
   • De cache slaat de basissleutel en de zijkant-code $z_j$ op.
3. Cache Uitlezen (Inferentie):
   • De huidige query $q$ wordt gebruikt om de zijkant-code te dequantiseren.
   • Een additieve logit correctie wordt toegepast:
     $$\hat{s}_j = \frac{q^\top \hat{k}^{base}_j}{\sqrt{d}} + \frac{q^\top U_h \hat{z}_j}{\sqrt{d}}$$
   • Deze correctie richt zich direct op de scorefout, waarbij de noodzaak om de volledige sleutelvector te reconstrueren wordt omzeild.

Belangrijk Onderscheid: HeadQ slaat de huidige query niet op; het slaat een residual-code op in een kalibratie-geleerde basis die de runtime query kan decoderen tot een logit-correctie.

4. Belangrijkste Bijdragen

1. Formulering van Model-Zichtbare Distorsie: Het artikel definieert sleuteldistorsie als "score-zichtbare fout modulo constante verschuivingen" en waardedistorsie als "uitlees-zichtbare fout met een $A^2$ diagonaal surrogaat". Dit verklaart waarom ruwe MSE een slechte proxy is voor gedragsprestaties.
2. HeadQ Algorithmus: Een praktisch cache-schrijf/uitlees-algoritme dat een laag-rang residual zijkant-code opslaat in een query-basis en deze toepast als een additieve logit-correctie.
3. Valificeerbare Bewijsketen:
   • Metrische Superioriteit: Fisher/score-ruimte fout voorspelt attention KL-divergentie aanzienlijk beter (Spearman $\rho \approx 0.62\text{--}0.69$ winst) dan ruwe sleutel-MSE over zes modellen.
   • Tegenvoorbeelden: Gekoppelde experimenten tonen aan dat een quantizer met lagere opslag-MSE hogere attention KL en outputfout kan hebben, wat bewijst dat opslag-MSE geen voldoende proxy is.
   • Null-Ruimte Interventies: Het herstellen van een "softmax-null" gemiddelde component (die dense attention negeert) verhoogt de Fisher-schade scherp voor niet-equivariante vast-grid quantizers, wat aantoont dat opslag-coördinaat verspilling gedragsprestaties breekt.
   • Implementeerbaarheid: Query-PCA-bases geleerd uit kalibratiegegevens benaderen de actieve residual succesvol, terwijl willekeurige bases van gelijke begroting falen.
   • Directionaliteit: Het omkeren van het teken van de HeadQ-correctie ("verkeerd-teken" controle) veroorzaakt catastrofale falen, wat bevestigt dat de methode werkt door het specifieke score-object in de juiste richting te corrigeren.

5. Experimentele Resultaten

Experimenten werden uitgevoerd op WikiText-103 met zes modellen (GPT-2, Pythia, Qwen2, Qwen2.5, TinyLlama, Mistral).

• K-Enige Experimenten (Dense Waarden):
  • Op de meest beschadigde 2-bit rijen verwijdert HeadQ (rang 4) ongeveer 84–94% van de excessieve perplexiteit (PPL) vergeleken met een KIVI-achtige baseline voor modellen zoals Qwen2, Pythia en GPT-2.
  • TinyLlama toont materiële verbetering; Mistral presteert in dit protocol al bijna op dense-niveau.
• Volledige-KV Compositie (2-bit):
  • Het combineren van HeadQ (sleutelzijde) met een $A^2$-gewogen waarde-beleid verbetert de prestaties over alle zes modellen bij 2-bit precisie.
• Anomalie Analyse:
  • Pythia-160M vertoonde een zwakkere correlatie tussen Fisher en KL in specifieke lage-entropie "route-flip" regimes (waar kleine verstoringen de top-1 token omdraaien). Deze anomalie was gelokaliseerd tot een specifieke familiestafgrens en verzwakte in grotere Pythia checkpoints (410M, 1B).

6. Betekenis en Claims

Het artikel beweert dat operator matching de kritieke ontwerpleer is voor KV-cache quantisatie:

• Sleutels zijn score-objecten; compressie moet query-zichtbare residuals behouden en vermijden om resolutie te verspillen aan softmax-null verschuivingen.
• Waarden zijn uitlees-objecten; compressie moet prioriteit geven aan $A^2$-gewogen fout.

De auteurs stellen expliciet dat dit werk mechanistisch en diagnostisch is, en geen claim van een productieklaar serversysteem. Zij claimen niet geoptimaliseerde gepakte kernen, latentie of doorvoer te hebben. De betekenis ligt in:

1. Het corrigeren van het theoretische doel van KV-cache quantisatie van opslagherstel naar model-zichtbare distorsie.
2. Het bieden van een implementeerbare methode (HeadQ) die aanzienlijke PPL-herstel bereikt in low-bit regimes door te richten op de juiste foutgeometrie.
3. Het demonstreren dat standaard opslagmetrieken (MSE) gedragsmatig misleidend kunnen zijn, wat een verschuiving naar score-ruimte proxies noodzakelijk maakt.

Het artikel concludeert dat de volgende stap is om dit operator-gematchte perspectief uit te breiden naar een volledig geoptimaliseerde volledige-KV codec, in plaats van opslag-MSE als standaarddoel te behandelen.