The Loss Does Not See the Basis, but Adam Does — part 5/12

Cleaned-text code points [60903, 76576). Bibliography omitted here; source.md includes it.

The resulting behavior on the zoo task is shown below, with improvements in both parameter recovery and effective rank as p p decreases towards p = 0 p{=}0 , without any change to the algorithm itself, simply by varying p p in the update rule. Two experimental setups were used to illustrate this result: a fixed-step version, where all learning rates were tested for all values of p p with the other hyperparameters fixed as well, and an envelope version, where for each p p the best learning rate was chosen individually based on recovery in the interpolating regime. In the fixed-step version, for every learning rate whose sweep both interpolates and resolves the flow limit (for the zoo task, the range from 10 − 3 10^{-3} to 3 × 10 − 2 3\times 10^{-2} ), performance was averaged over three random seeds, and parameter recovery consistently improved with decreasing p p : from 0.576 0.576 to 0.201 0.201 at 10 − 3 10^{-3} , from 0.570 0.570 to 0.256 0.256 at 3 × 10 − 3 3\times 10^{-3} , from 0.573 0.573 to 0.357 0.357 at 10 − 2 10^{-2} , and from 0.581 0.581 to 0.493 0.493 at 3 × 10 − 2 3\times 10^{-2} , with every run meeting the 10 − 7 10^{-7} interpolation threshold. Effective rank follows the same trend (Figure 3 ). 4 4 4 At the grid’s largest rate ( 10 − 1 10^{-1} ) every p p still interpolates, but recovery sits at 0.55 0.55 – 0.62 0.62 and effective rank at ≈ 13.8 {\approx}13.8 – 15.3 15.3 for all p p : a step this coarse leaves every endpoint far from its flow limit (Remark A.5 ), so the sweep flattens rather than restores. The depth of the restoration is rate-dependent; its monotonicity, at rates that resolve the flow limit, is not.

In the envelope setup, for each p p , the learning rate was selected which produced the best possible recovery in the interpolating regime, according to the standard procedure for the zoo task (Appendix D.2 ). The resulting envelope also shows monotonic improvement and matches the fixed-step results up to three decimal places at every p ≤ 0.75 p\leq 0.75 when the learning rate is 10 − 3 10^{-3} , indicating that the per- p p learning rate selection does not introduce hidden effects. Figure 3 displays the envelope sweep and provides tabulated results. Its p = 1 p{=}1 endpoint reads 0.570 0.570 where the zoo’s Adam row reads 0.5734 0.5734 , because the dial’s shared grid selects a different rate; Adam is flat across that range (Appendix C ).

This refines the finding of Wilson et al. (2017) into a continuous, adjustable relationship: along this tuning axis, the coordinate-wise anisotropy of the preconditioner directly explains the erosion of factored-model bias. The same anisotropy is also what buys Adam’s speed—switching to p = 0 p{=}0 increases the number of steps needed for interpolation by roughly a factor of eight (from 400 400 to 3200 3200 on this task, comparing each end at its selected rate). The balancedness drift of Proposition A.4 falls with the dial, and the endpoints bracket it: the zoo’s Adam row carries ‖ B T ‖ F = 5.37 \left\lVert B_{T}\right\rVert_{F}=5.37 and its scalar-Adam row 0.14 0.14 (Table 2 ), which ties the empirical adjustment to the conserved quantity. The monotonic trend is not tied to a specific problem scale: on H100, the pattern holds at n = 128 n{=}128 (recovery improves from 0.67 0.67 to 0.40 0.40 , effective rank from 36.6 36.6 to 12.0 12.0 as p p goes from 1 1 to 0 0 , across five seeds) and replicates at n = 40 n{=}40 with ten seeds (recovery improves from 0.52 0.52 to 0.14 0.14 on an extended flow-limit grid, both ends selected at 3 × 10 − 4 3\times 10^{-4} , under every rate the zoo grid carries, which is why the p = 1 p{=}1 end sits below the 0.56 0.56 – 0.58 0.58 of Remark A.5 ; the n = 128 n{=}128 endpoint reads 0.40 0.40 against the ladder’s 0.345 0.345 for the grid reason above).

## 8 Inside the equivariant class: the spectral-tail phase diagram

