Attention Residuals expliqués : repenser la profondeur des Transformers

Les LLM modernes sont devenus plus profonds, plus larges et plus gourmands en calcul, mais empiler des couches de Transformer supplémentaires n’apporte pas toujours des gains proportionnels. L’une des raisons est que les connexions résiduelles standard agrègent les sorties de couche avec des poids unitaires fixes : chaque couche hérite ainsi d’une somme uniforme de tout ce qui la précède. Cela peut diluer les représentations précoces, amplifier la magnitude des états cachés et compliquer la réutilisation sélective des caractéristiques intermédiaires les plus utiles.

Plutôt que de traiter la profondeur comme une récurrence additive fixe, l’article de l’équipe Kimi sur les Attention Residuals permet à chaque couche d’« attendre » sur les sorties des couches antérieures via des poids appris par softmax. Dans cet article, nous verrons pourquoi l’agrégation résiduelle standard devient un goulot d’étranglement, comment fonctionnent les Attention Residuals, pourquoi la variante par blocs compte, et ce que les résultats suggèrent réellement sur le passage à des modèles de langue plus profonds.

Si vous souhaitez approfondir les concepts derrière les architectures Transformer modernes, nous vous recommandons le cours DataCamp Transformer Models Tutorial in PyTorch.

Le problème : le signal noyé

Les connexions résiduelles sont un pilier de l’apprentissage profond depuis des années. Dans les Transformers, la mise à jour résiduelle standard est :

hl=hl−1​+fl−1​(hl−1​)

Cela favorise la circulation des gradients dans des réseaux très profonds. Mais les résiduels ne sont pas qu’une astuce de gradient : ils déterminent aussi comment l’information s’agrège en profondeur. Si l’on déroule la récurrence, on obtient :

hl​=h1​+∑i=1l−1​fi​(hi​)

L’état caché à la couche l est donc l’embedding plus une somme pondérée uniformément de toutes les sorties des couches précédentes. Autrement dit, chaque contribution reçoit en pratique le même poids.

À l’échelle, cela pose problème. L’article montre que dans les architectures PreNorm, l’accumulation résiduelle non pondérée fait croître la magnitude des états cachés avec la profondeur, en O(L) environ. 

À mesure que le flux résiduel enfle, les sorties des premières couches se diluent dans une somme courante de plus en plus grande. Une fois un signal précoce mélangé à cette accumulation, les couches profondes ne peuvent plus le récupérer sélectivement ; elles n’opèrent alors que sur l’état agrégé.

C’est ce que l’article appelle l’effet de « signal noyé ». Un signe empirique fort de cette inefficacité vient des études d’élagage de couches : une fraction significative des couches de modèles entraînés peut souvent être supprimée avec un impact minime sur les performances. Cela suggère que, même si l’on augmente la profondeur, les modèles manquent d’un mécanisme efficace pour l’exploiter pleinement. Au lieu de former une chaîne de raisonnement hiérarchique, le réseau agit comme un relais redondant où les signaux précoces se diluent progressivement.

Il existe aussi un compromis en passe avant. À mesure que le flux résiduel grandit, les couches ultérieures doivent produire des sorties de plus grande magnitude pour influencer de manière significative l’état accumulé. L’article relie cela au comportement PreNorm, où la magnitude des états cachés augmente de façon monotone avec la profondeur sous agrégation résiduelle standard.

Les connexions résiduelles aident les gradients à circuler, et le gradient par rapport à un état caché intermédiaire peut s’écrire :

Le terme identité préserve ici un chemin de gradient direct. Mais les résiduels imposent tout de même, en passe avant, d’agréger chaque couche antérieure avec un poids fixe de 1,0. C’est la limitation structurelle qu’AttnRes cherche à corriger.

Des résiduels additifs à l’attention selon la profondeur

L’idée conceptuelle la plus intéressante de l’article est la dualité temps–profondeur. Les connexions résiduelles compressent l’information en profondeur, comme les réseaux récurrents la compressent dans le temps. En modélisation séquentielle, l’attention a remplacé la récurrence en permettant à chaque position d’accéder sélectivement aux positions antérieures. AttnRes applique la même transition à la profondeur du réseau. Au lieu de définir le prochain état caché comme une somme fixe des couches précédentes, AttnRes permet à chaque couche d’« attendre » sur les sorties antérieures :

où les poids αi→l sont des poids d’attention softmax le long de la profondeur et somment à 1. Ces poids sont calculés comme suit :

avec, dans l’article :

Chaque couche reçoit un pseudo-vecteur requête appris wl, qui « attend » sur des clés et des valeurs construites à partir des sorties des couches précédentes. On applique RMSNorm aux clés afin que les couches produisant naturellement des sorties de grande magnitude ne dominent pas le softmax du seul fait de leur échelle.

Le changement clé proposé est de considérer les résiduels standard comme une forme d’attention linéaire selon la profondeur, et d’« upgrader » AttnRes en attention softmax selon la profondeur. Au lieu d’une accumulation uniforme, on obtient une récupération sélective en profondeur.

Un petit détail d’implémentation important concerne l’initialisation. Les auteurs recommandent d’initialiser à zéro tous les pseudo-vecteurs requêtes, de sorte que les poids d’attention initiaux soient uniformes ; le modèle commence ainsi l’entraînement comme une moyenne à poids égaux, plutôt qu’avec un mécanisme d’attention biaisé aléatoirement.

Figure 1 : vue d’ensemble des Attention Residuals (article Attention Residuals)

Full Attention Residuals

Avec Full Attention Residuals, chaque couche « attend » sur toutes les sorties des couches précédentes, ce qui lui permet d’agréger sélectivement l’information sur toute la profondeur. Le modèle gagne en flexibilité : une couche plus profonde peut privilégier sa prédécesseure immédiate, l’embedding d’origine, ou toute couche antérieure là où se trouve le signal le plus utile.

En entraînement standard, la surconsommation mémoire de Full AttnRes est moindre qu’il n’y paraît, car de nombreuses sorties de couche sont déjà conservées pour la rétropropagation. Mais à grande échelle, la donne change. Avec la recomputation des activations et le pipeline parallelism, ces sorties antérieures doivent être préservées et transmises explicitement pour que les couches ultérieures puissent y « attendre », ce qui devient coûteux.

Full AttnRes est donc la meilleure voie pour comprendre l’idée de base, mais ce n’est pas la version à privilégier pour un déploiement à l’échelle.

Block Attention Residuals

L’article introduit les Block Attention Residuals pour réduire l’empreinte mémoire et les coûts de communication. Les couches du modèle sont partitionnées en blocs : à l’intérieur d’un bloc, les sorties sont combinées via une accumulation résiduelle additive standard ; entre blocs, le modèle réalise une attention sur des résumés au niveau bloc, plutôt que sur chaque couche précédente individuellement.

Si Bn est l’ensemble des couches du bloc n, alors la représentation de bloc est donnée par :

Le modèle « attend » ensuite sur l’embedding b0=h1 , sur les résumés des blocs antérieurs, et sur la somme partielle du bloc courant au fur et à mesure des calculs. On réduit ainsi mémoire et communication de O(Ld) à O(Nd), où N est le nombre de blocs.

L’article rapporte qu’avec environ huit blocs, on retrouve l’essentiel du gain de la version complète, et que l’écart de performance entre Full AttnRes et Block AttnRes diminue à mesure que l’échelle augmente. Des tailles de bloc telles que S=2,4,8 restent proches de la version complète, tandis que des groupements beaucoup plus grossiers reviennent vers le comportement du baseline.

La leçon la plus utile est que vous n’avez pas besoin d’une attention complète selon la profondeur sur chaque couche pour obtenir l’essentiel du gain. Si Block AttnRes rend l’attention selon la profondeur viable en pratique, la déployer efficacement à l’échelle exige une conception système soignée. L’article présente aussi plusieurs optimisations rendant AttnRes utilisable en entraînement et en inférence réels.

Figure 2 : exemple de communication en pipeline avec cache, avec 4 rangs physiques et 2 étapes virtuelles par rang ; les cases hachurées marquent la fin des blocs AttnRes (article Attention Residuals)

Voici trois optimisations systèmes clés qui font d’AttnRes un remplacement « plug-and-play » des résiduels standard :

