Technical Summary: Symmetry-Compatible Principle for Optimizer Design

1. Problem Statement

A fundamental geometric disparity exists in modern deep learning practice. While neural network architectures (embeddings, LM heads, SwiGLU MLPs, MoE routers) possess rich intrinsic symmetry and equivariance properties, dominant optimization methods like Adam, AdamW, and RMSprop operate coordinate-wise. These methods treat matrix parameters as flattened vectors of independent coordinates, thereby ignoring the matrix structure and breaking the natural equivariance of the parameter space.

This "geometry-blind" approach fails to distinguish between the distinct geometries of different layer types. Consequently, popular optimizers may discard spectral structure, induce mismatched geometries, and fail to respect the invariance and equivariance structures inherent to specific parameter blocks (e.g., permutation invariance in vocabularies or expert interchangeability in MoEs). Current theoretical understanding and empirical comparisons between coordinate-wise adaptive methods and spectral optimizers remain limited, often relying on large-scale benchmarks without deep analysis of geometric behavior.

2. Methodology: The Symmetry-Compatible Principle

The paper proposes a layerwise symmetry-compatible principle for optimizer design. The core thesis is that the optimizer update rule $\mathcal{U}$ for a parameter block $W$ should be equivariant under the symmetry group $G$ acting on that specific block.

2.1 Theoretical Framework

Let $W \in \mathbb{R}^{m \times n}$ be a parameter matrix. If the coordinates of the input and output spaces are transformed by invertible matrices $P$ and $Q$, the linear map transforms as $\tilde{W} = P W Q^{-1}$. The gradient transforms contravariantly/covariantly accordingly. The optimizer update map $\mathcal{U}$ must satisfy:
$$\mathcal{U}(g \cdot D) = g \cdot \mathcal{U}(D)$$
where $D$ is an update direction (e.g., gradient or momentum) and $g$ is an element of the symmetry group.

The paper derives specific optimizer classes based on the symmetry groups of different layer types:

1. Ordinary Matrix Layers (Linear/Attention):

   • Symmetry: Bi-orthogonal ($O_m \times O_n$).
   • Optimizer Class: Spectral Optimizers. The update must be a spectral operator acting on singular values while preserving singular vectors. This unifies Stochastic Spectral Descent (SSD), Muon, Scion, and Polar Gradient methods. The update is typically the orthogonal polar factor of the momentum direction.

2. Embeddings and LM Heads:

   • Symmetry: Left-permutation ($P_v$) and Right-orthogonal ($O_d$). Rows index vocabulary items (discrete permutation), columns index hidden features (continuous rotation).
   • Optimizer Class: Left-Permutation Right-Orthogonal (LPRO) Equivariant updates.
     • Right-Spectral: Acts globally via the right Gram matrix ($D^\top D$).
     • Row-Norm: Acts locally on row norms (vocabulary items).
     • Hybrid: Compositions of row-norm and right-spectral updates.

3. SwiGLU MLP Projections:

   • Symmetry: Intermediate-neuron permutation ($P_{d_{ff}}$). The gate and up projections have rows corresponding to intermediate neurons; the down projection has columns corresponding to them.
   • Optimizer Class: Row-aware and Column-aware variants. Gate/Up projections use LPRO-compatible updates (row-norm/right-spectral); Down projections use the transposed geometry.

4. MoE Routers:

   • Symmetry: Expert-permutation ($P_e$) and Shared-logit-shift invariance (adding a constant to all logits does not change the softmax distribution).
   • Optimizer Class: Router-compatible updates. These operate on the centered direction $D_c = \Pi^\perp D$ (removing the shared shift).
     • Centered Left-Spectral: Acts on the centered expert Gram matrix.
     • Centered Row-Norm: Acts on centered expert row norms.

2.2 Practical Implementation

The authors instantiate these principles into practical momentum-based optimizers:

