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

Tim Tsz-Kit Lau; Weijie Su

Source-linked AI summary

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

Tim Tsz-Kit Lau, Weijie Su

arXiv:2605.18106v4May 18, 2026math.OCcs.AIcs.LGstat.ML

TL;DR

Coordinate-wise optimizers overlook the distinct symmetries of neural-network parameter blocks. This paper assigns symmetry-compatible updates layerwise and finds consistently better validation loss across dense and sparse MoE pre-training, with improved MoE load balance and training stability.

Problem
Coordinate-wise optimizers ignore matrix structure and layer-specific symmetries, while evidence on optimizers’ geometric behavior and language-model training dynamics remains limited.
Source: §1 · p. 2
Method
The paper assigns each major matrix-valued parameter class an optimizer whose update equivariance matches its layerwise symmetry group.
Source: §1 · p. 2, §5 · p. 31
Results
Across dense and sparse MoE pre-training, symmetry-compatible assignments consistently improve final validation loss over AdamW and persist across multiple model architectures.
Source: §1 · p. 4, §4.5 · p. 30
Takeaways & Limitations
The results support treating optimizer design as an architecture-aware problem rather than selecting one universal coordinate-wise update rule.
Source: §5 · p. 32
Hybrid row-norm/right-spectral updates achieve the best validation loss but incur higher computational overhead than cheaper row-norm-only alternatives.
Source: §4.2 · p. 27

Abstract

from arXiv · show

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.

1 Introduction

The paper argues that optimizers should be designed layerwise so each matrix-valued parameter block receives updates equivariant to its own symmetry and geometry. It develops corresponding spectral and architecture-specific optimizers and reports improved or more stable language-model pre-training outcomes in representative dense and sparse MoE settings.

Motivation

Coordinate-wise adaptive methods treat parameters as independent vector coordinates, ignoring matrix structure and the differing geometries of embeddings, projections, MLPs, and routers.The paper identifies this geometry-blind treatment as a mismatch with the natural equivariances of modern neural-network parameter blocks.Source: §1 · pp. 1–2, Figure 1 · p. 3

Symmetry-compatible principle

The central principle matches each major matrix-valued parameter class with an optimizer whose equivariance reflects that block’s layer-dependent symmetry and quotient geometry.This replaces a single coordinate-wise optimizer with architecture–optimizer co-design for embeddings, LM heads, dense layers, attention projections, SwiGLU projections, and MoE routers.Source: §1 · pp. 2, 4

Spectral optimizers

Bi-orthogonal equivariance unifies spectral optimizers including stochastic spectral descent, Muon, Scion [124], and polar gradient methods [91].These methods use the polar factor of a gradient or momentum update, preserving its singular-vector geometry; Muon’s spectral-norm steepest-descent interpretation follows from unitary invariance.Source: §1 · pp. 2–3

Architecture–optimizer co-design

The paper derives distinct updates for embeddings and LM heads, SwiGLU projections, and MoE routers from their respective permutation, orthogonal, neuron, and shared-shift symmetries.The resulting classes include one-sided spectral, row-norm, hybrid, row-aware, column-aware, projected centered row-norm, and left-spectral updates, with momentum variants such as RightPolarGradM and LeftPolarGradM.Source: §1 · p. 3

End-to-end pre-training evidence

Layerwise symmetry-compatible optimizer assignments consistently improve validation loss and MoE behavior across representative dense and sparse language-model pre-training experiments.Benefits include reduced expert load imbalance, improved router stability, fewer training-loss spikes, and reduced vocabulary-logit growth in Gemma 3 1B-style pre-training.Source: §1 · p. 4

Scope and limitations

The empirical claims are evidence that the principle is useful in representative pre-training regimes, not proof that equivariant optimizers dominate coordinate-wise methods universally.The paper explicitly does not present an exhaustive large-scale optimizer benchmark.Source: §1 · p. 4

2 Preliminaries and Related Work

This section establishes the notation and mathematical foundations for the paper, situates spectral and matrix-gradient optimizers within prior work, and frames equivariance as a principle for optimizer design. It distinguishes the paper’s parameter-space update-map perspective from architecture-level equivariance and curvature-based preconditioning.

2 Preliminaries and Related Work

The paper defines notation for matrix norms, singular values, inner products, vectorization, symmetric and positive-semidefinite matrices, orthogonal and permutation matrices, projections, and optimization domains.These preliminaries support the later treatment of spectral operators, symmetry groups, and matrix-valued updates.Source: §2 · p. 5

2.1 Matrix-Gradient Optimizers

Muon and stochastic spectral descent are related through spectral-norm steepest descent, whose polar-factor direction is bi-orthogonally equivariant; this paper extends that logic to embeddings, LM heads, SwiGLU projections, and MoE routers.The paper makes equivariance explicit and broadens the class of symmetry-based spectral optimizers beyond fully bi-orthogonal matrix layers.Source: §2.1 · pp. 5–6
Unlike K-FAC [115] [46], Shampoo [64] [7] [142] [45], SOAP [154], and related preconditioners, the framework derives update geometry from a layer’s symmetry group rather than primarily from curvature approximation.Weight constraints, variance reduction, and low-rank projection methods address complementary aspects of matrix-gradient optimization rather than the paper’s central symmetry-based design principle.Source: §2.1 · p. 6

