Résumé Technique : BHyT – Une alternative stable et efficace à la Normalisation Pré-Couche (Pre-Layer Normalization)

1. Énoncé du Problème

Les grands modèles de langage (LLM) reposent largement sur la Normalisation Pré-Couche (Pre-LN), généralement implémentée via RMSNorm, pour stabiliser l'entraînement et permettre un transfert d'apprentissage efficace. Cependant, à mesure que la profondeur du modèle augmente, la conception standard de la Pre-LN souffre de deux limitations critiques :

1. La Malédiction de la Profondeur : Des analyses récentes indiquent que l'interaction entre les connexions résiduelles et la Pre-LN entraîne une croissance rapide de la magnitude et de la variance des états cachés à travers les couches. Cette escalade pousse le Jacobien du bloc vers la matrice identité, faisant en sorte que les couches plus profondes se comportent comme des applications identités coûteuses et déstabilisant l'entraînement.
2. Compromis Efficacité vs Stabilité : Les solutions existantes ne traitent souvent qu'un seul aspect. Les méthodes orientées vers la stabilité comme la Peri-LN (normalisation avant et après les sous-couches) ou la Mise à l'Échelle de la LayerNorm (LNS) suppriment la croissance de la variance mais entraînent un surcoût computationnel important dû aux calculs statistiques répétés. À l'inverse, les méthodes sans normalisation orientées vers l'efficacité comme la Dynamic Tanh (DyT) éliminent le surcoût statistique mais ne parviennent pas à contrôler explicitement la croissance résiduelle en profondeur, ce qui conduit à une saturation potentielle et à une instabilité dans les réseaux profonds.

Le défi central est de concevoir un mécanisme qui combine le faible surcoût des transformations bornées avec un contrôle explicite de l'échelle et de la variance du flux résiduel afin de prévenir la dérive en profondeur sans sacrifier le débit d'entraînement.

2. Méthodologie : Bounded Hyperbolic Tanh (BHyT)

Les auteurs proposent BHyT, un remplacement direct pour la Pre-LN qui couple une non-linéarité tangente hyperbolique avec un bornage d'entrée explicite et dépendant des données.

Mécanisme Central

Contrairement à DyT, qui repose sur un scalaire global apprenable pour remettre à l'échelle les entrées, BHyT borne explicitement l'argument de la fonction $\tanh(\cdot)$ en utilisant des statistiques dépendantes des données.

• Bornage Probabiliste de l'Entrée : Pour éviter la saturation, BHyT contraint l'argument pré-$\tanh$ à se situer dans un intervalle prédéfini $[-\lambda, \lambda]$ avec une haute probabilité. En utilisant l'inégalité de Chebyshev, la méthode dérive une règle de mise à l'échelle basée sur la moyenne ($\mu_x$) et l'écart-type ($s_x$) de l'entrée.
• Forme Idéale ($BHyT^*$) : Théoriquement, la transformation remet à l'échelle l'entrée $x$ comme suit :
  $$BHyT^*(x) = \gamma \odot \tanh\left(\frac{\lambda}{\kappa s_x + |\mu_x|} x\right)$$
  où $\gamma$ est un paramètre d'échelle apprenable, $\lambda$ définit la plage cible, et $\kappa$ contrôle le conservatisme de la mise à l'échelle en fonction d'une probabilité cible $p$. Cela garantit que l'entrée reste dans la zone linéaire non saturante de $\tanh$ avec une haute probabilité.

Implémentation Pratique : Approximation de la Variance

Calculer la moyenne et la variance exactes à chaque site de Pre-LN au sein d'un bloc Transformer annulerait les gains d'efficacité. Pour y remédier, BHyT introduit une approximation de la variance au niveau du bloc :

