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

arXiv: https://arxiv.org/abs/2605.18106 · HTML: https://mlsanity.com/paper/2605.18106/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Placed in MoE theory and gating functions, one of 12 directions in Mixture-of-experts.
- Builds on 65 of its 182 references that are in the map.
- First posted on arXiv 18 May 2026.

## Placed in

- Mixture-of-experts / [MoE theory and gating functions](https://mlsanity.com/moe/moe-theory-and-gating-functions/)

## Builds on (most cited first, 30 of 65)

- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Gemma 3 Technical Report](https://mlsanity.com/paper/2503.19786/) (2025)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [gpt-oss-120b & gpt-oss-20b Model Card](https://mlsanity.com/paper/2508.10925/) (2025)
- [DeepSeek-V3 Technical Report](https://mlsanity.com/paper/2412.19437/) (2024)
- [GLM-5: from Vibe Coding to Agentic Engineering](https://mlsanity.com/paper/2602.15763/) (2026)
- [Kimi K2: Open Agentic Intelligence](https://mlsanity.com/paper/2507.20534/) (2025)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Models](https://mlsanity.com/paper/2508.06471/) (2025)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2402.01739/) (2024)
- [An Introduction to Vision-Language Modeling](https://mlsanity.com/paper/2405.17247/) (2024)
- [Understanding Gradient Orthogonalization for Deep Learning via Non-Euclidean Trust-Region Optimization](https://mlsanity.com/paper/2503.12645/) (2025)
- [A Note on the Convergence of Muon](https://mlsanity.com/paper/2502.02900/) (2025)
- [Practical Efficiency of Muon for Pretraining](https://mlsanity.com/paper/2505.02222/) (2025)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [Dion: Distributed Orthonormalized Updates](https://mlsanity.com/paper/2504.05295/) (2025)
- [Benchmarking Optimizers for Large Language Model Pretraining](https://mlsanity.com/paper/2509.01440/) (2025)
- [AdaMuon: Adaptive Muon Optimizer](https://mlsanity.com/paper/2507.11005/) (2025)
- [JetMoE: Reaching Llama2 Performance with 0.1M Dollars](https://mlsanity.com/paper/2404.07413/) (2024)
- [Benchmarking Neural Network Training Algorithms](https://mlsanity.com/paper/2306.07179/) (2023)
- [PolarGrad: A Class of Matrix-Gradient Optimizers from a Unifying Preconditioning Perspective](https://mlsanity.com/paper/2505.21799/) (2025)
- [Isotropic Curvature Model for Understanding Deep Learning Optimization: Is Gradient Orthogonalization Optimal?](https://mlsanity.com/paper/2511.00674/) (2025)
- [Step 3.5 Flash: Open Frontier-Level Intelligence with 11B Active Parameters](https://mlsanity.com/paper/2602.10604/) (2026)
- [On the Convergence of Muon and Beyond](https://mlsanity.com/paper/2509.15816/) (2025)
- [When do spectral gradient updates help in deep learning?](https://mlsanity.com/paper/2512.04299/) (2025)

## Built on by (most cited first, 11 of 11)

- [Dead-Direction Conditioners: Gauge-Equivariant Preconditioning for Deep Networks](https://mlsanity.com/paper/2606.29176/) (2026)
- [Dead Directions: Geometric Singular Learning](https://mlsanity.com/paper/2606.05957/) (2026)
- [Ember: A Kilobyte-State Optimizer for Token Interfaces](https://mlsanity.com/paper/2607.01455/) (2026)
- [Scaling Muon for Diffusion Transformers](https://mlsanity.com/paper/2608.20818/) (2026)
- [The Loss Does Not See the Basis, but Adam Does](https://mlsanity.com/paper/2608.05136/) (2026)
- [Toward a First-Principles Update Geometry for the Language-Model Head](https://mlsanity.com/paper/2608.22253/) (2026)
- [Blog: Survey of Optimizers](https://mlsanity.com/paper/2608.28557/) (2026)
- [AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models](https://mlsanity.com/paper/2610.01395/) (2026)
- [MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads](https://mlsanity.com/paper/2610.02705/) (2026)
- [Optimizing Large Language Models with Chained LMOs](https://mlsanity.com/paper/2610.10975/) (2026)
- [Taming Softmax at the Limits of Muon: A Unified Approach through Hilbert Geometry](https://mlsanity.com/paper/2610.17588/) (2026)

## Abstract (arXiv)

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.