QK-Normed MLA: QK normalization without full key caching

Sintesi Tecnica: QK-Normed MLA

Definizione del Problema
I moderni modelli linguistici di grandi dimensioni affrontano due pressioni contrastanti: la necessità di un decoding efficiente dal punto di vista della memoria negli scenari a lungo contesto e la necessità di stabilità nell'addestramento.

• Efficienza della Memoria: La standard Multi-Head Attention (MHA) memorizza interi vettori di chiave (key) e valore (value) per ogni token, creando elevati costi di memoria e di banda. La Multi-Head Latent Attention (MLA) affronta questo problema memorizzando stati latenti a bassa dimensionalità ($c_t$) invece di chiavi complete, recuperando l'effetto della chiave completa tramite una up-proiezione ($W^K_g$) durante il decoding.
• Stabilità dell'Addestramento: Scale incontrollate in query e chiavi possono portare alla saturazione dell'attenzione e a instabilità nell'addestramento. La normalizzazione QK (specificamente RMSNorm post-proiezione) stabilizza l'addestramento normalizzando i vettori prima del prodotto scalare. Tuttavia, la standard post-projection QK RMSNorm appare incompatibile con MLA perché sembrerebbe richiedere la chiave completamente proiettata ($k_{t,g} = c_t W^K_g$) per ogni token memorizzato al fine di calcolare la statistica di normalizzazione. Un'implementazione ingenua richiederebbe la materializzazione di chiavi complete, annullando i vantaggi di memoria di MLA.

Metodologia
Il documento propone QK-Normed MLA, una formulazione che integra la QK RMSNorm post-up-projection nel percorso di decoding latente di MLA senza materializzare chiavi complete. L'intuizione centrale è la decomposizione algebrica di RMSNorm:
$$\text{RMSNorm}(z) = z \odot \frac{\gamma}{S(z)}$$
dove $\gamma$ è un peso affine statico e $S(z)$ è una statistica RMS dinamica scalare.

Gli autori derivano una formulazione "assorbita" che separa questi componenti:

1. Assorbimento del Peso Statico: Il peso statico lato chiave $\gamma_k$ viene assorbito nella matrice di proiezione lato query. Invece di proiettare la query $q$ e poi normalizzare la chiave, la query viene pre-moltiplicata per $\gamma_k$ (e il peso statico $\gamma_q$) prima del prodotto scalare latente.
2. Caching della Statistica Scalare Dinamica: La statistica dinamica $S(z)$ si riduce a un singolo scalare inverso-RMS ($\alpha_{t,g} = 1/S_{t,g}$) per token e per gruppo KV.
3. Layout della Cache: La cache di decoding conserva lo stato latente originale a bassa dimensionalità $c_t$ e aggiunge una piccola cache scalare $A^k_{cache}$ contenente i valori inversi-RMS. Non memorizza le chiavi normalizzate complete.
4. Normalizzazione Blockwise: Per mantenere l'efficienza, il metodo normalizza separatamente i componenti "content" e "RoPE". Il percorso content utilizza la formulazione latente assorbita, mentre il percorso RoPE utilizza vettori standard materializzati. Ciò evita punti di sincronizzazione che complicherebbero le implementazioni di kernel fusi.

Durante il decoding, i logit dell'attenzione sono calcolati come:
$$\ell_{i,t,h} = (\tilde{q}_{i,h} \cdot c_t) \cdot \alpha^q_{i,h} \cdot \alpha^k_{t,g(h)}$$
dove $\tilde{q}$ include i pesi assorbiti, e i fattori scalari sono applicati come uno step di scaling riga-colonna prima della softmax.

Contributi Chiave

1. Formulazione Esatta: Il documento deriva una formulazione algebricamente esatta di post-up-projection QK RMSNorm che è compatibile con il percorso di decoding latente di MLA.
2. Design Efficiente della Cache: Propone un layout della cache che preserva la cache KV latente come stato dominante, aggiungendo solo una statistica RMS scalare per token e gruppo KV (ad esempio, aumentando la dimensione della cache di circa l'1,56% nelle configurazioni testate).
3. Validazione Empirica: Gli autori dimostrano che questo approccio ottiene una perdita di addestramento costantemente inferiore e una migliore accuratezza nei task downstream rispetto a QK clipping, introducendo un overhead di latenza trascurabile.

Risultati
Gli esperimenti sono stati condotti su modelli MLA da 400 milioni di parametri addestrati fino a 100 miliardi di token su GPU H800:

• Stabilità dell'Addestramento: QK-Normed MLA ha ottenuto una perdita di addestramento costantemente inferiore rispetto a QK-Clip durante l'intera esecuzione di 100 miliardi di token. Nei test di stress con alto tasso di apprendimento ($2 \times 10^{-2}$), QK-Clip è divergente con gradienti NaN, mentre QK-Normed MLA è rimasto stabile.
• Performance Downstream: Su una suite di task downstream, QK-Normed MLA ha migliorato l'accuratezza in 7 task su 8, alzando il punteggio medio da 44,75 a 46,33, e ha ridotto la perplexity di LAMBADA da 16,28 a 14,18.
• Efficienza di Decoding: I benchmark su contesti da 4k a 256k token hanno mostrato che il metodo aggiunge meno del 2% di overhead di latenza (media 1,40%). La cache scalare ha aumentato la dimensione totale della cache di solo ~2,8% per layer.

Significatività e Rivendicazioni
Il documento sostiene che l'apparente incompatibilità tra la normalizzazione QK e MLA sia un artefatto di implementazione piuttosto che un vincolo architettonico. Dimostrando che i pesi statici possono essere assorbiti e le statistiche dinamiche ridotte a scalari, gli autori rendono la normalizzazione QK un meccanismo pratico di stabilizzazione per i modelli MLA senza sacrificare l'efficienza della memoria del decoding latente.

Gli autori rimangono modesti riguardo all'ambito delle loro rivendicazioni:

• I risultati si basano su modelli da 400 milioni di parametri; non pretendono di stabilire un comportamento su scala frontier per modelli da miliardi di parametri, sebbene l'equivalenza algebrica valga indipendentemente dalla scala.
• L'analisi della scala dei logit è presentata come un diagnostico a livello di meccanismo piuttosto che come una teoria di ottimizzazione completa.
• La velocità pratica dipende dall'integrazione dei kernel (fondendo i fattori scalari nel percorso di attenzione) per evitare traffico di memoria non necessario.