1. Premier Site (Avant l'Attention) : Calcule la variance exacte par instance de l'entrée.
2. Second Site (Avant le MLP) : Au lieu de recalculer la variance exacte de la sortie résiduelle ($x + h_{Attn}$), il approxime la variance en utilisant la variance de l'entrée précédemment calculée et une estimation de la variance de la sortie d'attention dépendante du modèle.
   • La variance de la sortie d'attention est approximée par $\tilde{s}^2_{hAttn} \approx \frac{1}{Td} \|W_V W_O\|_F^2 \cdot \frac{\lambda_{Attn}^2}{\kappa^2}$, en supposant des poids d'attention uniformes et des entrées gaussiennes de moyenne nulle.
   • Cette approximation permet à la seconde normalisation de réutiliser les statistiques ou d'utiliser des paramètres de modèle précalculés, évitant ainsi une seconde opération de réduction dans le chemin de propagation critique.

Stabilité Théorique

L'article fournit une borne de propagation de la variance à profondeur finie. Le Théorème 3.5 stipule que pour un réseau de profondeur $L$, si les hyperparamètres satisfont $\lambda/\kappa < 1/\sqrt{L}$, la variance de sortie sous BHyT est strictement inférieure à celle sous la mise à l'échelle de la LayerNorm (LNS) pour chaque couche jusqu'à la profondeur $L$. Cela offre une garantie théorique que BHyT supprime la croissance de la variance en profondeur plus efficacement que la LNS.

3. Principales Contributions

1. Proposition de BHyT : Un remplacement direct pour la Pre-LN qui combine une transformation $\tanh(\cdot)$ bornée avec un bornage d'entrée dépendant des données pour contrôler la croissance des activations en profondeur.
2. Efficacité Théorique et Pratique :
   • Dérivation d'une borne de variance à profondeur finie prouvant que BHyT produit une borne de propagation de la variance plus petite que la LNS sous certaines conditions.
   • Introduction d'une approximation de la variance au niveau du bloc qui réduit les calculs répétés de la variance exacte, diminuant ainsi considérablement la latence.
3. Validation Empirique : Évaluation complète sur les modèles Llama-374M, Llama-1B et Llama-3B. L'étude démontre que BHyT améliore le compromis stabilité-efficacité tout en maintenant de solides performances en aval.

4. Résultats Expérimentaux

Les expériences ont été menées sur des modèles de type Llama (374M, 1B et 3B) pré-entraînés sur le corpus C4 avec des budgets de jetons proportionnels à la taille du modèle (environ 20 jetons par paramètre).

• Stabilité : BHyT a réussi à supprimer la croissance en profondeur de la magnitude et de la variance des activations, surpassant RMSNorm et DyT. Les statistiques d'activation sont restées contrôlées à travers les couches, comparables aux bases orientées vers la stabilité comme la Peri-LN.
• Performance :
  • Pré-entraînement : BHyT a obtenu la perte de pré-entraînement et la perplexité les plus basses pour toutes les échelles de modèles. Il a également atteint la précision moyenne la plus élevée sur les benchmarks à 3 coups (ex: ARC, MMLU, Winogrande), surpassant la deuxième meilleure méthode de 0,6 point sur le modèle 1B.
  • Post-SFT : Après le réglage fin supervisé (SFT) sur Lima1K, les modèles BHyT ont montré une perplexité plus faible sur C4 et WikiText-2 (indiquant moins d'oubli et une meilleure généralisation) et une précision en aval plus élevée par rapport à RMSNorm, Peri-LN, LNS et DyT.
  • Échelle : Dans une comparaison de 20 milliards de jetons sur Llama-3B, BHyT a surpassé la forte base de référence orientée stabilité Peri-LN de 3,4 points de pourcentage en précision moyenne en aval (pré-entraînement) et de 2,5 points (post-SFT).
• Efficacité :
  • Débit d'entraînement : BHyT a atteint un débit d'entraînement moyen de 1,6 % plus rapide que RMSNorm.
  • Débit de génération : BHyT a atteint 1 199,9 jetons/s sur Llama-1B, dépassant RMSNorm de 1,77 % et Peri-LN de 18,0 %. Bien que DyT soit resté le plus rapide (en raison de l'absence de statistiques), BHyT offrait un meilleur compromis entre vitesse et stabilité.
  • Ablation : Le mécanisme d'approximation de la variance a permis un accélération de 14,9 % par rapport à la forme exacte de BHyT* sans dégrader matériellement la convergence ou les performances en aval.

5. Signification et Revendications

L'article affirme que BHyT offre une voie pratique vers des modèles de fondation plus profonds et plus efficaces. En unifiant le faible surcoût des transformations bornées avec un contrôle explicite de la variance du flux résiduel, BHyT répond à la "malédiction de la profondeur" sans les pénalités de calcul des schémas de double normalisation.

Les auteurs soutiennent qu'un entraînement scalable peut améliorer l'efficacité sans sacrifier la stabilité. BHyT est présenté non pas simplement comme une alternative théorique, mais comme un remplacement robuste et direct pour RMSNorm qui permet un pré-entraînement stable et de fortes performances de transfert à travers différentes échelles de modèles (de 374M à 3B de paramètres) tout en réduisant la latence tant lors de l'entraînement que de l'inférence.