Résumé technique : Principe de compatibilité symétrique pour la conception d'optimiseurs

1. Énoncé du problème

Une disparité géométrique fondamentale existe dans la pratique actuelle de l'apprentissage profond. Alors que les architectures de réseaux de neurones (embeddings, têtes de langage, MLPs SwiGLU, routeurs MoE) possèdent de riches propriétés de symétrie et d'équivariance intrinsèques, les méthodes d'optimisation dominantes comme Adam, AdamW et RMSprop opèrent de manière coordinate par coordinate. Ces méthodes traitent les paramètres matriciels comme des vecteurs aplatis de coordonnées indépendantes, ignorant ainsi la structure matricielle et brisant l'équivariance naturelle de l'espace des paramètres.

Cette approche « aveugle à la géométrie » échoue à distinguer les géométries distinctes des différents types de couches. Par conséquent, les optimiseurs populaires peuvent rejeter la structure spectrale, induire des géométries inadaptées et ne pas respecter les structures d'invariance et d'équivariance inhérentes à des blocs de paramètres spécifiques (par exemple, l'invariance par permutation dans les vocabulaires ou l'interchangeabilité des experts dans les MoE). La compréhension théorique actuelle et les comparaisons empiriques entre les méthodes adaptatives coordinate par coordinate et les optimiseurs spectraux restent limitées, s'appuyant souvent sur des benchmarks à grande échelle sans analyse approfondie du comportement géométrique.

2. Méthodologie : Le principe de compatibilité symétrique

L'article propose un principe de compatibilité symétrique au niveau des couches pour la conception d'optimiseurs. La thèse centrale est que la règle de mise à jour $\mathcal{U}$ de l'optimiseur pour un bloc de paramètres $W$ doit être équivariante sous le groupe de symétrie $G$ agissant sur ce bloc spécifique.

2.1 Cadre théorique

Soit $W \in \mathbb{R}^{m \times n}$ une matrice de paramètres. Si les coordonnées des espaces d'entrée et de sortie sont transformées par des matrices inversibles $P$ et $Q$, l'application linéaire se transforme en $\tilde{W} = P W Q^{-1}$. Le gradient se transforme de manière contravariante/covariante en conséquence. La carte de mise à jour de l'optimiseur $\mathcal{U}$ doit satisfaire :
$$\mathcal{U}(g \cdot D) = g \cdot \mathcal{U}(D)$$
où $D$ est une direction de mise à jour (par exemple, le gradient ou la quantité de mouvement) et $g$ est un élément du groupe de symétrie.

L'article dérive des classes spécifiques d'optimiseurs basées sur les groupes de symétrie des différents types de couches :

1. Couches matricielles ordinaires (Linéaires/Attention) :

   • Symétrie : Bi-orthogonale ($O_m \times O_n$).
   • Classe d'optimiseur : Optimiseurs spectraux. La mise à jour doit être un opérateur spectral agissant sur les valeurs singulières tout en préservant les vecteurs singuliers. Cela unifie la Descente Spectrale Stochastique (SSD), Muon, Scion et les méthodes de gradient polaire. La mise à jour est typiquement le facteur polaire orthogonal de la direction de la quantité de mouvement.

2. Embeddings et têtes de langage (LM Heads) :

   • Symétrie : Permutation à gauche ($P_v$) et Orthogonale à droite ($O_d$). Les lignes indexent les éléments du vocabulaire (permutation discrète), les colonnes indexent les caractéristiques cachées (rotation continue).
   • Classe d'optimiseur : Mises à jour équivariantes à permutation gauche et orthogonale droite (LPRO).
     • Spectral à droite : Agit globalement via la matrice de Gram droite ($D^\top D$).
     • Norme de ligne : Agit localement sur les normes des lignes (éléments du vocabulaire).
     • Hybride : Compositions de mises à jour de norme de ligne et spectrales à droite.

3. Projections MLP SwiGLU :

   • Symétrie : Permutation des neurones intermédiaires ($P_{d_{ff}}$). Les projections de porte (gate) et de montée (up) ont des lignes correspondant aux neurones intermédiaires ; la projection de descente (down) a des colonnes correspondant à ceux-ci.
   • Classe d'optimiseur : Variantes conscientes des lignes et des colonnes. Les projections de porte/montée utilisent des mises à jour compatibles LPRO (norme de ligne/spectrale à droite) ; les projections de descente utilisent la géométrie transposée.

4. Routeurs MoE :

   • Symétrie : Permutation des experts ($P_e$) et invariance par décalage des logits partagés (ajouter une constante à tous les logits ne change pas la distribution softmax).
   • Classe d'optimiseur : Mises à jour compatibles avec le routeur. Elles opèrent sur la direction centrée $D_c = \Pi^\perp D$ (éliminant le décalage partagé).
     • Spectral à gauche centré : Agit sur la matrice de Gram des experts centrée.
     • Norme de ligne centrée : Agit sur les normes de lignes des experts centrées.

2.2 Mise en œuvre pratique

Les auteurs instancient ces principes dans des optimiseurs pratiques basés sur la quantité de mouvement :

