Optimizer Geometry Sets the Pace: Spectral Learning Dynamics in Matrix Factorization

Mahalakshmi Sabanayagam; Simon Lucey

cs.LGOct 3, 2026

Optimizer Geometry Sets the Pace: Spectral Learning Dynamics in Matrix Factorization

Authors: Mahalakshmi Sabanayagam, Simon Lucey

Organizations: Australian Institute for Machine Learning, Adelaide University

Abstract

Recent successes of matrix- and curvature-based optimizers have renewed interest in how update geometry shapes learning. These methods normalize or precondition updates, changing how different components progress during training. In deep matrix factorization, the geometry that slows gradient descent (GD) favors low-rank solutions by delaying the emergence of small singular modes. This raises the question of what remains of that spectral bias when normalization or curvature correction weakens or removes the slowdown. We study how optimizer geometry and depth jointly govern this behavior through a common framework for singular mode dynamics. Under explicit balance and alignment assumptions, we derive birth, saturation, and decay laws for Euclidean GD, coordinate-wise updates of SignGD and an instantaneous Adam approximation, spectral updates of Muon and cumulative Shampoo, and a block curvature model of K-FAC. The resulting picture is not a simple ordering from stronger to weaker low-rank bias. To highlight, SignGD and ideal Muon eliminate the divergent birth barrier and drive unsupported modes to zero in finite time. Cumulative Shampoo initially retains GD's depth-dependent barrier, then accumulated gradients produce a catch-up phase while making previously active modes increasingly persistent. Undamped K-FAC cancels the factorization-induced slowdown while preserving the ordering of the target singular values, whereas positive damping introduces a spectral threshold below which the slow GD phase laws reappear. These results give normalization, accumulated state, damping, and depth a direct interpretation as controls determining when modes emerge, persist, and disappear during training.

Figures & tables

Panel of Figure 1 — Figure 1: Optimizer geometry sets the pace of spectral learning. In balanced aligned deep matrix factorization, GD and Muon produce distinct spectral learning trajectories (experimental details in App. F ). (a) Under GD, larger target modes are learned earlier, while smaller modes remain suppressed for longer. (b) Under Muon, factor amplitudes grow at a common speed, allowing smaller targets to reach full recovery earlier. (c) From the same initial amplitude $0.01$ , Muon accelerates the growth of a positive target mode (red) and, importantly, removes the zero target mode in finite time (blue). GD instead exhibits delayed growth and slow decay. Thus, faster emergence of weak modes does not imply weaker removal of unsupported modes. Optimizer geometry determines which spectral components emerge, persist, and disappear. Other optimizers are discussed in later sections.

Optimizer	${Z_{\ell}}$	Update category
GD	$H_{\ell}$	Euclidean
SignGD ( Bernstein et al., 2018 )	$\sign(H_{\ell})$	Coordinate-wise
Adam ( Kingma and Ba, 2014 )	${m_{\ell}}\oslash({\sqrt{v_{\ell}}+\epsilon})$ with $\tau_{m}\dot{m}_{\ell}=H_{\ell}-m_{\ell},\,\,\tau_{v}\dot{v}_{\ell}=H_{\ell}^{\odot 2}-v_{\ell}$	Coordinate-wise (adaptive)
Ideal Muon ( Jordan et al., 2024 )	$U_{\ell}V_{\ell}^{\top}$ where $H_{\ell}=U_{\ell}\Sigma_{\ell}V_{\ell}^{\top}$	Spectral
Shampoo ( Gupta et al., 2018 )	$C_{\ell}^{-1/4}H_{\ell}D_{\ell}^{-1/4}$ with $\lambda_{\mathrm{Sh}}>0,$ $C_{\ell}(t)=\lambda_{\mathrm{Sh}}I+\int_{0}^{t}H_{\ell}(u)H_{\ell}(u)^{\top}\,du,$ $D_{\ell}(t)=\lambda_{\mathrm{Sh}}I+\int_{0}^{t}H_{\ell}(u)^{\top}H_{\ell}(u)\,du$	Spectral
K-FAC ( Martens and Grosse, 2015 )	$\operatorname{vec}(Z_{\ell})=(\mathcal{F}_{\ell}+\lambda_{\mathrm{KF}}I)^{-1}\operatorname{vec}(H_{\ell}),$ $\mathcal{F}_{\ell}=(A_{\ell}A_{\ell}^{\top})\otimes(B_{\ell}^{\top}B_{\ell})$ , $\lambda_{\mathrm{KF}}\geq 0$	Curvature based

