PaperlayerView exact revision on arXiv

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.18106v4math.OCcs.AIcs.LGstat.ML

TL;DR

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

  • Problem

    Coordinate-wise optimizers treat matrix entries independently, overlooking the distinct geometries and symmetries of embeddings, projections, LM heads, and MoE routers.

  • Method

    The paper assigns each matrix-valued parameter class an optimizer whose equivariance matches its layerwise symmetry, including spectral, row-norm, hybrid, and permutation-aware updates.

  • Results

    Symmetry-compatible assignments consistently improve final validation loss across dense and sparse MoE pre-training, while also reducing expert load imbalance in sparse MoE models.

  • Takeaways & Limitations

    The results support treating optimizer design as architecture-aware, with geometry-compatible updates offering a more principled alternative to one universal coordinate-wise rule.

  • Takeaways & Limitations

    Hybrid updates achieve the best validation loss in the Gemma experiment but incur higher computational overhead than the cheaper row-norm alternative.

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 layerwise and symmetry-compatible rather than coordinate-wise across all parameters. It unifies spectral methods and derives equivariant optimizer classes for diverse language-model matrix blocks, supported by dense and sparse MoE pre-training evidence.

  • Motivation: Coordinate-wise optimizers treat matrix entries independently, ignoring matrix structure and the differing geometries of embeddings, LM heads, projections, and routers.This geometry-blind treatment can break natural equivariances of matrix-valued parameters.
  • Symmetry-compatible principle: The central principle assigns each major matrix-valued parameter class an optimizer whose update equivariance matches that layer’s symmetry.The symmetry-aware view derives updates from layerwise geometry rather than vectorizing all parameters.
  • Spectral optimizers: Bi-orthogonal equivariance unifies spectral optimizers including stochastic spectral descent, Muon, Scion, and PolarGrad.These methods use exact or approximate polar updates that preserve the singular-vector geometry of the update direction.
  • Layerwise optimizer classes: Layer-specific constructions include one-sided spectral, row-norm, and hybrid row-norm/spectral updates for embeddings and LM heads, with projected updates required by shared-logit-shift invariance.The paper frames these constructions as architecture–optimizer co-design based on layerwise equivariance and quotient geometry.
  • Empirical evidence and scope: End-to-end dense and sparse MoE pre-training evaluates a stack assigning symmetry-compatible updates to all major matrix-valued parameter classes.The authors present the results as evidence that the principle is useful, not as an exhaustive benchmark or universal dominance claim.

2 Preliminaries and Related Work

This section establishes notation and situates the work within matrix optimization, spectral methods, and equivariance research. Its central distinction is extending equivariance from neural architectures and training dynamics to optimizer update maps acting on parameter tensors.

  • Notation: The paper defines notation for matrix norms, inner products, singular values, matrix products, vectorization, positive semidefinite matrices, orthogonal matrices, and permutation matrices.These definitions support the subsequent symmetry and matrix-optimization framework.
  • Related optimization methods: Unlike curvature-approximation and preconditioning methods, spectral and polar updates can enforce equivariance of the update map itself, with optimizer geometry determined by the layer’s symmetry group.Related approaches include K-FAC, Shampoo, BFGS/L-BFGS, SOAP, KL-Shampoo, KL-SOAP, and PSGD.
  • Matrix optimization: Matrix optimization is treated as distinct from vectorized optimization because matrices encode singular values, ranks, invariant subspaces, eigenvalues, and unitary symmetries.The section places the framework within convex matrix analysis, variational analysis, eigenvalue optimization, and spectral optimization.
  • Spectral optimizers: Spectral operators transform singular values while preserving singular vectors, providing the operator-theoretic structure underlying stochastic spectral descent, Muon, Scion, and polar gradient methods.Muon’s spectral-norm steepest-descent interpretation yields an orthogonal polar factor and an implicitly bi-orthogonally equivariant update.
  • Equivariance and optimizer design: The work extends equivariance from architecture design and existing training dynamics to optimizer design for parameter tensors, emphasizing the internal geometry of parameter space.This perspective makes the symmetry principle explicit for Muon and related methods and extends it to layers with different symmetries.

3 Equivariant Optimizers from Layerwise Symmetry

