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

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 LLMs, including Qwen3-0.6B-style, Gemma 3 1B-style, OLMoE-1B-7B-style, and downsized gpt-oss architectures. Across these experiments, symmetry-compatible updates consistently improve final validation loss, and in several cases training stability, over corresponding AdamW updates.

• The paper introduces a framework where optimizers are tailored to layer-specific symmetry groups in matrix-valued parameters, enabling geometry-compatible updates.
• It demonstrates that spectral and row-norm based updates outperform standard AdamW, particularly for embeddings, LM heads, SwiGLU MLPs, and MoE routers.
• Empirical results show improved training stability and lower validation loss in large-scale pretraining when applying symmetry-aware optimization.

Symmetry-Compatible Optimizer Design for Matrix-Valued Parameters in Deep Learning

Introduction

The explored paper presents a symmetry-compatible principle for optimizer design in deep learning, explicitly connecting the symmetries inherent in modern neural architectures with the geometry of update rules for matrix-valued parameters. The central thesis is that each parameter block in a neural network is endowed with a particular symmetry group induced by its functional and architectural role (e.g., permutation symmetry for embeddings, orthogonal symmetry for linear layers, and shift-permutation symmetry for routers in MoE architectures). The paper formalizes and operationalizes the notion that gradient update maps should be chosen to be equivariant with respect to these layerwise symmetries, rather than universally applying coordinate-wise adaptive optimizers such as AdamW.

Symmetry-Equivariant Optimizer Framework

The authors introduce a framework where optimizer update maps $U$ satisfy $U(g \cdot D) = g \cdot U(D)$ for a layer-specific symmetry group $G$ and update direction $D$ (often a gradient or momentum buffer). This generalizes conventional practices by asserting that the optimizer must respect the equivariance structure not of the scalar loss or input/output spaces, but rather of the parameter tensor considered as an intrinsic functional object—typically a matrix or tensor whose axes each carry a meaningful indexing set.

Concretely, the following symmetries and corresponding optimizer classes are derived:

• Fully-connected/attention projections: Bi-orthogonal symmetry (invariance/equivariance under left/right orthogonal transformations), motivating spectral or polar-gradient optimizers.
• Input embeddings, LM heads: Left-permutation (vocabulary) $\times$ right-orthogonal (feature) symmetry, yielding "LPRO-equivariant" optimizers: row-norm, right-spectral, and hybrid row-norm/right-spectral updates.
• SwiGLU MLP projections: Intermediate-neuron permutation symmetry, which prescribes row/column-aware variants of the above.
• MoE routers: Expert-permutation and shared-logit-shift invariance, leading to centered row-norm, left-spectral, and hybrid optimizers.

For all classes, the optimizer-architecture assignment is governed by the transformation properties under the appropriate group, formalized via update laws (see Section 3 in the source).

Technical Construction and Examples

Bi-Orthogonal Equivariance and Spectral Updates

A main result is the structural characterization that the only continuous direction-wise update maps $U$, equivariant under left and right orthogonal transformation ($U(PDQ^T) = P U(D) Q^T$), are spectral operators—i.e., maps that act in the singular vector frame of the update direction. This class subsumes gradient descent, spectral descent, polar updates, MUON, and related methods. Such updates preserve subspace geometry and are strictly compatible with the symmetry of general matrix blocks.

LPRO-Equivariant Updates: Embeddings and LM Heads

For parameters with left-permutation/right-orthogonal symmetry, the paper develops:

• Right-spectral updates: Act via the right Gram matrix (hidden features), e.g., $UR(D) = D(D^T D + \epsilon I)^{-1/2}$.
• Row-norm updates: Local rowwise normalizations: $U_{row}(D) = \text{Diag}(n(\| D_{i:} \|^2)) D$ for some function $n$.
• Hybrids: Sequential or composed normalization and spectral transformations.

The rationale is that for these parameter classes (tall, vocabulary-indexed matrices), row permutations or scalings are physically meaningful; thus row-aware or Gram-aware optimizers are both symmetry-preserving and practical.

SwiGLU: Permutation Symmetry in Neurons