2.2 Matrix Optimization Problems, Löwner Operators, and Spectral Operators

Spectral operators act on singular values while preserving singular vectors, providing the operator-theoretic structure shared by stochastic spectral descent, Muon, Scion, and polar gradient methods.This connects the framework to established matrix optimization, spectral functions, and unitarily invariant operator theory.Source: §2.2 · p. 6

2.3 Symmetry and Equivariance in Deep Learning

The paper transfers equivariance from architecture design to optimizer design by imposing it on parameter-tensor update maps, with the relevant symmetry arising from internal parameter-space geometry.This complements work on equivariant architectures and on the optimization or generalization implications of neural-network symmetries.Source: §2.3 · pp. 6–7

3 Equivariant Optimizers from Layerwise Symmetry

This section develops a framework for designing optimizers whose updates are equivariant to the symmetry groups of different parameter blocks, including ordinary matrices, embeddings and LM heads, routers, and SwiGLU projections. It derives corresponding quotient, row-norm, hybrid, and spectral update classes, while characterizing representative memoryless and stateful equivariant methods.

3.1 A General Symmetry-Induced Optimizer Geometry

The optimizer update should transform consistently with its input direction under the symmetry group of the parameter block, rather than requiring arbitrary coordinate-change invariance.This principle makes optimization independent of representations invisible to the model while allowing layer-specific symmetry groups.Source: §3 · p. 7, §3.1 · p. 8

3.3 Optimizers for Embeddings and LM Heads via Left-Permutation Right-Orthogonal Equivariance

Embeddings and LM heads use left-permutation/right-orthogonal equivariance, enabling row-norm, right-spectral, and hybrid updates instead of assuming full orthogonal symmetry.Untied LM heads additionally require horizontal updates because shared row-vector shifts leave softmax probabilities unchanged; centering and reprojection enforce this quotient constraint.Source: §3 · pp. 7–8, §3.3 · pp. 9–10

3.2 Bi-Orthogonal Equivariance for Ordinary Matrix Layers

For ordinary linear and attention matrices, orthogonal coordinate changes require bi-orthogonal equivariance, motivating polar and spectral updates; Nesterov momentum preserves this structure.The transformed momentum and polar direction satisfy fM_k = PM_kQ⊤, eN_k = PN_kQ⊤, and polar(eN_k) = P polar(N_k)Q⊤.Source: §3.2 · pp. 8–9

3.3.1 Right-Spectral Optimizers

Right-spectral embedding and LM-head updates act through the smaller feature Gram matrix D⊤D, matching low-dimensional gradient geometry when vocabulary size greatly exceeds hidden dimension.These updates are the one-sided analogue of spectral or polar-gradient methods and require only the d×d right Gram matrix.Source: §3.3.1 · p. 11
Right-spectral maps are computationally attractive but challenging for tall, ill-conditioned vocabulary matrices, while Adam is commonly used despite lacking LPRO equivariance.More robust polar decomposition algorithms such as QDWH and ZOLO-PD can compute these updates more accurately [91].Source: §3.3.1 · p. 11

3.3.2 Row-Norm and Hybrid LPRO-Equivariant Optimizers

Because vocabulary symmetry is discrete, LPRO equivariance permits local row-norm and hybrid row-norm/right-spectral maps; LM-head rowwise operations require projection back to the horizontal quotient.Right-spectral maps provide a global Gram-based subclass, whereas row-norm maps provide local row adaptation; composition preserves equivariance.Source: §3.3.2 · pp. 11–12

3.4 Optimizers for SwiGLU MLP Projections

SwiGLU projections have intermediate-neuron permutation symmetry: LPRO-compatible updates apply to gate and up rows, while the down projection uses the transposed, column-permutation geometry.The symmetry holds because permuting gate and up rows and correspondingly permuting down columns leaves the SwiGLU output unchanged; arbitrary orthogonal mixing generally fails with elementwise nonlinearities.Source: §3.4 · p. 13
For feed-forward blocks, hidden-unit permutations preserve the function when corresponding input rows and output columns are permuted, with gated projections permuted together.This symmetry applies to GLU, GeGLU, ReGLU, and SwiGLU when gate and value projections are jointly permuted with down-projection columns.Source: §3.4 · p. 13

3.5 Optimizers for MoE Routers

MoE router updates must be expert-permutation equivariant, shared-row-shift invariant, and horizontal, motivating centered left-spectral and projected centered row-norm families.The final projection is essential because row-dependent rescaling can reintroduce a shared-row component; hybrids combine global expert mixing with local expert-wise normalization.Source: §3.5 · p. 15

3.6 Symmetry-to-Optimizer Principle and Architecture–Optimizer Co-Design

Optimizer geometry should match each layer’s symmetry action, with horizontal projection removing redundant quotient directions such as shared softmax-logit shifts.The resulting classes include spectral, one-sided spectral, row-norm, hybrid, and projected centered updates selected according to the parameter block’s symmetry.Source: §3.6 · pp. 16–17

3.7 Practical Optimizers for Embeddings, LM Heads, SwiGLU MLP Projections, and MoE Routers

