The Loss Does Not See the Basis, but Adam Does
Devender Singh
cs.LG, math.OC, stat.ML
2026-08-06
Gradient descent finds low-rank solutions on factored models and Adam does not. The cause is gauge symmetry, which Adam's coordinate-wise second moment breaks.
Factoring a matrix into two thin ones, W = UV⊤, is the standard setup for low-rank modeling: matrix sensing, recommender systems, LoRA fine-tuning all do it. It comes with an old puzzle. From the same small initialization, gradient descent (GD) spontaneously finds a low-rank solution, as Gunasekar et al. proved in 2017. Adam, from the same initialization, does not, and lands on a higher-rank, worse-fitting solution. The phenomenon has been observed for eight years without a mechanistic explanation.
This paper traces the answer to a geometric property that had been overlooked: gauge symmetry. The loss L(U,V) depends only on the product UV⊤, so for any orthogonal matrix Q, replacing (U,V) with (UQ,VQ) leaves the product and the loss unchanged. The loss "does not see" which orthogonal basis you picked for the k-dimensional latent space. Adam does see it. It computes its second moment coordinate by coordinate, so changing the basis changes its update direction. The title, "The Loss Does Not See the Basis, but Adam Does," says exactly this.
The authors turn this thread into a complete formal framework.
The first step is to define gauge equivariance: an optimizer is gauge-equivariant if it behaves identically on all gauge-equivalent factorizations, depending only on invariants like the product or the Gram, not on the chosen basis. GD is naturally equivariant: its gradient transforms covariantly under the basis change, so the trajectory tracks only invariants. Adam is not. Its coordinate-wise second moment treats the basis as special.
Next comes a structure theorem: the memoryless equivariant update rules are exactly the "Gram-determined left preconditioners" of the form Φ(G) = H(GG⊤)G, where the update can depend on the current factors only through the invariant GG⊤. This turns equivariance from a property into a complete characterization. A companion transfer theorem says that shared-scalar preconditioned dynamics equal GD under a rescaled clock, so GD's pathwise properties carry over to common-scalar flows.
The cleanest causal isolation is a knob called Adam-p. At p = 1 it is standard Adam (coordinate-wise); at p = 0 it is shared-scalar Adam (equivariant), with a continuous interpolation between. The only thing this knob changes is how anisotropic the second-moment denominator is. The result is monotone: as p moves from 1 to 0, recovery error and effective rank both improve monotonically. The quantity that makes Adam non-equivariant is precisely the coordinate-wise anisotropy.
The central evidence is an "optimizer zoo" table. On a 40×40, rank-3, fully overparameterized (k = n = 40) matrix-sensing task, nine update rules are run until full interpolation (training loss all near zero) and ranked by recovery error against the true matrix. They split cleanly into two classes:
| Class | Optimizer | Recovery ‖W−X∗‖F/‖X∗‖F | Effective rank |
| Equivariant | Muon | 6.8×10⁻⁶ (near-exact) | 2.95 |
| Equivariant | GD | 0.131 | 4.51 |
| Equivariant | scalar-Adam | 0.201 | 5.43 |
| Equivariant | Shampoo | 0.286 | 6.95 |
| Coordinate-wise | Lion | 0.425 | 10.58 |
| Coordinate-wise | signum | 0.445 | 7.83 |
| Coordinate-wise | RMSProp | 0.527 | 12.80 |
| Coordinate-wise | Adafactor | 0.543 | 10.72 |
| Coordinate-wise | Adam | 0.573 | 14.37 |
Every equivariant method sits at recovery ≤ 0.286, every coordinate-wise method at ≥ 0.42, leaving a 0.14 gap between the classes. All methods drive training loss to the same level, so the difference can only come from which interpolating solution was selected. The true matrix has rank 3; Muon lands at effective rank 2.95 (exact recovery) while Adam lands at 14.37, far above the true rank. The table reports each method's best learning rate on its own grid, a choice that favors the coordinate-wise methods, so the class gap is a conservative lower bound.
On a transformer, the authors set up two gauge-equivalent initializations. Adam separates them on the very first step, while the equivariant optimizers stay indistinguishable to machine precision for several steps. The two end up with per-head invariants WQ⊤WK differing by 56% in relative Frobenius distance, and the disagreement sits in the invariant itself, so no per-head rotation can close it.
On real data, two hyperspectral image-completion tasks show GD cutting held-out RMSE by up to 43–44% at matched training loss, at the lowest sampling density and lowest effective rank.
For Muon, a "spectral schedule" phase diagram reconciles two opposing reports. When the true spectrum is clean (tail energy below about 4%), Muon's equal-rate updates recover the low-rank target exactly; once the spectral tail grows, the equal-rate schedule loses its edge and GD's greedy schedule wins. This is why some observe a strong low-rank bias in Muon and others see it disappear.
This paper supplies the mechanistic answer to why Adam and GD behave differently on factored models, and the answer has a practical corollary. Any parameterization that relies on a low-rank inductive bias, like LoRA, low-rank attention, or matrix completion, is quietly eroding that bias when paired with Adam. To keep it, you need an equivariant optimizer, or you need to replace Adam's coordinate-wise second moment with a shared scalar.
For optimizer design, the paper offers a clean criterion: gauge equivariance. It is a necessary (not sufficient) condition for preserving GD's implicit bias along that path. Structured optimizers like Muon and Shampoo work partly because they happen to be equivariant.
One thing to state plainly: this is a mechanistic explanation, not a tuning recipe. The authors write that "what we provide is a mechanistic explanation for an observed phenomenon," and the bias is mild at common scales, with no benchmark wins promised.
The authors' stated limits, plus what stays doubtful after reading: