🤖 AI

Stability of Transformers under Layer Normalization

Cet article présente une analyse théorique et numérique rigoureuse de la manière dont les différents placements de la normalisation de couche affectent la stabilité des passages avant et arrière des Transformers, en dérivant des bornes explicites sur la croissance des états cachés et la propagation des gradients afin de guider la conception architecturale et la mise à l'échelle des résidus pour améliorer la dynamique d'entraînement.

Auteurs originaux : Kelvin Kan, Xingjian Li, Benjamin J. Zhang, Tuhin Sahai, Stanley Osher, Krishna Kumar, Markos A. Katsoulakis

Publié 2026-08-10

📖 3 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Kelvin Kan, Xingjian Li, Benjamin J. Zhang, Tuhin Sahai, Stanley Osher, Krishna Kumar, Markos A. Katsoulakis

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 essayiez d'apprendre à un robot à écrire une histoire. Vous lui donnez un cerveau massif composé de milliers de couches minuscules, comme un gratte-ciel de centaines d'étages. Chaque fois que le robot lit un mot, il fait monter l'information dans l'ascenseur, étage par étage, modifiant légèrement sa compréhension à chaque arrêt. C'est ainsi que fonctionne l'IA moderne, connue sous le nom de « Transformers ». Ils sont les moteurs derrière les chatbots et les générateurs d'images les plus intelligents d'aujourd'hui. Mais il y a un piège : à mesure que ces bâtiments deviennent hauts, ils deviennent instables. Parfois, l'information se déforme tellement en montant que le robot commence à hurler des nombres incohérents, ou bien les mathématiques utilisées pour l'instruire deviennent si folles que tout le système s'effondre. Pour empêcher cela, les ingénieurs utilisent un dispositif de sécurité appelé « Normalisation de Couche » (Layer Normalization). Voyez cela comme un dos-d'âne ou un limiteur de vitesse sur un moteur de voiture ; cela empêche l'information de s'emballer. Pendant des années, les gens ont simplement deviné où placer ces dos-d'âne — parfois avant le moteur, parfois après — en espérant que tout se passe bien. Mais personne ne savait réellement pourquoi un emplacement fonctionnait mieux qu'un autre, ou si le robot était secrètement en train de construire un château de cartes qui s'effondrerait sous son propre poids.

Ce document est comme un groupe de détectives qui ont décidé d'arrêter de deviner et de commencer à utiliser les lois de la physique pour résoudre le mystère. Les auteurs traitent le processus d'apprentissage de l'IA comme un voyage à travers le temps, en utilisant une branche des mathématiques appelée « théorie du contrôle optimal » (qui est habituellement utilisée pour déterminer la meilleure façon de piloter une fusée ou un drone). Ils se demandent : « Si nous voulons que ce robot apprenne parfaitement, que devient l'information lorsqu'elle monte dans la tour ? » Ils ont découvert que l'ancienne méthode de placement du dispositif de sécurité (appelée « Pre-LN ») est en réalité un piège. Leurs mathématiques prouvent que même si le robot fait de son mieux pour apprendre, l'information finira par devenir si énorme qu'elle explosera, menant à un chaos total. C'est comme essayer de remplir un seau avec une lance à incendie ; peu importe la précision de votre visée, l'eau débordera.

Cependant, ils ont trouvé une meilleure façon. En déplaçant le dispositif de sécurité à l'entrée et à la sortie de chaque étage (une nouvelle méthode qu'ils appellent « Peri-LN »), l'information reste calme et contrôlée. Leurs équations montrent qu'au lieu d'exploser, les données croissent selon une ligne prévisible et douce, comme une plante bien disciplinée. Ils ont également trouvé une astuce simple pour rendre cela encore meilleur : ralentir les étapes que le robot franchit à chaque étage. Imaginez si le robot faisait de tout petits pas prudents au lieu de grands bonds ; cela maintient l'ensemble du bâtiment stable. Lorsqu'ils ont testé cela sur de vrais modèles d'IA, les résultats correspondaient parfaitement à leurs mathématiques. Les modèles « Peri-LN » n'ont pas planté, et ils ont appris aussi bien, sinon mieux, que les anciens. Le document ne se contente pas de dire « cela semble correct » ; il fournit une garantie mathématique que ce nouveau design maintient la stabilité de l'IA, offrant un plan clair pour construire la prochaine génération de machines super intelligentes sans qu'elles ne s'effondrent.