Practical implementations use momentum directions, GPU-friendly inverse-square-root or polar-factor approximations, float32 Gram calculations, and centered updates for LM heads and routers.Newton–Schulz iterations with Polar Express coefficients provide the inverse-square-root computations, while centered momentum removes vocabulary-logit or expert shared-shift directions.Source: §3.7 · p. 18, §3.7.1 · p. 18

3.7.1 One-Sided Spectral Optimizers

One-sided spectral optimizers exploit tall-skinny matrices by forming the smaller Gram matrix, using right-spectral updates for embeddings, LM heads, and SwiGLU projections and left-spectral updates for routers.LeftPolarGradM and RightPolarGradM compute the same exact polar direction when both are defined, but differ computationally and in intended layer symmetry.Source: §3.7.1 · p. 18

3.7.2 Row-Norm-Based and Hybrid Variants

Row-norm and hybrid variants provide cheaper alternatives, combining row-wise scaling with one-sided spectral steps while preserving the relevant permutation, orthogonal, and centered-subspace symmetries.For routers, HybridPolarGradM combines centered left-spectral updates with projected expert-wise normalization; for LM heads, projections follow rowwise nonlinear stages and the final update.Source: §3.7.2 · pp. 19–20
For LM heads and MoE routers, row-norm and hybrid updates are projected onto horizontal quotient subspaces to remove shared-row directions; routers additionally use centered expert geometry.The router alternatives include projected centered row-norm, centered left-spectral, and projected hybrid left-spectral/row-norm updates.Source: §3.7.2 · p. 20

3.8 Spectral Optimizers as the Bi-Orthogonally Equivariant Class

Bi-orthogonally equivariant continuous direction-wise update maps are exactly spectral operators that transform singular values while preserving singular-vector geometry.The characterization concerns optimizer updates under rotated gradient or momentum directions, not invariance of the layerwise loss itself.Source: §3.8 · p. 22

3.8.1 Examples of Spectral and Equivariant Matrix Optimizers

Spectral optimizers are memoryless direction-wise maps, whereas stateful methods such as Shampoo [64, 7] can remain equivariant through transforming auxiliary states without generally being spectral operators of the current gradient alone.SOAP [154] is likewise generally non-spectral because it uses history-dependent eigenspaces and coordinate-wise adaptive scaling, becoming spectral only in special aligned cases.Source: §3.8 · p. 22, §3.8.1 · p. 23
The spectral class includes gradient descent, momentum variants, stochastic spectral descent, Muon, Scion [124], and PolarGrad [91], while unitarily invariant normalization preserves spectrality.Row-wise or column-wise normalization generally breaks full bi-orthogonal equivariance but can suit layers with one-sided or permutation symmetries.Source: §3.8.1 · p. 23

4 Numerical Experiments

Across dense and sparse language-model pre-training, a layerwise symmetry-compatible optimizer stack consistently improves final validation loss over AdamW on large vocabulary-indexed matrices, with additional benefits for SwiGLU projections and MoE routers. The experiments support applying geometry-matched updates per parameter class rather than globally replacing AdamW.

4 Numerical Experiments

The experiments instantiate a layerwise stack: spectral updates for attention, row-norm or hybrid updates for embeddings and LM heads, geometry-aware SwiGLU updates, router-specific updates, and coordinate-wise optimization only for scalar or vector parameters.Fused tensors are interpreted according to functional intermediate-neuron geometry rather than storage layout, and untied embeddings and LM heads receive separate assignments.Source: §4 · pp. 23–24

4.2 Gemma 3 1B-Style Pre-Training

Gemma 3 1B-style training achieves final validation losses of 4.0516, 4.0435, and 4.0862 with RowNormM, HybridPolarGradM, and AdamW assignments, respectively, making HybridPolarGradM best.RowNormM remains substantially better than AdamW and is computationally cheaper because it avoids spectral inverse-square-root computations; projected LM-head updates also reduce final vocabulary-logit growth.Source: §4.2 · pp. 26–27, Figure 4 · p. 26

4.3 OLMoE-1B-7B-Style Pre-Training

In OLMoE-1B-7B-style training, final validation losses are 4.0955, 4.0815, 4.1187, and 4.1240, with symmetry-compatible configurations (i) and (ii) outperforming AdamW-router and fully AdamW configurations.The AdamW-all baseline initially progresses faster but is overtaken; geometry-compatible router variants have similar performance, and the fully coordinate-wise setup shows slightly larger training-loss spikes.Source: §4.3 · p. 28, Figure 5 · p. 28

4.4 Downsized gpt-oss Pre-Training

Downsized gpt-oss achieves final validation losses of 4.3080, 4.3136, 4.3388, and 4.3687, with RowNormM-based configurations substantially outperforming fully coordinate-wise AdamW.Configuration (i), using RowNormM for embeddings, LM heads, and routers, is best; symmetry-compatible router updates also reduce load imbalance and router z-loss relative to AdamW routers.Source: §4.4 · p. 29, Figure 6 · p. 29

4.5 Cross-Model Comparison