• RightPolarGradM / LeftPolarGradM: One-sided spectral updates using Newton-Schulz iterations for matrix inverse square roots.
• RowNormM: Row-wise normalization updates.
• HybridPolarGradM: Compositions of row-norm and one-sided spectral steps.
• Centering: For MoE routers, updates are applied to the centered momentum $M_{k,c} = \Pi^\perp M_k$.

3. Key Contributions

1. Symmetry-Compatible Principle: A formal argument that coordinate-wise optimizers are geometrically mismatched for matrix-valued parameters. The paper establishes that update rules should match the symmetry group of the parameter block.
2. Unifying View of Spectral Optimizers: Provides a symmetry-based interpretation of spectral optimizers (Muon, SSD, etc.) as the unique class of bi-orthogonally equivariant update maps.
3. New Equivariant Optimizer Classes: Derives specific update rules for layers with non-standard symmetries:
   • LPRO-equivariant updates for embeddings and LM heads.
   • Row/column-aware updates for SwiGLU MLPs.
   • Centered expert-permutation updates for MoE routers.
4. End-to-End Pre-training Evidence: Constructs a full layerwise optimizer stack where every major matrix class receives a symmetry-matched update. This is evaluated on dense and sparse MoE language models (Qwen3-0.6B, Gemma 3 1B, OLMoE-1B-7B, downsized gpt-oss).

4. Experimental Results

The authors conducted pre-training experiments on 10B-token subsets of FineWeb-Edu across four architectures.

• Dense Models (Qwen3-0.6B, Gemma 3 1B):

  • Replacing AdamW on vocabulary-indexed matrices (embeddings/LM heads) with RowNormM or HybridPolarGradM consistently improved final validation loss.
  • The gains were modest for the smaller Qwen3 model but more pronounced for the larger Gemma 3 model, suggesting geometry-aware updates become more critical as vocabulary matrices grow larger and gradients become more anisotropic.
  • Using HybridPolarGradM for SwiGLU MLP projections further improved performance, likely due to correcting row-scale imbalance in tall-skinny matrices.

• Sparse MoE Models (OLMoE-1B-7B, downsized gpt-oss):

  • Configurations using symmetry-compatible updates for embeddings, LM heads, and routers (e.g., RowNormM + LeftPolarGradM) outperformed configurations using AdamW for these components.
  • The fully coordinate-wise baseline (AdamW everywhere) performed the worst.
  • Symmetry-compatible router updates improved over coordinate-wise router updates and, in some cases, reduced training loss spikes, suggesting improved stability.

• General Trend: Across all models, replacing AdamW on large vocabulary-indexed matrices with symmetry-compatible optimizers consistently lowered final validation loss. The benefits persisted across different architectures and sparsity levels.

5. Significance and Claims

The paper claims to offer a principled shift in optimizer design from a "geometry-blind" coordinate-wise view to a "symmetry-aware" matrix view.

• Architecture-Optimizer Co-Design: The work argues that optimizer design should be an architecture-aware problem. Different parameter classes (embeddings, routers, linear layers) have distinct natural symmetries and require distinct optimizer geometries.
• Theoretical Grounding: It provides a rigorous mathematical characterization (via equivariance) for why spectral optimizers work for linear layers and extends this logic to derive new optimizers for other layers.
• Practical Utility: The results demonstrate that symmetry-compatible assignments are not just theoretically sound but practically beneficial, improving final loss and potentially training stability in large-scale pre-training.
• Modesty: The authors explicitly state they do not claim equivariant optimizers dominate coordinate-wise methods in all regimes. Instead, they present evidence that the layerwise symmetry-compatible principle leads to competitive and often beneficial assignments in representative pre-training settings. They view the results as evidence for the principle's usefulness rather than an exhaustive benchmark.

The paper concludes that introducing symmetry and equivariance into optimizer design opens a path toward more stable, parameterization-consistent, and theoretically grounded large-scale model pre-training.