Table 1: Layer update directions in the flow

\dot{W}_{\ell}=-Z_{\ell}

, where

H_{\ell}=B_{\ell}^{\top}GA_{\ell}^{\top}

. SignGD and Adam applies coordinate-wise updates, Muon uses polar normalization, Shampoo uses cumulative statistics, and K-FAC uses additive block damping. Detailed account of these optimizers is in App. B .

Figure 2: Verification of the phase laws in Theorem 2 at depth

L=4

. (a) The birth exponent

\alpha_{\mathrm{b}}

separates Euclidean, normalized, history adapted, and curvature corrected dynamics. (b) The decay exponent

\alpha_{\mathrm{d}}

distinguishes finite extinction of Muon (

\alpha_{\mathrm{d}}\!=\!q\!<\!1

) from the other optimizers. (c) The saturation coefficient

\beta

reveals target dependent mode ordering.

Panel of Figure 3 — Figure 3: Optimizer state parameters control the local birth regime. Adam’s stability parameter $\epsilon$ shifts the GD-to-SignGD crossover; Shampoo damping delays the transition from $p$ to $p-1$ ; and K-FAC damping restores the GD exponent below its curvature threshold.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Optimizer	Learning rate	Main state parameter	State value	Depth
GD	$5\times 10^{-2}$	none	none	$L=\{2,4\}$
GD	$8\times 10^{-2}$	none	none	$L=\{8\}$
Adam	$1\times 10^{-3}$	$(\beta_{1},\beta_{2},\epsilon)$	$(0,0,0.02)$	$L=\{2,4\}$
Adam	$2\times 10^{-3}$	$(\beta_{1},\beta_{2},\epsilon)$	$(0,0,0.02)$	$L=\{8\}$
Muon	$5\times 10^{-4}$	momentum	$0$	$L=\{2,4,8\}$
Shampoo	$2\times 10^{-4}$	$\lambda_{\mathrm{Sh}}$	$10^{-6}$	$L=\{2,4,8\}$

Appendix

Table 2: Optimizer settings used for the full observation matrix factorization experiments using depths

L=\{2,4,8\}

. The learning rates are chosen to keep the discrete trajectories close to the corresponding continuous flows while allowing all methods to reach comparable final training loss. Each optimizer is trained for a maximum of

5000

steps.

Optimizer	Learning rate	Main state parameter	State value
GD	$3$	none	none
Adam	$0.05$	$(\beta_{1},\beta_{2},\epsilon)$	$(0,0,0.02)$
Muon	$0.00025$	momentum	$0$
Shampoo	$0.003$	$\lambda_{\mathrm{Sh}}$	$10^{-6}$
K-FAC	$0.0005$	$\lambda_{\mathrm{KF}}$	$10^{-4}$

Appendix

Table 3: Optimizer settings used for underdetermined linear regression with depth

L=4

. The learning rates were selected over a different grid for different optimizers to obtain stable training and comparable final data loss.

Figure 4: Optimizer geometry sets the pace of spectral learning evaluation in matrix factorization. (a) Evolution of a weak target mode with

\mu=0.18

, shown as the learned fraction

\sigma(t)/\mu

. GD exhibits the factorization induced small mode barrier, whereas normalization and preconditioning substantially accelerate emergence. (b) Birth time to the common amplitude

b=0.03

as the target mode

\mu

varies. While ideal Muon is independent of the target mode ordering during birth, other optimizers are target sensitive. (c) Decay trajectories for an unsupported mode with target

\mu=0

. Ideal Muon removes the mode in finite time, undamped K-FAC produces exponential decay, while GD, instantaneous Adam with fixed stabilization, and cumulative Shampoo retain increasingly slow tails near zero. Optimizer geometry therefore determines which spectral components emerge, persist, and disappear.

Panel of Figure 5 — Figure 5: Optimization and rank dynamics in matrix factorization. All optimizers train the same depth $L=4$ deep linear factorization toward the same rank eight target and reach comparable final training loss. Left: training loss over optimization steps. Middle: stable rank of $W$ . Right: effective rank of $W$ . The optimizers exhibit markedly different transient rank dynamics despite converging toward similar final spectra. GD retains a lower rank representation for longer, whereas Adam, Muon, cumulative Shampoo, and K-FAC alter the rate at which spectral mass is distributed across modes.

