How Muon’s Spectral Design Benefits Generalization: A Study on Imbalanced Data
Authors: Bhavya Vasudeva
Puneesh Deora
Yize Zhao
Vatsal Sharan
Christos Thrampoulidis
Venue: ICLR 2026 Poster (2026)
Keywords: Muon
Shampoo
Spectral Gradient Descent
Generalization
Primary area: interpretability and explainable AI

Abstract: The growing adoption of spectrum-aware matrix-valued optimizers such as Muon and Shampoo in deep learning  motivates a systematic study of their generalization properties and, in particular, when they might outperform competitive algorithms. We approach this  question by introducing appropriate simplifying abstractions as follows: First, we use imbalanced data as a testbed. Second, we study the canonical form of such optimizers, which is Spectral Gradient Descent (SpecGD)—each update step is $\mathbf{U}\mathbf{V}^T$ where $\mathbf{U}\mathbf{\Sigma}\mathbf{ V}^T$ is the truncated SVD of the gradient. Third, within this framework we identify a canonical setting for which we precisely quantify when SpecGD outperforms vanilla Euclidean GD. For a Gaussian mixture data model and both linear and bilinear models, we show that unlike GD, which prioritizes learning dominant principal components of the data first, SpecGD learns all principal components of the data at equal rates. We demonstrate how this translates to a growing gap in class balanced loss favoring SpecGD early in training and further show that the gap remains consistent even when the GD counterpart uses adaptive step-sizes via normalization. By extending the analysis to deep linear models, we show that depth amplifies these effects. We empirically verify our theoretical findings on a variety of imbalanced datasets. Our experiments compare practical variants of spectral methods, like Muon and Shampoo, against their Euclidean counterparts and Adam. The results validate our findings that these spectral optimizers achieve superior generalization by promoting a more balanced learning of the data's underlying components.

Published as a conference paper at ICLR 2026
How MUON'S SPECTRAL DESIGN BENEFITS GENERALIZATION: A STUDY ON IMBALANCED DATA
Bhavya Vasudeva" " Puneesh Deora
+ Yize Zhao Vatsal Sharan " Christos Thrampoulidis
UUniversity of Southern California $ University of British Columbia
bvasudev@usc.edu :puneeshdeora, ethrampo dece. .uc.cubccc
ABSTRACT
The growing adoption of spectrum-aware matrix-valued optimizers such as Muon and Shampoo deep learning motivates a systematic study of their generalization properties and, in particular, when they might outperform competitive algorithms. We approach this question by introducing appropriate simplifying abstractions as follows: First, we use imbalanced data as testbed. Second, we study the canonical form of such optimizers, which is Spectral Gradient Descent (SpecGD) each update step is UV where UEVT is the truncated SVD of the gradient. Third, within this framework we identify a canonical setting for which we precisely quantify when SpecGD outperforms vanilla Euclidean GD. For a Gaussian mixture data model and both linear and bilinear models, we show that unlike GD, which prioritizes learning dominant principal components of the data first, SpecGD learns all principal components of the data at equal rates. We demonstrate how this translates to growing gap in class balanced loss favoring SpecGD early in training and further show that the gap remains consistent even when the GD counterpart uses adaptive step-sizes via normalization. By extending the analysis to deep linear models, we show that depth amplifies these effects. We empirically verify our theoretical findings on variety of imbalanced datasets. Our experiments compare practical variants of spectral methods, like Muon and Shampoo, against their Euclidean counterparts and Adam. The results validate our findings that these spectral optimizers achieve superior generalization by promoting more balanced learning of the data's 's underlying components.
INTRODUCTION
Spectrum-ave aware optimizers such as Shampoo (Gupta et al., 2018), Muon (Pethick al., 2025; Jordan et al., 2024) have recently gained significant traction in the deep learning community, delivering substantial training speedups for deep classifiers (Jordan et al., 2024) and transformer language models (Vyas et al., 2025; Liu et al., 2025) compared to standard methods like SGD (Robbins & Monro, 1951) with momentum or even Adam (Kingma & Ba, 2014). The key distinction lies in how these methods treat neural network parameters: while SGD and Adam operate entry-wise on vectorized parameters, Shampoo and Muon work directly with matrix-valued parameters (e.g.. weight matrices and attention matrices) at the layer level (Carlson et al., 2015; Large al., 2024). This matrix-level approach intuitively enables optimization trajectories that entry-wise methods cannot achieve. Despite their empirical success, a fundamental question remains unanswered: when do spectrum-aware optimizers generalize better than standard methods?
The challenge is substantial. Even establsabl optimizers like Adam, despite more than a decade of practical dominance, remain less well understood than SGD, whose Euclidean GD trajectory is well-characterrized both statistically and algorithmically. Recent theoretical progress has begun to illuminate these methods through the lens of implicit bias. For instance, while the implicit bias of GD toward --ror max-margin classifiers is well-established (Soudry et al., 2018: 8; Ji & Telgarsky, 2018), recent work proved that Adam converges toward -nom max-margin classifiers in linear settings (Zhang et al., 2024; Xie & Li, 2024). Conceptually, this difference can be understood by noting that Adam reduces to GignGD when momentum and preconditioning histories are ignored (i.e., 3 parameters set to zero). This reduction has often been leveraged before to study Adam's 's properties, since SignGD is simpler to analyze, e.g., Kunstner al. (2024); Vasudeva et al. (2025b).

---

