Symmetriekompatibles Prinzip für das Optimierer-Design: Einbettungen, LM-Köpfe, SwiGLU-MLPs und MoE-Router

Zusammenfassung

Eine auffällige geometrische Diskrepanz besteht seit Langem in der Praxis des Deep Learning. Während moderne neuronale Netzwerkarchitekturen von Natur aus reichhaltige Symmetrie- und Äquivarianzeigenschaften aufweisen, arbeiten gängige Optimierer wie Adam und seine Varianten inhärent koordinatenweise, sodass sie nicht in der Lage sind, die Äquivarianzstrukturen des Parameterraums zu respektieren. Wir adressieren diese Diskrepanz durch die Einführung eines symmetriekompatiblen Prinzips für das Design von Optimierern: Die Gradientenaktualisierungsregel sollte unter der auf den entsprechenden Gewichtsblock wirkenden Symmetriegruppe äquivariant sein. Diesem Prinzip folgend liefern wir zunächst eine einheitliche Perspektive auf bi-orthogonal äquivariante Aktualisierungen für allgemeine Matrixschichten, wie sie von Stochastic Spectral Descent, Muon, Scion und polaren Gradientenverfahren verwendet werden. Noch wichtiger ist, dass wir durch den Übergang von orthogonalen Gruppen zu Permutations- und gemeinsamen Shift-Symmetrien symmetriekompatible Optimierer für Parameterblöcke ableiten, deren Symmetrien sich von denen allgemeiner Matrixschichten unterscheiden: Einbettungs- und LM-Kopf-Matrizen, SwiGLU-MLP-Projektionen und MoE-Router-Matrizen. Diese Konstruktionen umfassen einseitig spektrale, Zeilen-Norm-, hybride Zeilen-Norm/Spektral-, zeilenbewusste, spaltenbewusste, zentrierte Zeilen-Norm- und linksspektrale Aktualisierungen. Sie ergeben einen durchgängigen schichtweisen Optimierer-Stack, in dem jeder wichtigen matrixwertigen Parameterklasse eine Aktualisierung zugewiesen wird, deren Äquivarianz mit ihrer Symmetriegruppe übereinstimmt. Wir untermauern dieses Prinzip durch Vor-Trainingsexperimente mit dichten und spärlichen MoE-Sprachmodellen, darunter Architekturen im Stil von Qwen3-0,6B, Gemma 3 1B, OLMoE-1B-7B und verkleinerten gpt-oss-Modellen. In diesen Experimenten verbessern symmetriekompatible Aktualisierungen durchgängig den endgültigen Validierungsverlust und in mehreren Fällen auch die Trainingsstabilität im Vergleich zu entsprechenden AdamW-Aktualisierungen.

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.