The Loss Does Not See the Basis, but Adam Does — part 1/12
Cleaned-text code points [0, 15859). Bibliography omitted here; source.md includes it.
# The Loss Does Not See the Basis, but Adam Does
###### Abstract
Gradient descent on a factored model W = U V ⊤ 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 ) ↦ ( U Q , V Q ) (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 ⊤ W K W_{Q}^{\top}W_{K} 56 % 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 43 – 44 % 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.
## 1 Introduction
The implicit regularization effect of gradient descent (GD) in factored models is by now well established: training an overparameterized factorization to convergence, i.e. to a global minimum of an underdetermined matrix problem starting from small initializations, does not yield an arbitrary interpolating solution. In well-studied regimes it yields approximately the minimum-nuclear-norm one, a convex proxy which strongly promotes low rank ( Gunasekar et al., 2017 ; Arora et al., 2019 ; Li et al., 2021 ) . By contrast, empirical results show that adaptive methods have different generalization properties—Adam-family optimizers, while achieving a similar train loss, can end up measurably worse on the test set ( Wilson et al., 2017 ; Keskar & Socher, 2017 ) . Recent work showed that per-coordinate adaptivity is basis-dependent ( Xie et al., 2025 ) , while a preconditioner shared across coordinates restores rotational equivariance ( Ling et al., 2022 ) . We share this concern about basis dependence. To the best of our knowledge, its implication for interpolant selection on a factored model has not been characterized: whether there is a criterion by which one can determine that an optimizer preserves or breaks the low-rank inductive bias, in such a way that a range of deployed methods can be sorted by that criterion and the sorting verified against ground truth.
This is particularly consequential in the context of individual attention heads, which carry the same symmetry: two transformers that compute the identical function at initialization start behaving differently from the first Adam update, and ultimately exhibit head-level invariants 56 % 56\% apart at matched validation accuracy (§ 6 ).
We answer this question by introducing a criterion according to which an optimizer can be said to preserve the implicit low-rank bias of factored models. The objective function L ( U , V ) = f ( U V ⊤ ) L(U,V)=f(UV^{\top}) remains unchanged under the transformation ( U , V ) ⟼ ( U Q , V Q ) , Q ∈ O ( k ) , (U,V)\;\longmapsto\;(UQ,\;VQ),\qquad Q\in\mathrm{O}(k), meaning that rotating the latent factor bases does not alter the resulting matrix W = U V ⊤ W=UV^{\top} or any model prediction. This property is referred to as the gauge symmetry. Optimization takes place in the space of ( U , V ) (U,V) factorizations, and different algorithms handle this symmetry to varying extents. A central question is whether an optimizer is equivariant under orthogonal transformations of the gauge. The ones that are can exhibit the low-rank behavior typically seen in gradient flow on factored models.
Equivariance allows us to reproduce gradient-flow dynamics, but does not in general lead to low-rank solutions (see Remark 4.7 ). For equivariant optimizers, the key to low-rank convergence is the spectral schedule: whether the dynamics it induces on the singular values of the factored weights during training lead to attraction to a low-rank manifold within the specified training budget. For optimizers whose updates respect equivariance—gradient descent, momentum, shared-scalar-preconditioned Adam, Muon, Shampoo—we can say that their behavior is consistent across all gauge-equivalent factorizations, while for coordinate-wise algorithms—Adam, RMSProp, signSGD, Lion, Adafactor—each optimizer reads a preferred basis, so it can select a different interpolant depending on where in the gauge orbit it starts. In particular, in the optimization problems considered in this paper, the coordinate-wise optimizers failed to reach the low-rank attractor, settling instead on interpolants of higher rank and worse recovery (Figure 1 , Table 2 ). Thus only equivariant algorithms were able to access the low-rank solutions found by gradient flow. Note, however, that this does not necessarily mean that a coordinate-wise optimizer cannot converge to a low-rank solution by some other route, just that it has less direct access to it. One example of such a case is given in Appendix D.10 .
##### Contributions.
We develop this criterion in five steps. 1. Classification and first-step defect analysis (§ 4 ). We characterize those standard optimization updates that are compatible with the gauge action. Gradient descent (GD), Polyak momentum, Adam with a shared scalar second moment, Muon’s matrix sign update (and all finite Newton–Schulz iterates of it), and Shampoo’s Kronecker preconditioner with damping are all equivariant. By contrast, any optimizer that applies a fixed entrywise nonlinearity at zero optimizer state (Adam, RMSProp, signSGD, Lion) is not. Similar distinctions have been noted previously for unfactored parameterizations, where the rotation acts on the whole parameter space (§ 2 , Table 1 ), but not for the internal gauge of a factored model. For Adam with zero-initialized state, we obtain an exact expression for the first-step defect. We then present a transfer theorem: optimization dynamics with a shared scalar preconditioner correspond to gradient flow under a rescaled time parameterization, hence previous results for gradient flow ( Gunasekar et al., 2017 ; Arora et al., 2019 ; Li et al., 2021 ) carry over directly whenever their own conditions are met by the same objective. A structure theorem turns the classification into a full characterization: a memoryless update is equivariant exactly when it takes the form Δ = H ( G G ⊤ ) G \Delta=H(GG^{\top})\,G for some preconditioner H H , whose choice defines the optimizer (Theorem 4.5 , stated for full-column-rank G G with k ≤ n k\leq n ). The defining property is equivariance, not balancedness—which both Muon and Shampoo violate while still preserving the bias.
2. Optimizer comparison (§ 5 ). In matrix sensing without weight decay, all optimizers are driven to interpolation, erasing residual training error as a source of variation. Under this controlled condition, the theory cleanly separates nine update rules (eight deployed optimizers plus a control, scalar-Adam) into two classes by their recovery error at this computational budget (equivariant 0.00 0.00 – 0.29 0.29 ; coordinate-wise 0.42 0.42 – 0.57 0.57 ; Figure 1 ). This distinction is robust to the various controls explored in Appendix D : symmetrized schedules (cosine decay for all methods), full learning-rate scans, and variable initialization amplitude. The predicted ranking of GD, scalar-Adam, and Muon at lower error than the coordinate-wise group is observed up to n = 256 n{=}256 ( 10 10 seeds, Appendix D.9 ), with two flagged exceptions to the strict two-class split (untuned Shampoo and long-anneal signum; § 11 ).
3. A dial that isolates the mechanism (§ 7 ). We construct a parameterized family of update rules which smoothly interpolates between preconditioner geometries as p p ranges over [ 0 , 1 ] [0,1] . The interpolation connects Adam’s standard coordinate-wise denominator ( p = 1 p{=}1 ) to a single coordinate-independent scalar ( p = 0 p{=}0 ), which is gauge-equivariant. Along this interpolation, recovery and effective rank improve monotonically. The only thing that changes with p p is the denominator’s coordinate-wise anisotropy—the part of the update that reads the gauge—plus the overall scaling that comes with it; adaptivity and momentum are kept constant throughout. The same trends hold when the learning rate is re-optimized at each p p (§ 7 ).
4. The same gauge lives in attention (§ 6 ). Because each head’s logits depend on W Q W_{Q} and W K W_{K} only through W Q ⊤ W K W_{Q}^{\top}W_{K} , each head has the same internal gauge symmetry as the factored model. We train two gauge-equivalent copies of a small transformer—the same function at initialization, written in different bases—to compare the behavior of Adam and heavy-ball SGD. Adam splits them at the first step, and the two copies end with head invariants W Q ⊤ W K W_{Q}^{\top}W_{K} that differ by 56 % 56\% . Heavy-ball SGD and scalar-Adam, by contrast, are indistinguishable to floating-point precision at that step and stay orders of magnitude below Adam’s split afterwards, and an A = I A{=}I twin confirms that the pipeline itself is exactly deterministic. A sanity check demonstrates that it is not some imprecision in the calculation that produces Adam’s divergence: perturbing the weights by 10 − 7 10^{-7} in the same basis leaves a step- 1 1 drift four orders of magnitude below the gauge split. Only for Muon does the apparent gauge divergence come from numerical instability in the matrix sign calculation, and in particular from the presence of small singular values. This is not specific to the small transformer or to deterministic training: the same effect can be seen when training a six-layer character-level language model with d model = 256 d_{\text{model}}{=}256 under mini-batch noise, where the Adam twins again diverge at step one while the equivariant ones stay at machine precision.
5. Inside the equivariant class, and on real data (§ 8 –§ 10 ). The optimizer sweep here has essentially two knobs: equivariance, which sets class membership and is what the dial of § 7 moves along, and the aggressiveness of the spectral schedule, which sets behavior inside the class. We characterize the latter by a phase diagram in the tail energy, the fraction of the target’s energy carried outside its low-rank part. Across 10 10 seeds this locates a boundary near 4 % 4\% tail energy, above which Muon’s equal-rate schedule ( Kang et al., 2026 ) stops helping and GD’s greedy schedule wins. This reconciles the exact recovery we see for Muon in § 5 with the removal of the simplicity bias reported for the same optimizer by Dragutinović et al. (2026) . A closed-form solution based on a decoupled two-timescale approximation recovers a timing mechanism compatible with that boundary (Proposition 8.1 ). The same consideration applies to practical tasks: on two hyperspectral image-completion benchmarks at equivalent train loss, GD achieves a 43 43 – 44 % 44\% better held-out error than Adam under higher underdetermination, at a much lower effective rank, and in all four seeds (§ 9 ; learning rates chosen throughout by a train-only criterion). Most importantly, our framework is constructive: the coordinate-wise clip on the injected flow velocity in our own previously published optimizer is by itself enough to break equivariance, quietly suppressing the very bias it was supposed to inject. By contrast, a clip based on a global norm restores equivariance of the injected velocity and achieves the best bias restoration among the Adam-type methods we test—though the hybrid becomes fully equivariant only once Adam’s diagonal preconditioner is taken to its shared-scalar end as well.
##### Scope.
This work deliberately avoids advocating any particular optimizer as a new state-of-the-art, and the mechanism does not promise benchmark wins. The bias it encodes is mild, and we highlight the conditions under which it is beneficial in § 10.1 : it wins over tuned explicit regularization only where the optimal regularization is itself mild and aligned in direction, and it falls short, in ways the theory anticipates, on tasks that demand stronger regularization or the kind of per-coordinate adaptation Adam provides. Like Gunasekar et al. (2017) and Arora et al. (2019) , what we provide is a mechanistic explanation for an observed phenomenon, stated as a testable claim and evaluated in settings where it is expected to help as well as where it is not.
## 2 Related work
##### Implicit bias of gradient methods on factored models.
Gunasekar et al. (2017) conjectured, and proved for commuting measurement matrices, that gradient flow on symmetric factorizations U U ⊤ UU^{\top} from infinitesimal initialization converges to minimum-nuclear-norm interpolants of underdetermined matrix sensing. Arora et al. (2019) extended the phenomenon to deep matrix factorizations and argued for a dynamical (rather than purely norm-based) picture, sharpened by Razin & Cohen (2020) and by the greedy low-rank learning dynamics of Li et al. (2021) . Much of this analysis leans on the quantity U ⊤ U − V ⊤ V U^{\top}U-V^{\top}V , which gradient flow conserves ( Arora et al., 2018 ; Du et al., 2018 ) . That conservation is bought by the same symmetry we study here: a continuous symmetry of the loss hands the flow a conserved quantity, in the Noether sense made precise for learning dynamics by Kunin et al. (2021) .
This entire line of work studies gradient descent and gradient flow. Our transfer theorem extends its pathwise conclusions to memoryless common-scalar flows under the original assumptions, and our experiments address a question this literature leaves open: what do the optimizers people actually deploy do to this bias?
##### Implicit bias of Adam and adaptive methods.
Wilson et al. (2017) first exhibited adaptive methods finding worse-generalizing solutions on constructed and real problems. The modern theory of Adam’s implicit bias is largely margin-based on separable classification, following Soudry et al. (2018) for gradient descent: ℓ ∞ \ell_{\infty} -norm-constrained characterizations of AdamW ( Loshchilov & Hutter, 2019 ) by Xie & Li (2024) , Karush–Kuhn–Tucker convergence for Adam ( Zhang et al., 2024 ; Cattaneo et al., 2024 ) , per-sample refinements ( Baek et al., 2026 ) , and a unified steepest-descent view assigning Adam the ℓ ∞ \ell_{\infty} and Muon the spectral geometry ( Gronich & Vardi, 2026 ; Fan et al., 2025 ; Bernstein & Newhouse, 2024 ) . Li et al. (2025) characterize a distinct sharpness measure that Adam minimizes near minimizer manifolds, driving it to qualitatively different solutions from SGD on overparameterized models. This is an adaptive-vs-SGD selection gap adjacent to ours, but it operates through a sharpness-reduction mechanism rather than a gauge mechanism, and it is not studied on factored-model recovery.