The Loss Does Not See the Basis, but Adam Does
Devender Singh
cs.LG, math.OC, stat.ML
Submitted: 2026-08-05
Comments: 22 pages main text + appendices, 5 figures. Code, seeds, and raw run records: https://github.com/idevender/loss-basis-adam
Code: https://github.com/idevender/loss-basis-adam
Project page: http://karpathy.github.io/2015/05/21/rnn-effectiveness
License: http://creativecommons.org/licenses/by/4.0/
The gist: Gradient descent on a factored model W = UV is implicitly biased toward low-rank solutions, while Adam, starting from the same small initialization, is not.
Terminology
Abstract
Gradient descent on a factored model W = UV 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) (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 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.
Sources
- Understanding and inverse design of implicit bias in stochastic learning: a geometric perspective
- Implicit Bias of Per-sample Adam on Separable Data: Departure from the Full-batch Regime
- The Spectral Dynamics and Noise Geometry of Muon
- Old Optimizer, New Norm: An Anthology
- Benchmarking Neural Network Training Algorithms
- How do simple rotations affect the implicit bias of Adam?
- Muon$^p$: Muon with Fractional Spectral Powers
- To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters
- Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data
- The Implicit Bias of Adam and Muon on Smooth Homogeneous Neural Networks
- Uniform Spectral Growth and Convergence of Muon in LoRA-Style Matrix Factorization
- Improving Generalization Performance by Switching from Adam to SGD
- Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers
- Muon is Scalable for LLM Training
- Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets
- Towards Understanding the Power and Limits of the Muon Optimizer: A River-Valley Perspective
- Dead-Direction Conditioners: Gauge-Equivariant Preconditioning for Deep Networks
- Symmetry Breaking in Transformers for Efficient and Interpretable Training
- FlowAdam: Implicit Regularization via Geometry-Aware Soft Momentum Injection
- How the Optimizer Shapes Learned Solutions in Equivariant Neural Networks
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks