Résumé technique : Couches différentes, variétés différentes

Énoncé du problème

La géométrie de l'espace des poids est un facteur critique dans l'optimisation des réseaux de neurones, pourtant les méthodes actuelles d'entraînement sous contrainte de variété appliquent généralement une contrainte géométrique unique et uniforme (par exemple, l'orthogonalité ou les contraintes de Stiefel) à toutes les matrices de poids d'un Transformer. Ce travail étudie si les différents modules de Transformer — spécifiquement les couches d'attention et les blocs MLP (Perceptron Multicouche) — préfèrent des géométries d'espace de poids distinctes. Les auteurs postulent que les rôles computationnels et les mécanismes non linéaires de ces modules diffèrent considérablement, suggérant qu'une assignation de variété uniforme peut être sous-optimale, voire déstabilisante.

Méthodologie

L'étude utilise Manifold Muon, un optimiseur qui applique des mises à jour normalisées par matrice aux poids des couches cachées tout en les contraignant à des variétés matricielles spécifiques. Les auteurs comparent deux familles de contraintes principales :

1. Contrainte de Stiefel : Impose l'orthonormalité des colonnes ($W^\top W = I$), encadrant strictement les valeurs singulières et fournissant un contrôle spectral fort.
2. Contrainte DGram : Une contrainte d'espace de Gram plus faible exigeant que les entrées hors-diagonale de la matrice de Gram soient nulles ($\text{Off}(W^\top W) = 0$). Elle préserve l'orthogonalité des colonnes mais permet aux normes des colonnes (et donc aux valeurs singulières) de varier, offrant ainsi des degrés de liberté de mise à l'échelle supplémentaires.

Les auteurs évaluent cinq stratégies d'assignation distinctes sur un modèle GPT-2 small (124M de paramètres) pré-entraîné sur OpenWebText :

• UNCONSTRAINED (Sans contrainte) : Aucune contrainte de variété.
• ALL-STIEFEL : Contraintes de Stiefel sur les poids d'Attention et de MLP.
• ALL-DGRAM : Contraintes DGram sur les poids d'Attention et de MLP.
• HETERO : Contraintes de Stiefel pour l'Attention, DGram pour le MLP.
• HETERO-INV : Contraintes DGram pour l'Attention, Stiefel pour le MLP.

Toutes les configurations partagent des architectures de modèles, des programmes de données, des calendriers d'entraînement et des hyperparamètres d'optimiseur identiques afin d'isoler l'effet de l'assignation de la variété.

Résultats clés

Les expériences révèlent une asymétrie claire dans la manière dont les différents modules répondent aux contraintes de variété :

• Configuration optimale : L'assignation HETERO (Stiefel pour l'Attention, DGram pour le MLP) a obtenu la perte de validation stable la plus basse (3,3544), surpassant à la fois la ligne de base non contrainte et la configuration uniforme ALL-STIEFEL.
• Instabilité de la DGram-Attention : Les configurations assignant des contraintes DGram aux couches d'attention (ALL-DGRAM et HETERO-INV) sont devenues instables sous le réglage d'hyperparamètres partagé. Ces exécutions ont présenté une croissance explosive de la perte de validation et n'ont pas réussi à s'en remettre.
• Stabilité de la DGram-MLP : Inversement, l'assignation de contraintes DGram aux couches MLP (lorsqu'elles sont couplées à une attention Stiefel) est restée stable et a amélioré les performances par rapport aux contraintes Stiefel uniformes.

Analyse mécaniste

Le document fournit une explication mécaniste de l'instabilité observée dans l'attention contrainte par DGram :

1. Croissance des valeurs singulières : Contrairement aux contraintes de Stiefel, la contrainte DGram permet aux valeurs singulières des matrices de poids de croître arbitrairement. Dans les couches d'attention, les projections de requête ($W_Q$) et de clé ($W_K$) se multiplient pour former les logits d'attention ($Z = X W_Q W_K^\top X^\top / \sqrt{d}$).
2. Amplification des logits : La croissance incontrôlée des normes spectrales de $W_Q$ et $W_K$ amplifie directement la magnitude des logits d'attention.
3. Saturation Softmax : Les grandes magnitudes de logits poussent la fonction softmax dans un régime de saturation, où la sortie approche un vecteur "one-hot". Cela provoque une dégénérescence du Jacobien de la softmax, menant à une dégénérescence du gradient et à une instabilité de l'entraînement.
4. Tolérance du MLP : Les couches MLP tolèrent les contraintes DGram car leurs fonctions d'activation (ex: GELU) opèrent de manière élément par élément et ne couplent pas les scores des jetons à travers un simplexe de probabilité normalisé. Par conséquent, l'amplification d'échelle dans les poids du MLP ne force pas l'ensemble du module dans un régime de routage saturé.

Contributions

Les auteurs identifient trois contributions principales :

1. Évaluation empirique : Ils évaluent les assignations de variété par module pour Manifold Muon dans le pré-entraînement de GPT-2, démontant que les contraintes uniformes sont insuffisantes.
2. Découverte de l'asymétrie : Ils montrent que les couches d'attention et MLP préfèrent des géométries d'espace de poids différentes, l'attention nécessitant un contrôle spectral fort (Stiefel) et le MLP bénéficiant d'une liberté de mise à l'échelle préservant la structure (DGram).
3. Mécanisme d'instabilité : Ils identifient la croissance des valeurs singulières et la saturation subséquente de la softmax comme un mécanisme plausible de l'instabilité de l'attention contrainte par DGram.

Signification et limites

Le papier soutient que l'optimisation sensible à la symétrie et à la géométrie pour les Transformers doit être spécifique aux modules plutôt qu'uniforme. Cela suggère que la géométrie de l'espace des poids est un axe de conception pratique pour améliorer la stabilité et la dynamique d'entraînement dans les réseaux modernes.

Les auteurs maintiennent une position modeste concernant leurs affirmations :

• Ils ne prétendent pas que la DGram-Attention est intrinsèquement instable dans toutes les conditions, mais qu'elle expose un chemin instable sous le réglage d'hyperparamètres partagé testé sans contrôle d'échelle supplémentaire.
• Les expériences sont limitées à GPT-2 small et à un ensemble spécifique de graines (seeds) ; l'écart de perte de validation entre la meilleure configuration hétérogène et la configuration uniforme est modeste.
• L'analyse mécaniste sert d'explication à l'instabilité observée plutôt que de preuve complète que la DGram-Attention ne peut pas être stabilisée (par exemple, via un réajustement des taux d'apprentissage ou une régularisation explicite de l'échelle).

Des travaux futurs sont suggérés pour développer des méthodes automatiques d'assignation des contraintes de variété à travers différents modules, architectures et échelles.