🤖 machine learning

QK-Normed MLA: QK normalization without full key caching

Cet article démontre que la normalisation QK peut être intégrée de manière transparente à l'attention latente multi-têtes (Multi-head Latent Attention - MLA) en reformulant mathématiquement l'opération pour éviter le cache complet des clés, atteignant ainsi une meilleure stabilité d'entraînement et une précision accrue sur les tâches en aval avec un surcoût de latence négligeable.

Auteurs originaux : Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun

Publié 2026-06-16

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous dirigez une bibliothèque immense où un bibliothécaire (l'IA) doit trouver le livre parfait (l'information) pour répondre à une question. Pour y parvenir efficacement, le bibliothécaire utilise deux astuces principales :

L'astuce du « Résumé Latent » (MLA) : Au lieu de garder une copie géante et lourde de chaque livre jamais écrit dans la pièce (ce qui prendrait trop de place), le bibliothécaire garde une minuscule « fiche de résumé » compressée pour chaque livre. Lorsqu'une question arrive, le bibliothécaire développe la fiche de résumé à la volée pour trouver la réponse. Cela permet d'économiser énormément d'espace.
L'astuce du « Contrôle du Volume » (Normalisation QK) : Parfois, le bibliothécaire s'emballe. Le « volume » de la connexion entre une question et un livre devient si fort qu'il couvre tout le reste, ce qui provoque une panique ou des erreurs chez le bibliothécaire. Pour corriger cela, nous mettons habituellement un « bouton de volume » sur les livres pour maintenir le niveau sonore stable.

Le Problème :
Pendant longtemps, ces deux astuces ne s'entendaient pas. L'astuce du « Contrôle du Volume » semblait nécessiter que le bibliothécaire ait la copie complète et lourde du livre en main pour ajuster le volume. Or, l'astuce du « Résumé Latent » était spécifiquement conçue pour éviter d'avoir ces copies lourdes. Il semblait qu'il faille choisir : soit économiser de l'espace (Résumé Latent), soit maintenir le volume stable (Contrôle du Volume), mais pas les deux.

La Solution (QK-Normed MLA) :
Les auteurs de cet article ont découvert un hack mathématique ingénieux qui permet d'avoir le beurre et l'argent du beurre. Ils ont réalisé que le bouton du « Contrôle du Volume » possède en fait deux parties :