SwiGLU blocks (and their relatives in MLP architectures) admit a permutation symmetry over intermediate neurons, which dictates that updates should be equivariant to permutations over the relevant axis. The optimizer assignments mirror those for embeddings, assigned along the correct axes.

MoE Routers: Centering and Expert-Permutation

Routers in mixture-of-experts models have discrete permutation and shared-shift symmetry. The update must act invariantly under relabeling and global logit shifts, leading to the default of centering over the expert axis and then applying either a left-spectral, centered row-norm, or hybrid update. This is practical and maintains the scientific structure of the optimization dynamics.

Empirical Results

The layerwise symmetry-compatible optimizer stack is instantiated and evaluated in end-to-end LLM pre-training experiments for multiple architectures:

• Dense: Qwen3-0.6B-style, Gemma 3-1B-style.
• Sparse/MoE: OLMoE-1B-7B, downsized gpt-oss.

In each case, parameter classes (embeddings, LM head, SwiGLU projections, routers) are assigned their symmetry-compatible optimizer. The following empirical findings are robust:

• Geometry-compatible optimizers consistently outperform ADAMW for vocabulary-indexed matrices in both dense and MoE settings, with the effect increasing as matrix size/anisotropy increases.
• In large models, symmetry-compatible updates both improve final validation loss and frequently confer better training stability.
• For SwiGLU blocks, hybrid row-norm/right-spectral updates are advantageous in the tall-skinny regime.
• For routers, left-spectral and centered row-norm updates are competitive, with hybrid variants offering further flexibility.

For example, in Gemma 3 1B settings, using HYBRIDPOLARGRADM for embeddings/LM-heads yields a final validation loss of 4.0461, compared to 4.0552 for RowNORMM and 4.0862 for ADAMW (exact numbers in Section 4.2).

Theoretical and Practical Implications

The symmetry-to-optimizer principle exposes the geometric misalignment of coordinate-wise optimizers with the functional role of high-level parameter blocks. It extends the longstanding, input-focused equivariance paradigm to the internal model geometry, advocating for a new regime of optimizer-architecture co-design. This is distinct from methods that merely use norm-based steepest descent or learned preconditioning: only update geometries compatible with the symmetry group avoid the issues of parameterization dependence, non-invariance, and high-rank noise injection seen in coordinate-wise methods.

Numerous open-source and industry efforts have begun to adopt matrix-aware optimizers at scale (e.g., MUON), and implementation challenges (efficient polar decompositions, distributed computation) are an active area. The proposed perspective harmonizes with, but is distinct from, works on modular norm theory, LMO-based optimization, and rotation-based updates—the critical addition is the symmetry-driven optimizer selection at the parameter-class level.

Outlook and Future Directions

The symmetry-compatible framework invites further development in several directions:

• Efficient implementation and numerical linear algebra (e.g., fast polar decomposition, Gram root approximations).
• Distributed and large-scale parallel training, memory trade-offs.
• Richer hybrid optimizers combining symmetry-compatible geometry with adaptive scaling, preconditioning, and regularization.
• Systematic benchmarking across model families, architectures, and scale regimes.
• Extension to more general parameter tensor symmetries (e.g., convolution, attention heads, multimodal blocks).
• Theoretical investigation into convergence, generalization, and scaling behavior for symmetry-compatible designs.

As the field moves towards heterogenous and increasingly modular architectures, the principle delineated here provides a coherent foundation for principled, scalable, and theoretically justified optimizer assignment.

Conclusion

The symmetry-compatible principle for optimizer design formalizes the notion that gradient updates for matrix-valued parameter blocks in deep neural networks must respect the natural symmetries of those blocks. By mapping parameter class to optimizer geometry via the layer's symmetry group, the framework generalizes and unifies a wide range of recent advances in matrix-gradient optimization, encompassing spectral, one-sided, row-norm, and hybrid approaches. Empirically, the approach delivers improved optimization outcomes in pre-training LLMs, especially as vocabularies and model sizes increase. The practical and theoretical consequences advocate for routine adoption of geometry- and symmetry-aware optimization, leading toward a more mature and architecture-aware science of large-scale neural network training.

Reference: "Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers" (2605.18106)