💻 computer science

Multi-Head Attention Residuals

Ce document introduit les Multi-Head Attention Residuals (MHAR), une modification architecturale sans paramètres qui permet à différents sous-espaces de caractéristiques d'accéder indépendamment à l'historique de profondeur via plusieurs têtes de routage, résolvant ainsi le désaccord de sous-espace et améliorant considérablement la perte de validation ainsi que les performances sur les tâches en aval à travers diverses échelles de modèles par rapport aux Transformers standards.

Auteurs originaux : Cheng Luo, Zefan Cai, Junjie Hu

Publié 2026-07-31

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Cheng Luo, Zefan Cai, Junjie Hu

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 construisez une immense bibliothèque à plusieurs étages, où chaque nouvel étage ajoute une nouvelle salle de savoir. Dans le monde de l'intelligence artificielle, ces « bibliothèques » sont appelées des Transformers, et elles sont les cerveaux derrière de nombreux outils intelligents que nous utilisons aujourd'hui. Pour donner du sens à une histoire ou à un problème mathématique, la bibliothèque doit se souvenir de ce qui s'est passé au premier étage tout en travaillant sur le dixième. Traditionnellement, ces bibliothèques possèdent un seul couloir étroit reliant chaque étage. Lorsqu'une pièce au dixième étage a besoin de regarder en arrière, elle ne peut jeter un coup d'œil qu'à la pièce située immédiatement en dessous d'elle. Elle doit faire confiance au fait que la pièce du dessous a parfaitement résumé tout ce qui s'est passé du premier étage jusqu'à ce point. C'est comme essayer de se souvenir de tout un film en demandant seulement à la personne assise juste devant vous ce qu'il s'est passé ; si elle a manqué un détail, vous le manquez aussi.

