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

As a meticulous researcher operating under extreme scrutiny, I have thoroughly analyzed both provided texts (A and B) concerning the paper on Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers.

The synthesis below integrates the high-level narrative from Text A with the granular technical details and structural organization from Text B to produce a comprehensive, detailed summary suitable for a high-stakes review.

Comprehensive Research Summary: Symmetry-Compatible Principle for Optimizer Design

This paper addresses a fundamental geometric disparity in deep learning optimization: while modern neural network architectures inherently possess rich symmetry and equivariance properties, popular coordinate-wise optimization methods (like Adam and its variants) operate independently of these structures, leading to suboptimal performance. The central thesis is that optimizer design must be layerwise and symmetry-compatible, meaning the update rule for a specific parameter block should respect the symmetry group acting on that block.

1. Core Theoretical Framework: The Symmetry-Compatible Principle

The paper introduces a guiding principle: the gradient update rule should be equivariant under the symmetry group acting on the corresponding weight block, and it must remove symmetry-redundant directions when the parameterization possesses a quotient structure. This shifts the paradigm from selecting a single universal update rule to deriving an architecture-aware stack of optimizers.

The unifying perspective is established by showing that bi-orthogonal equivariance naturally leads to the class of spectral optimizers, which are employed by methods such as stochastic spectral descent, Muon, Scion, and polar gradient methods.

2. Derivation of Layerwise Optimizer Classes (Architecture-Optimizer Co-Design)

The authors move beyond standard orthogonal group symmetries to incorporate permutation and shared-shift symmetries to derive tailored optimizers for specific parameter blocks:

• Embeddings and LM Heads: Driven by left-permutation and right-orthogonal equivariance, this leads to a family of updates including one-sided spectral, row-norm, and hybrid row-norm/spectral updates.

• SwiGLU MLP Projections: The intermediate-neuron permutation symmetry motivates the development of specialized row-aware and column-aware variants.

• MoE Routers: Leveraging expert-permutation symmetry combined with shared-logit-shift invariance, this yields novel projected centered row-norm and left-spectral updates.

This systematic derivation results in an end-to-end layerwise optimizer stack where the chosen update law precisely matches the symmetry group of its respective parameter class.

3. Convergence Analysis and Mathematical Rigor

The paper provides rigorous convergence analysis for these novel classes of optimizers, moving beyond simple empirical observation:

• General Update Scheme: The analysis begins with a general first-order iteration, establishing standing assumptions (L-smoothness and mu-Polyak–Łojasiewicz condition).

• Spectral Optimizers: Full spectral optimizers are shown to converge sublinearly to stationarity, while normalized polar-type spectral methods achieve linear convergence under the PŁ condition.

• One-Sided Spectral Optimizers: Analysis of right-spectral and left-spectral updates confirms descent and linear convergence results under the PŁ condition.

• Row-Norm Optimizers: Studies on row-norm variants (including smoothed versions like RowNormM) establish sublinear convergence to stationarity, with hybrid methods (e.g., HybridPolarGradM) demonstrating distinct geometric effects based on the order of normalization steps.

The analysis rigorously proves that these symmetry-compatible updates possess desirable convergence properties, contrasting them implicitly with the coordinate-wise dependence of AdamW.

4. Empirical Validation and Performance Superiority

The theoretical framework is corroborated through extensive pre-training experiments conducted on a diverse suite of language models, including dense and sparse Mixture-of-Experts (MoE) architectures such as Qwen3-0.6B, Gemma 3 1B, OLMoE-1B-7B, and downsized GPT variants.

Key Empirical Findings:

• Validation Loss Improvement: Across all tested architectures, symmetry-compatible update rules consistently improve final validation loss compared to the standard AdamW baseline.

• Load Balancing in MoE Models: The paper demonstrates that these updates effectively reduce expert load imbalance, a critical issue in sparse MoE models. Specifically, certain left-polar variants show the strongest load-balancing performance diagnostics (e.g., router z-loss).

• Logit Control: For LM head matrices, projected updates are shown to control final vocabulary-logit growth, which is complementary to techniques like final logit soft-capping and improves overall training stability.

Conclusion: The Architectural Imperative

The paper concludes that optimizer design for modern deep learning systems must be treated as an architecture-aware problem rather than the selection of a single universal update rule. By assigning an optimizer whose equivariance matches the layerwise symmetry, the resulting geometry-compatible stack demonstrably improves optimization trajectories and final validation loss across various model architectures.

Researcher's Note: The structure is robust. The transition from theoretical derivation (Section 3) to rigorous convergence proofs (Section F) and finally to empirical validation across complex MoE systems (Section G) provides a complete, high-fidelity argument for the proposed methodology. The distinction between orthogonal group symmetries and permutation/shared-shift symmetries is clearly articulated as the key driver for deriving novel update classes. This paper represents a significant methodological advance in optimizing neural network training dynamics based on geometric principles.

1. No longer rely on a single coordinate-wise optimizer for all parameters; instead, implement an architecture–optimizer co-design principle where each major matrix-valued parameter class is assigned an update whose equivariance matches its symmetry group. This ensures that updates are layerwise and symmetry-compatible, leading to improved final validation loss across dense and sparse MoE language models.

2. For vocabulary-indexed matrices (embeddings and LM heads), use left-permutation/right-orthogonal equivariant updates, such as row-norm or hybrid row-norm/right-spectral updates, which are tailored to their specific geometries rather than the general AdamW approach. This specifically improves optimization trajectories for these matrices in Qwen3 and Gemma 3 pre-training.

3. For SwiGLU MLP projections, utilize row-aware and column-aware variants motivated by intermediate-neuron permutation symmetry, enabling researchers to choose between row-norm or hybrid updates based on whether the bottleneck is tall or skinny ("dff >> dmodel").

4. For MoE router matrices, deploy projected centered row-norm and left-spectral updates which respect expert-permutation symmetry together with shared-logit-shift invariance, resulting in reduced expert load imbalance and improved router stability as measured by router z-loss.

5. Implement momentum variants like LeftPolarGradM or RightPolarGradM for ordinary matrix layers, which are derived from bi-orthogonal equivariance, ensuring that the update rule transforms consistently with the representation of its input direction.

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.