Equivariance is binary, but behavior inside the equivariant class is not—for memoryless spectral-separable maps, Proposition A.1 gives a coordinate: GD has h ⁡ ( σ ) = η ​ σ h(\sigma)=\eta\sigma , whereas the exact polar map has h ≡ η h\equiv\eta . Recent work highlights two seemingly contradictory trends for Muon: a strong simplicity (spectral) bias ( Fan et al., 2025 ; Gronich & Vardi, 2026 ) , and, for deep-linear and linear-attention models, a loss of GD’s sequential simplicity bias that makes Muon prone to fitting spurious features ( Dragutinović et al., 2026 ) . Our sensing and real-data experiments show both. In this section we explain how the spectral tail of the target reconciles these trends.

##### Design.

Plant X = 1 − τ 2 ​ X 3 + τ ​ E ⟂ X=\sqrt{1-\tau^{2}}\,X_{3}+\tau E_{\perp} where X 3 X_{3} is the rank-3 matrix and E ⟂ E_{\perp} a dense rank- ( n − 3 ) (n{-}3) component supported on the orthogonal complements of X 3 X_{3} ’s column and row spaces, with ‖ X 3 ‖ F = ‖ E ⟂ ‖ F = 1 \left\lVert X_{3}\right\rVert_{F}=\left\lVert E_{\perp}\right\rVert_{F}=1 . The support condition gives ⟨ X 3 , E ⟂ ⟩ F = 0 \left\langle X_{3},E_{\perp}\right\rangle_{F}=0 , hence ‖ X ‖ F = 1 \left\lVert X\right\rVert_{F}=1 , so that τ ∈ [ 0 , 1 ] \tau\in[0,1] and τ 2 \tau^{2} is exactly the fraction of target energy in the tail. Sense with m = 2 ​ dof m=2\,\mathrm{dof} counted at the rank- 3 3 target, as before. All methods interpolate, and we sweep τ \tau . The τ = 0 \tau{=}0 target is the zoo task of Table 2 , but this sweep is a separate 10 10 -seed run on its own two-rate grid, so its τ = 0 \tau{=}0 column reproduces that table’s ordering rather than its digits (GD 0.112 0.112 here vs. 0.131 0.131 there, and so on).

##### The two axes in the sweep.

Table 4 tabulates the sweep and Figure 4 plots it. On the symmetry axis, Adam has the highest recovery error at every τ \tau in Table 4 . Along the schedule dimension of the equivariant class, the degree of schedule aggressiveness (with Muon the most aggressive, followed by Shampoo and then GD) trades accurate low-rank recovery, which the aggressive schedules win, against robustness to the tail, which the moderate ones win. This creates a crossover region around 4 % 4\% tail energy (the difference between GD’s and Muon’s performance at τ = 0.2 \tau{=}0.2 is within the range of the seed-to-seed variance). Thus, Muon’s preference for uniform growth across modes ( Kang et al., 2026 ) is optimal if the signal is truly low-rank, but hurts performance on signals that carry a tail which should not be fit. Hence, there is a spectral regime in which each of the schedules is preferable, depending on the properties of the target signal.

This preference boundary does not depend on the type of decay schedule used: re-running the experiment with cosine-decayed schedules instead of constant ones keeps each method’s performance curve as a function of τ \tau unchanged, with the crossings left where their constant-schedule counterparts put them. That control ran on the 3 3 -seed grid, whose crossing is τ ∗ ≈ 0.35 \tau^{\ast}\approx 0.35 under both schedules; the estimate then became more precise, shifting to approximately 0.2 0.2 on a 10 10 -seed grid (see Appendix D.1 ).

##### An analytically tractable model.

In the simplified framework of decoupled dynamics—commonly used in studies of greedy dynamics—it is possible to precisely characterize both sides of the trade-off and explain why the scale of initialization influences one learning schedule but not the other.

###### Proposition 8.1 (Greedy vs. equal-rate schedules: a solvable two-timescale scenario) .