The section proposes that optimizer update maps should commute with each parameter block’s symmetry group, with updates also constrained to quotient or horizontal subspaces when symmetries create redundant directions. This principle yields spectral, row-aware, row-norm, hybrid, and centered or left-spectral optimizers tailored to ordinary matrices, embeddings and LM heads, SwiGLU projections, and MoE routers.

  • Layerwise equivariance principle: Optimizer updates should transform equivariantly with the gradient or update direction under the parameter block’s symmetry group.The requirement concerns the update map’s representation consistency, not necessarily global invariance of the layerwise loss.
  • SwiGLU MLP projections: Intermediate-neuron permutation symmetry yields row- and column-aware SwiGLU projection updates, including hybrids that combine row normalization with one-sided spectral normalization.The construction derives this optimizer family from SwiGLU symmetry and is related in spirit to Aurora.
  • Embeddings and LM heads: Embedding and LM-head symmetries permit right-spectral, row-norm, and hybrid updates because vocabulary symmetry is discrete rather than fully orthogonal.LM heads additionally require projection to the quotient or horizontal space induced by softmax shared-logit-shift invariance.
  • MoE routers: Expert-permutation and shared-shift symmetries yield centered row-aware, left-spectral, and hybrid router optimizers that preserve equivariance, shift invariance, and update horizontality.Hybrid variants combine global expert mixing from left-spectral updates with local expert-wise normalization from row-norm updates.
  • Ordinary matrix layers: Bi-orthogonal equivariance completely characterizes continuous direction-wise matrix updates as spectral operators acting through singular values while preserving singular vectors.This places SSD, Muon, Scion, PolarGrad, and related methods in the canonical class for ordinary matrix layers.

4 Numerical Experiments

Numerical experiments across dense and sparse MoE models show that symmetry-compatible optimizer assignments improve final validation loss over AdamW, with benefits persisting across architectures. Hybrid and router-aware updates often perform best, while their spectral computations introduce practical overhead.

  • Dense-model experiments: In the smaller dense experiment, final validation losses were 4.2017, 4.2050, and 4.2084 for configurations (i)–(iii), versus 4.1950, 4.1955, and 4.2046 after changing SwiGLU updates.RowNormM and HybridPolarGradM eventually overtook configuration (iii), while HybridPolarGradM achieved a slightly lower final loss than RowNormM.
  • Dense-model experiments: In the Gemma 3 1B-style experiment, HybridPolarGradM achieved the lowest final validation loss, with values of 4.0516, 4.0435, and 4.0862 after applying it to SwiGLU projections.The hybrid update benefits from the tall-skinny geometry of the projections, whereas RowNormM avoids spectral computations and has lower overhead.
  • Sparse MoE experiments: In the OLMoE experiment, final validation losses for configurations (i)–(iv) were 4.0955, 4.0815, 4.1187, and 4.1240, with symmetry-compatible assignments overtaking AdamW after faster initial baseline progress.Configurations (i) and (ii) use compatible updates for embeddings, LM heads, and routers; configuration (iv) uses AdamW for all three.
  • Sparse MoE experiments: In the downsized gpt-oss experiment, final validation losses were 4.3080, 4.3136, 4.3388, and 4.3687, while symmetry-compatible router updates also reduced load imbalance and router z-loss.The fully coordinate-wise AdamW configuration performed worst, and the same qualitative pattern persisted despite architectural differences from OLMoE.
  • Cross-experiment findings: Across experiments, replacing AdamW on large vocabulary-indexed matrices consistently improved final validation loss, with larger gains in Gemma 3 1B-style models and persistent gains in sparse MoE models.HybridPolarGradM and other polar-based updates require more computation and potentially higher numerical precision than row-norm-only updates.

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

The paper argues for layerwise, geometry-aware, symmetry-compatible optimization because coordinate-wise methods mismatch heterogeneous parameter symmetries. It situates this framework among related optimizer theories and outlines numerical, systems, and co-design directions for extending symmetry-aware training.

  • 5 Discussion and Outlook: Coordinate-wise Adam and AdamW remain robust defaults, but treating matrices as independent coordinates ignores their intrinsic geometry and may create parameterization-sensitive optimization.The paper connects this mismatch to possible effects on robustness, stability, scalability, and interpretability.
  • 5 Discussion and Outlook: Symmetry-compatible optimization assigns update geometries to parameter blocks according to their distinct algebraic symmetries rather than applying one coordinate-wise rule universally.The framework covers ordinary matrix layers, embeddings and LM heads, SwiGLU projections, and MoE routers.
  • 5 Discussion and Outlook: Scaling geometry-aware optimizers requires fast, stable, GPU-friendly matrix decomposition, including efficient polar oracles and orthogonalization routines.The discussion highlights Newton–Schulz iterations, QDWH, and Polar Express as relevant directions.
  • Appendix: The appendix organizes further related work, matrix-analysis background, geometric misalignment examples, omitted proofs, implementation details, and convergence analysis.It also identifies a dedicated section for practical optimizer details and another for convergence analysis.
  • A.1 Non-Euclidean Norm-Based Steepest Descent and LMO-Based Frameworks: Non-Euclidean and LMO interpretations clarify Muon, but symmetry analysis supplies a criterion for choosing norms: orthogonally invariant matrix geometries require unitarily invariant norms.The ℓp →ℓq operator norm generally fails this condition except when p = q = 2, yielding the spectral norm.
  • A.2 Modular Norm Theory: The framework differs from modular norm theory by deriving right-spectral updates from left-permutation/right-orthogonal equivariance, while the two approaches may complement each other.Modular norm theory targets scale invariance across width and depth, whereas the spectral framework targets directional geometry under layer symmetries.
  • A.3 Rotation-Based Optimizers: Unlike rotation-based optimizers that use only left rotations, this framework derives different optimizer classes for different layer symmetries and includes full, right-, and left-spectral, row-norm, and hybrid updates.It therefore presents a broader symmetry-based taxonomy without requiring the layerwise loss itself to be rotationally invariant.

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.

  • Unitarily invariant norms: Unitarily invariant norms satisfy ∥A∥ = ∥UAV∥ for all orthogonal U and V, with matrix norms additionally required to be submultiplicative.Over the reals, unitary invariance is equivalent to invariance under left and right orthogonal transformations.
  • Singular-value characterization: A norm is unitarily invariant if and only if it can be written as ψ(σ(W)), where ψ is a symmetric gauge function and σ(W) lists singular values in descending order.Symmetric gauge functions are norms invariant under coordinatewise absolute values and permutation matrices.
  • 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 are presented as important examples of unitarily invariant matrix norms.

C Geometric Misalignment of Coordinate-wise Adaptive Gradient Methods

Coordinate-wise adaptive methods such as Adam are not bi-orthogonally equivariant, so their updates depend on arbitrary coordinate choices rather than intrinsic matrix geometry. This mismatch can inject high-rank noise into low-rank optimization and contrasts with spectral methods that preserve matrix structure.

  • Geometric misalignment: Coordinate-wise adaptive optimizers can change under rotations of the input space, producing different training dynamics for equivalent reparameterizations.Their iterates therefore depend on arbitrary coordinate choices rather than solely on the parameter-space geometry.
  • Geometric misalignment: Unlike spectral methods, coordinate-wise methods are generally equivariant only under signed permutations, not bi-orthogonal transformations.This smaller symmetry group fails to capture the intrinsic geometry of matrix-valued optimization.
  • Adam: Adam’s coordinate-wise normalization makes its update direction depend on how matrix entries are indexed rather than on the matrix’s intrinsic geometry.The update is not equivariant under W → PWQ^⊤ for orthogonal P and Q.
  • Geometric consequences: Coordinate-wise methods tend to inject high-rank coordinate noise even when the underlying gradient is low rank, whereas spectral methods remain aligned with its matrix structure.The contrast is especially relevant when deep-network optimization is governed by an intrinsically low-dimensional landscape.
  • Sign descent: Sign descent corresponds to spectral descent on the diagonal lifting Diag(vec(G)), inheriting the degenerate geometry of that representation instead of G’s intrinsic matrix geometry.For gradients without exact zero entries, the polar factor of the diagonal lifting is Diag(sgn(vec(G))).

D Proofs of Main Text

The proofs establish equivariance of momentum-based polar updates and characterize symmetry-compatible update families for matrix layers, SwiGLU projections, and MoE routers. They also show that spectral-form updates are exactly the bi-orthogonally equivariant matrix updates.

  • Momentum equivariance: EMA, Polyak, and Nesterov momentum sequences transform bi-orthogonally, and their polar factors inherit the same equivariance.The transformed sequences satisfy f M_k = P M_k Q^⊤; Nesterov additionally satisfies eN_k = P N_k Q^⊤.
  • Mixed matrix symmetries: The composition 𝒰2 ◦ 𝒰1 is left-permutation and right-orthogonal equivariant.This provides the symmetry property needed for composed updates acting on parameter blocks with mixed symmetries.
  • SwiGLU symmetry: SwiGLU remains equivariant under hidden-unit permutations because coordinatewise activation and elementwise products commute with permutation matrices.The proof uses σ(PWgatex) = Pσ(Wgatex) and (Pu) ⊙ (Pv) = P(u ⊙v).
  • MoE router symmetry: Centered left-spectral and projected centered row-norm updates are expert-permutation equivariant and shared-row-shift invariant, making both router-compatible.The centered construction uses the projector Π⊥, while row norms transform by the corresponding expert permutation.
  • Spectral characterization: An update is bi-orthogonally equivariant exactly when it has spectral form 𝒰(D) = U Diag(ψ(σ(D)))V ⊤.Equivariance reduces the action to diagonal singular-value matrices, with ψ symmetric under compatible permutations and sign changes.

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

The paper provides practical implementation details and a Polar Express algorithm for matrix inverse square roots, then analyzes convergence for symmetry-compatible optimizers. Under smoothness, spectral updates admit monotonic descent and sublinear stationarity guarantees, while an additional PŁ condition yields linear convergence.

  • E Implementation Details of Practical Optimizers: The paper introduces implementation details for its proposed practical optimizers and presents a Polar Express procedure for computing matrix inverse square roots.The algorithm is based on the connection between polynomial iterations for orthogonal polar factors and inverse square roots.
  • E.1 Numerical Algorithm for Matrix Inverse Square Root via Polar Express: Polar Express couples matrix iterations so that, with convergence order q, the iterates approach A^1/2 and A^-1/2 with the same order q.The coupled iteration starts from X_0 = A and Y_0 = I_n.
  • F Convergence Analysis of Symmetry-Compatible Optimizers: The convergence analysis studies full spectral, one-sided spectral, row-norm-based, and hybrid optimizer classes through a shared first-order iteration.The analysis specializes common assumptions to each optimizer geometry.
  • F.1 General Update Scheme and Standing Assumptions: The general framework assumes an L-smooth layerwise loss, optionally imposes the µ-Polyak–Łojasiewicz condition, and excludes momentum for analytical simplicity.Smoothness gives monotonic descent and sublinear convergence to stationarity, whereas the PŁ condition is needed for linear convergence.
  • F.1 General Update Scheme and Standing Assumptions: In the stochastic setting, nonlinear symmetry-compatible update maps mean unbiased stochastic gradients do not generally produce unbiased update directions.A rigorous stochastic analysis therefore requires assumptions on the expected alignment and norm of the transformed stochastic gradient and is left for future work.
  • F.2 Full Spectral Optimizers: Full spectral convergence requires positive singular-value alignment and bounded update norms for the spectral map T(G) = U Diag(ψ(σ(G)))V^⊤.These structural conditions yield the alignment and norm bounds used in the descent analysis.
  • F.2 Full Spectral Optimizers: For γ_k ∈ (0, 2c1/(Lc2)), spectral updates ensure monotonic loss descent; with constant γ in this interval, they achieve sublinear stationarity convergence and, under PŁ, linear convergence.The linear rate is expressed through ρ := 1 − 2µγ(c1 − Lc2γ/2) ∈ (0, 1), and the bound is optimized at γ⋆ = c1/(Lc2).

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