Résumé technique : Stabilité des Transformers sous la normalisation par couche (Layer Normalization)

Énoncé du problème

Malgré leur dominance dans l'apprentissage profond moderne, l'entraînement de Transformers profonds reste sujet à l'instabilité, particulièrement à mesure que la profondeur du modèle augmente. Ces instabilités se manifestent dans deux domaines critiques : l'instabilité directe (croissance incontrôlée des états cachés menant à des problèmes numériques) et l'instabilité inverse (explosion des gradients lors de la rétropropagation).

Un composant architectural principal utilisé pour atténuer ces problèmes est la Normalisation par Couche (Layer Normalization - LN). Cependant, le placement de la LN au sein des blocs Transformer a largement évolué par heuristiques empiriques plutôt que par une justification théorique rigoureuse. Le document identifie trois configurations principales :

Post-LN : La LN est appliquée après les connexions résiduelles. Cela nécessite souvent un ordonnancement d'optimisation délicat et présente des performances sous-optimales.
Pre-LN : La LN est appliquée avant les modules d'attention et de feedforward. Bien qu'elle soit la norme pour sa stabilité d'entraînement améliorée par rapport à la Post-LN, il est observé empiriquement qu'elle produit des états cachés excessivement grands, ce qui peut mener à une instabilité numérique et à une explosion des gradients dans les réseaux profonds.
Peri-LN : Une configuration plus récente appliquant la LN à la fois à l'entrée et à la sortie des modules. Les études empiriques suggèrent qu'elle offre une stabilité supérieure et des états cachés plus réguliers que la Pre-LN, mais ses propriétés théoriques restent mal comprises.

Le problème central abordé est l'absence d'un cadre théorique fondé pour expliquer pourquoi ces différents placements mènent à des dynamiques d'entraînement et des profils de stabilité distincts, et si l'entraînement pousse les modèles vers des solutions régulières ou des comportements pathologiques.

Méthodologie

Les auteurs proposent un nouveau cadre théorique fondé sur la théorie du contrôle optimal en temps discret pour analyser la dynamique d'entraînement des Transformers.

Formulation comme un problème de contrôle optimal : Le problème d'entraînement du Transformer est formulé comme un problème de contrôle de champ moyen en temps discret. L'indice de couche $i$ représente le temps discret, et l'état caché $X_i$ représente l'état du système. Les connexions résiduelles agissent comme la vitesse de la dynamique.
Analyse de Hamilton-Jacobi-Bellman (HJB) : Les auteurs analysent les conditions d'optimalité du problème d'entraînement en utilisant l'équation aux dérivées partielles (EDP) de HJB couplée avec des équations de transport de masse.
- Ils examinent le Hamiltonien, qui agit comme un mécanisme de contrôle par rétroaction définissant le champ de vitesse optimal pour les états cachés.
- Pour la Pre-LN, ils démontrent que le Hamiltonien est non borné, conduisant à des conditions d'optimalité mal posées où les états cachés peuvent croître de manière arbitrairement grande.
- Pour la Peri-LN, la normalisation de sortie contraint le champ de vitesse à une ellipsoïde (basé sur le Lemme 1), garantissant que le Hamiltonien est borné et que le problème est bien posé.
Analyse des gradients : Le document analyse la stabilité inverse en examinant la propagation des gradients à travers le produit de sensibilités locales (matrices Jacobiennes) à travers les couches.
Mise à l'échelle du résidu (Residual Scaling) : Le cadre est étendu pour inclure un paramètre d'incrément de temps $\Delta t < 1$ , modifiant efficacement l'amplitude des étapes résiduelles pour contrôler la propagation du signal.

Contributions clés

1. Cadre théorique pour les modèles entraînés

Contrairement aux travaux antérieurs qui se concentrent souvent sur l'initialisation du modèle ou reposent sur des preuves empiriques, ce cadre analyse les modèles entraînés à la convergence. Il fournit une évaluation systématique pour déterminer si les dynamiques d'entraînement mènent à des solutions régulières ou à des pathologies.