Across four models, symmetry-compatible optimizers consistently improve final validation loss over AdamW on large vocabulary-indexed matrices, with gains increasing in Gemma and persisting in sparse MoE models.The comparison is not a scaling law because architectures, vocabulary sizes, training lengths, and sparsity structures differ.Source: §4.5 · p. 30

5 Discussion and Outlook … A.3 Rotation-Based Optimizers

The paper argues that optimizer design should be layerwise, geometry-aware, and symmetry-compatible, assigning different parameter classes update rules that match their natural equivariances. Its discussion situates this principle among related optimization frameworks while identifying numerical-linear-algebra and distributed-systems challenges for scaling geometry-aware methods.

5 Discussion and Outlook

Modern heterogeneous architectures make a single coordinate-wise optimizer increasingly unnatural because distinct parameter blocks have different intrinsic geometries and natural symmetries.The paper frames geometry-aware optimization as an architecture-aware alternative to universally applying Adam or AdamW.Source: §5 · pp. 30–31
Symmetry-compatible optimization assigns ordinary matrices spectral updates, embeddings and LM heads left-permutation/right-orthogonal updates, SwiGLU projections row- or column-aware updates, and MoE routers expert-permutation-equivariant updates.These constructions support an architecture–optimizer co-design principle in which each parameter class receives an update matched to its layerwise symmetry.Source: §5 · p. 31
Scaling geometry-aware optimizers requires efficient matrix decompositions and distributed implementations addressing communication, synchronization, sharding, and parallelism; recent work suggests these obstacles are tractable.The discussion highlights Newton–Schulz-based methods and distributed systems such as Distributed Muon [107] [126], Dion [3] [2], Disco [49], and Parallel Muon [103].Source: §5 · p. 31

Appendix

The appendix organizes further related work, matrix-analysis background, geometric misalignment, omitted proofs, optimizer implementation, convergence analysis, and numerical-experiment details.This overview identifies the appendix materials supporting the paper’s theoretical and empirical development.Source: Appendix · p. 45

A.1 Non-Euclidean Norm-Based Steepest Descent and LMO-Based Frameworks

For norm-based steepest descent, the symmetry analysis favors unitarily invariant norms for orthogonally invariant matrix geometries; ℓp →ℓq operator norms are generally inappropriate except when p = q = 2.This provides a geometric criterion for norm selection and helps explain discrepancies between theoretical prescriptions and empirical practice.Source: §A.1 · p. 45

A.2 Modular Norm Theory

Modular norm theory targets scale transfer across width and depth, whereas right-spectral updates derive directional geometry from left-permutation/right-orthogonal equivariance; combining them could yield fully geometry-aware optimization.The proposed combination uses modular-norm global scaling with symmetry-dictated spectral, right-spectral, or left-spectral directions.Source: §A.2 · pp. 45–46

A.3 Rotation-Based Optimizers

Unlike rotation-based optimizers, this framework does not require rotationally invariant layerwise losses or a single update rule, and it covers full spectral, right-spectral, left-spectral, row-norm, and hybrid classes.Its broader taxonomy follows the differing symmetry groups of matrix-valued layer types, including embeddings and LM heads.Source: §A.3 · p. 46

B Supplemental Technical Background on Matrix Analysis

This section defines unitarily invariant norms as matrix norms invariant under left and right orthogonal transformations and characterizes them through symmetric gauge functions of singular values.

Norm definitions

Unitarily invariant norms satisfy ∥A∥ = ∥UAV∥ for all orthogonal U and V; unitarily invariant matrix norms additionally require submultiplicativity.Over the real numbers, unitary invariance is equivalent to invariance under left and right orthogonal transformations.Source: B Supplemental Technical Background on… · p. 46

Singular-value characterization

A norm is unitarily invariant if and only if it equals a symmetric gauge function applied to its singular values, ψ(σ(W)).The singular values are arranged in descending order, while the gauge function is a norm invariant under coordinatewise absolute values and permutations.Source: B Supplemental Technical Background on… · pp. 46–47

Examples

Examples include Schatten p-norms for p ∈ [1, ∞], including the nuclear, Frobenius, and spectral norms, as well as the Ky Fan k-norm.These norms provide standard instances of the unitarily invariant class used in matrix analysis.Source: B Supplemental Technical Background on… · p. 47

C Geometric Misalignment of Coordinate-wise Adaptive Gradient Methods

Coordinate-wise adaptive methods depend on arbitrary parameter coordinates and generally fail to respect the bi-orthogonal geometry of matrix optimization, unlike spectral methods. This mismatch can inject high-rank noise into low-rank gradients and becomes more consequential as model size grows.

Geometric mismatch

Coordinate-wise optimizers are not bi-orthogonally equivariant, so equivalent matrix reparameterizations can produce different update directions and training dynamics.Their behavior depends on how matrix entries are indexed rather than solely on the intrinsic geometry of the matrix.Source: C Geometric Misalignment of… · pp. 47–48
Coordinate-wise nonlinearities such as division, squaring, and square roots distort gradient or momentum directions, unlike updates built only from linear operations.Adam relies on elementwise normalization and is therefore typically equivariant only under signed permutations, whereas spectral descent and Muon respect bi-orthogonal equivariance.Source: C Geometric Misalignment of… · p. 47

