Symmetrie-compatibel principe voor optimalisatorontwerp: embeddings, LM-heads, SwiGLU MLP's en MoE-routers

Samenvatting

Een opvallende geometrische discrepantie bestaat al lange tijd in de praktijk van deep learning. Hoewel moderne neurale netwerkarchitecturen van nature rijke symmetrie- en equivariantie-eigenschappen vertonen, werken populaire optimalisatoren zoals Adam en zijn varianten inherent coördinaatsgewijs, waardoor ze de equivariantiestructuren van de parameterruimte niet kunnen respecteren. We pakken deze discrepantie aan door een symmetrie-compatibel principe voor optimalisatorontwerp te introduceren: de gradiënt-updateregel moet equivariant zijn onder de symmetriegroep die op het corresponderende gewichtsblok werkt. Volgens dit principe geven we eerst een uniform perspectief op bi-orthogonaal equivariante updates voor algemene matrixlagen, zoals gebruikt door stochastic spectral descent, Muon, Scion en polaire gradiëntmethoden. Belangrijker nog, door over te stappen van orthogonale groepen naar permutatie- en gedeelde-verschuivingssymmetrieën, leiden we symmetrie-compatibele optimalisatoren af voor parameterblokken waarvan de symmetrieën verschillen van die van algemene matrixlagen: embedding- en LM-head matrices, SwiGLU MLP-projecties en MoE-router matrices. Deze constructies omvatten eenzijdige spectrale, rij-norm, hybride rij-norm/spectrale, rij-bewuste, kolom-bewuste, gecentreerde rij-norm en links-spectrale updates. Ze leveren een end-to-end laagsgewijze optimalisatorstack op waarin elke belangrijke matrixwaardige parameterklasse een update krijgt waarvan de equivariantie overeenkomt met haar symmetriegroep. We bevestigen dit principe door pre-trainingsexperimenten op dichte en sparse MoE-taalmodellen, waaronder Qwen3-0.6B-stijl, Gemma 3 1B-stijl, OLMoE-1B-7B-stijl en verkleinde gpt-oss-architecturen. In al deze experimenten verbeteren symmetrie-compatibele updates consequent het uiteindelijke validatieverlies, en in verschillende gevallen de trainingsstabiliteit, ten opzichte van overeenkomstige AdamW-updates.

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.