Consider the factored dynamic system decomposed into independent, balanced scalar components ( w = u ​ v w=uv , with u = v = w u=v=\sqrt{w} ; aligned target; this is the idealization adopted in Arora et al. (2019) and Li et al. (2021) ). Suppose target modes satisfy s 1 > s 2 > 0 s_{1}>s_{2}>0 , all initialized at w i ​ ( 0 ) = w 0 w_{i}(0)=w_{0} where 0 < w 0 ≤ s 2 / 2 0<w_{0}\leq s_{2}/2 , and let T 1 T_{1} denote the first time the leading mode reaches ( 1 − δ ) ​ s 1 (1-\delta)s_{1} , with δ \delta fixed such that 0 < δ < 1 − s 2 / s 1 0<\delta<1-s_{2}/s_{1} . (Both lower bounds are strict and necessary: w 0 = 0 w_{0}=0 is a stationary point, δ ≤ 0 \delta\leq 0 would make the threshold unreachable, and the proof involves division by these terms.) (i) Greedy schedule ( h ⁡ ( σ ) = η ​ σ h(\sigma)=\eta\sigma ): Each mode follows logistic growth, w ˙ i = 2 ​ η ​ w i ​ ( s i − w i ) \dot{w}_{i}=2\eta\,w_{i}(s_{i}-w_{i}) , and at time T 1 T_{1} the smaller mode remains inactive: w 2 ​ ( T 1 ) ≤ C δ , ρ ​ s 1 ​ ( w 0 s 1 ) 1 − ρ → w 0 → 0 0 , ρ ≔ s 2 / s 1 , C δ , ρ = 2 ​ ( 1 − δ δ ) ρ . w_{2}(T_{1})\;\leq\;C_{\delta,\rho}\,s_{1}\Big(\tfrac{w_{0}}{s_{1}}\Big)^{1-\rho}\xrightarrow[\;w_{0}\to 0\;]{}0,\qquad\rho\coloneqq s_{2}/s_{1},\quad C_{\delta,\rho}=2\big(\tfrac{1-\delta}{\delta}\big)^{\rho}. Thus, reducing initialization enhances the timescale gap between leading and trailing modes indefinitely. (ii) Equal-rate schedule ( h ≡ η h\equiv\eta ): All modes grow at the same rate ( w i \sqrt{w_{i}} increases linearly with time) and saturate at their respective s i s_{i} ; by T 1 T_{1} , the smaller mode is already fully fitted, w 2 ​ ( T 1 ) = s 2 w_{2}(T_{1})=s_{2} , regardless of w 0 w_{0} . Here, the flat schedule does not let the initialization scale differentiate between the modes.

Proposition 8.1 (proved in Appendix B ) shows that the decoupled-mode model has a transient dynamics consistent with the observed phase behavior. In particular, if the target matrix has exactly low rank, uniformly fitting all the planted modes recovers the full signal in this model, as suggested by the equal-rate principle of Kang et al. (2026) for Muon. In the case of a target matrix with a less-determined tail (sensed with m = 2 ​ dof m=2\,\mathrm{dof} measurements counted at the head), the flat schedule biases the fit towards directions the measurements do not constrain, while the greedy schedule suppresses those unconstrained modes down to the level of w 0 1 − s 2 / s 1 w_{0}^{\,1-s_{2}/s_{1}} . Thus, the observed transition may be determined by a combination of the spectral properties of the target matrix, the measurement design, the initialization scale, the training length, the schedule, and the optimization algorithm.

## 9 Real data at matched training loss

While the synthetic data was constructed to have a known ground truth, real-world matrices are expected to have spectral tails. As discussed in § 8 , this suggests that the same mechanism should be at work in practice, with additional contributions from the tail components. We test this hypothesis on two hyperspectral image-completion benchmarks. The first is the Indian Pines dataset ( Baumgardner et al., 2015 ) , consisting of 2000 2000 randomly selected pixels across 200 200 spectral bands, derived from a corrected Airborne Visible/Infrared Imaging Spectrometer (AVIRIS) cube with water-absorption bands excluded. The second is Pavia University, featuring 2000 2000 random pixels and 103 103 bands captured by the Reflective Optics System Imaging Spectrometer (ROSIS) sensor. In both cases, the spectral data are approximately low-rank due to physical properties, and both datasets were made publicly available by Graña et al. (2021) .

