QK-Normed MLA: QK normalization without full key caching

Technische Samenvatting: QK-Normed MLA

Probleemstelling
Moderne grote taalmodellen worden geconfronteerd met twee concurrerende drukfactoren: de noodzaak voor geheugenefficiënt decoderen in scenario's met een lange context en de noodzaak voor trainingsstabiliteit.

• Geheugenefficiëntie: Standaard Multi-Head Attention (MHA) cache volledige key- en value-vectoren voor elke token, wat hoge geheugen- en bandbreedtekosten met zich meebrengt. Multi-Head Latent Attention (MLA) adresseert dit door laag-dimensionale latente toestanden ($c_t$) te cachen in plaats van volledige keys, waarbij het volledige key-effect wordt hersteld via een up-projectie ($W^K_g$) tijdens het decoderen.
• Trainingsstabiliteit: Ongecontroleerde schalen in queries en keys kunnen leiden tot verzadiging van de aandacht (attention saturation) en instabiliteit tijdens de training. Query-Key (QK) normalisatie (specifiek RMSNorm) stabiliseert de training door vectoren te normaliseren vóór de dot product. Echter, standaard post-projectie QK RMSNorm lijkt incompatibel met MLA omdat het schijnbaar de volledig geprojecteerde key ($k_{t,g} = c_t W^K_g$) voor elke gecachte token vereist om de normalisatiestatistiek te berekenen. Een naïeve implementatie zou de volledige keys moeten materialiseren, wat de geheugenvoordelen van MLA tenietdoet.

Methodologie
Het artikel stelt QK-Normed MLA voor, een formulering die post-up-projectie QK RMSNorm integreert in het MLA latente decodepad zonder volledige keys te materialiseren. Het kerninzicht is de algebraïsche decompositie van RMSNorm:
$$\text{RMSNorm}(z) = z \odot \frac{\gamma}{S(z)}$$
waarbij $\gamma$ een statisch affien gewicht is en $S(z)$ een dynamische scalaire RMS-statistiek.

De auteurs leiden een "geabsorbeerde" formulering af die deze componenten scheidt:

1. Statische Gewichtabsorptie: Het statische key-zijde gewicht $\gamma_k$ wordt geabsorbeerd in de query-zijde projectiematrix. In plaats van de query $q$ te projecteren en vervolgens de key te normaliseren, wordt de query vooraf vermenigvuldigd met $\gamma_k$ (en het statische gewicht $\gamma_q$) vóór de latente dot product.
2. Dynamische Scalaire Caching: De dynamische statistiek $S(z)$ reduceert tot een enkele inverse-RMS scalaire ($\alpha_{t,g} = 1/S_{t,g}$) per token en per KV-groep.
3. Cache Layout: De decode-cache behoudt de oorspronkelijke laag-dimensionale latente toestand $c_t$ en voegt een kleine scalaire cache $A^k_{cache}$ toe die de inverse-RMS waarden bevat. Het slaat niet de volledige genormaliseerde keys op.
4. Blockwise Normalisatie: Om efficiëntie te behouden, normaliseert de methode de "content" en "RoPE" componenten afzonderlijk. Het content-pad gebruikt de geabsorbeerde latente formulering, terwijl het RoPE-pad standaard gematerialiseerde vectoren gebruikt. Dit voorkomt synchronisatiepunten die de implementatie van gefuseerde kernels zouden compliceren.

Tijdens het decoderen worden de attention logits berekend als:
$$\ell_{i,t,h} = (\tilde{q}_{i,h} \cdot c_t) \cdot \alpha^q_{i,h} \cdot \alpha^k_{t,g(h)}$$
waarbij $\tilde{q}$ de geabsorbeerde gewichten bevat, en de scalaire factoren worden toegepast als een rij-kolom schalingstap vóór de softmax.

Belangrijkste Bijdragen

1. Exacte Formulering: Het artikel leidt een algebraïsch exact exacte formulering af van post-up-projectie QK RMSNorm die compatibel is met het MLA latente decodepad.
2. Efficiënt Cache Design: Het stelt een cache layout voor die de latente KV-cache als de dominante toestand behoudt, waarbij slechts één scalaire key RMS-statistiek per token en KV-groep wordt toegevoegd (bijv. een toename van de cachegrootte met ~1,56% in geteste configuraties).
3. Empirische Validatie: De auteurs demonstreren dat deze aanpak consequent een lagere training loss en betere downstream nauwkeurigheid bereikt vergeleken met QK clipping, terwijl er een verwaarloosbare latentie overhead wordt geïntroduceerd.

Resultaten
Experimenten werden uitgevoerd op 400M-parameter MLA-modellen getraind voor maximaal 100B tokens op H800 GPU's:

• Trainingsstabiliteit: QK-Normed MLA behaalde consistent een lagere training loss dan QK-Clip gedurende de volledige 100B-token run. In high-learning-rate stress-tests ($2 \times 10^{-2}$) divergeerde QK-Clip met gradiënt NaNs, terwijl QK-Normed MLA stabiel bleef.
• Downstream Prestaties: Op een reeks downstream taken verbeterde QK-Normed MLA de nauwkeurigheid op 7 van de 8 taken, waarbij de gemiddelde score steeg van 44,75 naar 46,33, en verminderde de LAMBADA perplexiteit van 16,28 naar 14,18.
• Decode Efficiëntie: Benchmarks op contexten variërend van 4k tot 256k tokens toonden aan dat de methode minder dan 2% latentie overhead toevoegt (gemiddeld 1,40%). De scalaire cache vergrootte de totale cachegrootte met slechts ~2,8% per laag.

Betekenis en Claims
Het artikel claimt dat de ogenschijnlijke incompatibiliteit tussen QK normalisatie en MLA een implementatiefout is in plaats van een architecturale beperking. Door aan te tonen dat de statische gewichten geabsorbeerd kunnen worden en de dynamische statistieken gereduceerd kunnen worden tot scalairen, maken de auteurs QK normalisatie mogelijk als een praktische stabilisatiemechanisme voor MLA-modellen zonder het geheugenefficiëntievoordeel van latente decoding op te offeren.

De auteurs blijven bescheiden over de reikwijdte van hun claims:

• De resultaten zijn gebaseerd op 400M-parameter modellen; zij claimen geen frontier-schaal gedrag voor miljarden-parameter modellen vast te stellen, hoewel de algebraïsche equivalentie onafhankelijk van schaal geldt.
• De logit-scale analyse wordt gepresenteerd als een mechanisme-niveau diagnostiek in plaats van een volledige optimalisatietheorie.
• De praktische snelheid hangt af van kernel-integratie (het fuseren van scalaire factoren in het attention pad) om onnodig geheugentransport te voorkomen.