Sign descent

Sign descent uses coordinate-wise signs and is governed by the coordinate-wise ℓ∞ geometry rather than unitarily invariant matrix geometry.It can be interpreted as spectral descent on the diagonal lifting Diag(vec(G)), thereby inheriting the pathological geometry of that degenerate representation.Source: C Geometric Misalignment of… · pp. 47–48

Geometric mismatch

Coordinate-wise methods inject high-rank noise even when gradients are low rank, while spectral methods remain aligned with the low-dimensional matrix structure governing many deep-network optimization problems.The advantage of geometry-aware methods typically increases with model size.Source: C Geometric Misalignment of… · pp. 47–48

D Proofs of Main Text

The proofs establish equivariance of momentum-based polar updates, SwiGLU transformations, router-compatible updates, and characterize bi-orthogonally equivariant matrix updates through spectral forms.

Momentum and polar-factor proofs

EMA, Polyak, and Nesterov momentum preserve bi-orthogonal equivariance, and their associated polar factors transform equivariantly under orthogonal actions.The transformed momentum sequences satisfy f M_k = P M_k Q^⊤, while Nesterov’s sequence satisfies e N_k = P N_k Q^⊤.Source: D Proofs of Main Text · pp. 48–49

Matrix-layer equivariance

The composed update 𝒰2 ◦ 𝒰1 is left-permutation and right-orthogonal equivariant, supporting the required symmetry of the corresponding matrix-layer construction.Source: D Proofs of Main Text · p. 49

SwiGLU equivariance

Coordinatewise activation and elementwise multiplication commute with hidden-unit permutations, yielding the stated SwiGLU equivariance under permuted gate, up, and down projections.The proof uses σ(PW_gate x) = Pσ(W_gate x) and (Pu) ⊙ (Pv) = P(u ⊙ v).Source: D Proofs of Main Text · p. 49

Router-compatible updates

Centered left-spectral and projected centered row-norm updates are expert-permutation equivariant and invariant to shared row shifts, with horizontality ensured by centering.The proofs establish these properties separately for both update families and conclude that both are router-compatible.Source: D Proofs of Main Text · p. 50

Characterization of equivariant updates

A matrix update is bi-orthogonally equivariant exactly when it acts spectrally as 𝒰(D) = U Diag(ψ(σ(D)))V^⊤ for an absolutely symmetric ψ.Equivariance reduces the action to diagonal singular-value matrices, while sign and permutation symmetries force the spectral form and make it independent of the chosen SVD.Source: D Proofs of Main Text · pp. 50–51

E Implementation Details of Practical Optimizers … F.2 Full Spectral Optimizers

The paper supplies practical optimizer implementations, including Polar Express for matrix inverse square roots, and analyzes symmetry-compatible updates under smoothness and related assumptions. For full spectral optimizers, the analysis establishes descent, sublinear stationarity, and linear convergence under a PŁ condition.

E Implementation Details of Practical Optimizers

The paper provides implementation details for its proposed practical optimizers, connecting the numerical inverse-square-root procedure to the broader optimizer stack.The implementation discussion is distinct from the subsequent convergence analysis.Source: E Implementation Details of Practical… · p. 51

E.1 Numerical Algorithm for Matrix Inverse Square Root via Polar Express

Polar Express computes matrix inverse square roots through a coupled polynomial iteration linked to orthogonal polar-factor iterations, converging to A^1/2 and A^-1/2 with the same order q.The implementation section presents the numerical algorithm used for the inverse-square-root computation.Source: §E.1 · p. 51

F Convergence Analysis of Symmetry-Compatible Optimizers

The unified convergence framework controls the alignment ⟪G, T(G)⟫_F and update norm ||T(G)||_F, with bounds specialized to each optimizer geometry.Scale-compatible updates use gradient-norm scaling, whereas fully normalized directions are analyzed through geometry-dependent ratios.Source: F Convergence Analysis of… · p. 52, §F.1 · pp. 52–53

F.1 General Update Scheme and Standing Assumptions

The general analysis assumes an iteration without momentum and uses L-smoothness for descent; the PŁ inequality is needed only for linear convergence, while smoothness alone yields descent and stationarity convergence.The layerwise loss is analyzed under standard regularity conditions with a generic update map T.Source: §F.1 · p. 52
A rigorous stochastic extension remains open because nonlinear symmetry-compatible maps do not preserve unbiasedness of stochastic gradients; additional assumptions on expected alignment and norm are required.The stated limitation applies to spectral, one-sided spectral, row-norm-based, and hybrid updates.Source: §F.1 · p. 52

F.2 Full Spectral Optimizers

Full spectral updates monotonically decrease the loss when 0 < γ_k < 2c_1/(Lc_2), assuming L-smoothness and singular-value alignment and boundedness.The update map must provide positive gradient alignment and controlled update norm.Source: §F.1 · p. 52, §F.2 · p. 53
With a constant admissible learning rate, full spectral optimizers achieve sublinear convergence to stationarity; under the PŁ condition, they converge linearly with factor ρ = 1 − 2µγ(c_1 − Lc_2γ/2).The stationarity bound depends on the initial optimality gap and the descent coefficient γ(c_1 − Lc_2γ/2).Source: §F.1 · p. 52, §F.2 · pp. 53–54