We perform rank-48 factored matrix completion without weight decay, testing at two underdetermined sampling densities per dataset, calibrated so that m / dof 24 ≈ 1.15 m/\mathrm{dof}_{24}\approx 1.15 and 1.9 1.9 . Here, degrees of freedom (dof) are based on the intrinsic rank of 24 24 —where each singular-value spectrum stabilizes—rather than the higher model rank of 48 48 . This ensures the ratio reflects how underdetermined the true signal is (see Appendix C ). The intrinsic rank is determined directly from each scene’s singular spectrum, sets a fixed sampling density for all methods, and is not tuned using validation performance. We refer to this ratio as m / dof m/\mathrm{dof} henceforth.

A key challenge in comparing optimization methods is early stopping: slower algorithms may appear superior simply because they benefit from implicit regularization due to incomplete fitting. To address this, we compare methods along their full (training loss, test error, effective rank) trajectories at matched training loss levels. Specifically, for any given training loss ℓ \ell , we record each method’s held-out root-mean-square error (RMSE) at the first point it reaches ℓ \ell . Learning rates are selected using a training-only criterion (deepest convergence followed by fewest steps; see Appendix D.6 ), and results are averaged over four random seeds. That rule sets the size of every gap we report below. It selects for fast convergence within the budget, hence for larger rates, which puts GD in its rate-invariant region while handing Adam the worst rate on its own grid: giving each method the rate that minimizes its own held-out error would leave GD ahead by 13 % 13\% rather than 43 % 43\% . We keep the train-only rule because selecting and evaluating on the same held-out statistics would reintroduce the bias the protocol exists to avoid, and neither the ordering nor the rank profile changes under the other one (Appendix D.6 ).

At matched training loss ≤ 3 × 10 − 5 \leq 3\times 10^{-5} , at an observed entry density of 0.15 0.15 ( m / dof ≈ 1.15 m/\mathrm{dof}\approx 1.15 ), gradient descent (GD) achieves a held-out RMSE of 0.0150 ± 0.0001 0.0150\pm 0.0001 at an effective rank of 11 11 , while Adam reaches 0.0268 ± 0.0005 0.0268\pm 0.0005 at rank 28 28 . Thus GD’s held-out error is 44 % 44\% lower than Adam’s at equivalent fit, in all four seeds, aligning with the rank behavior predicted by the mechanism.

The full learning trajectories offer more insight than final performance alone (Figure 5 ). GD’s test error decreases steadily as training progresses, consistent with a greedy, head-first fitting pattern. In contrast, Adam’s test error increases during interpolation (from 0.0251 0.0251 at training loss 3 × 10 − 4 3\times 10^{-4} to 0.0268 0.0268 at 3 × 10 − 5 3\times 10^{-5} ), as its effective rank rises from 22 22 to 28 28 . This shows that, on real data, deeper fitting with an optimizer that breaks gauge symmetry can actually degrade generalization. Muon shows little improvement in held-out error until very late stages, with effective rank remaining close to the model limit (dropping from 46 46 to 36 36 )—a real-data manifestation of the tail behavior described in § 8 .

The performance gap narrows as sampling density increases, as predicted by the selection mechanism: the advantage of GD reduces from + 44 % +44\% at m / dof = 1.15 m/\mathrm{dof}=1.15 to + 28 % +28\% at m / dof = 1.9 m/\mathrm{dof}=1.9 (consistent across all four seeds), and appeared to fall to nearly zero by m / dof ≈ 3 m/\mathrm{dof}\approx 3 in an earlier interpolation-based evaluation, which we do not tabulate because the matched-loss protocol supersedes it. As more data become available, the solution becomes better constrained, reducing the impact of interpolant selection—a transition noted in § 10.1 .

##### Validation and a second benchmark.

The observed effect is neither an artifact of where the comparison is taken nor specific to one dataset. Read at the stricter threshold ( ≤ 10 − 5 \leq 10^{-5} , the level Table 5 reports), the same Indian Pines runs give + 43.0 % +43.0\% at m / dof ≈ 1.15 m/\mathrm{dof}\approx 1.15 and + 27.6 % +27.6\% at 1.9 1.9 , against the + 44 % +44\% and + 28 % +28\% quoted above at 3 × 10 − 5 3\times 10^{-5} : the gap narrows slightly as the fit deepens, and does not turn on the choice of stopping level.