Principe compatible avec la symétrie pour la conception d'optimiseurs : plongements, têtes LM, MLP SwiGLU et routeurs MoE

Résumé

Une disparité géométrique frappante persiste depuis longtemps dans la pratique de l'apprentissage profond. Alors que les architectures modernes de réseaux de neurones présentent naturellement de riches propriétés de symétrie et d'équivariance, les optimiseurs populaires tels que Adam et ses variantes fonctionnent intrinsèquement par coordonnées, ce qui les empêche de respecter les structures d'équivariance de l'espace des paramètres. Nous abordons cette disparité en introduisant un principe compatible avec la symétrie pour la conception d'optimiseurs : la règle de mise à jour du gradient doit être équivariante sous l'action du groupe de symétrie agissant sur le bloc de poids correspondant. Suivant ce principe, nous proposons d'abord une perspective unifiée sur les mises à jour bi-orthogonalement équivariantes pour les couches matricielles générales, telles qu'employées par la descente spectrale stochastique, Muon, Scion et les méthodes de gradient polaire. Plus important encore, en passant des groupes orthogonaux aux symétries de permutation et de décalage partagé, nous dérivons des optimiseurs compatibles avec la symétrie pour les blocs de paramètres dont les symétries diffèrent de celles des couches matricielles générales : les matrices d'embedding et de tête LM, les projections SwiGLU du MLP et les matrices de routeur MoE. Ces constructions incluent des mises à jour spectrales unilatérales, de norme par ligne, hybrides norme par ligne/spectral, conscientes de la ligne, conscientes de la colonne, de norme par ligne centrée et spectrales à gauche. Elles produisent une pile d'optimiseurs couche par couche de bout en bout dans laquelle chaque classe majeure de paramètres matriciels se voit attribuer une mise à jour dont l'équivariance correspond à son groupe de symétrie. Nous corroborons ce principe par des expériences de pré-entraînement sur des modèles de langage denses et MoE parcimonieux, incluant les architectures Qwen3-0.6B-style, Gemma 3 1B-style, OLMoE-1B-7B-style et gpt-oss réduite. Dans ces expériences, les mises à jour compatibles avec la symétrie améliorent systématiquement la perte de validation finale, et dans plusieurs cas la stabilité de l'entraînement, par rapport aux mises à jour AdamW correspondantes.

English

A striking geometric disparity has long persisted in the practice of deep learning. While modern neural network architectures naturally exhibit rich symmetry and equivariance properties, popular optimizers such as Adam and its variants operate inherently coordinate-wise, rendering them unable to respect the equivariance structures of the parameter space. We address this disparity by introducing a symmetry-compatible principle for optimizer design: the gradient update rule should be equivariant under the symmetry group acting on the corresponding weight block. Following this principle, we first provide a unified perspective on bi-orthogonally equivariant updates for general matrix layers, as employed by stochastic spectral descent, Muon, Scion, and polar gradient methods. More importantly, by moving from orthogonal groups to permutation and shared-shift symmetries, we derive symmetry-compatible optimizers for parameter blocks whose symmetries differ from those of general matrix layers: embedding and LM head matrices, SwiGLU MLP projections, and MoE router matrices. These constructions include one-sided spectral, row-norm, hybrid row-norm/spectral, row-aware, column-aware, centered row-norm, and left-spectral updates. They yield an end-to-end layerwise optimizer stack in which each major matrix-valued parameter class is assigned an update whose equivariance matches its symmetry group. We corroborate this principle through pre-training experiments on dense and sparse MoE language models, including Qwen3-0.6B-style, Gemma 3 1B-style, OLMoE-1B-7B-style, and downsized gpt-oss architectures. Across these experiments, symmetry-compatible updates consistently improve final validation loss, and in several cases training stability, over corresponding AdamW updates.