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

• URL: http://arxiv.org/abs/2605.18106v2
• Date: Mon, 25 May 2026 18:11:39 GMT
• Title: Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers
• Abstract summary: A striking geometric disparity has long persisted in the practice of deep learning.<n>We introduce a symmetry-compatible principle for design: the gradient update rule should be equivariant under the parameter group acting on the corresponding weight block.
• Score: 3.433766572511366
• License: http://creativecommons.org/licenses/by/4.0/
• 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 load imbalance in sparse MoE models, and in several cases improve training stability over the corresponding AdamW updates.

Related papers

• Aligning Network Equivariance with Data Symmetry: A Theoretical Framework and Adaptive Approach for Image Restoration [71.51826063983434]
  Image restoration is an inherently ill posed inverse problem.<n>We use non strict symmetry as a constraint to formulate the restoration inverse problem.<n>We show that the equivariance error of the optimal restoration operator is strictly bounded by the data symmetry error and the discretization mesh size.
  arXiv  Detail & Related papers  (2026-05-13T16:22:19Z)
• Uncovering Symmetry Transfer in Large Language Models via Layer-Peeled Optimization [6.787340999768117]
  Large language models (LLMs) are convex convex by minimizing crossentropy loss for next-shift prediction.<n>We show that when the target tokens exhibit a cyclic symmetry, the optimal logit matrix is exactly circulant and the Gramtoken both form circulant output as well.
  arXiv  Detail & Related papers  (2026-05-12T21:10:34Z)
• Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation [30.86228838205224]
  We introduce Pion, a spectrum-preserving for large language model (LLM) training based on orthogonal equivalence transformation.<n>We derive the Pion update rule systematically examine its design choices, and analyze its convergence behavior along with several key properties.<n> Empirical results show that Pion offers a stable and competitive alternative to the standard LLMs for both pretraining and finetuning.
  arXiv  Detail & Related papers  (2026-05-12T17:59:34Z)
• Muon-OGD: Muon-based Spectral Orthogonal Gradient Projection for LLM Continual Learning [15.604900995606451]
  We propose Muon-OGD, a spectral-norm-aware continual learning framework.<n>We evaluate the proposed method on standard continual learning benchmarks.
  arXiv  Detail & Related papers  (2026-05-09T13:42:08Z)
• Rethinking Diffusion Models with Symmetries through Canonicalization with Applications to Molecular Graph Generation [56.361076943802594]
  CanonFlow achieves state-of-the-art performance on the challenging GEOM-DRUG dataset, and the advantage remains large in few-step generation.
  arXiv  Detail & Related papers  (2026-02-16T18:58:55Z)
• Reinforcement Learning Using known Invariances [54.91261509214309]
  This paper develops a theoretical framework for incorporating known group symmetries into kernel-based reinforcement learning.<n>We show that symmetry-aware RL achieves significantly better performance than their standard kernel counterparts.
  arXiv  Detail & Related papers  (2025-11-05T13:56:14Z)
• Generalized Linear Mode Connectivity for Transformers [87.32299363530996]
  A striking phenomenon is linear mode connectivity (LMC), where independently trained models can be connected by low- or zero-loss paths.<n>Prior work has predominantly focused on neuron re-ordering through permutations, but such approaches are limited in scope.<n>We introduce a unified framework that captures four symmetry classes: permutations, semi-permutations, transformations, and general invertible maps.<n>This generalization enables, for the first time, the discovery of low- and zero-barrier linear paths between independently trained Vision Transformers and GPT-2 models.
  arXiv  Detail & Related papers  (2025-06-28T01:46:36Z)
• Beyond the Permutation Symmetry of Transformers: The Role of Rotation for Model Fusion [43.299430093251736]
  We introduce rotation symmetry, a novel form of parameter space symmetry for transformers.<n>Unlike permutation symmetry, rotation symmetry operates in a continuous domain, thereby significantly expanding the equivalence set for transformers.<n>We propose a theoretically optimal matching algorithm as a plug-and-play module to enhance model fusion.
  arXiv  Detail & Related papers  (2025-02-01T01:44:55Z)
• Incorporating Arbitrary Matrix Group Equivariance into KANs [69.30866522377694]
  Kolmogorov-Arnold Networks (KANs) have seen great success in scientific domains.<n>We propose Equivariant Kolmogorov-Arnold Networks (EKAN) to broaden their applicability to more fields.
  arXiv  Detail & Related papers  (2024-10-01T06:34:58Z)
• Learning Layer-wise Equivariances Automatically using Gradients [66.81218780702125]
  Convolutions encode equivariance symmetries into neural networks leading to better generalisation performance. symmetries provide fixed hard constraints on the functions a network can represent, need to be specified in advance, and can not be adapted. Our goal is to allow flexible symmetry constraints that can automatically be learned from data using gradients.
  arXiv  Detail & Related papers  (2023-10-09T20:22:43Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.