• RightPolarGradM / LeftPolarGradM : Mises à jour spectrales unilatérales utilisant des itérations de Newton-Schulz pour les racines carrées inverses de matrices.
• RowNormM : Mises à jour de normalisation ligne par ligne.
• HybridPolarGradM : Compositions d'étapes de norme de ligne et spectrales unilatérales.
• Centrage : Pour les routeurs MoE, les mises à jour sont appliquées à la quantité de mouvement centrée $M_{k,c} = \Pi^\perp M_k$.

3. Contributions clés

1. Principe de compatibilité symétrique : Un argument formel selon lequel les optimiseurs coordinate par coordinate sont géométriquement inadaptés pour les paramètres à valeur matricielle. L'article établit que les règles de mise à jour doivent correspondre au groupe de symétrie du bloc de paramètres.
2. Vue unifiée des optimiseurs spectraux : Fournit une interprétation basée sur la symétrie des optimiseurs spectraux (Muon, SSD, etc.) comme la classe unique de cartes de mise à jour équivariantes bi-orthogonales.
3. Nouvelles classes d'optimiseurs équivariants : Dérive des règles de mise à jour spécifiques pour les couches ayant des symétries non standard :
   • Mises à jour équivariantes LPRO pour les embeddings et les têtes de langage.
   • Mises à jour conscientes des lignes/colonnes pour les MLPs SwiGLU.
   • Mises à jour de permutation d'experts centrées pour les routeurs MoE.
4. Preuves de pré-entraînement de bout en bout : Construit une pile complète d'optimiseurs au niveau des couches où chaque grande classe de matrices reçoit une mise à jour adaptée à sa symétrie. Cela est évalué sur des modèles de langage denses et MoE épars (Qwen3-0.6B, Gemma 3 1B, OLMoE-1B-7B, gpt-oss réduit).

4. Résultats expérimentaux

Les auteurs ont mené des expériences de pré-entraînement sur des sous-ensembles de 10 milliards de tokens de FineWeb-Edu à travers quatre architectures.

• Modèles denses (Qwen3-0.6B, Gemma 3 1B) :

  • Le remplacement d'AdamW sur les matrices indexées par le vocabulaire (embeddings/têtes de langage) par RowNormM ou HybridPolarGradM a systématiquement amélioré la perte de validation finale.
  • Les gains étaient modestes pour le modèle Qwen3 plus petit mais plus prononcés pour le modèle Gemma 3 plus grand, suggérant que les mises à jour conscientes de la géométrie deviennent plus critiques à mesure que les matrices de vocabulaire grandissent et que les gradients deviennent plus anisotropes.
  • L'utilisation de HybridPolarGradM pour les projections MLP SwiGLU a encore amélioré les performances, probablement en raison de la correction du déséquilibre d'échelle des lignes dans les matrices hautes et minces.

• Modèles MoE épars (OLMoE-1B-7B, gpt-oss réduit) :

  • Les configurations utilisant des mises à jour compatibles avec la symétrie pour les embeddings, les têtes de langage et les routeurs (par exemple, RowNormM + LeftPolarGradM) ont surpassé les configurations utilisant AdamW pour ces composants.
  • La référence entièrement coordinate par coordinate (AdamW partout) a obtenu les pires résultats.
  • Les mises à jour de routeur compatibles avec la symétrie ont amélioré les mises à jour de routeur coordinate par coordinate et, dans certains cas, ont réduit les pics de perte d'entraînement, suggérant une stabilité améliorée.

• Tendance générale : Pour tous les modèles, le remplacement d'AdamW sur les grandes matrices indexées par le vocabulaire par des optimiseurs compatibles avec la symétrie a systématiquement abaissé la perte de validation finale. Les bénéfices ont persisté à travers différentes architectures et niveaux de sparsité.

5. Importance et revendications

L'article revendique un changement de principe dans la conception d'optimiseurs, passant d'une vue « aveugle à la géométrie » coordinate par coordinate à une vue « consciente de la symétrie » matricielle.

• Co-conception Architecture-Optimiseur : L'ouvrage soutient que la conception d'optimiseurs doit être un problème conscient de l'architecture. Différentes classes de paramètres (embeddings, routeurs, couches linéaires) possèdent des symétries naturelles distinctes et nécessitent des géométries d'optimiseurs distinctes.
• Fondement théorique : Il fournit une caractérisation mathématique rigoureuse (via l'équivariance) expliquant pourquoi les optimiseurs spectraux fonctionnent pour les couches linéaires et étend cette logique pour dériver de nouveaux optimiseurs pour d'autres couches.
• Utilité pratique : Les résultats démontrent que les affectations compatibles avec la symétrie ne sont pas seulement théoriquement solides mais pratiquement bénéfiques, améliorant la perte finale et potentiellement la stabilité de l'entraînement dans le pré-entraînement à grande échelle.
• Modestie : Les auteurs déclarent explicitement ne pas revendiquer que les optimiseurs équivariants dominent les méthodes coordinate par coordinate dans tous les régimes. Au lieu de cela, ils présentent des preuves que le principe de compatibilité symétrique au niveau des couches conduit à des affectations compétitives et souvent bénéfiques dans des contextes de pré-entraînement représentatifs. Ils considèrent les résultats comme une preuve de l'utilité du principe plutôt que comme un benchmark exhaustif.

L'article conclut que l'introduction de la symétrie et de l'équivariance dans la conception d'optimiseurs ouvre la voie à un pré-entraînement de modèles à grande échelle plus stable, cohérent avec la paramétrisation et fondé sur des bases théoriques.