F.2.1 Specialization to Normalized Polar-Type Spectral Methods … G Experimental Details

This section develops convergence guarantees for spectral, polar-type, row-norm, and hybrid updates under conditions involving scaling, alignment, rank, and related bounds. It also details the architectures, numerical routines, optimizer hyperparameters, learning-rate schedules, and implementation choices used in the four pre-training experiments.

F.2.1 Specialization to Normalized Polar-Type Spectral Methods

Normalized polar updates admit monotone descent when 0 < γk < 2RT(Gk)/L, while linear convergence additionally requires ratio and gradient-alignment lower bounds.For T(G) = polar(G), Rpolar(G) = ||G||nuc/rank(G), and stronger rates require controlling ||G||nuc relative to ||G||2; flatter spectra improve the ratio bound.Source: §F.2.1 · p. 55

F.2.2 Specialization to PolarGrad with Nuclear-Norm Scaling

Nuclear-norm-scaled PolarGrad has exact alignment and norm identities, guarantees descent for γk < 2/(Lrk), and improves the one-step decrease when gradients have nontrivial stable rank.With bounded gradient rank, it achieves O(1/T) stationarity and linear convergence under the PŁ condition; local PŁ or KŁ conditions yield corresponding eventual or generalized rates.Source: §F.2.2 · pp. 56–57, Since |||Gk|||2 · p. 57

Since |||Gk|||2

Under bounded gradient rank, nuclear-norm-scaled PolarGrad converges to stationarity at O(1/T) and converges linearly under PŁ with γ ∈ (0, 2/(Lr)).Smoothness alone suffices for monotonic descent and the stationarity rate, while local PŁ or KŁ assumptions can recover local linear or sublinear convergence.Source: Since |||Gk|||2 · pp. 57–58

F.3 One-Sided Spectral Optimizers

One-sided spectral optimizers use right- or left-Gram spectral operators, and the same smoothness framework yields descent, sublinear stationarity, and PŁ linear convergence under alignment and norm control.Right updates have the form TR(G) = G Ψ(G⊤G), whereas left updates have the form TL(G) = Φ(GG⊤)G.Source: §F.3 · p. 58

F.3.1 Right-Spectral Optimizers

Right-spectral updates achieve descent for γk ∈ (0, 2cR,1/(LcR,2)) and linear convergence under PŁ with contraction factor ρR = 1−2µγ(cR,1 −LcR,2γ/2).Their alignment is lower-bounded by cR,1||G||2 under the eigenvalue alignment and boundedness assumptions.Source: §F.3.1 · pp. 58–59

F.3.2 Left-Spectral Optimizers

Left-spectral updates have analogous descent and PŁ convergence guarantees, while nuclear-norm-scaled left and right PolarGrad variants inherit the full PolarGrad descent, stationarity, and linear-convergence results.The one-sided variants recover the same closed-form alignment and norm identities after nuclear-norm scaling.Source: §F.3.2 · pp. 59–60

F.4 Row-Norm-Based Optimizers

Row-norm updates scale each gradient row independently and, under uniform scaling bounds, guarantee descent for γk ∈ (0, 2η/(Lη2)), with sublinear stationarity and PŁ linear convergence.This structure is especially suited to matrices whose row axis carries the relevant symmetry, including embeddings, LM heads, and MoE routers.Source: §F.4 · pp. 60–61

F.4.1 Specialization to Smoothed Row Normalization

Smoothed row normalization remains bounded at zero, interpolates toward full normalization, and inherits descent, stationarity, and PŁ linear-convergence guarantees under a uniform row-norm bound.For η(t)=1/(t+ε), the scaling lies between 1/(M+ε) and 1/ε; constant γ∈(0, 2ε^2/(L(M+ε))) guarantees monotonic descent.Source: §F.4.1 · p. 62

F.5 Nuclear-Norm-Scaled Right-Spectral/Row-Norm Hybrid Optimizers

The right-spectral/row-norm hybrid preserves the original feature geometry before local row normalization, with descent controlled by alignment Ahyb(G) and active row support srow(G).Its update satisfies a descent condition when γk∈(0, 2Ahyb(Gk)/(Lsrow(Gk))); uniform lower alignment and bounded support yield sublinear stationarity and PŁ linear convergence.Source: §F.5 · pp. 62–64

F.6 Nuclear-Norm-Scaled Row-Norm/Right-Spectral Hybrid Optimizers

The row-norm/right-spectral hybrid instead computes its polar factor after row reweighting, so convergence depends on Arowhyb(G), effective rank rrowhyb(G), and a row-normalization comparability bound.Under these assumptions it has sublinear stationarity and PŁ linear-convergence results; the smoothed choice η(t)=1/(t+ε) gives a concrete learning-rate range γ∈(0, 2ε/(Ld)) under a uniform row-norm bound.Source: §F.6 · pp. 65, 66, 68
For the smoothed row-norm/right-spectral hybrid, a uniform row bound R verifies κrow=1/(R+ε), alignment arow=ε, and rank bound rrow=d.These substitutions establish the assumptions needed for the specialized convergence corollary.Source: §F.6 · pp. 67–68

