Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers

Abstract

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 update rules consistently improve final validation loss, reduce expert load imbalance in sparse MoE models, and in several cases control final vocabulary-logit growth, improve router stability, and overall training stability over the corresponding AdamW updates.

A Degeneracy-Aligned and Discrete Maximum-Principle- Compatible Peridynamic Differential Operator for the Heston Stochastic Volatility Model
Dan Qu, Yanbo Ma
Mathematics2 Oct 2026Peer-reviewedCC BY 4.0General mathematicsCrossref
Weak Asymptotic Symmetry of Quasisymmetric Embeddings on Quasilines
Katsuhiko Matsuzaki, Fei Tao
arXiv17 Sep 2026PreprintCC BY 4.0Full textAnalysis & PDEsarXiv
Minimal output entropy for the channels that add or remove a box of a Young diagram, and a Pauli principle for every permutation symmetry
Robin Reuvers
arXiv6 Oct 2026PreprintarXiv licenceFull textMathematical physics Quantum physics & quantum informationarXiv
Task- and Session-Level Model Routing: A Common-Interface Hybrid Evaluation of Four Open-Source Routers Across Four Benchmarks
Kiran N. Kumar, Santhosh K. Saminathan
arXiv1 Aug 2026PreprintCC BY 4.0Full textStatisticsarXiv
Multiplicative vs. Additive State Coupling in Learned Optimizer Schedules
M. Harper Langston, Pierre-David Letourneau, Dalton Jones +7
arXiv23 Oct 2026PreprintCC BY-NC-ND 4.0Full textNumerical analysis & scientific computing Optimisation & controlarXiv

Cite

BibTeX RIS

BibTeX

@misc{lau2026symmetrycompatible,
  title = {{Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers}},
  author = {Tim Tsz-Kit Lau and Weijie Su},
  howpublished = {arXiv},
  year = {2026},
  month = oct,
  doi = {10.48550/arxiv.2605.18106},
  eprint = {2605.18106},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2605.18106},
  note = {CC BY 4.0},
}

RIS

TY  - UNPB
TI  - Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers
AU  - Tim Tsz-Kit Lau
AU  - Weijie Su
PY  - 2026
DA  - 2026/10/14
PB  - arXiv
DO  - 10.48550/arxiv.2605.18106
UR  - https://arxiv.org/abs/2605.18106
AB  - 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 update rules consistently improve final validation loss, reduce expert load imbalance in sparse MoE models, and in several cases control final vocabulary-logit growth, improve router stability, and overall training stability over the corresponding AdamW updates.
ER  -

Abstract

Related

Cite