• Calcul en deux phases : en découplant le pseudo-requête de l’état caché, on peut regrouper les requêtes de toutes les couches d’un bloc. La phase 1 calcule en parallèle l’attention inter-blocs, tandis que la phase 2 gère les dépendances intra-bloc séquentielles via un Online Softmax Merge. Cela réduit fortement l’I/O par couche par rapport aux approches multi-flux comme mHC.
• Mise en cache inter-étapes : en pipeline parallelism, nous éliminons les transferts redondants en mettant en cache localement les blocs déjà reçus. Cela procure un gain de V× (où V est le nombre d’étapes virtuelles, voir Figure 2) sur le coût de pic par transition, permettant de recouvrir complètement la communication par le calcul.
• Prefilling économe en mémoire : en shardant les représentations de blocs sur des dispositifs en tensor parallelism, l’empreinte mémoire pour les longues séquences (128K+) est réduite de plusieurs ordres de grandeur (par ex. de 15 Go à 0,3 Go par dispositif).

Résultats et analyse des performances

L’article évalue AttnRes à plusieurs niveaux : lois d’échelle, ablations, dynamiques d’entraînement et benchmarks applicatifs.

Figure 3 : courbes de lois d’échelle pour les Attention Residuals (article Attention Residuals)

Le résultat phare vient des expériences de lois d’échelle. Sur cinq tailles de modèle, Full AttnRes et Block AttnRes obtiennent systématiquement une perte de validation inférieure au baseline. D’après les courbes ajustées, Block AttnRes atteint la même perte qu’un baseline entraîné avec environ 1,25× plus de calcul. AttnRes n’est donc pas qu’une amélioration théorique : il est aussi plus économe en calcul.

Sur les tâches aval, AttnRes dépasse le baseline sur tous les benchmarks évalués. Les gains les plus marqués concernent les tâches de raisonnement multi-étapes, les maths et le code, comme le montre le tableau 1 :

Tableau 1 : comparaison des performances d’AttnRes avec le baseline, après la même recette de pré-entraînement (article Attention Residuals)

Les gains sont particulièrement nets sur le raisonnement en plusieurs étapes, les mathématiques et le code, ce qui corrobore l’hypothèse centrale : si les couches tardives peuvent récupérer sélectivement des représentations précoces au lieu d’hériter d’un agrégat flou, le raisonnement compositionnel s’améliore.

Comparé au baseline, AttnRes affiche une perte de validation plus faible tout au long de l’entraînement, des magnitudes de sortie plus bornées en profondeur et des magnitudes de gradient plus uniformes entre couches. Le baseline souffre du schéma de dilution PreNorm habituel, où la magnitude des états cachés croît monotoniquement avec la profondeur, tandis que Block AttnRes produit un profil plus contrôlé et borné grâce à l’agrégation sélective aux frontières de blocs.

On note aussi, à l’échelle de l’architecture, que pour un budget de calcul et un nombre de paramètres fixes, la configuration optimale avec AttnRes se déplace vers un modèle plus profond et plus étroit que le baseline. Cela suggère qu’AttnRes rend la profondeur additionnelle plus utile qu’avec l’agrégation résiduelle standard.

Conclusion

Les résiduels sont souvent vus comme une astuce d’entraînement nécessaire : un raccourci identité qui préserve les gradients. AttnRes montre que cette vision est trop réductrice. Les chemins résiduels sont aussi le mécanisme d’acheminement de l’information en profondeur, et l’accumulation additive fixe est peut-être trop primitive pour des modèles très profonds.

Les Attention Residuals ne visent pas tant à colmater un problème identifié qu’à rehausser un choix de conception négligé. La modélisation séquentielle est passée de la récurrence à l’attention parce que la récurrence fixe était trop restrictive. AttnRes soutient que la profondeur est prête pour la même transition.

Le modèle conserve une forte localité : les couches « attendent » souvent surtout sur leurs voisines proches, mais apprennent aussi des schémas de sauts non triviaux, maintiennent un poids persistant sur l’embedding et conservent des comportements distincts entre les couches pré-attention et pré-MLP. 

La conclusion pratique n’est pas que tous les modèles doivent remplacer les résiduels par une attention complète selon la profondeur. Elle consiste plutôt à reconsidérer l’agrégation résiduelle comme un espace de conception architectural à part entière.