ArXiv

The Loss Does Not See the Basis, but Adam Does

Authors
Devender Singh
Categories
cs.LG
arXiv
https://arxiv.org/abs/2608.05136v1
PDF
https://arxiv.org/pdf/2608.05136v1

Brief

The paper analyzes why gradient descent on factored matrices W = U V^T is biased to low-rank interpolants while Adam is not, attributing the gap to gauge symmetry ((U,V)→(UQ,VQ)) and whether an optimizer is gauge-equivariant. It proves characterization and transfer results for equivariant, memoryless preconditioners, and validates them across matrix sensing, transformers (56% per-head invariant gap), and hyperspectral tasks (43–44% lower held-out error for GD).

Why it matters

Gradient descent on factored models W = U V^T is implicitly biased toward low-rank solutions, while Adam, RMSProp and other coordinate-wise methods are not when started from the same small initialization.

Key details

  • The difference traces to gauge-equivariance (invariance under (U,V) → (UQ,VQ)): equivariant optimizers (gradient descent, momentum, shared-scalar Adam, Muon, Shampoo) can inherit gradient-flow's low-rank mechanism. The paper proves a structure theorem (memoryless equivariant rules = Gram-determined left preconditioners) and a transfer theorem to common-scalar flows.
  • Empirically, interpolating from coordinate-wise to shared-scalar preconditioning restores the low-rank bias monotonically; Adam separates two gauge-equivalent initializations at the first step producing per-head invariants W_Q^T W_K 56% apart (relative Frobenius), and on two hyperspectral datasets gradient descent cuts held-out error by 43–44% at the lowest sampling density and yields lower effective rank.
Source evidence

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 $WQ^\top WK$ 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.

Comment: 22 pages main text + appendices, 5 figures. Code, seeds, and raw run records: https://github.com/idevender/loss-basis-adam