Résumé technique : Profitez de votre normalisation de couche avec l'efficacité computationnelle de RMSNorm

Énoncé du problème

La normalisation de couche (Layer Normalization, LN) est un composant fondamental des architectures d'apprentissage profond modernes, en particulier les Transformers, où elle stabilise l'entraînement et améliore l'optimisation. Cependant, lors de l'inférence, la LN nécessite un centrage par échantillon (soustraction de la moyenne) et une mise à l'échelle, ce qui introduit une surcharge computationnelle non négligeable par rapport aux alternatives sans normalisation.

RMSNorm a été proposé comme une alternative plus efficace en supprimant l'opération de centrage, ne conservant que la mise à l'échelle. Bien que RMSNorm offre des accélérations significatives (réductions rapportées de 7 % à 64 % du temps d'exécution dans des travaux antérieurs), il élimine l'opération de centrage. Cette suppression peut entraîner des plages de sortie moins stables et des variations d'activation plus importantes, risquant ainsi de perdre les avantages associés au centrage établis dans la littérature antérieure (LeCun et al., 1990 ; Schraudolph, 1998 ; Ba et al., 2016). Les tentatives précédentes d'équivalence entre LN et RMSNorm (par exemple, Jiang et al., 2023) étaient limitées à des structures spécifiques de Transformers pré-normalisés, laissant ouverte la question de l'existence d'un cadre général permettant à des réseaux de neurones profonds (DNN) arbitraires de remplacer la LN par RMSNorm sans altérer la fonction du modèle ni la dynamique d'entraînement.

Méthodologie

L'article propose un cadre général pour déterminer si une couche LN dans un DNN arbitraire peut être remplacée par RMSNorm tout en maintenant une équivalence mathématique à la fois en inférence et en entraînement. La stratégie centrale consiste à « plier » l'opération de centrage de la LN dans les couches linéaires en amont.

1. Contrainte de centrage par colonne (CCC) et centrage des poids par colonne (CBWC)

Les auteurs analysent d'abord le cas où la LN suit immédiatement une couche linéaire. Ils observent que si la sortie de la couche linéaire a une moyenne nulle à travers les neurones, l'étape de centrage dans la LN devient redondante.

• Contrainte de centrage par colonne (CCC) : Une couche linéaire satisfait la CCC si les éléments de chaque colonne de sa matrice de poids somment à zéro. Sous cette contrainte, la sortie de la couche linéaire est garantie d'avoir une moyenne nulle pour toute entrée.
• Centrage des poids par colonne (CBWC) : Pour imposer la CCC pendant l'entraînement sans restreindre la capacité du modèle, les auteurs introduisent une technique de reparamétrisation. Une matrice de paramètres proxy $W$ est utilisée, qui est transformée en une matrice de poids effective $V$ via une projection : $V = (I - \frac{1}{m}\mathbf{1}_m\mathbf{1}_m^\top)W$. Cette transformation garantit que les poids effectifs satisfont la contrainte de somme nulle par colonne tout en préservant la dynamique d'optimisation du modèle original.

Proposition : Un modèle utilisant une couche linéaire avec CBWC suivie de RMSNorm est mathématiquement équivalent à un modèle utilisant la même couche linéaire (sans reparamétrisation) suivie de la LN standard.

2. Généralisation aux DNN arbitraires : LN pliables

Le cadre est étendu aux structures de réseau arbitraires en utilisant une analyse basée sur les graphes.

• LN pliable : Une LN est définie comme « pliable » si son entrée peut être garantie d'avoir une moyenne nulle grâce à des contraintes sur les couches en amont.
• Graphe à moyenne nulle : Les auteurs construisent un graphe orienté représentant le flux de données vers une LN. En remontant depuis la LN, ils identifient les « nœuds feuilles » (couches en amont).
  • Si tous les nœuds feuilles correspondent à des couches linéaires générales (y compris les convolutions et les RNN avec poids partagés) ou à des couches garantissant intrinsèquement des sorties à moyenne nulle, la LN est pliable.
  • Les opérations scalaires (par exemple, dropout en mode inférence, mise à l'échelle de la température) préservent la propriété de moyenne nulle et ne brisent pas la pliage.
  • Les connexions résiduelles (addition) préservent la pliage si chaque branche est individuellement pliable.
  • Les opérations non linéaires (par exemple, ReLU, Softmax) ou les concaténations brisent généralement la propriété de moyenne nulle, empêchant le pliage sauf si un centrage auxiliaire est inséré.

3. Algorithme de détection

L'article présente un algorithme (Algorithme 1) qui effectue un parcours avant du graphe de calcul pour détecter automatiquement les LN pliables et identifier les couches linéaires générales spécifiques en amont qui nécessitent un CBWC. Ce processus est léger en termes de calcul et ne nécessite qu'un seul passage avant.

Contributions clés

1. Cadre général : Un cadre théorique pour déterminer la remplaçabilité de la LN par RMSNorm dans des DNN arbitraires, dépassant les limites des travaux antérieurs restreints aux Transformers pré-normalisés.
2. Équivalence mathématique : L'introduction du CBWC, qui permet le pliage exact de l'opération de centrage de la LN dans les poids en amont, garantissant que le schéma CBWC+RMSNorm produit des sorties et des gradients identiques au modèle original basé sur la LN, tant pendant l'entraînement que pendant l'inférence.
3. Détection de LN pliables : La définition des « LN pliables » et le développement d'un algorithme de détection basé sur les graphes qui identifie quelles LN dans une architecture donnée peuvent être converties.
4. Extension pratique : Une discussion sur la gestion des cas où les conditions de pliage exact ne sont pas remplies (par exemple, en raison du dropout pendant l'entraînement), suggérant que CBWC+RMSNorm reste compétitif par rapport à la LN classique même lorsque l'équivalence exacte est partiellement rompue.

Résultats expérimentaux

Les auteurs ont validé leur méthode sur plusieurs architectures (GPT-2, BERT, OPT, Phi, ViT, Swin-Tiny) et tâches (classification de texte, traduction, classification d'images, affinage).

• Accélération de l'inférence :

  • La méthode atteint des accélérations d'inférence de bout en bout de 2 % à 12 % sur des modèles représentatifs (par exemple, GPT-2, BERT, OPT) sans modifier les prédictions du modèle.
  • L'accélération est plus prononcée avec des longueurs de séquence plus longues (jusqu'à 24,95 % d'accélération pour BERT à 16 384 jetons).
  • Le débit augmente considérablement (par exemple, +1 119 à +3 385 jetons/s sur des GPU A100).
  • La surcharge de l'algorithme de pliage est minimale (amortie après quelques étapes d'inférence) et n'augmente pas l'empreinte mémoire.

• Performance d'entraînement et d'affinage :

  • Dans les scénarios où les conditions de pliage exact sont remplies (par exemple, inférence), la méthode est identique à la référence.
  • Dans les paramètres d'entraînement pratiques où des conditions comme le dropout brisent l'équivalence exacte, CBWC+RMSNorm se comporte systématiquement entre la LN standard et RMSNorm classique.
  • Sur la classification d'images (ImageNet-100 avec Swin-Tiny), la méthode a obtenu la meilleure précision de test et la meilleure perte parmi les trois variantes (LN, RMSNorm, CBWC+RMS).
  • Dans les tâches de génération de texte et d'affinage (GPT-3 sur Alpaca), la méthode a obtenu la meilleure perte de test et la meilleure perplexité tout en offrant des gains d'efficacité.

Importance et affirmations

L'article prétend fournir une manière principielle de réduire le coût computationnel de la normalisation de couche tout en préservant autant que possible ses avantages théoriques et pratiques.

• Équivalence exacte : Pour les LN « pliables », la méthode offre un remplacement mathématiquement exact, permettant le déploiement de RMSNorm dans des modèles initialement entraînés avec LN sans réentraînement ni dégradation des performances.
• Efficacité : Elle délivre des accélérations d'inférence de bout en bout mesurables (2 %–12 %) sur des architectures largement utilisées, ce qui est significatif pour les déploiements à grande échelle de grands modèles de langage (LLM) et de modèles vision-langage (VLM).
• Robustesse : Même lorsque les conditions de pliage exact ne sont pas pleinement remplies (par exemple, pendant l'entraînement avec dropout), la méthode reste compétitive par rapport à la LN standard, offrant un compromis favorable entre fidélité et efficacité.
• Généralité : Contrairement aux approches précédentes, ce cadre s'applique aux DNN arbitraires, pas seulement à des variantes spécifiques de Transformers, et inclut un algorithme de détection pour automatiser le processus.

Les auteurs reconnaissent que leur analyse couvre des architectures représentatives mais n'épuise pas tous les modèles modernes à grande échelle. Ils notent que l'équivalence exacte peut échouer dans des architectures avec des modules comme le dropout pendant l'entraînement, suggérant que des travaux futurs pourraient affiner les traitements pour de tels modules.