Panel of Figure 6 — Figure 6: Singular mode trajectories reveal optimizer dependent spectral clocks. GD exhibits strongly sequential learning;Adam reduces these delays, while Muon makes the modes evolve on nearly the same clock until they approach their respective targets. Cumulative Shampoo shows a faster history adapted growth of weaker modes. K-FAC further reduces the factorization induced separation between mode time scales through curvature correction. The common final amplitudes are determined by the unique full observation target, whereas the transient trajectories expose the distinct inductive biases of the optimizers.

Figure 7: Depth changes both the spectral exponents and the optimizer state transitions. Left: for instantaneous Adam, increasing depth changes both the GD like exponent

p=2-2/L

and the normalized exponent

q=1-1/L

. Middle: for cumulative Shampoo, depth changes the initial GD like exponent

p

, the history adapted exponent

p-1

, and the amplitude at which accumulated history begins to dominate. Right: for damped KFAC, deeper factorizations strengthen the small mode GD like barrier while curvature correction still drives the exponent toward zero above the damping controlled mobility threshold.

Figure 8: Effect of depth on decay laws and Shampoo saturation order. Left: empirical decay exponents

\alpha_{d}

for unsupported modes as depth varies. All optimizers closely follow the theoretical predictions in general. Right: Shampoo saturation coefficient

\beta

as a function of the positive target mode

\mu

for

L\in\{2,4,8\}

. Shampoo is singled out here because its saturation phase is where mode order reversal can happen theoretically. At

L=2

\beta

decreases with

\mu

, so smaller target modes saturate faster. At

L=8

, the ordering reverses and larger target modes saturate faster, while

L=4

lies near the transition. These results show that depth influences not only the size of the phase law exponents, but also the ordering of mode fitting under optimizer state dynamics.

Panel of Figure 9 — Figure 9: Optimization and rank dynamics in full observation matrix factorization. All optimizers train the same depth $L=2$ deep linear factorization toward the same rank eight target and reach comparable final training loss.

Panel of Figure 10 — Figure 10: Singular mode trajectories show optimizer dependent spectral clocks for $L=2$ .

Panel of Figure 11 — Figure 11: Optimization and rank dynamics in full observation matrix factorization. All optimizers train the same depth $L=8$ deep linear factorization toward the same rank eight target and reach comparable final training loss.

Panel of Figure 12 — Figure 12: Singular mode trajectories show optimizer dependent spectral clocks for $L=8$ .

Figure 13: Birth, decay and normalized saturation dynamics of Adam. Dotted and solid lines mark the theoretical and empirical values for

L=4

, respectively. These values align in both the decay and saturation rate plots across different values of

\epsilon

Figure 14: Phase transitions induced by cumulative Shampoo state. Solid curves show empirical dynamics and dashed curves show the corresponding aligned theory for depth

L=4

. Together, the panels show that cumulative Shampoo changes its effective spectral geometry during training while retaining explicit memory of earlier gradient activity.

Figure 15: Additive K-FAC damping controls the spectral phase transitions. Solid curves show empirical dynamics and dashed curves show the corresponding aligned theory for depth

L=4

. It shows that

\lambda_{\rm KF}

defines a spectral mobility threshold separating curvature corrected dynamics from GD-like behavior.

Panel of Figure 16 — Figure 16: Optimizer dependent dynamics persist in underdetermined linear regression. We train the same deep linear predictor on an underdetermined, nonwhitened regression problem using GD, Adam, Muon, cumulative Shampoo, and KFAC. All methods reach comparable final loss and final spectra of $W$ despite substantially different transient dynamics shown by the stable rank and effective rank throughout training. The optimizer geometry strongly changes how quickly spectral mass concentrates.

Panel of Figure 17 — Figure 17: Singular value trajectories in nonaligned underdetermined linear regression. Although all methods approach similar final singular spectra, their transient spectrum differ substantially. The qualitative optimizer dependent dynamics predicted by the aligned phase analysis remain visible even when the empirical data covariance couples singular directions and the exact scalar mode reduction no longer holds.

Explore similar work

Jun 29, 2026

Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics

Matrix factorization (i.e., problems of the form

\min_{\mathbf{P},\mathbf{Q}} \|\mathbf{M}^\star - \mathbf{P}^\top\mathbf{Q}\|_\mathrm{F}^2

) is a minimal learning problem that exhibits both nonlinear parameter dynamics and representation learning. In this setting, we study how parameter trajectories under the Muon optimizer differ from those of gradient descent. We identify three main dynamical differences: 1) Muon avoids the slow saddle-to-saddle dynamics from small initialization. Muon instead learns all the top modes of

\mathbf{M}^\star

at the same rate, with the smaller modes converging first. 2) Muon remains stable even when the learning rate exceeds the critical threshold set by the local loss sharpness. This frees the learning rate from the condition number of the problem, enabling rapid convergence via exponential learning rate annealing. 3) Once the weights are aligned with each other and the target, Muon flow conserves the matrix quantity

\sqrt{\mathbf{P}^\top \mathbf{P}}-\sqrt{\mathbf{Q}^\top \mathbf{Q}}

, while gradient flow is known to conserve the matrix

\mathbf{P}^\top\mathbf{P} - \mathbf{Q}^\top\mathbf{Q}

. Despite having distinct conserved quantities, both optimizers find the so-called balanced solution from vanishing initialization. When training from small random initialization, the weights spontaneously align early in training. We derive the alignment rates in simple settings and show that they predict the empirical alignment rates in general. Finally, we exploit structural properties of Muon to construct a learning rate schedule that achieves near-perfect alignment in only two optimization steps.

Mark Rhee, Jamie Simon, Dhruva Karkada

UC Berkeley · Imbue

Muon Optimizer Matrix Optimization Matrix Factorization Gradient Descent Dynamics

Ask PDF

Jul 14, 2026

Reassessing Muon for Matrix Factorization

Muon has recently emerged as a strong optimizer for large-scale deep learning, where it reshapes gradient updates through approximate orthogonalization and has been reported to outperform Adam and AdamW in large language model training. Its empirical success has motivated a growing body of theoretical work that interprets Muon as steepest descent under the spectral norm. Yet it remains unclear which of Muon's advantages stem from its update rule itself and which are artifacts of the scale, architecture, and data of modern deep networks. In this work, we isolate the optimizer from these confounding factors by studying Muon on a simple, well-understood, and spectrally structured problem: low-rank matrix factorization. Through a controlled comparison against carefully tuned adaptive baselines, we find that Muon does not consistently outperform AdamW in this setting and that several previously reported advantages are sensitive to hyperparameter choices. Our results provide a more nuanced picture of when spectrum-aware orthogonalization is beneficial and argue for evaluating modern optimizers on controlled problems in addition to end-to-end benchmarks.

Ali Parviz, Gal Mishne, Alex Cloninger

†. Halicioğlu Data Science Institute, UC San Diego. · ∗. Department of Mathematics, UC San Diego.

Deep Learning Optimization Muon Optimizer Adaptive Gradient Methods

Ask PDF

Jun 7, 2026

The Spectral Dynamics and Noise Geometry of Muon

Muon replaces a matrix gradient

G=UΣV^\top

by its polar factor

UV^\top

. This keeps the singular directions selected by the gradient, but makes the update spectrum flat. We study the optimization bias created by this operation. Under explicit alignment assumptions, we prove that the polar update is the one-step entropy-maximizing choice among bounded updates that use the gradient singular directions and do not adapt to the current weight spectrum. In an underdetermined regression model, we derive exact singular-value dynamics for continuous-time Muon and identify a measurement-dependent condition under which the normalized spectrum moves toward equal nonzero singular values. This geometry also rules out a common low-rank interpretation: at fixed Frobenius norm, Muon's distinguished state has a flat spectrum, whereas nuclear-norm minimization favors spectral concentration. Controlled matrix-sensing experiments separate the effect from simple gradient rescaling, show that norm-matched gradient descent does not reproduce Muon, and recover the predicted flattening trend across broad ablations. In small NanoGPT pretraining, Muon preserves stable rank, has a broad learning-rate plateau, and improves validation loss relative to AdamW; in a matched small-ViT control, the ranking reverses. The resulting picture is regime-dependent: Muon is not universally superior, but its flat-spectrum bias can help when many spectral directions need to remain active.

Pierfrancesco Beneventano, Mahmoud Abdelmoneum, Tomaso Poggio

Massachusetts Institute of Technology

Muon Optimizer Matrix Optimization Spectral Methods Implicit Regularization Gradient Descent Dynamics

Ask PDF