这篇论文解决了深度学习优化器忽视架构对称性的根本问题,做大规模语言模型预训练的团队值得关注——它可能成为替代AdamW的新范式,实验显示能稳定提升损失和训练稳定性。
这篇论文提出了一种对称性兼容的优化器设计原则,要求梯度更新规则在参数块的对称群作用下保持等变性。作者首先统一了双正交等变更新(如Muon、Scion等方法)的视角,然后针对嵌入层、语言模型头、SwiGLU MLP投影和MoE路由器矩阵等具有不同对称性的参数块,推导了相应的对称性兼容优化器(如单边谱、行范数、混合行范数/谱等更新)。通过在Qwen3-0.6B、Gemma 3 1B、OLMoE-1B-7B等架构上的预训练实验,验证了该方法相比AdamW能持续改善验证损失和训练稳定性。这项工作为不同参数类型提供了端到端的层级优化器栈,使更新规则与参数块的对称性相匹配。
Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers
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.