G Experimental Details

Experiments reduce hidden layers and experts relative to the source architectures, initialize 2D weights with Gaussian standard deviation 0.02, and use specified polar-factor routines, smoothing, and optimizer-dependent schedules.Polar Express and Gram Newton–Schulz use five inner steps with εNS=10^-7; ε=10^-8 is used for row scaling, while scalar/vector AdamW uses 100-step warmup with cosine decay and other parameters use stable decay.Source: G Experimental Details · p. 69
Training configurations and optimizer hyperparameters are reported for all four model pre-training experiments.Source: G Experimental Details · p. 69
The gpt-oss SwiGLU implementation is unconventional because it includes clamping and a residual connection.Source: G Experimental Details · p. 69

G.1 Qwen3-0.6B-Style Pre-Training

Qwen3-0.6B-style pre-training configurations had similar wall-clock times within each setup, while setup (b) was slower than setup (a) because HybridPolarGradM adds computational overhead for SwiGLU MLP projections.

G.1 Qwen3-0.6B-Style Pre-Training

Configurations (i)–(iii) in setup (a) took 7.360, 7.498, and 7.369 hours, respectively.Source: §G.1 · p. 70
Configurations (i)–(iii) in setup (b) took 7.765, 7.886, and 7.771 hours, respectively.Source: §G.1 · p. 70
Setup (b) required more time than setup (a) because HybridPolarGradM for SwiGLU MLP projection matrices has additional computational overhead over Muon.Source: §G.1 · p. 70

G.2 Gemma 3 1B-Style Pre-Training · G.2.1 Gemma 3 1B-Style Pre-Training Learning Rate Sweep

Gemma 3 1B-style pre-training shows that symmetry-compatible embedding and LM-head updates outperform AdamW across reasonable learning-rate choices. This ordering persists whether SwiGLU projections use Muon or HybridPolarGradM, indicating robustness beyond simultaneous optimizer replacement.

G.2 Gemma 3 1B-Style Pre-Training

Wall-clock training time is similar across configurations, ranging from 8.132 to 8.807 hours in setting (a) and from 8.469 to 9.101 hours in setting (b).For (a), configurations (i)–(iii) take 8.341, 8.807, and 8.132 hours; for (b), they take 8.675, 9.101, and 8.469 hours.Source: §G.2 · p. 71
The Gemma 3 1B-style pre-training experiment defines multiple training configurations and corresponding optimizer hyperparameters for comparing these optimizer assignments.The supplied tables identify the configuration and hyperparameter summaries but do not provide their cell values.Source: Table G.4 · p. 71, Table G.5 · p. 71

G.2.1 Gemma 3 1B-Style Pre-Training Learning Rate Sweep

The learning-rate sweep varies equal embedding and LM-head base learning rates while holding scalar/vector and other matrix learning rates fixed.More delicate tuning of the two vocabulary-indexed matrix rates is possible, but the experiment keeps them equal for simplicity.Source: §G.2.1 · p. 72
Across swept γ0,emb = γ0,head values, RowNormM and HybridPolarGradM remain below AdamW in later validation loss, despite AdamW’s comparable or faster initial progress.The ordering holds with either Muon or HybridPolarGradM on SwiGLU projection matrices.Source: §G.2.1 · p. 72, Figure G.1 · p. 73
The observed validation-loss improvement is attributed to symmetry-compatible vocabulary-indexed updates rather than a single learning-rate choice or simultaneous replacement of Muon on SwiGLU projections.This conclusion follows from the stable ordering across the sweep and both SwiGLU optimizer settings.Source: §G.2.1 · p. 72
Figure G.1 compares the sweep under Muon and HybridPolarGradM assignments for the SwiGLU MLP projection matrices, with RowNormM and HybridPolarGradM consistently outperforming AdamW.The two panels differ only in the SwiGLU projection optimizer.Source: Figure G.1 · p. 73

G.2.2 Gemma 3 1B-Style Pre-Training Across Random Seeds

Across two additional random seeds, symmetry-compatible optimizers preserve their advantage over AdamW for Gemma 3 1B-style pre-training, with HybridPolarGradM achieving the lowest mean final validation loss.

Random-seed robustness

HybridPolarGradM achieves the lowest mean final validation loss, improving over AdamW by approximately 0.0441, while RowNormM improves by approximately 0.0375.The qualitative ordering remains stable across random seeds, indicating the main experiment’s gains are not an artifact of a single initialization.Source: Table G.7 · p. 74, §G.2.2 · p. 74

Experimental setup

Across two additional random seeds, RowNormM and HybridPolarGradM consistently outperform AdamW on the vocabulary-indexed input embedding and LM head matrices.The three assignments differ only in the optimizer used for those matrices; SwiGLU MLP projections use HybridPolarGradM throughout.Source: §G.2.2 · p. 74, Figure G.2 · p. 74

G.2.3 Vocabulary-Logit Growth and Final Logit Soft-Capping

Projected symmetry-compatible LM-head updates control shared-shift drift in final vocabulary logits while preserving softmax-relevant centered logits. In Gemma 3 1B-style pre-training, they improve validation loss relative to AdamW, and complement final-logit soft-capping as a direct stabilizer.

