为不同层设计对称兼容的优化器,提升训练稳定性和模型性能。
Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers
- 基于对称性原则设计更新规则,匹配各类参数块的几何结构。
- 在多个语言模型上验证,显著降低专家负载不均衡并改善稳定性。
- 适用于嵌入、LM头、SwiGLU和MoE路由器等特殊结构,适合研究高效优化器者。
深度学习实践中长期存在显著的几何不一致性:现代神经网络架构具有丰富的对称性与等变性,但主流优化器如Adam及其变体本质上是坐标轴独立的,无法尊重参数空间的等变结构。本文提出一种对称兼容的优化器设计原则——梯度更新规则应与作用于相应权重块的对称群保持等变。基于此,我们统一了矩阵层的双正交等变更新(如随机谱下降、Muon、Scion、极坐标梯度法)。更重要的是,通过从正交群转向排列与共享位移对称性,推导出适用于非通用矩阵层的对称兼容优化器:包括嵌入与语言模型头矩阵、SwiGLU MLP投影、MoE路由器矩阵。具体构造涵盖单边谱、行范数、混合行范数/谱、行感知、列感知、中心化行范数、左谱更新等。这些方法构成端到端的逐层优化器栈,每类矩阵参数均配备与其对称群匹配的更新方式。在密集与稀疏MoE语言模型(含Qwen3-0.6B、Gemma 3 1B、OLMoE-1B-7B及缩减版gpt-oss)的预训练实验中,该原则一致提升最终验证损失,减少稀疏MoE中的专家负载不均衡,并在部分情况下控制词汇表逻辑值增长、增强路由器稳定性与整体训练稳定性,优于对应AdamW更新。
原文摘要 · Abstract (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 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。