The Loss Does Not See the Basis,
but Adam Does

Devender Singh
arXiv:2608.05136
Paper Code BibTeX
Ground-truth recovery error for nine optimizers. The four gauge-equivariant methods fall between 0.000 and 0.286 in blue; the five coordinate-wise methods fall between 0.425 and 0.573 in vermilion, with an empty gap between the groups.
Nine optimizers, two clusters, no overlap. Ground-truth recovery on underdetermined matrix sensing. All nine run to interpolation, so nothing here is explained by residual training error. Every gauge-equivariant method lands at or below $0.286$; every coordinate-wise one at or above $0.42$. The space between is unoccupied.

Abstract

Gradient descent on a factored model $W = UV^\top$ is implicitly biased toward low-rank solutions, while Adam, starting from the same small initialization, is not. We trace the difference to the gauge symmetry of the loss, its invariance under $(U, V) \mapsto (UQ, VQ)$. Gradient flow's low-rank mechanism is available to an optimizer only if that optimizer is gauge-equivariant, a condition necessary for the transfer but not sufficient for low-rank recovery. Gradient descent, momentum, “shared-scalar” Adam, Muon, and Shampoo satisfy it. Adam, RMSProp, and the other coordinate-wise methods do not. A structure theorem characterizes the memoryless equivariant rules as exactly the Gram-determined left preconditioners, and a transfer theorem carries gradient flow's pathwise properties to common-scalar flows. We then sort nine update rules on underdetermined matrix sensing by recovery error against the planted ground truth. A one-parameter family from coordinate-wise to shared-scalar preconditioning restores the bias monotonically, isolating anisotropy as the cause. A “spectral schedule” reconciles two opposing reports about Muon: equal-rate updates recover exactly low-rank targets but lose their edge as the spectral tail grows. In transformers, Adam separates two gauge-equivalent initializations at the first step, where the equivariant optimizers stay at float precision, and ends with the per-head invariants $W_Q^\top W_K$ $56\%$ apart in relative Frobenius distance, a gap no per-head rotation can close. On two hyperspectral datasets at matched training loss, gradient descent cuts held-out error by $43$–$44\%$ at the lowest sampling density, and at lower effective rank. Basis choice is therefore not a tuning detail but a decision about which interpolant the optimizer selects.

The gauge

The objective $L(U,V) = f(UV^\top)$ is unchanged by

$$(U,V) \;\longmapsto\; (UQ,\; VQ), \qquad Q \in O(k),$$

because rotating the latent bases alters neither $W = UV^\top$ nor any prediction. The loss cannot see $Q$. The question is whether the optimizer can.

The optimizer zoo

methodrecovery ↓ erankbal train lossnotes
min-nuclear-norm0.03353.30——convex reference
gauge-equivariant
Muon0.00002.951.655.7e−8near-exact (6.8e−6 unrounded)
GD0.13124.510.066.5e−8Gunasekar anchor
scalar-Adam ($p{=}0$)0.20105.430.143.5e−8equivariant; flow proxy
Shampoo0.28566.951.474.5e−8
coordinate-wise
Lion0.424810.587.358.0e−8
signum0.44547.831.796.5e−8low-rank-but-wrong
RMSProp0.526612.804.746.0e−8needs decay
Adafactor0.543010.723.405.6e−8factored diagonal still breaks it
Adam0.573414.375.371.2e−11Wilson anchor

Matrix sensing, $40\times40$, rank-3 ground truth, no weight decay, three paired seeds, all runs continued to interpolation. Recovery is $\|W - X^\ast\|_F/\|X^\ast\|_F$; bal is $\|U^\top U - V^\top V\|_F$. The separation claim itself rests on the ten-seed ladder in the appendix, where it holds at every size up to $n{=}256$.

The same gauge lives in attention heads

Log-scale plot of relative logit distance versus training step. Adam's gauge twins jump four orders of magnitude at step one and saturate; SGD and scalar-Adam gauge twins stay at float noise; Muon's gauge and noise twins trace the same curve.
Adam is basis-dependent in attention. Each head's logits depend on $W_Q$ and $W_K$ only through $W_Q^\top W_K$, so each head carries the same symmetry. Two transformers are initialized as the same function in different bases. Adam's gauge twins split at step one, four orders of magnitude above its own noise twin, then saturate; the head invariants finish $56\%$ apart, a gap no per-head rotation can close. SGD and scalar-Adam stay at float noise. For Muon the gauge and noise twins trace one curve, identifying its separation as numerical chaos in the matrix sign rather than basis dependence.

Isolating preconditioner anisotropy with a dial

Algorithm 1, Adam-p, the preconditioner-anisotropy dial. Pseudocode in which a pooled scalar s-bar is computed as the square root of the mean of the bias-corrected second moments over both factors, and the denominator D-theta interpolates between the per-coordinate and shared-scalar forms with exponent p.
Algorithm 1, reproduced from the paper. Only the denominator moves with $p$; the gradient, moment estimates, bias corrections, step size, and (zero) weight decay are the same throughout. The pooling step in blue is the only inter-factor cross-talk, and it is precisely what makes the update equivariant at $p{=}0$: by pooling $\hat v$ over all entries of both factors, $\bar s$ depends only on $\|G_U\|_F^2 + \|G_V\|_F^2$, which is gauge-invariant. Setting $p{=}1$ recovers standard Adam.
Effective rank of the solution falling monotonically as the dial parameter p decreases from 1 to 0, approaching gradient descent's value.
Anisotropy is the cause, not a correlate. As the preconditioner changes shape from coordinate-wise ($p{=}1$) toward isotropic ($p{=}0$), recovery and effective rank improve monotonically toward gradient descent's. Adaptivity and momentum are held fixed. An envelope arm (learning rate re-optimized at each $p$) and a fixed-step arm agree closely, which is what pins the effect on anisotropy rather than on step-size rescaling. The same anisotropy is what buys Adam its speed: $p{=}0$ needs roughly eight times as many steps to interpolate.

Inside the equivariant class

Recovery error versus tail-amplitude parameter tau for Adam, GD, Shampoo and Muon over ten seeds. Muon is exact at tau equals zero and crosses above GD near tau equals 0.2. Adam is worst at every tau.
A second axis: the spectral schedule. Equivariance is necessary, not sufficient. Within the equivariant class, Muon's equal-rate schedule is exact at $\tau{=}0$ and cedes to GD near $\tau^\ast \approx 0.2$, about $4\%$ tail energy. This reconciles “Muon recovers exactly” with “Muon removes the simplicity bias” — two readings of one schedule at different points on the same axis. Bands are $\pm1$ s.d. over ten seeds.

Real data at matched training loss

Two panels plotting held-out error and effective rank against training loss on Indian Pines. GD sits below Adam at every matched training loss; Adam's effective rank climbs to 28 while GD settles between 5.6 and 13.2.
Gradient descent cuts held-out error by $43$–$44\%$. Two hyperspectral image-completion benchmarks, four seeds, learning rates chosen by a train-only criterion. The optimizers are compared only at equal training loss, so no fitted run is ever set against an unfitted one. GD settles at effective rank $5.6$–$13.2$; Adam climbs to $28$ against the scene's intrinsic rank of $24$.

BibTeX

@article{singh2026lossbasis,
  title         = {The Loss Does Not See the Basis, but Adam Does},
  author        = {Singh, Devender},
  year          = {2026},
  eprint        = {2608.05136},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2608.05136}
}

Send feedback and questions to devenders@mun.ca — Devender Singh.