Un réglage statique (un cadran fixe qui ne change jamais).
Une lecture dynamique (un coup d'œil rapide pour voir à quel point le livre actuel est bruyant).

Ils ont démontré que :

Le réglage statique peut être déplacé avant même que le bibliothécaire ne regarde la fiche de résumé. Il est intégré directement dans la question elle-même.
La lecture dynamique est si simple qu'au lieu d'avoir besoin du livre entier, le bibliothécaire a seulement besoin de noter un seul petit nombre (un scalaire unique) pour chaque fiche de résumé.

Le Résultat :
Désormais, le bibliothécaire peut continuer à utiliser les minuscules « fiches de résumé » (économisant ainsi un espace massif) tout en ayant un contrôle parfait du volume. Il n'a plus besoin de porter les livres lourds.

Ce que les expériences ont montré :
L'équipe a testé cette nouvelle méthode sur un modèle de 400 millions de paramètres (une IA de taille moyenne) entraîné sur une quantité massive de texte (100 milliards de mots).

Meilleur apprentissage : Le modèle doté de ce nouveau « Contrôle du Volume » a appris plus rapidement et a commis moins d'erreurs que les modèles utilisant une méthode plus grossière appelée « clipping » (qui consiste simplement à crier « STOP ! » quand les choses deviennent trop fortes).
Meilleures réponses : Lors de tests sur diverses tâches, la nouvelle méthode a donné des réponses plus précises.
Vitesse : Le seul inconvénient est la vérification de ce petit nombre, ce qui a ajouté moins de 2 % au temps nécessaire pour lire une longue histoire. C'est si peu que cela est à peine perceptible.

En bref : L'article prouve que vous n'avez pas à sacrifier l'efficacité de la mémoire pour maintenir la stabilité de votre IA. Vous pouvez utiliser un petit tour mathématique ingénieux pour contrôler le « volume » sans jamais avoir besoin de stocker les données lourdes et volumineuses que vous essayiez justement d'éviter.

Résumé Technique : QK-Normed MLA

Énoncé du Problème
Les modèles de langage de grande taille modernes sont confrontés à deux pressions divergentes : la nécessité d'un décodage efficace en termes de mémoire pour les scénarios à contexte long et le besoin de stabilité de l'entraînement.

Efficacité Mémoire : L'attention multi-têtes standard (MHA) met en cache les vecteurs clés et valeurs complets pour chaque jeton, ce qui crée des coûts élevés en mémoire et en bande passante. L'attention latente multi-têtes (MLA) répond à cela en mettant en cache des états latents de faible dimension ( $c_t$ ) plutôt que des clés complètes, récupérant l'effet de la clé complète via une up-projection ( $W^K_g$ ) lors du décodage.
Stabilité de l'Entraînement : Des échelles incontrôlées dans les requêtes (queries) et les clés (keys) peuvent entraîner une saturation de l'attention et des instabilités d'entraînement. La normalisation QK (plus précisément RMSNorm) stabilise l'entraînement en normalisant les vecteurs avant le produit scalaire. Cependant, la norme post-projection QK RMSNorm semble incompatible avec MLA car elle semble nécessiter la clé entièrement projetée ( $k_{t,g} = c_t W^K_g$ ) pour chaque jeton mis en cache afin de calculer la statistique de normalisation. Une implémentation naïve nécessiterait de matérialiser les clés complètes, annulant ainsi les avantages de MLA en matière de mémoire.

Méthodologie
Le papier propose QK-Normed MLA, une formulation qui intègre la QK RMSNorm post-up-projection dans le chemin de décodage latent de MLA sans matérialiser les clés complètes. L'idée centrale est la décomposition algébrique de RMSNorm :
$\text{RMSNorm}(z) = z \odot \frac{\gamma}{S(z)}$
où $\gamma$ est un poids affine statique et $S(z)$ est une statistique RMS dynamique.

Les auteurs dérivent une formulation « absorbée » qui sépare ces composants :

Absorption du Poids Statique : Le poids statique côté clé $\gamma_k$ est absorbé dans la matrice de projection côté requête. Au lieu de projeter la requête $q$ puis de normaliser la clé, la requête est pré-multipliée par $\gamma_k$ (et le poids statique $\gamma_q$ ) avant le produit scalaire latent.
Mise en Cache de la Statistique Dynamique : La statistique dynamique $S(z)$ se réduit à un seul scalaire d'inverse-RMS ( $\alpha_{t,g} = 1/S_{t,g}$ ) par jeton et par groupe KV.
Configuration du Cache : Le cache de décodage conserve l'état latent de faible dimension original $c_t$ et ajoute un petit cache scalaire $A^k_{cache}$ contenant les valeurs d'inverse-RMS. Il ne stocke pas les clés normalisées complètes.
Normalisation par Blocs : Pour maintenir l'efficacité, la méthode normalise séparément les composantes « contenu » et « RoPE ». Le chemin de contenu utilise la formulation latente absorbée, tandis que le chemin RoPE utilise les vecteurs matérialisés standards. Cela évite les points de synchronisation qui compliqueraient les implémentations de noyaux (kernels) fusionnés.

Pendant le décodage, les logits d'attention sont calculés comme suit :
$\ell_{i,t,h} = (\tilde{q}_{i,h} \cdot c_t) \cdot \alpha^q_{i,h} \cdot \alpha^k_{t,g(h)}$
où $\tilde{q}$ inclut les poids absorbés, et les facteurs scalaires sont appliqués comme une étape de mise à l'échelle ligne-colonne avant le softmax.

Contributions Clés

Formulation Exacte : Le papier dérive une formulation algébriquement exacte de la QK RMSNorm post-up-projection qui est compatible avec le chemin de décodage latent de MLA.
Conception de Cache Efficace : Il propose une configuration de cache qui préserve le cache KV latent comme état dominant, ajoutant seulement un scalaire de statistique RMS de clé par jeton et par groupe KV (par exemple, augmentant la taille du cache d'environ 1,56 % dans les configurations testées).
Validation Empirique : Les auteurs démontrent que cette approche permet d'obtenir une perte d'entraînement plus faible et une meilleure précision sur les tâches en aval par rapport au QK-Clipping, tout en introduisant un surcoût de latence négligeable.

Résultats
Des expériences ont été menées sur des modèles MLA de 400 millions de paramètres entraînés jusqu'à 100 milliards de jetons sur des GPU H800 :

Stabilité de l'Entraînement : QK-Normed MLA a obtenu une perte d'entraînement systématiquement plus faible que QK-Clip tout au long de la course de 100 milliards de jetons. Lors de tests de stress à taux d'apprentissage élevé ( $2 \times 10^{-2}$ ), QK-Clip a divergé avec des gradients NaN, tandis que QK-Normed MLA est resté stable.
Performance en Aval : Sur une suite de tâches en aval, QK-Normed MLA a amélioré la précision sur 7 tâches sur 8, faisant passer le score moyen de 44,75 à 46,33, et a réduit la perplexité de LAMBADA de 16,28 à 14,18.
Efficacité du Décodage : Les benchmarks sur des contextes allant de 4k à 256k jetons ont montré que la méthode ajoute moins de 2 % de latence (moyenne de 1,40 %). Le cache scalaire n'a augmenté la taille totale du cache que d'environ 2,8 % par couche.

Signification et Revendications
Le papier affirme que l'incompatibilité apparente entre la normalisation QK et MLA est un artefact d'implémentation plutôt qu'une contrainte architecturale. En montrant que les poids statiques peuvent être absorbés et que les statistiques dynamiques peuvent être réduites à des scalaires, les auteurs permettent la normalisation QK comme un mécanisme de stabilisation pratique pour les modèles MLA sans sacrifier l'efficacité mémoire du décodage latent.

Les auteurs restent modestes quant à la portée de leurs affirmations :

Les résultats sont basés sur des modèles de 400 millions de paramètres ; ils ne prétendent pas établir un comportement à l'échelle des frontières pour des modèles de plusieurs milliards de paramètres, bien que l'équivalence algébrique soit indépendante de l'échelle.
L'analyse de l'échelle des logits est présentée comme un diagnostic de niveau mécanisme plutôt que comme une théorie d'optimisation complète.
La vitesse pratique dépend de l'intégration des noyaux (fusion des facteurs scalaires dans le chemin d'attention) pour éviter les transferts de mémoire inutiles.

Articles similaires