Pour corrão cela, des chercheurs ont récemment testé une nouvelle idée appelée « Attention Residuals » (résidus d'attention). Au lieu de simplement regarder la pièce du dessous, la bibliothèque a doté chaque pièce d'un télescope magique capable de regarder n'importe quel étage précédent. Cependant, il y avait un piège : le télescope n'avait qu'un seul objectif. Chaque détail de la pièce — chaque couleur, chaque son, chaque nombre — devait passer par ce même objectif unique pour décider sur quel étage se concentrer. C'était comme forcer un groupe d'amis aux goûts totalement différents (l'un adore les films d'action, un autre les documentaires, un troisième les dessins animés) à tous s'entendre sur une seule et même chaîne à regarder ensemble. Ils devaient faire des compromis, et souvent, ce compromis signifiait que personne ne voyait vraiment ce dont il avait besoin.

Ce document présente une mise à jour ingénieuse appelée Multi-Head Attention Residuals (MHAR). Les chercheurs ont réalisé que, tout comme un cinéma possède plusieurs écrans pour diffuser différents genres en même temps, le télescope de la bibliothèque devrait avoir plusieurs objectifs. Ils ont divisé cet objectif unique et partagé en plusieurs petits objectifs indépendants. Désormais, la partie « film d'action » de la pièce peut regarder les étages riches en action, tandis que la partie « documentaire » regarde les étages historiques, le tout sans se disputer ou faire de compromis.

L'équipe a testé cette idée en construisant des bibliothèques de différentes tailles, allant de petits modèles de 100 millions de paramètres jusqu'à un modèle massif d'un milliard de paramètres. Ils ont constaté que pour les petites bibliothèques, l'ancien système à objectif unique fonctionnait bien. Mais à mesure que les bibliothèques devenaient plus grandes et plus complexes, l'objectif unique devenait un goulot d'étranglement, rendant l'IA moins performante que le système de couloir standard. Le nouveau système à objectifs multiples, en revanche, devenait de plus en plus performant à mesure que les bibliothèques grandissaient. À la plus grande échelle, la nouvelle méthode a réduit la confusion de l'IA (mesurée par la perte de validation) d'environ 0,063 par rapport au modèle standard, tandis que l'ancienne méthode à objectif unique aggravait les choses de 0,140.

Les chercheurs ont également montré que cette mise à jour ne nécessitait pas de construire un plus gros télescope ni d'ajouter du poids à la bibliothèque ; il s'agissait simplement d'une manière plus intelligente de disposer le verre existant. Ils ont même construit des outils spéciaux, ultra-rapides, pour s'assurer que la bibliothèque ne soit pas ralentie par l'observation supplémentaire, maintenant la vitesse presque aussi rapide que le design original. Enfin, ils ont prouvé que cette astuce fonctionne même sur une immense bibliothèque pré-construite de 8 milliards de paramètres, renforçant sa capacité à résoudre des problèmes mathématiques et à répondre à des questions scientifiques complexes sans avoir besoin d'être reconstruite de zéro. La découverte fondamentale est qu'à mesure que les modèles d'IA deviennent plus larges et plus complexes, ils cessent de s'entendre sur ce qu'ils doivent mémoriser, et donner plusieurs façons indépendantes de regarder en arrière est la clé pour libérer leur plein potentiel.

Résumé Technique : Multi-Head Attention Residuals (MHAR)

1. Énoncé du Problème

Les Transformers standards propagent l'information à travers la profondeur via un flux résiduel additif unique, où chaque sous-couche ne lit que l'état immédiatement précédent ( $h_\ell = h_{\ell-1} + f_\ell(h_{\ell-1})$ ). Bien que les caractéristiques utiles calculées dans les couches précédentes restent présentes dans la somme courante, elles deviennent individuellement inaccessibles.

Les Attention Residuals (Kimi, 2025) relaxent cette contrainte en traitant toutes les sorties des sous-couches antérieures comme une mémoire, permettant à chaque sous-couche de récupérer une combinaison apprise et pondérée par softmax de l'ensemble de l'historique de profondeur. Cependant, ce mécanisme souffre d'un goulot d'étranglement structurel : il utilise une requête de routage unique partagée ( $q \in \mathbb{R}^d$ ) pour évaluer l'historique de profondeur. Par conséquent, les $d$ coordonnées de caractéristiques sont forcées de lire les mêmes couches dans les mêmes proportions via une distribution $\alpha$ unique.

L'article soutient que cela crée un « compromis forcé ». À mesure que la largeur du modèle augmente, les différents sous-espaces de caractéristiques développent probablement des préférences divergentes concernant les couches historiques les plus pertinentes. Une requête unique ne peut pas satisfaire simultanément ces préférences divergentes, ce qui fait s'effondrer les distinctions de préférences de profondeur en une lecture unique. Le coût de ce compromis croît avec la largeur du modèle, rendant potentiellement le routage à requête unique préjudiciable à grande échelle.

2. Méthodologie : Multi-Head Attention Residuals (MHAR)

Les auteurs proposent les Multi-Head Attention Residuals (MHAR), une généralisation stricte des attention residuals qui élimine le compromis forcé sans ajouter de paramètres ou de calcul significatif.

Mécanisme Central

Au lieu d'une requête unique $q$ , MHAR reforme la requête de routage en $H$ têtes par sous-espace, chacune possédant son propre softmax indépendant sur l'historique de profondeur.

Reformage (Reshaping) : Le vecteur de requête unique $q \in \mathbb{R}^d$ est remodelé en un tenseur de forme $(H, d/H)$ .
Routage Indépendant : Pour chaque tête $h$ , le modèle calcule une distribution de profondeur distincte $\alpha^{(h)}$ basée sur la tranche du vecteur source correspondant à cette tête.
Lecture Bloc-Diagonale : Le mélange récupéré est bloc-diagonal ; chaque sous-espace lit l'historique de profondeur de manière indépendante.
Récupération : Fixer $H=1$ permet de retrouver exactement les single-head attention residuals originaux.

Détails d'Implémentation

Efficacité des Paramètres : La transformation est un remodelage, ajoutant zéro paramètre par rapport à la ligne de base à tête unique.
Efficacité de Calcul : Le surcoût est négligeable, impliquant seulement $H$ petits softmaxes sur l'axe de la profondeur (longueur $\le 2L+1$ ) au lieu d'un seul grand.
Optimisation Système : Reconnaissant que le routage de profondeur est limité par la mémoire (nécessitant un trafic de $O(L^2 B T d)$ ), les auteurs fournissent des noyaux Triton fusionnés. Ces noyaux effectuent un softmax en ligne avec un RMSNorm fusionné, éliminant la matérialisation intermédiaire des clés et des sources empilées. Cela réduit le temps de routage de $2\times$ par rapport à torch.compile et de $6,4\times$ par rapport aux implémentations de référence en mode eager. Cela porte le débit d'entraînement à $0,55\text{--}0,88\times$ de la ligne de base.

Réglage des Hyperparamètres

Le nombre de têtes de routage $H$ est fixé égal au nombre de têtes KV du modèle. Les auteurs constatent que cela constitue une valeur par défaut quasi optimale et sans hyperparamètre, alignant la granularité du routage sur la granularité de consommation du modèle.

3. Contributions Clés

Architecture MHAR : Une modification sans paramètre qui divise la requête de routage en têtes indépendantes, supprimant le compromis forcé du routage de profondeur à requête unique.
Noyaux Efficaces : Des noyaux Triton fusionnés qui rendent le routage de profondeur pratique à grande échelle, réduisant considérablement la surcharge mémoire et la latence.
Conversion Préservant l'Identité : Une méthode pour greffer MHAR sur des modèles pré-entraînés en utilisant des « delta attention residuals » (ajoutant le terme routé au flux existant plutôt que de le remplacer), permettant une phase de milieu d'entraînement sans pics de perte.
Insight Mécaniste : Le sondage direct des requêtes entraînées confirme que le désaccord des sous-espaces croît avec la largeur et est le principal moteur de la dégradation du routage à tête unique.

4. Résultats Expérimentaux

Pré-entraînement (From Scratch)

Entraîné sur FineWeb-Edu à trois échelles (100M, 350M, 1B) :

Perte de Validation : MHA améliore systématiquement la perte de validation par rapport aux Transformers standards et aux single-head attention residuals.
- 100M : $-0,049$ nats par rapport à la ligne de base.
- 350M : $-0,080$ nats par rapport à la ligne de base.
- 1B : $-0,063$ nats par rapport à la ligne de base.
Comportement de Mise à l'Échelle (Scaling) : L'avantage de MHA sur le routage à tête unique s'élargit avec l'échelle. Alors que le routage à tête unique aide à 100M ($-0,039$), il devient préjudiciable à 350M ( $+0,055$ ) et 1B ( $+0,140$ ).
Équivalence de Calcul : MHA atteint un gain équivalent en calcul de $1,3\text{--}1,5\times$ par rapport à la ligne de base.

Milieu d'Entraînement (8B)

Utilisant un contrôle apparié sur un modèle de 8B (Marin-8B) poli sur un corpus orienté vers les capacités :

Préservation de l'Identité : La conversion n'introduit aucun choc d'optimisation ; les courbes de perte coïncident avec la ligne de base pour les 100 premières étapes.
Gains en Aval : MHA produit des améliorations statistiquement significatives sur les tâches de raisonnement :
- GSM8K : $+3,2$ points (McNemar apparié $p=0,004$ ).
- GPQA : $+3,1$ points (McNemar apparié $p=0,038$ ).
- Les gains sur MMLU, MATH et les tâches de code n'étaient pas statistiquement significatifs dans cette configuration spécifique.

Études d'Ablation

Nombre de Têtes : Le nombre de têtes optimal $H$ suit le nombre de têtes KV. Dans les régimes sous-entraînés, plus de têtes aident toujours ; dans les régimes convergés, $H=KV$ est le « coude » de la courbe.
Largeur vs Profondeur : Le sondage des requêtes entraînées révèle que le « désaccord de largeur » (préférences de routage divergentes entre les sous-espaces) croît de manière monotone avec la largeur du modèle, validant l'hypothèse du compromis forcé.
Alignement des Têtes : Aligner les têtes de routage sur les groupes KV de l'attention ne produit aucun bénéfice mesurable par rapport au routage de sous-espace arbitraire, suggérant que le bénéfice provient de la multiplicité des têtes plutôt que de leur alignement spécifique aux structures d'attention.

5. Signification et Revendications

L'article affirme que les Multi-Head Attention Residuals résolvent une limitation architecturale fondamentale dans la manière dont les Transformers accèdent à l'historique de profondeur. En permettant à différents sous-espaces de caractéristiques d'être attentifs à différentes couches de manière indépendante, MHA passe efficacement à l'échelle là où le routage à requête unique échoue.

Les auteurs soulignent que :

L'amélioration n'est pas due à l'ajout de paramètres ou de calcul, mais plutôt à l'élimination d'un goulot d'étranglement structurel.
La méthode est pratique à grande échelle grâce à des noyaux fusionnés qui atténuent la nature limitée par la mémoire du routage de profondeur.
L'approche est polyvalente, fonctionnant dès le départ et via un entraînement de milieu de parcours préservant l'identité.
La nécessité du routage multi-têtes est une conséquence directe de la largeur du modèle ; à mesure que les modèles croissent, le « compromis forcé » d'une requête unique devient un handicap, faisant de MHA un composant critique pour le passage à l'échelle des architectures conscientes de la profondeur.