Published as a conference paper at ICLR 2026
Sigrum
Maon
Eigenvaluss of F(xx')
Figenvactors Heatmap
Reconstructions using spectral componints
1.0
16
Row Index
channd
7
7
Meortt
3 Orignal
10
15
Index
Elgerveccor Index
F 75 32
Epech
Figure 1: (Left) Test accuracy for training an MLP on ColOred-MNIST with 99% digit-colour correlation using NMD, Signum, and Muon (see text for details). All optimizers achieve near- perfect accuracy on majority groups (same digit-colour labels), but Muon outperforms the others on minority groups (opposite digit-colour labels), especially early in training. (Middle) Eigenvalues and eigenvectors of the empirical moment matrices for all training images show that the dominant spectral components correspond to colour. (Right) Reconstructions of sample images reveal that digit-shape information lies less dominant components. This suggests that Muon's 's spectral design promotes balanced learning of these components, leading to superior generalization.
A similar reduction exists for spectrum-aware optimizers: Shampoo and Muon without accumulations and with exact matrix operations reduce to Spectral Gradient Descent (SpecGD) (Pethick et al., 2025; Bernstein & Newhouse, 2024), where each update step is UV with UEVT denoting the truncated SVD of the gradient (see Sec. for details). Thus, just as SignGI serves as the canonical form for understanding Adam, SpecGD serves as a canonical form for understanding Shampoo and Muon. Moreover, both SignGD and SpecGD are instances of normalized steepest descent with respect to max and spectral norms, respectively (Bernstein, 2009; Carlson et al., 2015; Pethick et al., 2025).
Recently, Fan et al. (2025); Tsilivis et al. (2024) charac- terized SpecGD's ' optimization bias in linear multi-class and homogeneous binary classification, showing that it drives weights toward max-margin classifier with re- spect to the spectral norm. However, these results have two limitations. First, they only describe the algorithm's 's behavior in the terminal phase of training, which may not reflect practical deep learning scenarios that often employ early stopping. Second, and more importantly, implicit bias results provide no direct guarantees about general- ization performance-the the ultimate objective in machine learning. For instance, even in linear settings, the mini- mum spectral norm solution may not be unique. Recent work has also studied Muon and Shampoo, but primarily from the perspective of optimization properties (Pethick et al., 2025; Morwani et al., 2024; Vyas et al., 2025; Li & Hong, 2025; Chang et al., 2025; Chen et al., 2025) scalability (Liu et al., 2025; Boreiko et al., 2025), leaving their generalization behavior and underlying mechanisms comparatively underexplored.
38
2.0
0.0
462
3.800
Wortt dlas
Class belonced
1.0
0.4
Time (1)
Figure 2: Dynamics of the iterates (top) and loss (bottom) for GD, NGD, and SpecGD on a linear model with class im- balance (priors pı > P2 > pa). SpecGD learns all spectral components at the same rate, whereas (N)GD prioritizes more dominant components. Although all con- verge to the same solution, SpecGD sig- nificantly outperforms in worst-class and loss early-on in training. See Sec. 3.2 for a precise characterization of SpecGD dynamics and a proof of its superiority over (N)GD.
Motivated by the apparent lack of understanding of the generalization properties of SpecGD, we ask: Can we identify concrete settings where SpecGD generalizes better than standard (Euclidean) GD?
Imbalanced Data as Playground. We introduce imbal- anced data as a testbed for studying SpeeGD's 's potential generalization advantages. Fig. provides a concrete demonstration: we train an MLP on Colored-MNIST (Arjovsky et al., 2020; Pezeshki et al., 2021) under severe group imbalance, where each digit appears in its majority color 99% of the time during training, creating strong spurious digit-color correlation (see App. D. for details). At test time, we evaluate on majority and minority groups with same and opposite digit-color associations, re- spectively. The results revealing: whereas Muon improves on both groups early in training, its Frobenius and max norm counterparts, namely normalized momentum descent (NMD) and Signum (Bersttein al., 2018), respectively, take much longer to improve on the minority group. We observe
2

---

Published as a conference paper at ICLR 2026
similar effect on the CIFAR dataset (Krizhevsky et al., 2009) with class imbalance (see Sec. 2 for details).
Theoretical Comparison in Class-ibbalanced Data. To demystify this behavior and analyze the implicit regularization of different update rules, we analyze in detail a linear model trained with squared-loss under class imbalance. We derive closed-form expressions for the training trajectories of (Euclidean) GD and SpecGD, and use them to show that, with early stopping, SpecGD achieves a lower worst-class and class-balanced risk than GD, whereas both approach the same risk asymptoti- cally in time. We also show that introducing adaptive step-size in GD via normalization does not equalize the gap: SpeeGD outperforms normalized GD (NGD) as well. This illustrated in Fig.
Extending to deep networks, we derive the closed-form training trajectory for SpecGD with L 2 1 layers. Our findings reveal two key effects of depth: i) it accelerates the learning of all spectral components, and ii) it narrows the time gap between the saturation points of different components.
Experimental Results. We empirically validate our theoretical insights through experiments on datasets with class or group imbalance. We demonstrate that practical variants of SpecGD, such as Muon and Shampoo, achieve superior generalization over SGD by promoting a more balanced learning of the spectral components of the data. While our theory focuses on uncovering the distinct learning mechanisms of SpecGD and GD, our experiments also include Adam to provide a comprehensive benchmark against commonly-used baseline.
2 RESULTS ON CIFAR WITH STEP-IMBALANCE
In this section, we conduct a preliminary exper- iment to investigate whether Muon's superior- ity observed in the group-imbalaneed setting in Fig. also persists under class imbalance. We consider step-imbalanced CIFAR-10 and CIFAR- 100 (Krizhevsky et al., 2009) with imbalance ratio 20 (majorities contain 20x more samples than minorities). We train ResNet- 1 8/50 on CIFAR- 10/100 with three optimizers: NMD, Signum, and Muon (see App. D.2 for details and additional re- sults with lower imbalance ratios). As shown Fig. 3, Muon achieves superior minority-class per- formance compared to NMD and Signum. While all optimizers achieve similar majority-class ac- curacy, Muon more effectively reduces the perfor- mance gap between minority and majority classes, particularly early on in training. This result fur- ther suggests that spectral methods may offer ad- vantages in imbalanced settings, motivating our subsequent theoretical investigation into the mech- anisms underlying this behavior.
NMD
Signum
Muon
Minonty Majority A0
1.0
Minonty Acc Majontty AcE:
0.8
0.6
0.4
0.2
0.0
0.8
0.4
0.2
0.0
0.0
10"
10a
10"
102
Epochs
Epochs
(a) CIFAR-I0
(b) CIFAR-100
Figure 3: Test accuracy dynamics for training ResNet-based models (a) CIFAR-10, and (b) CIFAR-100, both with STEP imbalance (20:1 majority- tomiintttto-mmnoritt class ratio) using NMD, Signum, and Muon. All optimizers achieve com- parable accuracy on majority classes, but Muon outperforms the others on minority-class perfor- mance, especially early in training.
3 LINEAR MODEL ON CLASS-IMBALANCE DATA
Notation. We denote matrices, vectors and scalars by A, and respectively. We denote the (i.j)-th entry of matrix A as Ali,j]. Let I1- ll. Imaxx and [|-112 denote the Frobenius, max and spectral norms, respectively, where || |||maax == max,j J|. Let ||a||2 denote the l norm of 1 [ denotes the indicator function, e.g., la 2 bj = 1f a 2 and 0 otherwise. For matrices A, B. the inner product (A, B) tr(AB
Let W& E RK denote the weight matrix of a linear model, and L(W) denote the loss. Let W, and V. VL(WW) denote the iterate and gradient at time I,, respectively. Normalized steeptst-descent ||-1|, with step-size 1) > 0, are (Boyd & Vandenberghe, 2009): W = W.- m where A, := argmax A S (V,, A).
(NSD) updates with respect to norm
As discussed in Sec. 1, NGD, SignGD and SpecGD are instances of normalized steepest descent Eq. (1) with respect to Frobenius, max and spectral norm, respectively. Normalized momentum steepest-descent (NMD) updates are defined similarly as Eq. (1) but use momentum M, -M--1
3

---

Published as a conference paper at ICLR 2026
Table Normalized steepest descent (NSD) and normalized momentum descent (NMD) updates for different norms. Here, V denotes the gradient, M, = - + (1 B) V denotes the momentum term, and V. U and M, Ü,E,V V denote their SVDs. Setting 3 = in the NMD updates yields the corresponding NSD updates.
Norm
NSD update A
argmax S1 V A
NMD update A := argmaxx |All (M, A)
|H||F
NGD: A Ville
NMD: A ||
M,
SignGD: A sign V
Signum: A, -M.
1:112
SpecGD: A, U.V.
Muon: A, U.
(1 3) v in place of V. Table summarizes the updates for NSD and NMD for the three norms. Note that Muon with exact matrix operations is NMD with spectral norm, and reduces to SpecGD when 3 0. Also see App. A, where we discuss how Shampoo and Adam without accumulations reduce to SpecGD and SignGD, respectively.
Data Model. Let y {1, (l,.... k} denote the class labels, and let the corresponding one-hot labels be denoted as E {ee where e, is the e-th standard basis vector in R* Class probabilities given by Pe == Pr(y Me E Rd, and samples for class are generated as isotropic Gaussians with mean Mec Finally, we assume that the means μhe are orthogonal. Put together, the data model we study is such that:
such that Σ Pe 1. Each class has associated mean vector
Pr(y = c) Pe, CE [k],
æly ~ N (Ny, o'laa) and L...L pi | || = μ..
(DM)
Define minority class index m == arg mince[s] Pe with class prior Pm, and majority class index
M :==ary arg maX Pe with class prior PM. Let SNR:= e
oz denote the signal hignal-to-oois to-noise ratio (SNR).
3.1 EXPERIMENTAL RESULTS
For the experiments, sider a heavy-tailed class- imbalanced setting by choos- ing Pe & 1. and we sample each Me independently from zero- cero-mean isotropic Gaus- sian distribution and normal- ize it. We use 20 classes, 100 samples, d = 200 and Tz 0.1. We initialize the weight matrix W, by sampling each entry indepen- dently from (0, 1). We use learning rates 0.025, 0.005 and 5 x 10-4 for NGD, SignGD and SpecGD, respectively. These choices are made such that train loss curves of the algorithms comparable. The training is stopped when the norm of the gradient for SpecGD falls below 10-6 as small gradients can introduce numerical imprecision in the SVD computation.
NGD
SIgnGD
Spec3D
Normiized Margin
100
Training Loss
10"
Tosa Accurayy
'Wel
10"
10"
50
10"
10"
10"
10
13"
10"
10
10
13"
Heraaoro
leralions
Baralions
Figure 4: Results for linear model trained with NGD, SignGD and SpecGD on heavy-tailed Mlass-imbllane data using cross-entropy loss. Early-stopped SpecGD attains higher class-balanced and worst- class accuracy compared to other update rules or stopping points.
Fig. 4 shows results of training linear model in this setting using NGD/SignGD/SpecGD to minimize cross-entropy loss. We observe that for all three update rules, as train loss approaches 0, the iterates W, converge to solution that maximizes the margin defined with respect to the corresponding norm (Fan al., 2025). However, comparing test performance, early-stopped SpecGD attains higher and worst-class test accuracies compared to NGD or SignGD at any stopping point.
3.2 THEORETICAL ANALYSIS
In this section, we analyze and compare the dynamics of GD and SpecGD. For tractability, we consider squared-loss, and population setting. Specifically, let L(W)= y W where the expectation
is over the joint distribution of in (say) (DM). In addition, define C. (W) )- ( [[alpp=lly W to be the class-conditional loss for class e [k) and let Laai Eeel (W( be the balanced
4

---

Published as a conference paper at ICLR 2026
loss. Define the population moment matrices Ewx: := 2x and Eye E= Our analysis holds when the (full) SVDs of these moment matrices are jointly diagonalizable.
Condition (Joint Diagonalizability). There exist orthonormal matrices E RR V E Ra and matrices S IR Saz E Raxe with non-zero entries only along their main diugonals" 0 and s1x 2 82 ... 2 d 2 0, respectively), such that, Eyz= US ya V
{81 2 yx 2 %%² 2
2 and Exg-
V S
Under this condition applied to the empirical moment matrices, Saxe et al. (2013); Gidel et al. (2019) derive closed-food form training dynamics of two-layer linear networks. Here, we instead apply this condition on population statistics making it possible to study test statistics. In the lemma below, we show that our data model (DM) satisfies this condition. See App. B.1 for the proof.
Lemma 3.1. The population moment matrices of data model (DM) satisfy Condition 1, with s for E [k:], se = 1² prpe + o² for 1 S c S k, and -ex = 0 k < c S d.
This result shows that the data model (DM) is one such setting where Condition 1 strictly holds. Notably, in Gidel al. (2019), the authors validate whether weaker version of this condition is satisfied for some datasets used in practice. Specifically, they check whether there exist orthonormal matrices U, V such that the empirical moment matrices are jointly diagonalizable Uy V and V(E + B)VT, respectively, where || is much smaller than |Em=|| |l- They find that on datasets like MNIST, CIFAR, etc., the ratio is indeed small (see Table 1 in Gidel et al. (2019)), 9)), so this can be considered as a reasonable condition. We consider Condition I, where || BB| O, mainly to make the analysis more tractable. See App. C for further discussion, and experimental results that examine the effect of relaxing this condition.
Evolution of W.. We compare the evolution of the weight matrix W, for GD and SpecGD over iterations t. For each iteration t. define W. := == U W.V and recall that W clc, c] denotes the ch main-diagonal entry of We for c E [k:.
Under Condition 1, Saxe et al. (2013) (see also App. B.2) shows that when initialized zero and run with sufficiently small step size 1) mineggs (s**)=,, GD iterates W, such that We is diagonal at each iteration with diagonal entries evolving as (and approximation becoming accurate for gradient flow (GF) 1) 0):
W [.(c,c] 3
1-(1 (1-
-000 (1-
(1 ---e-1i
The following result characterizes the dynamics of W, for SpecGD. See App. B.2 for the proof. Proposition 1. Assume zero initialization Wo 0 and Condition I holds. Then, for SpecGD with a step size 17 < min sea, at each iteration, W. UW, V where W, & is zero except its main diagonal along which the entries evolve as follows for E [k]]
W.lc.d = nt
Comparing the above two displays, which contrast GD's and SpecGD's iterate evolution, shows the following. Although both methods asymptotically converge the same solution, their trajectories differ significantly. GD learns component с at rate proportional to s²x meaning more dominant spectral components are learned faster. In contrast, SpeeGD learns all components at the same rate until each individual value saturates and converges to its terminal value. This difference is also demonstrated in Fig. 2 (top). which tracks W Wc, c for GD and SpeeGD in a class-imbalanend (data model (DM)) with r) 0.01, d = k- = μ 1, o² 0.125, and p1; P2, P3 set 0.5, 0.3, 0.2, respectively.
setting
In App. C, we compare the dynamics of GD and SpecGD derived under Condition I to more general setting using the same data model but with finite samples and random initialization, and show that the theoretically derived dynamics closely match those observed empirically (see Fig. 11). Additionally, in App. C, we also compare the dynamics of Muon (with 3 = 0.9) to those of SpecGD, both when Condition 1 holds (see Fig. 10) and in the finite sample setting (see Fig. 12). In both cases, we find that similar to SpecGD, Muon promotes learning different spectral components at a similar rate.
"With slight abuse of terminology, we will refer to rectangular matrices having non-zero entries only on their main diagonals simply as 'diagonal' hereon.
5

---

Published as a conference paper at ICLR 2026
Generalization of SpecGD vs. (N)GD. We now show that this property of SpecGD to learn concepts at equal rate translates to superior generalization compared to GD an imbalanced setting where least-significant spectral components of the moment matrices are associated to minority classes. Concretely, under data model (DM), the k first eigenvectors of En align with the class-mean vectors, ordered in decreasing class prior probability (see proof of Lemma 3.1). Intuitively, then, learning spectral components earlier during training should translate to generalization gains. Indeed, in Fig. 2 (bottom), we observe that SpecGD attains lower minority-class and class-balanced loss compared to GD. The following theorem formalizes this intuition. We use continuous- time versions of GD and SpecGD here, i.e., gradient flow (GF) and spectral GF (SpecGF). The dynamics for GF are characterized by the approximation in Eq. (2), while the SpecGF dynamics follow Proposition 1, both with = (see App. B.2 for details). We use W (t) and W W(t) to denote the continuous- time iterates and the corresponding singular value matrix, respectively.
Theorem (Simplified). Assume data model (DM) and zero initialization. Let
be the first
time SpecGF fits the minority class Further, assume that p. 2 1, 3ps, and pm S 5588+66k"
Then, for every t E (0, )"), SpecGF outperforms GF with minority-class and balanced loss gaps between the algorithms growing linearly with time as follows:
C (t) c (t) Z pt/4,
C Saa (t) 2 pl//2
For clarity, we state the above theorem under some simplifications on the condition on minority class prior Pm- See App. B.3 for the complete version and the proof.
Next, we investigate whether these gains of SpecGD arise because uses normalized updates while GD does not. To account for this difference, we compare SpecGD with NGD, the normalized version of GD. We first simulate the dynamics of NGD in the setting of Fig. 2, and find that although NGD converges much faster than GD, its iterates W. follow the same trajectory as GD (Fig. (top)). Note that at the end of training. all three optimizers eventually converge to the same solution. However, importantly, we find that early on in training, SpecGD still outperforms NGD in terms of both minority-class and class-balanced loss (Fig. 2 (bottom)). We formalize and prove this observation in the theorem below. See App. B.4 for the complete version and the proof.
Theorem 2 (Simplified). Assume data model (DM), zero initialization, and gradient flow algorithms NGF and SpecGF. Further, assume that Pm 25NR-4k- and k 2 Par-P... Then, for every t E (0,t*), SpecGF outperforms NGF with minority-class and balanced loss gaps between the algorithms growing linearly with time as follows:
CNGF (t) SSpce (t) 2 2t/2,
NGF (t)
C.pee (t) 2 ut/2
Proof Sketch. We now provide brief proof sketch for Theorem highlighting the key ideas. From Proposition I, since every SpecGF iterate (l) remains in the same basis U, i.e.. W(1)= = V, the singular values in (tt) are the key quantitics of interest for bounding the loss.
Concretely, for any class e E [k). we can relate its loss to the tt singular value W(]fe, c| =: (1):
(t) = 1/2 (1-- (ao() + + a(())
(3)
For any given I, the class-balanced loss gap is A ALbal = Ee Ee (CD(L) - CSpe CCSp0e()) (1). In order to bound this gap, we start by computing its derivative
A (1) - Ec na (A) (1) + 02 2. ((())(F())
a
+ 1. pa Spae
Spec (t)) Spec Spec( (t) 02 E. ac
Spec (t) Spec (t).
Next, from the expression of a (t) in Eq. (2) and a (t) = ppme sia we use Gf ( GF (t) 2
Spec and 0 (1) Supm- Combining with am (1) from Proposition and Spec (I) = for S 1", we have:
(t - + (1 μt) o
Next, we use 1 1" Om, and using the assumption on Pm in the definition of am, get A (tat(() /1/.. This is then integrated over / to obtain the final result.

---

Published as a conference paper at ICLR 2026
(biinenn:
1]
W[2:2]
(33.3]
SpecOD
SpecC2 (lowr)
Spec120 (Niiwar)
Linear
Bilinear
1.5
1,0
1.0
1.0
B 0.5-
S: 0.5
0.5
0.0
10 10
10
10
10'
10
10"
Time int)
Time (ni))
Time (nt)
Figure 5: Comparison of the dynamics of the squared norm of singular value matrix of the iterates, and the individual singular values, for GD and SpecGD across linear (I=1) 1) and bilinear (L 2) models for -mbalancedd data (priors p1 > P2 > Ps). In both cases, GD learns more dominant spectral components first, whereas SpecGD learns all components at the same rate. The gap between saturation of different components for SpecGD decreases as model depth increases from to 2, as proved in Proposition 2.
The proof of Theorem 2 based on a similar idea. The key distinction is that we don't 't have the precise dynamics of NGD(F) iterates. Specifically, the NGF update for e [k] is written as
NGF de (t)
Te(() R(J
where ro(t) == S.e Se a
NGF (t) and R(t)
Σ c==1 r2(().
Then, we have that 0 < (t) < sxx and using this with R(t) > max. Te(t) gives d (t) S 1, and hence a NGF S Using these we lower bound A (t) and Aa and integrate over
3.3 EFFECT OF DEPTH
To investigate how increasing model depth L affects SpecGD dynamics, consider deep linear model
W : L-1 WL-1-i with total loss L(W", WL-1) _ - WL-1-ix
Bilinear Model. We begin with the case L 2, which is particularly instructive due to its connection to the Unconstrained Features Model (UFM). The UFM is widely used to study the geometry of learned representations in deep and nonlinear networks under finite-sample regimes (Yang et al., 2018; Mixon et al., 2020). serves as a proxy for architectures with sufficient capacity to freely optimize both the feature representation (W*z) and the classification head (W'). The bilinear model recovers this framework when the input features æ are orthogonal, making it a natural setting for analyzing the interplay between feature learning and classifier dynamics.
The following result characterizes the SpecGD dynamics of W. = W.W.P
Proposition (Simplified). Suppose Condition holds. Let the weights be initialized as W xd1 Q"Iaa 0a-d; VV and Wi E 0a,= where Q E [d1xd1 is an orthonormal matrix, k < di < d, and
a constant. Then, for SpecGD with, at each iteration, W QW V W. UW.Q where W W are zero except their main diagonal along which, in the limit + oc, they evolve as follows for E [b]]
Wh l. mtt
For clarity, we consider - oo in the above proposition to state the SpeeGD dynamics. See App. B.5 for the complete version (for any > 0) and the proof.
We find that, in contrast to the linear model where SpecGD learns the c component in time 2 to g for the bilinear model, it saturates in time to R + (for large ö). This quite
different from the effect of increasing model depth on the dynamics of GD. As discussed in Saxe et al. (2013); Gidel et al. (2019), for GD, the time to reach e-close to the optimal (for the dh In(e)
Iu(1/c) (for small 17) for linear model, compared to to s 7p2
component) is to In )
In(1/e-1)+In 11 2n
1=(1/c)+5 (for small c and large 5) for bilinear model. To make the 2ne
comparison more concrete, recall from Lemma 3.1 that for class-imbalanced setting with data
7

---

Published as a conference paper at ICLR 2026
model (DM), when μ 1, se Pc + SNR and s Pe- The GD dynamics for linear vs. bilinear model depend on the SNR: when SNR is high, we observe the stage-wise learning dynamics in both cases, i.e., spectral components are learned in descending order of class priors. However, when SNR is low, these stage-wise dynamics may not be apparent for the linear model as l. X P- A SNR.
1/7.
In contrast, for SpecGD, t are qualitatively similar for both models; however, the gap between saturation times on different components is smaller for the deeper bilinear model.
where L {1,2}. This means the dynamics
Pc++5NR
We illustrate these differences between GD and SpecGD and the effect of increasing model depth in Fig. 5, where similar to Fig. 2. we consider a dlass-imbalances setting with 7) 0.05, d = k = 1, - 0.125, and P1.P2, set as 0.55, 0.3, 0.15, respectively. We set e = 0.05, 8 = 10.
Next, we investigate whether for SpecGD, the gap between saturation on different components becomes smaller as model depth is increased further.
Deep Linear Model. We next consider deeper models with L 2 2. The dynamics of W, I w for SpecGD can be characterized in similar way Proposition 2. Specifically,
+i== defined analogously remain diagonal with entries E [R]], in the limit 5 - oo, evolving as:
W/l, cl = nf t S
+
I
t
See App. B.6 for a formal statement and the proof. We can quantify the gap between saturation of the minority vs. majority
KG0 X.
MGD 41
1/L
component as AT: == Amax fmin
1. For the dass-imbalanced setting in Eq. (DM), using Lemma 3.1.AT AT
Emin
1/L
SNR++1/Pm
0.25
SNR+1/PM
10
Spech
We find that as L increases, the gap between the saturation of the different components becomes smaller. We empirically validate this result for SpecGD in Fig. 6 by comparing the test accuracy dynamics of 2- 2-layer and 4-layer MLPs on the Colored- MNIST dataset. A similar trend observed for NGD, where depth speeds up learning of the minority component. This is an interesting empirical parallel, given that precise theoretical dynamics for GD have primarily been characterized for depth (Gidel
Figure 6: Test accuracy dynamics of NGD and SpecGD for a 2-layer vs. 4-layer MLP on the Colored- MNIST dataset.
et al., 2019; Saxe et al., 2013).
EXPERIMENTAL RESULTS
In this section, we present experimental results on datasets with group and class imbalance, comparing the performance of Muon to SGD. Our analysis is guided by the spectral perspective developed in Sec. 3.2. For class imbalance, recall that spectral components correspond to class priors. We extend this analogy to group imbalance, where inputs contain both core and spurious features, and groups defined by different core-spurious feature combinations. As seen in our motivating Golored-MNIS, example (Fig. 1). spurious features (e.g.. color) act as the dominant spectral component, while core features (e.g., digit shape) are less dominant (see also (Ng et al., 2024), which demonstrates this for some other datasets with spurious correlations). Our experiments show that Muon promotes more balanced learning of different spectral components, thereby improving generalization relative to SGD, Finally, while our theory focused on the dynamics of SpecGD versus (N)GD, our empirical study also includes Adam, offering broader comparison with widely used adaptive optimizer.
Group Imbalance Datasets. Here, we consider three widely used datasets with group imbalance.
Dominoes Dataset. We consider the MNIST-CIFAR dataset from the Dominoes benchmark (Shah et al., 2020; Pagliardini et al., 2022), where each image is generated by stacking an MNIST digit from class (0, 1) on the top with CIFAR-10 10 image from class (automobile, truck) on the bottom. The MNIST part is spuriously correlated with the label 95% of the time, while the CIFAR part is 100% predictive the label. Example images are shown in Fig. (left).
8

---

Published as a conference paper at ICLR 2026
We train ResNet-34 model using SGD, Adam, Shampoo, and Muon and compare their generalization us- ing two metrics on the test set: the worst-group accuracy and the decoded worst-group accuracy (see App. D.3 for details). The latter obtained by freezing the learned representation of the trained model, re-training only the final linear layer with logistic re- valida-
9GU Adam
Werd Ward grop Decodedfee
Sharppoe
Maann
0.4
02
0.
5.
Eppers 025)
Tran Epodte (e30)
Figure 7: Example images from the MNIST-CIFAR dataset (left), and comparison of worst-grout group test accuracy (middle) and worst-group decoded test accuracy (right) dynamics of SGD, Adam, Shampoo and Muon (see text for discussion). Muon learns core features faster and has superior worst- group performance while SGD relies on spurious features.
gression on a group-balance tion set, and then evaluating its perfor- mance on the test set. This mcasures the extent to which the model has learned the core feature information in its representations. From Fig. 7, we see that Muon consistently outperforms SGD across both metrics throughout training. Compared to Adam, Muon shows greater gains early in training, eventually attaining similar perfor- mance. Muon, as well as Shampoo and Adam, achieve much higher decoded accuracy than SGD, which suggests that they successfully learn the core features (which correspond to less dominant spectral components), whereas SGD struggles to do so.
Subgroup Robustness Benchmarks. Next, we consider two benchmark subgroup robustness datasets: MultiNLI (Williams et al.. 2018) and CelebA (Liu et al., 2015). MultiNLI consists of sentence pairs, where the task is classify the relationship of the second sentence to the first as entailment, neutral, or contradiction. The spurious feature is the presence of negation words, which often indicate contradiction. CelebA contains images of celebrity faces, with the task of predicting whether the hair color is blondelnot blonde, and the celebrity's s gender (malelfemale) serves as the spurious attribute.
Shumpoo
Adam
Maoo
Wore-graa twee Aoe
1,0
Aug Koe
Ang Teat Aoc
0.9
Woretgron
0.5
3.7
FT (Epoohh))
FT Budgel (Epochss)
FT Budget (Epcohs)
FT Budgel (Epoors)
Figure& 8: Comparison of the worst-group and average test accuracy over the number of fine-tuning (FT) epoch budget (see text for details) for different optimizers on the (left) MultiNLI and (right) CelebA datasets. The average test accuracy is comparable across all optimizers. For the worst-group accuracy, spectrum aware optimizers (Muon and Shampoo) consistently outperform SGD. However, compared to Adam, for MultiNLI, Muon is slightly worse, while for CelebA, its performance becomes comparable for more FT epochs.
We fine-tune a pretrained BERT pert-baaseuncased model (Devlin et al., 2019) on the ResNet-50 model on CelebA, using different optimiz- balanced) test accuracy and the worst-group accuracy over the
MultiNLI dataset, and an ImageNet- ers and evaluate the average number of finetuning epochs. To ensure our comparisons are robust, we conduct extensive hyper- parameter sweeps for all optimizers (see App. D.4 for details). We select the best hyperparameter configuration for any given optimizer and number of fine-tuning epochs, T, based the highest average worst- worst-group accuracy achieved on the validation set, averaged across multiple initialization seeds, any point up to T epochs. We then report the performance for the optimal configuration for different FT epochs. From Fig. 8, we see that Muon outperforms SGD on both datasets. On MultiNLI, Adam attains better performance than Muon/Shampoo whereas CelebA, its perfor- mance is comparable to Muon. These experiments show that Muon promotes more balanced learning of spectral components than SGD, thereby improving generalization.
Datasets. Here, we extend our investigation of class imbalance to language modeling. While our previous experiments and theory focused on one-hot classification, we now shift our comparison of optimizers to a next-token prediction (NTP) task. This task presents a natural imbalance over the distribution of tokens, as word frequencies follow long-tail distribution described by Zipf's law (Piantadosi, 2014).

---

Published as a conference paper at ICLR 2026
TinyStories Dataset. To examine how dif- ferent optimizers compare on rare token learning, we train small Transformer mod- els on the TinyStories (Eldan & Li, 2023) corpus and compare the performance of Muon, SGD, and Adam. We group tokens into "frequent" and "rare" categories based on their corpus statistics and evaluate the token-level accuracy for each group sepa- rately on a held-out test set. Complete train- ing details are available App. D.5. The performance on rare and frequent tokens is reported in Fig. 9. Consistent with our findings in the one-hot daas--imbllaneed classification, we observe that Muon gener- alizes faster on both the majority (frequent) and minority (rare) components compared to SGD, and in this case, even compared to Adam.
Majonty
Minortty
Accuracy
Top-5 Accuracy
D.7
0.6
00
0.5
0.5
0.4
0
0.3
022
0.0
Esochs
Epochs
Figure 9: Performance comparison of SGD, Adam, and Muon on frequent (majority) vs. rare (minority) tokens, evaluated by Top-1 and Top-5 accuracy on a validation set of TinyStories. The results demonstrate that Muon achieves faster generalization on both frequent and rare tokens compared to SGD and Adam.
Attribute- Orgbbnii- Classification. To extend our findings from statistical imbalance to a setting with explicit conceptual structure, we conduct a toy experiment on hierarchical concept learning (Zhao & Thrampoulidis, 2025). In this task, model must classify organisms (c.g., "dog") based on their attributes, which are shared at different levels of a taxonomic tree (e.g., "needs oxygen" for all animals). We observe distinct learning dynamics: GD learns the coarse, high-level categories (plant vs. animal) first, while Muon exhibits more balanced learning progression across both coarse and fine-ge line-graieed concepts (see App. D.6 for details)
DISCUSSION AND OUTLOOK
Our work investigates how the spectral design of optimizers like Muon and Shampoo yields general- ization benefits, using imbalanced data as a testbed. The core mechanism we identify is that their canonical form, SpecGD, implicitly learns all principal components of the data an equal rate, in stark contrast to GD.
crucial test of this insight is how its practical variants, Muon and Shampoo, perform against strong, widely-used baselines like Adam. In our experiments in Sec. 4, we find that Adam's generalization performance is only slightly worse or sometimes even better-than that of Muon and Shampoo. This is consistent with recent large-scale benchmarks for LLM pretraining, which show that well-tuned Adam variants remain highly competitive (Wen et al., 2025; Semenov et al., 2025). Previous work has used different perspectives to explain Adam's 's superiority over (S)GD: i) Kunstner et al. (2024) show how heavy-tailed maas-immalaaceed settings lead to gains for Adam by focusing on their optimization dynamics; ii) Vasudeva et al. (2025b) show that Adam promotes richer feature learning than SGD in neural networks, which can lead to better generalization in group- moou-imbalanced settings. Together, these works underscore Adam's 's advantages in imbalanced regimes, albeit through different mechanisms. Our motivation for using this testbed stems from recognizing that spectral components of imbalanced data directly correspond majority and minority classes or groups, and our analysis reveals that SpecGD achieves its gains through a fundamentally different and interpretable bachaanimmmaalaneee learning of principal components.
A natural question is how the implicit balancing effect of spectral methods compares to explicitly designed techniques like loss re-weighting (Byrd & Lipton, 2019; Cao et al., 2019; Sagawa et al., 2019; Menon et al., 2020; Liu et al., 2021). In App. D.8, we show that while weighted cross-entropp (wCE) effectively encourages balanced learning of spectral components on NMD, Muon provides similar, though weaker, effect "for free" without requiring explicit group priors. This suggests spectral methods could useful when such information is unavailable for re-weighting.
Our work provides framework to understand the generalization of spectrum-awar optimizers within a controlled setting using population statistics and squared loss. A natural next step is to build upon this foundation by extending the analysis to more complex regimes using finite samples and cross-entropy (CE) loss. To bridge this gap, we conduct an initial empirical analysis of rate of learning spectral components for nonlinear models trained with CE loss, by tracking singular values of the logit matrices during training (see App. 1.10 for details). We outline future directions and their associated challenges in App. F.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
The authors acknowledge use the Discovery cluster by USC CARC and the Sockeye cluster by UBC Advanced Research Computing. This work was partially funded by the NSERC Discovery Grant No. 2021-03677 and the Alliance Grant ALLRP 58 1099-2. PD and YZ were also supported by the UBC 4YF Doctoral Fellowship. This work was also supported in part by NSF CAREER Award CCF-2239265, an Amazon Research Award, a Google Research Scholar Award and an Okawa Foundation Research Grant.
REFERENCES
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint 10XXiv:907.028900 2020.
Dennis Bernstein. Matrix mathematics: theory, facts, and formulas. Princeton university press, 2009.
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. arXiv preprint M:Xiv:2409.20325, 2024.
Jeremy Bernstein, Yu-Xiang Wang. Kamyar Azizzadenesheli, and Anima Anandkumar. signsgd: Compressed optimisation for non-convex problems. arXiv preprint arXiv: WVXiv:/802.04444 2018.
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions. arXiv preprint arXiv:2211.12316, 2023.
Valentyn Boreiko, Zhiqi Bu, and Sheng Zha. Towards understanding of orthogonalization muon. In High-dimensional Learning Dynamics 2025, 2025.
Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2009.
Jonathon Byrd and Zachary Lipton. What is the effect of importance weighting in deep learning? In International Conference Machine Learning, pp. 872-881. PMLR, 2019.
Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma. Learning imbalanced datasets with label-distrbutiinnaawar margin loss. In Advances in Neural Information Processing Systems, pp. 1567-1578, 2019.
Yuan Cao, Zhiying Fang, Yue Wu, Ding- Xuan Zhou, and Quanquan Gu. Towards understanding the spectral bias of deep learning. arXiv preprint 19Xiiv1192.0000 198, 2020.
David Edwin Carlson, Edo Collins, Ya-Ping Hsieh, Lawrence Carin, and Volkan Cevher. Precondi- tioned spectral descent for deep learning. In Neural Information Processing Systems, 2015.
Da Chang, Yongxiang Liu, and Ganzhao Yuan. On the convergence of muon and beyond. arXiv preprint arXiv:2509, 158111100911166.. 2025.
Lizhang Chen, Jonathan Li, and Qiang Liu. Muon optimizes under spectral norm constraints. arXiv preprint arXiv:2506.1500 2025.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4171 4188,, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 1e.11855/1/119/
Ronen Fldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint arXiv:2305.0775: 2023.
Chen Fan, Mark Schmidt, and Christos Thrampoulidis. Implicit bias of signgd and adam on multiclass separable data, XrXiv preprint arXiv:2500..44664 2025.

---

Published as a conference paper at ICLR 2026
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks, arXiv preprint arXiv:1810.0000 2018.
Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle: Fourier analysis sheds light on deep neural networks. Communications in Computational Physics, 18(5):1746-176,, June 2020. ISSN 1991-7120.
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https:
//kel lerjordan g thub. / ot/ poots/muon/.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14Xv.69412.6988. 2014.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66-71, 1, 2018.
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Hcavy-tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv:2402. 19449, 2024.
Tim Large, Yang Liu, Minyoung Huh, Hyojin Bahng, Phillip Isola, and Jeremy Bernstein. Scalable optimization in the modular norm. In The eigrth- Annual Conference on Neural Information Processing Systems, 2024.
Jiaxiang Li and Mingyi Hong. A note the convergence of muon and further. arXiv preprint 2502.22500.029000 2025.
Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Perey Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp- 6781-6792. PMIR, 18-24 Jul 2021.
Jingyuan Liu, Jianlin Su, Xingeheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training, arXiv preprint ar XrXiv22522.1600 2025.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Martin Marek, Sanae Lotfi, Aditya Somasundaram, Andrew Gordon Wilson, and Micah Goldblum. Small batch size training for language models: When vanilla sgd works, and why gradient accumulation is wasteful. arXiv preprint arXiv:2507.07101, 2025.
Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, and Sanjiv Kumar. Long-tail learning via logit adjustment. arXiv preprint arXiv:2007.077 07314, 2020.
Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. arXiv preprint ar 2rXiv:20/1.11661 2020.
12

---

Published as a conference paper at ICLR 2026
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on shampoo's 's preconditioner. arXiv preprint arXiv22406.17745, 2024.
Preetum Nakkiran, Gal Kaplun, Dimitris Kalimeris, Tristan Yang, Benjamin L. Edelman, Fred Zhang. and Boaz Barak. Sgd on neural networks learns functions of increasing complexity. arXiv preprint arXiv:/905. 1604, 2019.
Yeat Jeng Ng, Ainhize Barrainkua, and Novi Quadrianto. Understanding and addressing spurious correlation via neural tangent kernels: A spectral bias perspective, 2024.
Antonio Orvieto and Robert Gower. In search of adam's secret sauce. arXiv preprint arXiv:2505.21829, 2025.
Matteo Pagliardini, Martin Jaggi, François Fleuret, and Sai Praneeth Karimireddy. Agree to disagree: Diversity through disagreement for better transferability. arXiv preprint arXiv:2202.06 200.:2202.4444, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32, 2019.
Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with norm-cons norm-constrained Imos. arXiv preprint arXiv:2502.00529, 2025.
Mohammad Pezeshki, Eékou-Qumar Kaba, Yoshua Bengio, Aaron Courville, Doina Precup, and Guillaume Lajoie. Gradient starvation: learning proclivity in neural networks. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Steven T. Piantadosi. Zipl's word frequency law in natural language: critical review and future directions. Psychonomic Bulletin & Review, 21:1112- 1130, 2014.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred A. Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. arXiv preprint 1806:0/806007344 2019.
Harold Robbins and Sutton Monro. A stochastic approximation method. Annals of Mathematical Statistics, 22(3):400-400 1951.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint
2019.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 1312.611 2013.
Andrei Semenov, Matteo Pagliardini, and Martin Jaggi. Benchmarking optimizers for large language model pretraining. arXiv preprint arXiv:2509.014400 2025.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 9573-9585. Curran Associates, Inc., 2020.
Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data-parallel pytorch im- plementation of the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint arXiv:2309.006// 2023.
Daniel Soudry. Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1-57, 2018.
13

---

Published as a conference paper at ICLR 2026
Teodora Srećković, Jonas Geiping, and Antonio Orvieto. Is your batch size the problem? revisiting the adam-sgd gap in language modeling. arXiv preprint arXiv:2506. 2025.
Nikolaos Tsilivis, Gal Vardi, and Julia Kempe. Flavors of margin: Implicit bias of steepest descent in homogeneous neural networks. arXiv preprint arXiv:2410.22069 2024.
Bhavya Vasudeva, Deqing Fu, Tianyi Zhou, Elliott Kau, Youqi Huang, and Vatsal Sharan. Transformers learn low sensitivity functions: Investigations and implications. arXiv preprint arXiv:2403.06925, 2025a.
Bhavya Vasudeva, Jung Whan Lee, Vatsal Sharan, and Mahdi Soltamolkotabi. The rich and the simple: On the implicit bias of adam and sgd. arXiv preprint arXiv:2505. 240222 2025b.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham M. Kakade. SOAP: Improving and stabilizing shampoo using adam for language modeling. In The Thirteenth International Conference on Learning Representations, 2025.
Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them. arXiv preprini 2eX.i::559.0200 2025.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Tech- nologies, Volume 1 (Long Papers), pp- 1112-1122. Association for Computational Linguisties, 2018.
Shuo Xie and Zhiyuan Li. Implicit bias of adamw: L-infinity norm constrained optimization. arXiv preprint arXiv:2404.044: 2024.
Zhilin Yang, Zihang Dai, Ruslan Salakhutdinov, and William W. Cohen. Breaking the softmax bottleneck: A high-rank rnn language model. arXiv preprint
2018.
Chenyang Zhang, Difan Zou, and Yuan Cao. The implicit bias of adam on separable data. arXiv preprint arXiv:2406.10m: 2024.
Yize Zhao and Christos Thrampoulidis. On the geometry of semantics in next-token prediction. arXiv preprint arXiv:2505.08348 2025.
14

---

Published as a conference paper at ICLR 2026
APPENDIX
Optimizers
15
Proofs
16
B.1 Proof of Lemma 3.1
16
B.2 Proof of Proposition
17
B.3 Proof of Theorem
18
B.4 Proof of Theorem 2
20
B.5 Proof of Proposition 2.
22
B.6 SpecGD Dynamics for Deep Linear Model
23
Additional Results in the Linear Setting
24
D Additional Results and Details of Experimental Seftings
25
D.I Colored-MNIST ......T.........................
26
D.2
CIFAR-10/100 with Class Imbalance
................................
26
D.2.1 Effect of Class Imbalance Ratio
29
D.2.2
Fine-tuning from ImageNet Pretraining
................................
29
D.3 MNIST-CIFAR Dataset
30
D.4
Subgroup Robustness Benchmarks
30
D.5
TinyStories
D.6
Attribute-Organism Classification
32
D.7 Mixture of TinyStories and Add-k
32
D.8 Effect of Loss Re-weighting
34
D.9 MNIST with Heavy-tailed Class Imbalance
34
D.10 Spectral Analysis for Non-linear Models
34
E Related Work
35
Limitations and Open Questions
36
A OPTIMIZERS
In this section, we list the update rules for all the optimizers considered in the paper for completeness. We start with the update rules of NGD, SignGD, and SpecGD, then list their momentum versions (NMD, Signum, and Muon), and then write the updates for Shampoo and Adam.
Using the notation introduced in Sec. 3, updates are of the form
W = W. n,, where A == argmaa A <1 V. A).
(4)
NGD update. For NGD, the arg max in Eq. (4) uses [|,,|,, and we get A
15

---

Published as a conference paper at ICLR 2026
SignGD update. For signGD, the arg max in Eq. (4) uses max, and get A, sign Ve). where sign(r): = # and sign(0) = 0, applied element- wise on the matrix V..
SpecGD update. For SpecGD, the arg max in Eq. (4) uses the spectral norm. Let the truncated SVD of V. be U, where U, and V. are orthonormal matrices and Σ, is a diagonal matrix with positive diagonal entries. Using this, we have A, = U.V.
For NMD, Signum, and Muon (with exact matrix operations), the updates mirror NGD, SignGD, and SpecGD, respectively, except that they use the momentum term M; = 3M + (1 () V in place of the raw gradient.
Shampoo update. For Shampoo, first define the preconditioning matrices
V where the parameter 32 denotes the preconditioning accumulation parameter, and the momentum ma- trix M, M 81) V Using these preconditioners, the update is A, L, 1/4 M,
V. L, L + (1 V V
and R, 82 (1
1/4
It is easy to see that Shampoo with exact matrix operations reduces to SpecGD when we set 31 B2 0 as follows. Let the SVD of V, be U, Then, since B2 0, the precon- ditioners are L, = U,T,U,T and R. = V Using these, since 3 = 0, the update is
A, U,E, 1/2 (U, (V, V 1/2 1/2v, UVT, which is same the SpecGD update.
Adam update. For Adam, let M, Mi
1
(B.M: + 31
denote the bias-
1- 1-
corrected first-moment estimate, and Z.
Z1+1
1- B2 + (1 32
denote the
1-8
bias-corrected second (raw) moment estimate, where denotes the Hadamard product, and 31, 82 denote the momentum parameters.
Then, the update is A; = (Ž, + +11 ) @ Mr where (.) denotes the Hadamard power, e >> the numerical precision parameter, and 1 denotes the all-ones vector. It is easy see that Adam reduces to SignGD when we set 31 = 32 = c 0.
B
PROOFS
B. PROOF OF LEMMA 3.1
The covariance matrix is
Σn = [[: | = Ey,e (my + E + +
Eg Mylly + Eg My [Eg 1 + Ee|e|En|[u + E[e=]] E, + oila.
=/Er
since e and y are independent and E[e] = 0. Further,
En = Σ P.po = MPM := p Po [s] E Rdxk p p² diag(pi....,
Then, we can write Σn VS V where
Sez diag(u² + 02 o ... 12pp 0 2
02)
(5)
times
V M v1 p ph. Ugt1...., faa Radd Here, we used the assumption orthonormality of the means.
(urld 1 (Mel
The cross-covariance
E = Elyx Σ Pe Me E Elx 1 y==c Σ
upe e p..
16

---

Published as a conference paper at ICLR 2026
We write Zya
U Sya VT, where
i = = diagg(PP......... ........ Od-x].
(6)
Note that Syz is a k x d matrix, where class priors Pt, i E [k] denote entries along the main diagonal with zeros in the rest of the entries. For brevity, we will denote this as Syz = p diag(p1, ... Pke kxd from here on.
B.2 PROOF OF PROPOSITION
For completeness, we start by deriving the discrete-time evolution for GD.
We can write the gradient
VL(W,) = [V- -US V + W.V VT.
For GD, have
W4+1 = We - (V)(W)
-We We+nU WUU V W Sax V W:(1 1/V Sax VT) + +nUSs V
V
1+1
U yz (I 11 2 7 V
T=0
Since Wo 0, this gives
W W++1= :- U W.
=nSyp yz Σ (I 11
7=0
Assuming 7
Jman' we get
8 1-(1 111
W W4t1/ nsy"
Σ
For sufficiently small 7, this gives the approximation
W Witlii, (1- -00f((+11))
For SpecGD, note that the gradient can be written terms of W, as
VL(W,) U S yo V W.V Swm V
U(S W. Syz) VT.
Starting W. 0 W. gives Vo -
US ym VT.
Wiii = n f1
i i= j E [k]] i i
Thus, Wi = VUV
Here, Ve E Raxk denotes the first
columns V. Proceeding this way, we arrive at the following update rule for all t,
W411 W. +n
Σ
u.u
&W "W
This gives
W 41 W
W.+n
Σ e.e.T
i: Wi
Since Wo 0, we conclude with the desired:
W W (4 1)
[ (1+1)
(1+1)>
For the continuous time versions, aka gradient flow (GF) and spectral gradient flow (specGF)) the steps very similar:
17

---

Published as a conference paper at ICLR 2026
Gradient Flow (GF). We first state the GF equation
d
(t) = - VZ(((())
dt
Using this, we can write the evolution of ( U W as
d
(1) W (1)S + S yas
W(0) 0,
dt
(t) S Sax (I e 1Swa S yz
or entrywise
xx (1--
Spectral Gradient Flow (SpecGF),
We first state SpecGF equation
d (tt) -t()V(()
dt
where the gradient has SVD VC(W(()) (I)) Following the same steps as for SpecGD, we can show that (t) U. V, and W W(t) evolves as
d
(t)li, = {
[ W (t)lū, 11st < 8/"
dt
W (1) 1000 = 6%
with solution
W(t)[i, = t1 t <
B.3 PROOF OF THEOREM
We first state the full version of Theorem below followed by its proof.
Theorem 3. Assume data model (DM) and zero initialization. Let t"
2xx be the first time SpecGF
fits the minority class m. Further assume 2 and Pm. S 3S3R+IK' Then for every t E (0,t*) SpecGF outperforms GF with a growing minority-class loss gap CGF (t) - Spee (t) 2 pt/4. Moreover,
k-2p if Pm I 2psn8 2232 then the gap in the balanced loss also grows as C (t) Spec (t) 2 pt/2.
Proof. The population loss for class using iterate W, is written as
L.(1)= == _ y W. 2: 1| = [ 2y
W. W W. W.
1 2 W. Me W Me + W
1/ e. W. Me 1 fo W.
C. UW,V T Mell2 + 1 W..VVT
1
e. e. [c. c|ull + oo W W.ll
Define at(t) = Welc, c). Then the per-class loss in terms of the singular values of W, is written as
K
(t( = ((1-pa,(())) - to2
(t).
Using Prop. and Lem. 3.1, shows that the singular values of Gr (f) and W W- SpeeGF
evolve
of Tur (1)c, cl ac exp (-Re!t))
Spee
W SpooGF (t)c,, = +1 (t arl a ac),
(9)
18

---

Published as a conference paper at ICLR 2026
where 0i denotes the ratio of the singular values for class c, and using Eqs. (5) and (6) from the proof of Lem. 3.1, Og Pelt +o2"
Polt
The time derivatives of uF (1) and Spee (1) for L S 6" are given by
CF (4) =
dag o
Poflee to]+ ( = da Spe dt
1.
dB
Define the gaps
Aoa(1) == 0% Spec (1) oe C (1) 0,
Aác(1) = ( (1) -cc (1) = - PPofe (o2tpor?) (10)
Note that (t) is increasing in t.
Minority-class loss gap. For fixed t, consider the minority loss gap
Am (t) L (t) L. (t).
(11)
Using the per-elass loss from Eq. (7) and differentiating.
sêm (f) = - (1 - -1- a (t)) GF (t) - (1 μ am
Spec (t)) Spec (t)
Temm-1(0)
o² Σ aj (t)a (t)
La Spec (t)(a))) Spec)))
(12)
Term-2(VD)
Term-1 (Φ).
Add-subtract (1 μ CP (t)) Spec (f) Term- 1, and we have a
(1 / a (t)) OF (t) (1 Spec a Spee (t)
= (1 4 a (t)) Ddm (t) + μ Amm (t) Spee (t).
(13)
We know that
Da(t)= = (' dd((T) S tåà (increasing integrand).
Substituting in Eq. (13), get
+
< (1 ( t á (t)) Adm (t)
S (1 am t) Ddm (t)
S (1 2/1 akm) Adem (1).
(14)
The second and third inequalities assume that Adm (t) 2 O, which holds true under Pm S 1 1/2p. which is true as we assume μ 2 1. The third inequality also uses the fact that t < t" = am
Term-2 (V) Before any aj Spec saturates,
Σ
Spec a a;
kt.
which gives for all t" that
1 2 -kom
(15)
(1) 2 u(1 24 (m Adim - 0 kum
2 u(( 24 am)) (1 - pm)) - o² kaa
I 4.
(16)
Here, first Pmft 1/2 as μ 1 and Pm I 1/2 as it is miniorty class. Next, the Pmnl definition of O'm Pm toz (0, 1) gives the final bound.
and the assumption Pm S 3588 Since AL, = 0, integrating
19

---

Published as a conference paper at ICLR 2026
loss gap. For fixed t, the -laasbbllanced loss gap is
(t) Σ C (t)
C (t)
(17)
Using the per-class-loss from Eq. (7) and differentiating,
GF (t) + ox Σ ae (t)a. ACbal(t) - Σ (1 ua (t)
GF
+ 1: Σ (1 pac Spo (1))
1
Spec (1)
-00 Σ Spec (1)c (t)
2-
1 k
Σ udc (t) + E -oo
Spec (d)
Spec (t) - Σ
a
Spec (t) (t)
2 t/p² + (1 - pt) Lt
SKR+ 2/²² + μ PmH Pm S3R+1
The first inequality follows by using ac (t) (t) 2 For the second inequality, we use ā (t) 2 μ4P.. and substitute the expressions for de definition of SNR.
Spee (1) and d
Spec (t). In the last step, we use / S Om and the
k-2n As we assume Pm S + we have that
Dlbal!) 2 1 - tu²²
Pmitt SNR++
SNR+1
2 2.
Integrating this over (0, t), we get AL (t) 2 pt/2 for St.
B.4 PROOF OF THEOREM 2
We first state the full version Theorem 2 followed by its proof.
Theorem 4. Assume data model (DM), zero initialization, and gradient flow algorithms NGF and SpecGF. Further, assume that
Pm 2 2mm (pm SN8+1)²
[I-pm (1-Pm(0NN422R) (BNR+2k]) Pm SNR-4K Then, for every
t E (0,t"), SpecGF outperforms NGF with growing minority-clas. loss gap CNGF (t)
SSpee (f) 2
/2 Moreover, if PM Pm Z k(1 Pe, (SNR+ 2kJ). then the gap in the balanced loss also grows as
2(pm 1)²
CNGP (t) S (t) 2 pt/2.
Proof. The continuous time update of ae
Spec (t) for t t = am is given by
Spec 1,
Spec t.
In the continuous time limit, the NGF update can be written as
rt()) R(t)'
(N)
re(t)= = si s o
NGF (t),
R(t)
(t).
We have 0 5 r(()) < 8y up:, R(I) 2 ra)x which gives (N < and hence
a NGF (t) 5t t- an Using this, get R(t) 2 S sya, Using Lemma 3. have
R(t)= 2 PMU Pmp Pmfl 2 PM pm
pMp²+a² +a²
20

---

Published as a conference paper at ICLR 2026
Using Eq. (7), write the minority class loss gap as
Am (t) = C.D (t) Spec (t)
= (1 pa
NGF (t)) o² Σ (ac
NGF (t))² -(1 So (t)) o² 2 Ca (t))
(l) --(1- pam
NGF NGF (1) + 0x
(U))
a
NGP (1)a (FF((()
pamm Spec
(I))uc Spec (1) 0x Σ
Speer (1) (1)
2 Ritt le + (1 - pt) 22
(t)(t)
2 p am (1² ko
Pm (pm +
(2 Pmn) 16.
Since we assume that
Ppm (pm SNR+1) Pm < SN3+22
PM Pm z
1- (SNB+26)
we have
a2 +Pmm) (0² -D ( (1++2R02)))
(tm 2 1 am (1 ko+))
20 Pm 1Dml++022
1
Pmmlll++ae +ko?)
(o Pmu 2pm (1 (42 +ko
Pmnli +0
11 2(Pmn t.22)
=1/2.
Using this, we get CNCD(() (t) CSpec for St5s.
Similarly, for class-balanced loss, we have that
Abbal(t) == NGD (t)
Spec bal (t)
11 pace
NGF (1)) + 0² E C(a (o))2
E
ua
Soo (1)) +o² >5
(a Spec (1))
21
26
ALbai(1) Σ (1 m
NGF (1)) NGF
(1) + 0z Σ
NGF NGF (1)d (1)
ae
Σ p
Spec (I))
Spec (1) - ox Σ
Spec (1) (1)
2 * Σ
rell)
+ (1 - pt) o2kt
R(C
2 am (1² + ko
(pm#"+o7
ko k -p Pb
where we use E. roll) ps E. Pc μ. Since we assume that
2 SNR+117 1)2 P.n. (SNR+2k]]
PM Pm
Pm SER+2E
we have
ALbal 2 11 am am(u2 ko 23 (o2 2hpm.) Pmntt (pmpe +0?)
p
2n (pmpe to
(1-
Pm +
20 (1- (1-kpm) kpm) (Vmm 2 -
2(pm
1/2.
Using this, we get LNGD (1) Spec (1) 2 pl/2 for 0 5150.
21

---

Published as a conference paper at ICLR 2026
B.5 PROOF OF PROPOSITION 2
We first state the full version Proposition 2 followed by its proof.
Proposition 3. Suppose Condition holds. Let the weights be initialized as W [[a Oa-d,V and W = e U 0a, 0a,-a]Q where Q E fgdixdi is an orthonormal ma- trix, d1 < d, and 6 > a constant. Then, for SpecGD, at each iteration, W == QW; VT, w. UW, QT, where W t: W. are except their main diagonal along which the entries evolve as follows for i E [k]:
W w. (nt + 0 t t
( - (v
1 1
)
1
Proof. Here, assume that d > di k.' We can write the gradients
WwoL(W? W;)) = -(W,) TE [yv - W, w:WW
W x)) "] = (W?) E z + (W?) W W Σg:
WwI Wi) -El(V w. w (W,) -y (W9) T + W; (W
Let W: Q Q W and W. = U W, Q. Define W. = W. 1 2. Then, we can write
W W;!) -Q(W.) J Sy yu V + Q(W W W.SAEV V
W L(W), L. W.') = -USS yz W Q + UW,, S TQ.
This follows by writing Wi QW"VT w! UW,Q QT and substituting in the gradient expressions.
We now simplify the gradients the first iteration. Note that we have
W Id, 0a-d,l:
W = e e- [ Od,, Odl-
W 28 I Os-a].
Using these, we can get gradients first step
W 2 Wh W = - Q I Od, -tl (S - e-2 [ 0g Sam V (
W L (W o W = - (Syz e e-2 O 0u-x) [ O2 1T
Then, for sufficiently small initialization, the iterates after one step of SpecGD are written
W = e [T, 0 V + I1 (V =x 02 TT
W = e - Oa, - + (QU 0a, TT
Using this iterate after step 1, we can see that W 1, W W remain diagonal. Using the gradient expressions, that the update at step 1 remains in the same basis. This implies the update remains in the same basis. Proceeding similarly, we can see that with the given initialization, W and W remain diagonal, and the updates for W and W using SpecGD can be written as
W W = W +1
eie,
[
Ix
Oa-k 0a-
Oa.
W +1
W +n
Σ
e,e; [T Oct
[1 cm
Simplifying, we get
(nt
t t 1 n
(v==e
-
1
-
e
"The proof works for other cases well and we just use this for the purpose of illustrating the technique.
22

---

Published as a conference paper at ICLR 2026
B.6 SPECGD DYNAMICS FOR DEEP LINEAR MODEL
The following result characterizes the dynamics of SpecGD for deep L-layer linear model (L 2 2). Proposition 4. Assume k d, I. and suppose Ass. I holds. Let weights be initialized
W Q.v i W - UQ-
where {Qi) are orthonormal matrices and > is a constant. Then, for SpecGD, at each 1-
iteration, W = Q.W.v. w! Q1+1 W111 W -UW,-1Q Q where W for I [L] is zero except its main diagonal along which entries evolve as follows for E [k]:
W.l, (nt + )
( 1/- (+))--+))
1/L 1
1/L
Further, the gap between saturation of the minority majority component is AT ==
Emin
1/L
-.
Proof. We have that W T w' (Q+41) T and W wh-1 U WL WL-IQr.- Qr.-1.
where E (1....,74 2}. Also define
L-1
H L-1 (W'), w = [W). w
W = II w
1=
x=1+1
1=0
Then, we can write the gradients as
L( W! W?'F [(y W 1 We = W.? Eye Wel w, W
Note that through change of basis, write
W U VT, w Q4+1 II TU) = [ QT,
t=0
i=+11
and consequently, for <15 T, 2,
L-1
Vw.L(WP ..... wh-1) QQ+1 II 1=++1
(w))
S
II W
Smn I QT.
1=0
We can see that W SpeeGD update is
o is diagonal at initialization. Then, if for j E [d], 5/
II
W ool.. f1s5%, the
w w/ Q
Proceeding similarly, for any I [L - 1], we can write the SpecGD updates at any t > 0 as
W -- +
Σ
e,e].
# A1
Simplifying, we get
W c] = (nt+
1 ((#)
1/L -
c - 1)]
1/1.
1/I -
-- )
+
t>
23

---

Published as a conference paper at ICLR 2026
C ADDITIONAL RESULTS IN THE LINEAR SETTING
In this section, we present some additional experimental results for the synthetic setting with class imbalance considered in Sec. 3, where we use linear model and MSE loss. We consider the same setting and parameter values as in Fig. 2, unless stated otherwise.
In Fig. 10, we compare the dynamics of SignGD, Adam (with 31 0.9, 32 0.999) and Muon (with 3 = 0.9) with GD, NGD and SpecGD. We find (from the rightmost plot) that in contrast to (N)GD and SpecGD iterates that remain in the UV eigenbasis (defined as in Condition 1), SignGI) and Adam iterates take some time to align to that basis. We also find that SignGD and Adam learn the spectral components
more balanced way compared to (N)GD, and get slightly better worst-class loss early on in training. However, since SpecGD learns the spectral components at an equal rate, it attains the best worst-class and class-balanced loss early on. Notably, the dynamics of Muon closely follow those of SpecGD. All optimizers eventually lead to convergence to the same solution.
GD
NGD
SignCD
SpeeGD
Adam
Muon
W1
W22
W[3.3]
Woret-class
Cass-balanced
Residual Projected lterates
1.0
0.5
0.4
Si
- WW
0.4
Adam
0.5
0.3
0.2
0.0
0.0
10"
10'
10
10 10 10
10
10"
10"
10"
10"
Time
Time (t)
Time (t)
Figure 10: Dynamics of the iterates (left) and loss (middle) for GD, NGD, SignGD, SpecGD, Adam and Muon on a linear model with class imbalance (priors p1 > P2 > Ps). SpecGD and Muon learn all spectral components at the same rate, whereas (N)GD prioritizes more dominant components. SignGD and Adam also seem to learn the spectral components at the same rate; however, their iterates take time align to the UVT eigenbasis (as seen from the residual of projected iterates (right)). Although all converge to the same solution, SpecGD and Muon significantly outperform in worst-class and bllancee loss early on in training.
CD- Thaery
SpecCD Theary
CD-Expt
Spee
3]
R12.2
W.
tolanood
Theory-Eapt Gap
W- U.W
0.4
114 Waselll
0.5
Lose
0,05
10'
10'
Time (1)
Time (1)
Time (1)
Figure 1I: Comparison of the dynamics of iterates (left) and loss (middle) for GD and SpecGD on linear model with class imbalance (priors P1 > P2 > Ps) in the population limit with zero initialization (theory) and in the finite sample setting with random initialization (expt), and norm of the theoretical/expeeimeetaal difference between the iterates (right). We find that the dynamics characterizes by the theory closely match the experimental behaviour.
In Fig. 11, we compare the dynamics of GD and SpecGD under the theoretical setting (population limit and zero initialization) vs. the experimental setting (using 2000 samples with zero-mean Gaussian initialization with variance 10 2). We use a learning rate of 0.02 in this case. For comparing the iterates, we project them onto the UV eigenbasis from the population setting. We observe that the iterate and loss dynamics from theory closely predict the behaviour seen in the experimental setting.
In Fig. 12, we consider the same theoretical and experimental setting as in Fig. 11, and compare the dynamies of SpecGD in both cases to the dynamics of Muon (with / 0.9) in the experimental setting. We use a learning rate of 0.005 in this case. We observe that the iterate and loss dynamics for Muon closely follow those predicted theoretically for SpecGD.
24

---

Published as a conference paper at ICLR 2026
Opect--heecty
SpecsDDEttt
Muon Ches-bularood
WP3.1
Worsictta
Theory Expo Gap
0.10
UWWW
0.4
0.5
Loo8
0.05
10
Time 01)
Time 01)
Time
Figure 12: Comparison of the dynamics of iterates (left) and loss (middle) for SpecGD (both in the population limit with zero initialization (theory) and in the finite sample setting with random initialization (expt)) and Muon (expt) a linear model with class imbalance (priors p1 > P2 > Pa), and norm of the difference between the iterates (SpecCiD-theory vs SpecGD/Muon expt) (right). We find that the dynamics characterizes by the theory closely match the experimental behaviour.
In the next experiment, we examine the effect of relaxing Condition on the dy- namics GD and SpecGD iterates. Specif- ically, in the relaxed version, the joint diag- onalization holds with Sa= + B, instead of just diagonal matrix Swx. such that BB c. We note that while Condition satisfied for data model (DM) in the population limit, only the relaxed ver- sion holds in the finite sample setting. For instance, under the setting considered in Fig. 11, B| x 0.013. In Fig. 13, we com- pare the dynamics of iterates in the finite sample setting B|| = c) a baseline, where we the UV eigenbasis from
0D.(B=3)
SpeeCD (6=4)
GD ()
Spec [[][||
M11.1
(922.21 .... IR3.3]
1.0
0.10
[|0,
0.0
0.00
10'
10" 10
10'
10
10"
10"
(1)
Figure 13: Comparison of the dynamics of iterates (left)
for GD and SpecGD on linear model with class imbal- ance (priors P1 P2 > Ps) when B = O (Condition I holds) vs. when B = e (Condition relaxed), and
norm of the difference between the iterates (right). We
find that the dynamics in the latter setting closely match those under Condition 1.
this setting but set = O. We find that the
iterates in the two settings closely follow each other.
D ADDITIONAL RESULTS AND DETAILS OF EXPERIMENTAL SETTINGS
This section presents the details of experimental settings for our experiments with the Colored-MNIST dataset and those in Sec. and Sec. 4, well as additional results.
For our experiments with Shampoo, we rely on the distributed-shampo reference im- plementation of Shi al. (2023) and use matrix inverse stabilization parameter precondit wtt.ttiiuiorerueeerreerrm = 5. max mexx_ppecoodiiiionee __meeerrriimmi 8192, unless stated oth- ngCon erwise. To stabilize training, we also use SGDGraft SSDGraftingConf ig or AdamGra ft fnnnnn f i for
10 10-6
the update grafting. Here, $1 governs the exponential moving average of the raw gradients, while 3g governs the accumulation in the matrix preconditioners (see App. A). Further, is a small diagonal "jitter' that stabilises the matrix inverse, and SGD/Adam grafting simply rescales the Shampoo update so its overall magnitude matches that of a plain SGD/Adam step (with default 3 parameters).
For experiments with Muon, represents the momentum parameter, and we perform one Newton-Schulz iteration in each step. By default, the Adam optimizer is used update all vector parameters. For hyperparameter tuning, set 31 for Adam to be the same as 3 for Muon. We use the default value of 3, for Adam with training with Muon.
the implementation from Jordan al. (2024). Here, B
For our experiments with Adam and SGD, we use the standard implementations from the Py TyTrhh library Paszke et al. (2019). For SGD 8 refers the momentum parameter whereas for Adam 3 refers to momentum and Bg is the exponential moving average parameter (see App. A).
Following recent works (Srećković al., 2025; Marek et al., 2025), which show that the performance gap between SGD and adaptive optimizers like Adam is larger for larger batch sizes, we use larger batch sizes our experiments.
25

---

Published as a conference paper at ICLR 2026
D.1 COLORED-MNIST
Dataset. In Sec. 1, we consider a variant of the Colored-MNIST dataset, a benchmark used commonly in the literature on spurious correlations (Arjovsky et al., 2020; Pezeshki et al., 2021). The task is to classify each digit as either < or 2 5. The digits in each class are injected with a background color (red green) that is correlated with the label and acts a spurious feature. In setting, the spurious correlation in the train set is 99%.
Model Architecture and Training. We train a four-layer ReLU MLP with hidden layer dimensions 512, 128, and 32, using inputs from the standard Coloured-MNIS" dataset. Each image is 28 x 28 with 3 channels (RGB). Training is performed for 500 epochs the full-batch setting. We don't use weight decay or learning rate scheduler in this setting. We use preconditioner frequency 1 for Shampoo.
Results. Fig. 14 shows the set of hyperparameters we tried for different optimizers and effect on majority- and minority-group test accuracy. All results are averaged over five independent runs. We select the best hyperparameter setting based on the best average group-balanced accuracy at the end of training.
Building our initial comparison of optimizers from the Frobenius, max, and spectral norm families in Fig. 1, we present a broader comparison that includes additional optimizers in Fig. 15.
Skynu
8
8
0.090
&
Shampoo
Figure 14: Hyperparameter sweeps for different optimizers the Colored-MNIS: dataset, showing the majority- and minority-group test accuracies at the end of training.
D.2 CIFAR- WITH CLASS IMBALANCE.
Dataset and Preprocessing. We construct glass-imbalanced versions of CIFAR-I0 and CIFAR- 100 (Krizhevsky et al., 2009) using step imbalance with ratio R = 20. For CIFAR-10, we sample half the classes (5) to be majority classes containing 20x more samples (5000) than the rest of the (minority) classes (250). The same procedure is applied to CIFAR- 100 using all 100 classes.
We apply standard data augmentation including random horizontal flips, random crops, and normal- ization using dataset-specific channel means and standard deviations. The datasets are split into training and test sets following the original CIFAR splits.
Model Architecture and Training. We train ResNet- on CIFAR-10 10 and ResNer-50 on CIFAR-100 following standard architectures (He al., 2016). All models use random initialization.
Training is conducted with batch size 2048 for 200 epochs. We compare four optimizers representing different bosed approaches, namely NMD, Signum and Adam, and Muon. We use weight decay for all optimizers. All optimizers use cosine learning rate scheduler. We select the learning rate and momentum by grid search (see Fig 16, Fig 17). The hyperparameters with the highest minority-class accuracy are used to report the final results. We use the following learning rate (77) and momentum (3 or (1) configurations:
26

---

Published as a conference paper at ICLR 2026
NSD
SigneD Signum
SpacED
RMSSPrp
vShampoo
Shampoo
10
10
10
Epech
Figure 15: Comparison of optimizer test accuracy on the Colored-MNIST dataset. The top and bottom panels show performance the majority and minority groups, respectively. We compare several families of optimizers, starting with Normalized Steepest Desceni methods-NGD. SignGD, and SpecGD-whiel normalize the gradient using the Frobenius, max, and spectral norms, respectively. We then show their corresponding Normalized Mometum Descent variants: NMD, Signum, and Muon. Finally, we compare other adaptive methods, contrasting Adam (element-wise preconditioning) with Shampoo (full-matrix preconditioning). along with their non-momentum analogs, RMSProp and vShampoo, to isolate the effect of preconditioning. We observe that Muon and SpecGD significantly outperform their Frobenius (NMD, NGD) and max (SignGD, Signum) counterparts in terms of test accuracy on minority groups. However, we do not observe a discernible performance gap when comparing Shampoo with Adam or vShampoo with RMSProp.
NMD
Signum
Adam
F-0.9
8:0.8
8:0.95
8-0.5
55
Minority Class Acc.
Minoriy Class Aoc.
Minorty Class Acc.
Mimority Apc.
0.6
0.1
0.4
0.5
0.2
02
0.0
0.0
10
10
10
10
10
10"
10
10"
10"
10"
Figure 16: Hyperparameter sweeps for different optimizers on the CIFAR-10 dataset, showing minority-class test accuracies at the end of training.
NMD
Signum
Muon
Adam
p-0.55
$-0.9
8:035
Minority Closs Act.
0.20
Minority Class Abc.
Minonty Class Acc.
0.20
Minority Acc.
0.20
0.15
0.15
0.15
0.15
0.10
0.10
0.10
0.10
0.05
0.05
0.06
0.00
0.00
0.00
10
10"
10
10
10
10"
10"
10"
10"
7
Figure 17: Hyperparameter sweeps for different optimizers on the CIFAR-100 dataset, showing minority-class test accuracies at the end of training.
CIFAR-10: NMD (n 0.1, 8 0.95), Signum (7 0.001, 8 0.95). Adam (n 0.001, 31 = 0.9), Muon (17 = 0.1, B = 0.9).
CIFAR-100: NMD
0.1, 8 0.9), Signum (n 0.001, B 0.9), Adam (n 0.001, 31 0.9), Muon (7 = 0.1, 3 = 0.95).
Results. As shown in Fig. 18 (CIFAR-10) 10) and Fig. 19 (CIFAR-100), Muon consistently outperforms other optimizers on both majority- and minority-class accuracies for both experiments, with the performance advantage being most pronounced during the early stages of training.
27

---

Published as a conference paper at ICLR 2026
MM8
Sigmam
Muon
Adar
Loss
Minority Acc.
Majority Acc.
0.8
10
2.4
Training Loss
2.0
Accuracy
Accuracy
0.4
6
0.-
10a
10"
136
1.a
Eoochs
Epochs
Epocht
Figure 18: Comparison of training loss, and minority- and majority-class test accuracies when training a ResNet-18 18 model on CIFAR-I0 using different optimizers. Muon achieves superior minority-clase test accuracy while maintaining competitive majority- class performance, especially early on in training.
NME
Signam
Muon
Adaamn
Minority
Majority
Troining Less
3.0
Accurccy
Accurasy
1.5
0.0
100
10
10
0.0
10-
13P
10
0.0
10
135
10a
13
Epochs
Epochs
Epochs
Figure 19: Comparison of training loss, and minority- and majority-class test accuracies when training ResNet-50 model on CIFAR-100 using different optimizers. Muon achieves superior minority-class test accuracy while maintaining competitive majority-class performance, especially early on in training.
NMD
Signum Majority
Muon
Adam
Minority
10
R-20
Accuracy
Epachs
Eaachs
Epochs
Epochs
Figure 20: Test accuracy dynamics on CIFAR-10 with ResNet- across varying imbalance ratios R. Muon consistently outperforms other optimizers on minority classes, with the performance gap becoming more pronounced as R increases.
NMD
Signum Majority
Muon Minority
Adam
R=1
R-5
R-10
R=20
Accuracy
Epochs
Epochs
Epochs
Figure 21: Test accuracy dynamics on CIFAR- 100 with Re across varying imbalance ratios R. In contrast to the behaviour on CIFAR-10, we find that in this setting Muon maintains superior or comparable performance across all imbalance levels.
28

---

Published as a conference paper at ICLR 2026
D.2.1 EFFECT OF CLASS IMBALANCE RATIO
To investigate how the severity of class imbalance affects the relative performance of different optimizers, we conduct experiments varying the imbalance ratio R E {1, 5, 10, 20} on both CIFAR- 10 and CIFAR-100.
Experimental Setup. We use the same step-imbalanced construction as in Sec. 2, but vary the imbalance ratio R. which denotes the majority-to-minet balanced dataset, while higher values indicate progressively more severe imbalance. For CIFAR-10, we train ResNet-18, and for CIFAR-I00, 100, we train ResNet-50, following the same data augmentation and preprocessing pipeline described in App. D.2.
sample ratio. R = 1 corresponds to a
We use the optimal hyperparameters identified in our R = 20 experiments (App. D.2) for each optimizer across all imbalance ratios to ensure fair comparison. All models are trained for 200 epochs with batch size 2048 and cosine learning rate scheduling.
Results. Figures 20 and 21 show test accuracy dynamics across varying imbalance ratios. On CIFAR- 10, Muon's advantage grows with imbalance severity. This gap is particularly pronounced early in training: while all optimizers perform similarly on balanced data (R 1). at R 20 Muon achieves about 20% minority-class accuracy compared to < 10% for NMD and Signum at the end of epoch 1. On CIFAR-100, the effect of changing R is less pronounced and Muon maintains the best performance across all settings.
D.2.2 FINE-TUNING FROM IMAGENET PRETRAINING
In this section, we investigate whether the observations for (pre)training the model from random initialization are consistent with the case where we fine-tune ImaeeNet--retrrinedd model.
Experimental Setup. We fine-tune a ResNet- 18 model initialized with ImageNet-trained weights (IMAGENETIK, V1) CIFAR-10 with step imbalance (R 20). While retaining the pretrained weights, we modify three components for CIFAR-10 10: the first convolutional layer is replaced with a new layer adapted for 32 32 inputs (kernel size 3, stride 1, padding 1), the maxpooling layer is removed (replaced with identity), and the final fully connected layer is replaced for 10-class output. All layers remain trainable during fine-tuning. Models are fine-tuned for 200 epochs with batch size 2048 and cosine learning rate scheduling.
Following the protocol specified earlier, the hyperparameters used report the final results are selected based on the best minority-class accuracy. We compare NMD (1- = 10-2. 3 0.99), Signum (1) = 10-4, B 0.9), Adam (1) = 10-3 31 0.9. 5g2 0.999), and Muon (1 10 10-2 3 = 0.9). All optimizers use zero weight decay.
Results. As shown in Fig. 22, when fine-tuning from ImageNet initializatton, the optimizer dynamies differ from training from scratch (Fig. 18). Adam, Muon, and Signum achieve comparable training loss and majority-class accuracy, while NMD exhibits slower convergence. For minority-clo -laass accuracy, Adam demonstrates the best performance, followed by Muon and Signum, while NMD struggles to learn minority classes quickly.
NHK
Sinmum
Muon
Adam
Loss
Minority Acc.
Majority Acc.
2.4
Trining Loss
2.6
Accuracy
Accurscy
0.8
0.0
10
1.0
10F
10
13P
10'
107
0.
10'
Epochs
Epochs
Epochs
Figure 22: Fine-tuning ImageNet- all optimizers except NMD achieve similar majority-clas accuracy, Adam demonstrates superior minority-class performance, followed by Muon and Signum. NMD shows significantly delayed learning both majority and minority classes.
ResNet- 18 on imbalanced CIFAR- 10 (R= 20). While
29

---

Published as a conference paper at ICLR 2026
D.3 MNIST-CIFAR DATASET
We train the model for 500 epochs or until convergence (loss reaches 0.01). We use a batch size of 2048 and a cosine annealing learning rate scheduler. We use fixed weight decay of 10-5.
For SGD and Muon, we test learning rates 0.01, and momentum values 1 0.9. For Adam and Shampoo, we test learning rates 10 10 31 values 10- 0.5, 0.9 and 82 values 10 8 0.5, 0.999. The results are averaged over three independent runs. The results are reported for the best hyperparameter setting for each optimizer, selected based on the best seed-averaged worst-group accuracy on the validation set.
D.4 SUBGROUP ROBUSTNESS BENCHMARKS
For MultiNLI, we use a batch size of 512 and the results are averaged over four independent runs. For CelebA, we use batch size of 1024 and the results are averaged over five independent runs.
To ensure our comparisons are robust, we conduct extensive hyperparameter sweeps for all optimizers, as shown for the MultiNLI dataset in Fig. 23 and CelebA in Fig. 24. This selection protocol confirms that the consistent superior performance of spectral optimizers like Muon and Shampoo over SGD is not an artifact of a specific hyperparameter choice.
SGD
Adam B 10
Adam Bg 0.5
Adam 0.999
200 -
240,011
Ax2
1x00031
2x0
Je 2H00001
Worstgrou group Test Acc
0.0
3.5
Wordtgroop Teat Agg:
Word 2
0.4
4-4
4-
84
0.20
0.25
2.50
0.75
1.00
0.00
122
0.50
1.00
Muon
Shampoo = 10
Shampoo
Shampoo A. 0.999
240
1H00001
300
140000
Worstgroup AOC
Tees Ave
Wooetgroup Tees Aoc
gondggrou Teat Aoz
0.4
24
0.00
2.2
0.50
1.00
0.2
0.50
0.25
1.00
P1
P1.
Figure 23: Hyperparameter sweep for different optimizers on the MultiNLI dataset, showing worst- group test accuracy at the last fine-tuning epoch. The sweep varies the momentum parameters (B for SGD/Muon, 31, 32 for Adam/Shampoo), learning rate (1), and weight decay (A). The results indicate that the best performance is not necessarily obtained at the default momentum values for SGD (3- = 0.9) or Adam (3, = 0.9, B2 0.99). For Muon, a higher 3 is generally better for each learning rate, although performance tends to fall for B > 0.9, except in the case of small learning rates like 7 10 5 10 4 where performance remains high. The performance of Shampoo is highly stable across different 31 and 322 values for any learning rate 17, in sharp contrast to Adam.
Adam 10-1
A Adam
Muon
1-30000
300 0.25
3.50
8.75
230
330
1.2
8:30
0.75
310
010
0.5
028
8%
Figure 24: Hyperparameter sweep for different optimizers on the CelebA dataset, showing worst- group test accuracy at the last fine-tuning epoch. The sweep varies the momentum parameters (3 for SGD/Muon, 31, 32 for Adam/Shampoo), learning rate (1)), and weight decay (A).
30

---

Published as a conference paper at ICLR 2026
D.5 TINYSTORIES
Dataset and Preprocessing. We use a subset of the TinyStories corpus (Eldan & Li, 2023), randomly sampling 500 stories from the full dataset as training set. The selected stories are tokenized using SentencePiece (Kudo & Richardson, 2018) with unigram tokenization and vocabulary size 2000. We sampled another 50 stories as validation set.
We define token frequency buckets based on corpus statistics computed over the training set. Tokens are categorized as: (1) "rare" tokens: frequency I 50th percentile. (2) "frequent" tokens: frequency 2 80th percentile. See Fig. 26 (left) for a demonstration of token frequency distribution.
Evaluation Metrics. We evaluate models the validation set every epoch using cross-entropy loss on both training and validation sets, and top- 1/pp1/ accuracy on the validation set.
Model Architecture and Training. We train 4-layer Transformer models with the following specifica- tions: maximum sequence length 64, 4 heads, embedding dimension 256, feed-forward dimension 512, and dropout rate 0.1. The token embeddings and unembeddings are tied. All weights are initialized randomly.
Training conducted with batch size 32 for 100 epochs using three optimizers: SGD (1 0.1, 3 = 0.9, weight decay 10-4), Adam (1) 0.001, 31 = 0.9, 8n = 0.999, weight decay 10 10-*), and Muon (1) = 0.1, 3 = 0.95, no weight decay). Learning rates and momentum values are selected via grid search based on the lowest validation loss (see Fig. 25).
Results. As shown in Fig. 26 (right), Muon and Adam converge faster than SGD, with Muon showing rapid loss reduction around epoch 10. Critically, SGD exhibits a large gap between majority and minority token performance, while Muon significantly reduces this gap. achieving more balanced learning across token frequencies. Also see Fig. 9 which compares the validation set accuracy, where the effect is more pronounced.
SGD
Adam
Muon
Validation Loss
P:0.85
7.0
An0.8
p:0.85
0:0.9
p-0.85
8.50
6.6-
Validation
8.5
Validation
8.0
5.5
6.0
8.00
5.0
5.6
7.75
7.50
10
10
10
10
10
10
10
10
10
10
10
10"
10'
Figure 25: Effect of using different learning rates and momentum on the final validation loss of different optimizers on TinySso
SGD
Muon
Majoriyy (Frequent)
--- Minority (Rare)
50th percentie
80th persentile
Token Frequency Distribution
Training Loss
Validation Loss
500
15.0
Number of movess
200
10.0
150
7.5
2-
5.0
50
2.5
30
300
10
10a
10a
10a
10'
10han Frequency (kg scalo)
Epochs
Epochs
Figure 26: Left: Token frequency distribution in training set. The x-axis shows token rank ordered by frequency (most to least frequent), and the y-axis shows log frequency. Tokens at or below the 50th percentile (left of dashed line) are classified as rare tokens, while tokens or above the 80th percentile (right of solid line) are classified as frequent tokens. Right: Comparison of train and validation loss dynamics of various optimizers on Storis. Solid lines: frequent tokens (2 80th percentile), dashed lines: rare tokens (< 50th percentile). Muon leads to reduced performance gap between frequent and rare tokens on the train set compared to SGD and Adam.
31

---

Published as a conference paper at ICLR 2026
D.6 ATTRIBUTE- ORRANNSM CLASSIFICATION
We consider structured attribute prediction using an organism dataset to study hierarchical concept learning and the effect of using different optimizers.
Dataset and Task. The dataset contains 12 organisms across hierarchical categories: plants (apple, orange, rose, tulip), animals (dog. wolf, cat, tiger, sparrow, eagle, salmon, trout), mammals (dog. wolf, cat, tiger), and other animals. Each organism has 13 attributes with taxonomic structure: plant attributes ("uses sunlight for food," "grows roots"), general animal attributes ("needs oxygen," "can move around"), mammal attributes ("stays warm-blooded"), etc. The conditional probability matrix, given attribute and predict subject, shown in Fig. 27. This creates a hierarchical classification challenge where the model must learn both individual attribute predictions and the underlying taxonomic relationships.
Model Architecture and Training. We bilinear model with logits W W where W Rd and W E RRxxd with n = 12 (number of organisms). & = 13 (number of attributes), and 128 hidden dimensions. GD uses learning rate y 0.001 and Muon uses learning rate 0.01, with Muon using momentum 0.9.
Results. Fig. 27 shows confusion matrices tracking category classification for GD and Muon. The categories are organized from coarse to fine in these matrices. Each entry indicates the number of examples in the row category being classified in the column category. The results reveal a key difference in learning dynamics: GD learns the coarse categories (Animals and Plants) much faster than fine mra-ird categories, showing strong diagonal entries for these broad classes early in training. In contrast, Muon learns all hierarchical levels at similar rates, achieving more balanced progress across both coarse and fine categories. This differential learning behavior directly supports our theoretical finding that spectrum-awan optimizers like Muon learn principal components at equal rates, while GD prioritizes learning dominant patterns first.
Step
Step
True category
GD
Prodicted category
Step 20
Step
Step 70
Step 120
Thue categoryy
Muon
Prodicted categbry
Figure 27: Confusion matrix evolution attribute classification for GD and Muon. Categories are organized from coarse (top-left) to fine (hottom-right).. Each cell value indicates the number of attributes with true category (row) classified as predicted category (column). A diagonal matrix indicates perfect classification. GD learns coarse categories first (evident at step 20), while Muon learns all hicrarchical levels at similar rates (step 70). Both optimizers achieve perfect classification by step 120.
D.7 MIXTURE OF TINYIS AND ADD-k
In this section, we investigate whether our findings on imbalanced learning extend to the level of knowledge acquisition and skill learning. We introduce domain-imbalanced setting by constructing
32

---

Published as a conference paper at ICLR 2026
composite dataset where the majority component consists of natural language narratives (TinyStories), while the minority component requires learning a structured mathematical skill.
Dataset and Task. We train mixture of the TinyStories dataset (majority) and mathematical task (minority). specifically add-k. For TinyStories, we follow the same train/test process as in App. D.5. For add-k: example sequence is of the form "Q: 21, A: 1/, Q: 22, A: y2. Q: 23, A: ys", where U; =; + k. We consider two-digit natural numbers for both the inputs I; and the offsets k. where is drawn from a fixed set of values during both training and evaluation. We train using the next token prediction objective and evaluate the exact match accuracy to predict the last ys, given the rest of the sequence.
Model Architecture and Training. We use the same Transformer architecture as in our TinyStories experiments (see App. D.5 for details). We train on a composite dataset of 20, 000 samples. This includes 18, 000 TinyStories samples (majority, with 0.9 fraction) and 2. 000 add-k arithmetic sequences (minority, with 0.1 fraction). We evaluate on 500 held-out samples for each domain. We train for 200 epochs with batch size of 256 and compare three optimizers. SGD uses f = 5 10-? momentum 0.95, and weight decay 10-4. For Adam, we set = 10-3 -3,03,00 31 0.9, 32 0.999, and weight decay 10-2. For Muon, we set 7 10-2 and momentum 0.95.
Results. We compare the top-1 and top-5 accuracies on TinyStooies and the exact match accuracies on a for SGD, Adam, and Muon in Fig. 28. We observe that for Muon, the accuracies on both datasets increase at a more similar rate compared to SGD, and it generalizes better on the minority add-k task. Adam, the other hand, achieves performance comparable to Muon, similar to some of our results on subgroup robustness datasets in Sec. 4.
50
60C
Matth
Spach
Epec1
Cooch
Figure 28: Comparison of SGD, Adam and Muon when training a transformer on mixture of TinyStories (majority) and add-k (minority) datasets: train loss (left), top-1 and top-5 test accuracies on TinyStories (middle) and exact match test accuracy on add-k (right). Muon promotes more balanced learning compared to SGD while performing comparably Adam.
MID
MNDwCE
Muon
MacHiCE
1.00
CD cmmomennuum
Adam
Shampoo
Muon
0.75
10
Trening Loss
Total Training
0.5
5-
14norty
0.0
10
10
Ecoch
300
200
Epochs
Epochs
Figure 29: Effect of training with weighted cross-entropy (wCE) on Colored-MNIST using NMD and Muon. For NMD, using wCE causes majority- and minority-group test accu- racies to rise at the same rate, outper- forming Muon early in training; Muon with wCE further amplifies this effect.
Figure 30: Comparison of GD, Adam, Shampoo and Muon when training CNN on the Barcoded MNIST dataset from Kunstner et al. (2024), a variant of MNIST with heavy- tailed class imbalance. GD only drives the loss majority classes towards O and makes little progress on the minority classes. In contrast, Adam, Shampoo and Muon drive the loss on both majority and minority classes towards O.
33

---

Published as a conference paper at ICLR 2026
D.8 EFFECT OF LOSS RE- WEIIGGHTTNG
In Fig. 29, we compare training with weighted cross-entropy (wCE) loss on the Colored-MNIST dataset (same setting as Fig. 1) using NMD and Muon. We find that for NMD with wCE, the majority- and minority-group accuracies rise at the same rate, and it outperforms Muon (with CE) very early in training, while using Muon with wCE further amplifies this effect. Since Muon with CE already exhibits a weaker form of this effect 'for free", it may still be useful in settings where explicit group information for loss ne-weighting is unavailable.
D.9 MNIST WITH HEAVY-TAILED CLASS IMBALANCE
Dataset. In this section, we consider the Barcoded MNIST dataset, which is a variant of MNIST with heavy-tailed class imbalance introduced in Kunstner et al. (2024). Barcoded MNIST contains two types of classes: 10 classes with 5000 samples each from the original MNIST dataset (majority), and 10 21t additional classes with samples each (minority). The images the minority classes are generated by taking images from the original MNIST dataset, and encoding a 10-bit barcode into the top left corner of the images, for each of the 10 original classes.
Model Architecture and Training. Following Kunstner et al. (2024), train 2-layer CNN on this dataset in the full-batch setting. In Fig. 30, we compare the total train loss as well as train loss on the majority and minority classes (each comprising about 50% of the total samples) for the four optimizers; GD with momentum, Adam, Shampoo and Muon (we use learning rates 0.005, 10-4, 10-4, and 0.005, respectively). We use default momentum parameter values for all optimizers. For Shampoo, we use AdamGraft ingConfig, and set the stabilization parameter to 10-5.
Results. Consistent with the results in Kunstner et al. (2024), we observe while that GD only minimizes the loss on the majority classes, and makes negligible progress on the minority classes, Adam minimizes loss on both majority and minority classes. In addition, we find that training with Muon or Shampoo leads to a similar behaviour as Adam: they also minimizes the loss on both majority and minority classes, in contrast to GD.
D.10 SPECTRAL ANALYSIS FOR NON-LINEAR MODELS
In this section, we extend our analysis on the rate of learning of different spectral components beyond linear models. Specifically, we consider two settings with step imbalance (k// majority and k/2 minority classes) trained using cross-entropy (CE) loss: i) one-hidden-layer MLP (with fixed outer layer, as used in Fig. 1) trained on the data model (DM) with an imbalance ratio of 10, and ii) a ResNet- 18 trained on CIFAR-10 10 with an imbalance ratio of 20 (same setting App. D.2 for Fig. 3). In both settings, we track the singular values of the k x n logit matrix using all n train samples across the ls classes training progresses. Specifically, measuring the KL. divergence between the uniform probability vector g and the vector of normalized singular values diag(E,) over training time, where Σ, the singular value matrix at time t. tr.)..)
quantify the extent of balanced learning by
3,3
03
17
1
17
Cocch
Opech
Figure 31: Comparison of NGD, SignGD, and SpeeGD on a one-hidden ReLU MLP on Gaussian mixture data generated using the data model Eq. (DM). From left to right: train loss, minority- and majority-class test accuracies, and KL divergence between the normalized all-ones vector and the vector with normalized singular values of the logit matrix over time. SpecGD maintains a consistently low KL divergence throughout training. In contrast, NGD and SignGD exhibit an increase this metric, indicating period of spectral imbalance.
Gaussian Mixture Setting. We consider n = 1000, d = 200. Similar to Fig, 4, we compare optimizers NGD, SignGD and SpecGD in this case, with learning rates x 10-2, 8 x 10-1, and
34

---

Published as a conference paper at ICLR 2026
2 x 10-2, respectively. We use the full batch for training, and use held out set of 2000 test samples for reporting test accuracies. The results are reported after averaging over 5 independent seeds for initialization and data selection. As shown in Fig. 31. for NGD and SignGD, the KL divergence-based metric initially increases as training progresses, indicating more imhalanced singular values, and becomes smaller in later stages of training, as both majority and minority classes learned. In contrast, SpecGD maintains a consistently lower value for this metric throughout training, demonstrating that its normalized singular values remain closer uniform and that spectral components are learned at more balanced rate. These findings are supported by the corresponding majority- and minority-class accuracies.
CIFAR-10 with Class Imbalance. We repeat the experiment described above on the imbalanced CIFAR-10 10 dataset (same setting as App. D.2 for Fig. 3). In Fig. 32, we observe that Muon similarly demonstrates more balanced learning of the spectral components, This is supported by the dynamics of majority- and minority- accuracies (see Fig. 3).
NHD
Signum
Muon
KLLunifum normal 1o0d 0)
0.6
0.4
10
10a
10
102
Epochs
E RELATED WORK
Figure 32: Comparison of NMD, Signum, and Muon on imbalanced CIFAR-10 10 show- ing the KL divergence between the uniform probability vector and the vector of normal- ized singular values of the logit matrix. Muon maintains a lower KL divergence throughout training compared to NMD and Signum indi- cating more balanced learning of the spectral components.
Spectrum-awee -ware and Other Adaptive Optimizers. Our work is situated within the growing body of re- search on spectrum-aware and matrix-level optimiz- ers, which have demonstrated significant empirical success. The development of methods like Sham- poo (Gupta et al., 2018) and, more recently, Muon (Pethick et al., 2025; Jordan et al., 2024) has high- lighted the potential of optimizers that operate with matrices at the layer level (Carlson et al., 2015; Large et al., 2024), rather than vectorized parameters. Recent theoretical inquiries have sought to understand the mechanisms behind their performance. One line of work has focused on characteriz- ing the implicit bias of their canonical form, SpecGD. For instance, Fan et al. (2025); Tsilivis et al. (2024) have shown that SpecGD converges to a max- margin classifier with respect to the spectral norm, though this analysis primarily describes the terminal phase of training. Concurrently. other research has investigated the advantages of adaptive methods like Adam, particularly in settings with heavy-tailed class imbalance (Kunstner et al., 2024) or from the perspective of feature learning (Vasudeva et al., 2025b). While these works highlight Adam's advantages on different forms of imbalanced data through various mechanisms, we use this testbed to reveal a fundamentally different principle behind SpecGD's D's success: its balanced learning of the data's 's principal components, which correspond directly to majority and minority groups
Further, given that experiments benchmark spectral methods against strong adaptive baselines, it is worth noting that recent large-scale empirical studies on language model pretraining confirm that well-tuned Adam variants remain highly competitive. Extensive benchmarks have shown that Adam W is a robust and strong baseline (Wen et al., 2025; Semenov et al., 2025), which is consistent with our experiments in Sec. 4, where Adam is on par with or sometimes beats Shampoo and Muon. Orvieto & Gower (2025) showed, via extensive hyperparameter tuning, that 81 = 83 is a near-optimal choice for momentum parameters in Adam, contrary to conventional wisdom. They also showed that Signum recovers much of the performance gap between SignGD and Adam, a trend we also observe in our experiments (Fig. 15). key factor in its performance relative to SGD the batch size, as the Adam-SGD gap is most pronounced in bagge training regimes (Srećković et al., 2025; Marek et al., 2025). In line with these works, we tune hyperparameters such as the learning rate and momentum parameters for all our experiments (see App. D).
Training Dynamics. Our theoretical analysis of training dynamics builds directly upon the framework established by Saxe et al. (2013) and subsequently extended by Gidel et al. (2019). These works provided exact, closed-form solutions for the learning dynamics of deep linear networks trained with standard GD. Their key insight was to analyze the system under the condition that the input and output moment matrices are jointly diagonalizable, which simplifies the complex, non non-linear dynamics into a set of decoupled scalar equations. While their analysis revealed how GD prioritizes
35

---

Published as a conference paper at ICLR 2026
learning the dominant spectral components of the data, work adopts this lens and extends it to SpecGD. This extension allows for a direct, analytical comparison between the training trajectories of GD and SpecGD, enabling to precisely characterize the latter's more balanced learning rate across all spectral components. Furthermore, our analysis of the bilinear model (Sec. 3.3) also connects to the literature the Unconstrained Features Model (UFM), which is often used theoretical proxy to study the geometry of feature learning in deep networks, positing a model with enough capacity to freely optimize both the learned representations and the final classification layer (Yang al., 2018; Mixon et al., 2020).
Spectral Bias in Deep Learning. More broadly, our work connects the well-documented phe- nomenon of spectral bias or simplicity bias in deep learning, where networks trained with (S)GD show predisposition to learn simpler patterns first. Prior works have demonstrated that networks pri- oritize learning low-frequency functions before high-frequency ones (Rahaman et al., 2019; John Xu al., 2020), and that SGD) guides networks to learn functions of progressively increasing complexity (Nakkiran al., 2019; Cao et al., 2020). This principle extends even modern architectures, with recent work showing that Transformers also exhibit bias towards learning low-sensitivity functions (Bhattamishra et al., 2023; Vasudeva al., 2025a). However, this inherent bias towards simplicity can be detrimental, as it can lead models to rely on simple but spurious correlations in the data, which hinders generalization Our analysis complements this line of work by showing how spectrum-aware optimizers offer a mechanism to somewhat mitigate this bias by promoting more uniform learning across the spectral components of the data.
underrepresented groups (Shah al., 2020; Vasudeva al., 2025b).
F LIMITATIONS AND OPEN QUESTIONS
Our theoretical results are derived in population setting with squared loss, and extending them to the more practical snime- regime is key direction. In this setting, where many interpolating solutions may exist, can characterize the implicit bias of SpecGD and establish formal general- ization guarantees analogous to those in our population-level analysis? Another crucial extension to move beyond squared loss analyze commonly used losses like cross-entropy. A primary challenge here that our key technical assumption of joint diagonalizability may no longer hold, requiring new analytical tools. Furthermore, while experiments language modeling (Sec. 4) suggest that Muon's balanced learning persists (higher top-k validation accuracy on rare tokens), formal extension of our theory non-trivial. It remains open question what data structures, analogous to the moment matrices Exe and Σye in our multi-class classification setup, define the principal components when doing next-token prediction with sequential, language data. Answering these questions would provide a more complete picture of why spectrum-aware optimizers succeed in modern deep learning.
36


Reviews:
[{"review_id": "1n0HWhN9HZ", "reviewer": "Reviewer_sqYt", "summary": "This paper tries to dive deep into the effectiveness of the optimizers which consider Spectral properties of the parameters like specifically MuOn. The authors study SpecGD which can be seen as an approximation of Muon optimizer. Authors show that the SpecGD method is effective on the imbalanced datasets, as it tends to learn the all the data moments at almost equal rate in comparison to NGD which tends to learn more aggressively from the data moments of the majority classes. They show that this effect is also prominent for deep networks theoretically, and show experimental validations for the same.", "questions": "Could the authors explain more about the applicability of the data model they have considered with orthogonal means and specific variance’s generalizability in the real-world datasets?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1.\tThe paper is well-written and quite comprehensive.\n2.\tNice intuitive experimentation is done for each claim in theory which makes understanding easy and also shows real world applicability.\n3.\tThe results of the paper uncover the reasons for effectiveness of the spectral optimizers like MuOn, with faster convergence on imbalanced datasets.", "weaknesses": "There is a line of work that shows that the generalization on imbalanced data could be improved by escaping saddle points. Hence, considering algorithms like Perturbed Gradient Descent (PGD) or Sharpness Aware Minimization could be used for optimization and show some connection with.\n\n[R1] Jin, Chi, et al. \"How to escape saddle points efficiently.\" International conference on machine learning. PMLR, 2017.\n[R2] Rangwani, Harsh, et al. \"Escaping saddle points for effective generalization on class-imbalanced data.\" Advances in Neural Information Processing Systems 35 (2022): 22791-22805.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DiFVVpyULt", "reviewer": "Reviewer_HTPs", "summary": "This paper investigates when spectrum-aware optimizers like Muon and Shampoo generalize better than standard gradient descent methods by analyzing their canonical form, Spectral Gradient Descent (SpecGD). Using imbalanced data as a testbed under a mixture-of-gaussians data model, the authors prove for a linear model that SpecGD learns all principal components of the data at equal rates, unlike gradient descent which prioritizes dominant components. Paper shows that this leads to superior balanced accuracy early in training for linear and deep linear models under a Gaussian mixture data model. Authors also present experiments across image classification, natural language inference, and language modeling tasks demonstrate that Muon achieves faster generalization on minority classes and rare tokens compared to SGD, though Adam sometimes matches or exceeds Muon's performance.", "questions": "1. Since the authors are arguing that Muon/Shampoo perform differently than Euclidean methods, can they add Adam to Fig. 1 (left) as well? Also add it to Fig. 3 and others. I am not sure what to read from the eigenvalue and eigenvector figures.\n2. Is there any insight into what about the class imbalance/spectra dictates the speedup in generalizing on minority classes with Muon?\n3. In Fig. 4, what is the structure of the linear model—just $Wx$?\n4. In Proposition 1, there is a dependence on the iterate $t$ and step size $\\eta$. What are the implications of the step size value?\n5. A stretch question: why is momentum in Adam insufficient to leverage spectral information in the dataset? Experiments do show that Muon is not always better than Adam, e.g., in Fig. 8 at the end.\n\n### Typos and Editorial Suggestions\n1. The paper notes that Adam reduces to SignGD without momentum and preconditioning histories, which has been helpful for understanding Adam. Can the authors briefly mention here how Shampoo/Muon reduce to SpecGD and under what assumptions? This would make the writing parallel because in line 73, the authors argue for studying SpecGD as the canonical form for Shampoo and Muon.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a clear empirical observation (Fig. 1) to motivate imbalanced data as a testbed, which is an effective approach.\n2. The writing is clear, and the experiments cleanly motivate the theoretical arguments. The paper provides valuable insight in Theorem 1 on the gap between gradient descent and spectral GD (specifically between their continuous flow counterparts), showing that spectral GD outperforms GD.\n3. The analysis extends beyond a single-layer linear network to a sequence of matrices, creating a deep linear network framework.\n4. Experiments span several tasks (image classification, sentence relationship classification, next-token prediction), datasets (Dominoes, MultiNLI, TinyStories), and architectures, demonstrating that worst-group accuracy is generally learned faster with Muon. This is an interesting observation across various settings.", "weaknesses": "1. In the data model (DM), requiring orthogonal $\\mu_i$ seems like a strong assumption. It is unclear whether orthogonal means for Gaussian mixtures is a standard assumption. What breaks down if this assumption is removed? In fact, in line 173, experiments are done with $\\mu_i$ sampled uniformly from a Gaussian, so only $\\mathbb{E}[\\mu_i, \\mu_j] = 0$ holds.\n2. Why is joint diagonalizability a useful condition for the analysis? More importantly, are there other data models, theoretically or empirically, that would satisfy this condition? When does it break down? The paper does not discuss this condition extensively, but it would be briefly useful to discuss it in the text.\n3. I am slightly confused about the takeaways for neural network training with non-linear activations. The analysis would be significantly harder, but the paper has a disconnect between presenting closed-form training dynamics for a linear model on a special data model (DM) and the motivation and experiments for deep neural networks on general image classification datasets. SVD analysis works for MSE minimization in a linear model, but would this be useful even in two-layer networks with a nonlinear activation?\n4. In Proposition 2, the initialization for matrices is extremely specific, and it is unclear what all the values mean. This is unlike Proposition 1 and Theorem 1, where the initialization was zero. This makes me question the effect of initialization on the results.\n5. Like the spectral dynamics visualized and derived for linear networks, the paper would benefit from a similar analysis in the experiment section that goes beyond just worst-group accuracy with different optimizers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dpvOgvixYW", "reviewer": "Reviewer_v3RF", "summary": "This paper verified that spectral gradient descent learns all principal components of the data at equal rates, especially for imbalanced datasets. The conclusion extends to deep linear models and spectral optimizers including Muon and Shampoo, where the study concluded that these optimizers achieve better generalization through the balanced learning rates.", "questions": "1. The argument in Section 3.4 relies on assumptions of similar scaling between Hessian blocks (Theorem 6).  Could you comment on how violations of this assumption might affect the results? For instance, in real networks, do earlier layers and later layers exhibit vastly different spectral norms, and how would that impact the perturbation analysis?\n2. Is it possible to construct an artifical dataset with a small proportion of domain-specific text corpus, e.g. mathematics or coding, and investigate if Muon exhibits a stronger learning progression in the domain, and if the learning dynamics generalizes to a knowledge-level?\n3. While spectral optimizers may promote equal learning on imbalanced datasets, does the conclusion leads to a hypothesis that an ideal data preprocessing algorithm is equivalent to adopting Muon to Adam?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The study promotes the understanding of the success of spectral optimizers from a fundamental perspective, especially when the datasets are imbalanced in class or group.\n2. The studies in this paper convered a wide range of neural architectures. The experimental designs are trustworthy.\n3. Good story-telling from multiple perspectives.", "weaknesses": "1. This paper frames SpecGD as the canonical form for Muon and Shampoo. However, the practical implementations involve crucial components like preconditioner history accumulation  and various approximations/stabilization techniques. The paper could benefit from a more detailed discussion of how these practical elements might alter the \"equal-rate\" learning dynamics of the idealized SpecGD. \n2. The extended analytics to language modeling (next token prediction) is limited to rare and frequent tokens learning, while data imbalance on linguistics or knowledge and skills are investigated less in the study. \n3. Adam and other second-order preconditioning approximators should also be studied in simple linear models and image classification tasks, as its relation to second-order curvatures and simplicity in implementation and computation are still critical perspectives at present.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8VXryBXme6", "reviewer": "Reviewer_oR6c", "summary": "This paper studies a simplified algorithm spectral GD, as a canonical algorithm to understand algorithms such as Muon, Shampoo, and SOAP. They adopt a linear model under squared-loss, using imbalanced classes as their data, where these classes have orthogonal means, and employ continuous dynamics or discrete dynamics that mimic the continuous case. The results show that for worst-class and class-balanced risks, when considering early phase loss decay, specGD is faster than GD with or without adaptive step-size (normalized gradient). They also prove that depth can accelerate the learning of minor components. Experiments are also provided under practical settings to validate the theoretical findings.", "questions": "1. Please clarify the theoretical relationship between Spectral GD and practical adaptive methods such as Muon, Shampoo, and SOAP. In particular, does Spectral GD explicitly approximate their behavior along the dominant eigen-directions of the preconditioning matrix?\n\n2. Equation (1): The definition of the inner product is unclear. Since \\nabla L should have the same dimensions as W_t, the resulting inner product seems dimensionally inconsistent. Please provide a precise definition.\n\n3. It would be helpful to explicitly present the update rules (for NGD, specGD, GD, and their gradient-flow versions) used in Theorems 1 and 2 for clarity.\n\n4. The early phase and asymptotic behavior of specGD and (N)GD are different. It would be better if both can be included in the theorem to give more complete picture of these algorithms.\n\n5. Section 3.3: Proposition 2 and its counterpart for the deep linear model do not seem clear. \n\n5.1 The requirements on \\eta are not clearly stated. Theorems 1 and 2 analyze gradient flow, whereas Proposition 2 uses gradient descent. Please clarify the rationale and distinction between these settings.\n\n5.2 The results hold when \\delta\\to\\infty, which makes the current statement a bit wierd since you can then just assume zero initialization. The statement should be rewritten. \n\n5.3 The counterpart result for deep linear models is not formulated as a proposition. It may be clearer to combine these results into a single theorem and then discuss the role of depth.\n\n5.4 The variable t appears in the results, but as t \\to \\infty, \\bar{W} diverges.\n\n5.5 There is a \\Delta T that directly characterizes the time of learning different classes but this is not included in any theorems or propositions. \n\n5.6 Overall, section 3.3 needs to be reorganized.\n\n6. Experiments compare toy algorithms (canonical algorithms) for theory, and the practical algorithms for real applications. However, it is unclear whether specGD can adequately represent Muon and Shampoo in practice, given their distinct preconditioning mechanisms. It would be better to include additional experiments that directly compare specGD with Muon and Shampoo for validation.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper provides a relatively thorough analysis of Spectral GD under imbalanced datasets, characterizing both the loss decay and the impact of network depth. The problem they studied is important towards understanding real applications, although under simplied settings. There is both theoretical and empirical evidence to support the claim.", "weaknesses": "The overall concern is that the setting of this paper is too strong, More precisely, the authors adopt the specGD, which is a perfectly preconditioned algorithm, while in practice, the main difficulty of adaptive algorithms is to design the preconditioning, and different preconditionings lead to different behavior. \n\nRegarding the theory, the assumptions seem pretty strong. The authors use linear model, orthogonal data, joint diagonalizability (another strong “orghogonality” assumption on data), zero (or near-zero) initialization, and gradient flow (or near-gradient flow).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)