The convergence analysis develops ratio-based guarantees for normalized polar methods, extends them to one-sided, row-norm, and hybrid spectral updates, and specifies the implementation choices used in experiments. These results establish descent, stationarity, and PŁ linear-convergence conditions under explicit alignment, norm, rank, and support assumptions.

  • F.2.1 Specialization to Normalized Polar-Type Spectral Methods: Normalized polar updates require a spectral advantage ratio because their update norm does not vanish at zero gradients; descent holds when 0 < γk < 2RT(Gk)/L.For the polar update, Rpolar(G) = ||G||nuc/rank(G), so stronger convergence additionally requires controlling nuclear-norm alignment with ||G||F^2.
  • F.3 One-Sided Spectral Optimizers; F.3.1 Right-Spectral Optimizers; F.3.2 Left-Spectral Optimizers: Right- and left-spectral optimizers obtain descent and convergence from Gram-operator alignment bounds, with admissible learning rates 2cR,1/(LcR,2) and 2cL,1/(LcL,2), respectively.Nuclear-norm-scaled one-sided PolarGrad variants recover the full PolarGrad alignment identities and inherit its descent, stationarity, and PŁ linear-convergence results.
  • F.4 Row-Norm-Based Optimizers: Row-norm optimizers act locally on rows, matching embedding, LM-head, and MoE-router symmetries, and guarantee descent, O(1/T) stationarity, and PŁ linear convergence under bounded row scaling.Their monotonic-descent condition is γk ∈(0, 2η/(Lη2)).
  • F.4.1 Specialization to Smoothed Row Normalization: Smoothed row normalization η(t) = 1/(t + ε) avoids the singularity of η(t) = 1/t, and under row-norm bound M permits monotonic descent for γ ∈(0, 2ε2/(L(M + ε))) with the preceding convergence guarantees.It preserves local row adaptivity while remaining bounded at zero.
  • F.5 Nuclear-Norm-Scaled Right-Spectral/Row-Norm Hybrid Optimizers: The right-spectral/row-norm hybrid has explicit alignment and norm identities governed by Ahyb(G) and srow(G), yielding descent for γk ∈(0, 2Ahyb(Gk)/(Lsrow(Gk))) and PŁ linear convergence under uniform bounds.Right-polar-first preserves the original feature geometry before row normalization.
  • F.6 Nuclear-Norm-Scaled Row-Norm/Right-Spectral Hybrid Optimizers: The row-norm/right-spectral hybrid instead computes its spectral step from the row-normalized Gram matrix, requiring alignment, effective-rank, and comparability assumptions for sublinear and PŁ linear convergence.With η(t) = 1/(t + ε) and uniform row-norm bound R, the assumptions hold with arow = ε, rrow = d, κrow = 1/(R+ε), and γ ∈(0, 2ε/(Ld)).
  • G Experimental Details: Experiments reduce hidden layers and experts, initialize all 2D weights with N(0, 0.02), use five inner steps for Polar Express and Gram Newton–Schulz with εNS = 10−7, and set ε = 10−8 for the specified optimizers.AdamW uses warmup and cosine decay for scalar/vector parameters, while other parameters use stable decay; four pre-training configurations are reported.

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

Qwen3-0.6B-style pre-training compared two configuration groups with similar wall-clock times across configurations. The second group was slower because HybridPolarGradM adds computational overhead for SwiGLU MLP projection matrices relative to Muon.

  • Training time: Configurations (i)–(iii) in (a) required 7.360, 7.498, and 7.369 hours, respectively.
  • Training time: Configurations (i)–(iii) in (b) required 7.765, 7.886, and 7.771 hours, respectively.
  • Training time: Group (b) took longer than (a) because HybridPolarGradM incurs additional overhead for SwiGLU MLP projection matrices compared with Muon.

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

The Gemma 3 1B-style experiment evaluates training configurations and sweeps a shared embedding/LM-head learning rate. Symmetry-compatible RowNormM and HybridPolarGradM remain better than AdamW across the sweep, rather than only at one learning rate.

  • G.2 Gemma 3 1B-Style Pre-Training: Configurations (i)–(iii) require 8.341, 8.807, and 8.132 hours in setting (a), versus 8.675, 9.101, and 8.469 hours in setting (b).The wall-clock comparison covers both optimizer settings and all three configurations.
  • G.2 Gemma 3 1B-Style Pre-Training: Table G.5 specifies the optimizer hyperparameters used for Gemma 3 1B-style pre-training.These hyperparameters define the optimization settings for the configurations evaluated in the experiment.
  • G.2.1 Gemma 3 1B-Style Pre-Training Learning Rate Sweep: The experiment sweeps a shared base learning rate for embedding and LM-head matrices while keeping scalar/vector and other matrix learning rates fixed.The sweep enforces γ0,emb = γ0,head, although separate tuning could be more delicate.
  • G.2.1 Gemma 3 1B-Style Pre-Training Learning Rate Sweep: Across swept learning rates, validation-loss gaps between configuration (iii) and configurations (i)–(ii) remain substantial.This indicates that the observed separation persists across the tested base learning-rate values.
  • G.2.1 Gemma 3 1B-Style Pre-Training Learning Rate Sweep: Across both optimizer settings, RowNormM and HybridPolarGradM remain consistently better than AdamW for the input embedding and LM-head matrices.AdamW makes comparable or slightly faster initial progress but remains consistently above the symmetry-compatible alternatives in validation loss.

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

Across two additional random seeds, Gemma 3 1B-style pre-training preserved the optimizer ordering for vocabulary-indexed matrices. RowNormM and HybridPolarGradM consistently outperformed AdamW, with HybridPolarGradM achieving the lowest mean final validation loss.

  • Experimental setup: The experiment varied only the input embedding and LM head optimizers among RowNormM, HybridPolarGradM, and AdamW, while SwiGLU MLP projections used HybridPolarGradM throughout.The MLP assignment used a row-norm/right-spectral composition.
  • Results: The qualitative optimizer ordering remained stable across random seeds, with RowNormM and HybridPolarGradM consistently outperforming the AdamW baseline.These gains were observed on the vocabulary-indexed matrices.
  • Results: 0.0441: HybridPolarGradM improved mean final validation loss over AdamW, compared with 0.0375 for RowNormM over AdamW.HybridPolarGradM achieved the lowest mean final validation loss.

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