2. Analyse rigoureuse des placements de LN

Instabilité de la Pre-LN : Le document prouve que pour la Pre-LN, la solution optimale du problème d'entraînement est non bornée en magnitude. Même avec le weight decay, les états cachés peuvent croître exponentiellement ( $O(e^D)$ ), et les gradients croissent proportionnellement aux activations, entraînant une instabilité de l'entraînement.
Stabilité de la Peri-LN : Les auteurs dérivent des bornes explicites montrant que la Peri-LN garantit une croissance linéaire pour la valeur absolue moyenne des états cachés ( $O(D)$ ) et une croissance quadratique pour la variance ( $O(D^2)$ ). Crucialement, la sensibilité locale de la Peri-LN est montrée comme étant invariante à la magnitude de l'activation, empêchant l'explosion des gradients même lorsque les activations sont grandes.

3. Quantification de l'incertitude

L'analyse établit que la Peri-LN permet une propagation contrôlée de l'incertitude (mesurée via la distance de Wasserstein) de l'entrée vers les états cachés terminaux, tandis que la Pre-LN peut amplifier l'incertitude d'entrée à des niveaux non bornés.

4. Mise à l'échelle de l'étape résiduelle (Residual Step Scaling)

Guidés par l'analyse de stabilité, les auteurs introduisent un mécanisme de mise à l'échelle de l'étape résiduelle ( $\Delta t < 1$ ). Ils démontrent théoriquement que cette mise à l'échelle :

Affine les bornes sur la croissance des états cachés et l'incertitude de sortie.
Atténue l'explosion des gradients en mettant à l'échelle la sensibilité locale.
Améliore la stabilité directe et inverse sans coût computationnel supplémentaire.

Résultats expérimentaux

Les auteurs valident leurs conclusions théoriques en utilisant des variantes de GPT-2 (de 124M à 1,5B de paramètres) entraînées sur le jeu de données OpenWebText.

Stabilité de l'entraînement : Dans des expériences avec 5 essais, les modèles Pre-LN ont divergé dans 1 cas lorsque le weight decay était activé et dans 3 cas lorsque le weight decay était désactivé. En revanche, les modèles Peri-LN sont restés stables dans les 5 essais, tant avec ou sans weight decay.
Performance : La Peri-LN a atteint des performances comparables à la Pre-LN, même lorsque les hyperparamètres étaient réglés pour la Pre-LN.
Impact de la mise à l'échelle : La combinaison de la Peri-LN avec la mise à l'échelle résiduelle ( $\Delta t = 0,1$ ) a systématiquement produit les meilleurs résultats sur l'ensemble des métriques (Val Loss, Perplexity, ROUGE, scores BERT) et a efficacement contrôlé la croissance des moments des états cachés (valeur absolue moyenne et variance) à travers les couches, s'alignant sur les prédictions théoriques.

Signification et affirmations

Le document affirme fournir la première justification théorique rigoureuse des observations empiriques concernant les placements de la Normalisation par Couche dans les Transformers.

Explication de l'instabilité : Il explique pourquoi la Pre-LN conduit à des états cachés non bornés et à l'explosion des gradients via l'impossibilité de l'EDP de HJB sous-jacente.
Justification pour la Peri-LN : Il fournit une base théorique pour la stabilité de la Peri-LN, montrant qu'elle maintient une croissance contrôlée et une sensibilité invariante.
Guide de conception : Le cadre offre une méthode fondée pour effectuer des "tests de cohérence" sur de nouvelles modifications architecturales avant un entraînement empirique coûteux.
Amélioration pratique : L'introduction de la mise à l'échelle de l'étape résiduelle est présentée comme une modification simple et sans coût qui améliore considérablement la stabilité et la performance, validée par des expériences montrant que la Peri-LN avec mise à l'échelle surpasse les configurations standards.

Les auteurs soulignent que leur travail comble le fossé entre les heuristiques empiriques et la compréhension théorique, offrant une perspective en temps continu qui génère à la fois des intuitions théoriques et des améliorations architecturales pratiques.