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

Researchers introduce a symmetry-compatible principle for neural network optimizer design that aligns gradient updates with the geometric properties of different parameter types. The approach yields specialized update rules for embeddings, language model heads, SwiGLU MLPs, and mixture-of-experts routers, demonstrating improved validation loss and training stability across multiple language model architectures compared to standard AdamW optimization.

The fundamental challenge addressed in this research stems from a longstanding architectural mismatch in deep learning: modern neural networks exhibit inherent symmetries and equivariance properties, yet dominant optimizers like Adam operate coordinate-wise without respecting these geometric structures. This disconnect limits optimization efficiency and leaves performance gains on the table across diverse model architectures.

The work builds on prior innovations in equivariant optimization methods—including stochastic spectral descent, Muon, and Scion—which use bi-orthogonal updates for general matrix layers. The key contribution extends this framework beyond orthogonal groups to permutation and shared-shift symmetries, enabling specialized update rules for parameter blocks with unique geometric properties. The resulting layerwise optimizer stack assigns matched update rules to embeddings, language model heads, SwiGLU projections, and MoE routers, each respecting its corresponding symmetry group.

For the AI infrastructure and language model development community, this research carries practical significance. Experiments spanning dense models and sparse mixture-of-experts architectures—including Qwen3, Gemma 3, and OLMoE variants—consistently show symmetry-compatible updates outperforming AdamW baselines on validation loss while improving training stability. In sparse models, the approach additionally reduces load imbalance, a known efficiency problem in mixture-of-experts systems.

The implications extend beyond incremental improvements. Better optimization could reduce computational overhead in model training, lower energy consumption, and accelerate development cycles for large language models. Future work will likely explore whether these principles scale to state-of-the-art model sizes and whether they integrate with emerging training techniques like mixture-of-depths or dynamic routing strategies.

• →Symmetry-compatible optimizer design aligns gradient updates with geometric properties of different neural network parameter types.
• →Specialized update rules for embeddings, LM heads, SwiGLU MLPs, and MoE routers consistently improve validation loss over AdamW.
• →The approach reduces load imbalance in sparse mixture-of-experts models, addressing a known efficiency bottleneck.
• →Experiments span multiple architectures including Qwen3, Gemma 3, and OLMoE, demonstrating broad applicability.
• →Framework extends equivariant optimization beyond orthogonal groups to permutation and shared-shift symmetries.