The section evaluates whether projected, symmetry-compatible LM-head updates control final vocabulary-logit growth through the shared-logit-shift geometry. In Gemma 3 1B-style pre-training, these updates reduce raw-logit growth and improve validation loss relative to AdamW, with or without final-logit soft-capping.

  • Optimizer-side logit control: Projected LM-head updates remove the shared-logit-shift direction, targeting redundant logit drift while preserving centered logits that determine the softmax distribution.This quotient-geometry mechanism is intended to control final vocabulary logits without shrinking all logit variation.
  • Without final logit soft-capping: Without final-logit soft-capping, projected RowNormM and HybridPolarGradM LM-head updates achieve slightly lower validation loss than corresponding AdamW-LM-head configurations.AdamW variants exhibit substantially larger raw logit RMS, with RowNormM + AdamW showing a pronounced late-stage increase.
  • Without final logit soft-capping: AdamW-driven raw-logit growth is mainly shared-shift drift, because its centered logit RMS remains comparable to or smaller than that of the projected methods.The centered diagnostic measures quotient-space growth rather than the shared-logit-shift component.
  • With final logit soft-capping: With final-logit soft-capping enabled, AdamW raw-logit growth is largely suppressed, and maximum log-sum-exp remains in a comparable range across methods.This confirms soft-capping as an effective direct stabilization mechanism for final vocabulary logits.
  • With final logit soft-capping: Even with soft-capping, projected RowNormM and HybridPolarGradM updates achieve lower validation loss while producing larger centered logit RMS than AdamW variants.The result suggests projected updates suppress redundant shared-shift drift while allowing useful centered logit variation.

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

The OLMoE-1B-7B-style pre-training configurations required approximately 10.7–10.8 hours of wall-clock training time, with only minor differences across configurations.

  • Training time: Wall-clock training times were 10.79 hours for configuration (i), 10.82 hours for (ii), 10.71 hours for (iii), and 10.75 hours for (iv).The four configurations therefore had closely matched training durations.

G.4 Downsized gpt-oss Pre-Training

Downsized gpt-oss pre-training configurations (i)–(iv) have closely similar wall-clock training times, ranging from 18.91 to 19.34 hours.

  • Configurations (i) and (ii) take 19.08 hours and 18.91 hours, respectively.
  • Configurations (iii) and (iv) take 19.34 hours and 19.22 hours, respectively.
  • The four configurations therefore require approximately 19 hours of wall-clock training time.

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

The diagnostics evaluate router expert balance and logit-scale stability using load-balancing, z-loss, and assignment statistics. Across objectives with and without auxiliary penalties, symmetry-compatible router updates outperform AdamW on routing behavior without harming validation language-modeling performance.

  • Diagnostics: The evaluation logs expert assignment 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 L_LB = E Σ_i f_iP_i and is near 1 for a perfectly balanced router; router z-loss penalizes excessive logit scale.
  • Without auxiliary losses: Without auxiliary losses, all symmetry-compatible router updates achieve substantially lower measured load-balancing loss and router z-loss than AdamW on training and validation batches.The comparison varies only the router-matrix optimizer while keeping the other optimizer assignments fixed.
  • Without auxiliary losses: LeftPolarGradM and uncentered LeftPolarGradM nearly eliminate dead experts and produce the most balanced expert assignments, while RowNormM improves moderately and HybridPolarGradM is intermediate.AdamW has lower load entropy, higher load coefficient of variation, a larger dead-expert fraction, and higher maximum expert load.
  • Without auxiliary losses: These routing improvements do not hurt validation cross-entropy and slightly improve over AdamW, with similar centered and uncentered LeftPolarGradM behavior.The results suggest the dominant gain comes from left-polar router geometry, while centering supplies the shared-logit-shift-invariant quotient-space formulation.
  • With auxiliary losses: With L_LB and L_z weighted 0.01 and 0.001, respectively, auxiliary losses improve routing for every method, but RowNormM, LeftPolarGradM, and uncentered LeftPolarGradM remain best clustered.HybridPolarGradM improves over AdamW but remains less balanced than the strongest router-specific updates, showing that auxiliary losses and optimizer geometry are complementary.
Loading 2605.18106v4…