Vocabulary-Logit Growth and Final Logit Soft-Capping

The results support quotient geometry: projected LM-head updates suppress shared-logit-shift drift while preserving centered logits that determine the softmax distribution.Raw RMS includes the shared-shift component, whereas centered RMS removes it; AdamW’s growth is mainly in the shared-shift direction rather than relative vocabulary-logit differences.Source: §G.2.3 · pp. 75–76
Projected RowNormM and HybridPolarGradM LM-head updates achieve lower validation loss than AdamW while reducing raw-logit growth and preserving comparable or stronger centered-logit behavior.Without soft-capping, AdamW shows substantially larger raw-logit RMS, with RowNormM + AdamW exhibiting late-stage growth in raw RMS and maximum log-sum-exp; projected updates avoid this growth. With soft-capping, projected methods still achieve lower validation loss and larger centered-logit RMS.Source: Figure G.3 · p. 76, §G.2.3 · pp. 76–77, Figure G.4 · p. 77
Final-logit soft-capping largely suppresses AdamW’s raw logit-scale growth and keeps maximum log-sum-exp comparable across methods.Attention-logit soft-capping is disabled, isolating the final vocabulary logits; soft-capping therefore acts as a direct stabilizer of these logits.Source: §G.2.3 · pp. 75, 77, Figure G.4 · p. 77
Soft-capping and projected optimizer geometry are complementary: soft-capping bounds effective final logits, while projected updates control redundant shared-shift drift and permit useful centered vocabulary separation.With soft-capping enabled, projected methods have larger centered logit RMS than AdamW while maintaining better validation loss, showing they do not simply shrink all logit variation.Source: §G.2.3 · pp. 75, 77

G.3 OLMoE-1B-7B-Style Pre-Training

The OLMoE-1B-7B-style pre-training configurations required nearly identical wall-clock training times, all around 10.7–10.8 hours.

Training time

Configurations (i)–(iv) took 10.79, 10.82, 10.71, and 10.75 hours, respectively.The four runs differed by only 0.11 hours in total wall-clock training time.Source: §G.3 · p. 78

G.4 Downsized gpt-oss Pre-Training

Downsized gpt-oss pre-training configurations had nearly identical wall-clock training times, ranging from 18.91 to 19.34 hours. The section also provides tables for their training configurations and optimizer hyperparameters.

Training time

Wall-clock training times were 19.08, 18.91, 19.34, and 19.22 hours for configurations (i)–(iv), respectively.The four configurations therefore had closely matched training durations.Source: §G.4 · p. 79

Training setup

Table G.10 lists the downsized gpt-oss pre-training configurations, while Table G.11 lists their optimizer hyperparameters.Source: Table G.10 · p. 79, Table G.11 · p. 79

G.4.1 Load-Balancing Loss and Router z-Loss Diagnostics

This section diagnoses how router optimizers affect expert balance and router-logit scale, with and without auxiliary losses. Symmetry-compatible updates improve routing dynamics over AdamW while preserving or slightly improving validation language-model performance.

Diagnostics

The diagnostics track routed fractions, router probability mass, load-balancing loss, router z-loss, entropy, variation, dead experts, maximum load, and router-logit scale.The load-balancing loss is LLB = EΣ_i f_iP_i and is close to 1 for a perfectly balanced router; router z-loss penalizes excessive logit scale.Source: §G.4.1 · pp. 79–80

Training objective without auxiliary losses

Without auxiliary losses, all symmetry-compatible router updates substantially reduce measured load-balancing loss and router z-loss versus AdamW, with LeftPolarGradM variants strongest on load balancing.Figures G.5 and G.6 report training and validation diagnostics; HybridPolarGradM improves over AdamW but trails the purely left-polar variants.Source: §G.4.1 · p. 80, Figure G.5 · p. 81, Figure G.6 · p. 81
Router improvements do not hurt validation cross-entropy and slightly improve over AdamW, while centered and uncentered LeftPolarGradM perform similarly.The dominant empirical gain appears associated with left-polar router geometry, whereas centering provides the symmetry-compatible quotient-space formulation for shared-logit-shift invariance.Source: §G.4.1 · p. 80, Figure G.5 · p. 81
Without auxiliary losses, LeftPolarGradM variants nearly eliminate dead experts and achieve the most balanced assignments, while RowNormM and HybridPolarGradM provide intermediate improvements over AdamW.AdamW has lower load entropy, higher load coefficient of variation, larger dead-expert fraction, and higher maximum expert load; these gains arise without an explicit load-balancing objective.Source: §G.4.1 · p. 80, Figure G.7 · p. 82

Training objective with auxiliary losses

With auxiliary losses weighted 0.01 for LLB and 0.001 for Lz, all methods improve routing, but symmetry-compatible updates retain lower losses and better expert-assignment statistics than AdamW.RowNormM, LeftPolarGradM, and uncentered LeftPolarGradM cluster as the best-balanced methods, while HybridPolarGradM remains better than AdamW but less balanced.Source: §G.4.1 · pp. 80–81, Figure G.8 · p. 83, Figure G.9 · p. 83