🤖 machine learning

The Loss Does Not See the Basis, but Adam Does

This paper demonstrates that Adam's failure to recover low-rank solutions in factored models stems from its lack of gauge equivariance, a property possessed by gradient descent and other shared-scalar optimizers that preserves the implicit bias toward low-rank interpolants.

Original authors: Devender Singh

Published 2026-08-06

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Devender Singh

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a giant, messy puzzle where the pieces are numbers arranged in a grid. In the world of artificial intelligence, this is called "matrix factorization." You have a huge, complicated picture (the data), and you want to break it down into two smaller, simpler stacks of cards that, when multiplied together, recreate the original picture. The goal is to find the simplest possible stacks that still fit the picture perfectly. This is important because simpler solutions usually work better in the real world, avoiding the trap of memorizing noise instead of learning patterns.

To solve this puzzle, computers use a method called "gradient descent," which is like a hiker trying to find the bottom of a valley by always stepping downhill. For a long time, scientists noticed that if you start with very small cards, this hiker naturally finds the simplest solution. However, a newer, faster hiker called "Adam" (a popular tool in AI) often gets lost, finding a complicated, messy solution even when a simple one exists. The big question has been: Why does the fast hiker fail to find the simple path, and can we fix it without making it slow?

This paper, titled "The Loss Does Not See the Basis, But Adam Does," investigates why this happens. The author discovered that the problem isn't just about speed; it's about how the hiker sees the world. The puzzle has a hidden symmetry: you can rotate your stacks of cards in many different ways, and the final picture stays exactly the same. Think of it like spinning a globe; the continents move around, but the map is still the same. The traditional hiker (Gradient Descent) ignores the specific rotation and just looks at the shape of the valley, naturally finding the simplest solution. But the fast hiker (Adam) gets distracted by the specific rotation of the cards. It treats one direction as "special" and another as "different," even though they are mathematically identical. This distraction causes it to pick a complicated, high-rank solution instead of the simple, low-rank one.

The author proved that any optimizer (a tool for solving the puzzle) that respects this rotational symmetry will naturally find the simple solution, while those that break the symmetry will get stuck in complex ones. They tested nine different optimizers and found a clear split: the "symmetry-respecting" ones (like Gradient Descent, Muon, and a modified version of Adam) found solutions with errors as low as 0.000006, while the "symmetry-breaking" ones (like standard Adam and RMSProp) had errors over 0.42—a massive difference.

To prove it wasn't just a fluke, they built a "dial" that smoothly changes Adam from its chaotic, rotation-hating mode into a rotation-friendly mode. As they turned the dial, the solution got simpler and more accurate, showing that the specific way Adam looks at the data is the exact cause of the problem. They even tested this on real-world data, like hyperspectral images of the earth, and found that the rotation-friendly methods reduced errors by about 44% compared to the standard Adam.

Interestingly, the paper also found that being "rotation-friendly" isn't always a magic bullet. If the puzzle itself is messy and has a lot of random noise (a "spectral tail"), the super-fast, rotation-friendly hiker called Muon sometimes gets too eager and fits the noise, while the slower, steady hiker (Gradient Descent) does better. So, the best tool depends on the specific nature of the puzzle.

Finally, the author looked at how this affects modern AI models called Transformers (the brains behind chatbots). They found that if you start two identical AI models with the same math but just rotate their internal cards differently, the standard Adam optimizer makes them behave completely differently after just one step. They end up with completely different internal structures, even though they started as the same function. This means that the choice of optimizer isn't just a tuning detail; it fundamentally decides which version of the solution the AI learns. The paper concludes that to get the best, simplest results, we need optimizers that respect the hidden symmetries of the math, rather than getting distracted by the specific way we arrange the numbers.

Technical Summary: "The Loss Does Not See the Basis, But Adam Does"

Problem Statement
Gradient descent (GD) on factored models ( $W = UV^\top$ ) is known to exhibit an implicit bias toward low-rank solutions, often converging to minimum-nuclear-norm interpolants in underdetermined matrix sensing tasks. However, adaptive optimizers like Adam, despite achieving similar training losses, frequently fail to recover these low-rank structures, yielding solutions with higher effective rank and worse generalization. While recent work has identified that adaptive methods are basis-dependent (sensitive to coordinate rotations), the specific mechanism by which this dependence breaks the low-rank inductive bias in factored models has not been characterized. This paper addresses the question: Is there a criterion to determine whether an optimizer preserves or breaks the low-rank bias, and can this criterion sort deployed methods by their recovery performance?

Methodology and Theoretical Framework
The author introduces gauge equivariance as the central criterion for classifying optimizers. The loss function $L(U,V) = f(UV^\top)$ possesses a gauge symmetry: it is invariant under the transformation $(U,V) \mapsto (UQ, VQ)$ for any orthogonal matrix $Q \in O(k)$ . This symmetry implies that rotating the latent factor bases does not alter the model's predictions.

Classification of Update Rules:
- Equivariant Optimizers: An optimizer is gauge-equivariant if its update rule commutes with the gauge transformation. The paper proves that memoryless update rules are equivariant if and only if they take the form of a Gram-determined left preconditioner: $\Delta = H(GG^\top)G$ , where $H$ depends only on the Gram matrix $GG^\top$ .
- Non-Equivariant Optimizers: Optimizers that apply fixed entrywise nonlinearities to the gradient or optimizer state at initialization (e.g., standard Adam, RMSProp, Lion, SignSGD) break this symmetry. Specifically, Adam's coordinate-wise second-moment adaptation fails to commute with orthogonal rotations, leading to a "first-step defect" where the product matrix $W$ changes immediately upon rotating the initialization.
Transfer Theorem:
The author establishes a transfer theorem showing that optimizers with a shared scalar preconditioner (a single scalar denominator for all coordinates) correspond to gradient flow under a rescaled time parameterization. Consequently, theoretical results regarding gradient flow's low-rank bias (e.g., Gunasekar et al., 2017; Li et al., 2021) apply directly to these equivariant methods.
Spectral Schedule:
While equivariance is necessary for accessing the low-rank manifold, it is not sufficient for low-rank recovery. The specific behavior within the equivariant class is determined by the spectral schedule (how singular values grow). Aggressive equal-rate schedules (like Muon) can recover exact low-rank targets but may fail if the target has a significant "spectral tail" (energy outside the low-rank part), whereas greedy schedules (like GD) suppress these tails.

Key Contributions and Experiments

Optimizer Zoo Comparison: The author evaluated nine optimizers on an underdetermined matrix sensing task ( $40 \times 40$ , rank 3, $m=2 \times \text{dof}$ ) until full interpolation.
- Result: A clean separation emerged. Equivariant methods (GD, Scalar-Adam, Muon, Shampoo) achieved recovery errors $\le 0.286$ . Coordinate-wise methods (Adam, RMSProp, Lion, SignSGD, Adafactor) failed to reach the low-rank attractor, with errors $\ge 0.42$ .
- Muon's Performance: Muon achieved near-exact recovery ( $6.8 \times 10^{-6}$ error) on pure low-rank targets but degraded rapidly as the target's spectral tail increased, crossing over to GD's performance at approximately 4% tail energy.
The "Dial" Experiment: To isolate the cause of Adam's failure, the author constructed a one-parameter family of optimizers (Adam- $p$ ) that interpolates between standard Adam ( $p=1$ , coordinate-wise) and Scalar-Adam ( $p=0$ , shared scalar).
- Result: As $p$ decreased from 1 to 0, recovery error and effective rank improved monotonically. This confirmed that coordinate-wise anisotropy in the preconditioner is the specific mechanism destroying the low-rank bias, independent of momentum or adaptivity.
Attention Head Analysis: The paper demonstrates that individual attention heads in transformers share the same gauge symmetry ( $W_Q^\top W_K$ ).
- Result: Training two gauge-equivalent transformer copies (identical function, different bases) with Adam caused them to diverge structurally after a single step, resulting in a 56% relative Frobenius distance in their head invariants ( $W_Q^\top W_K$ ). Equivariant optimizers (SGD, Scalar-Adam) remained indistinguishable at machine precision. The author notes that while alignment techniques (like permutation re-basin) can reconcile the gauge, they cannot reconcile the invariants themselves if they differ; thus, basis-dependent training removes the agreement required for model merging, as no per-head rotation can close the gap between the divergent invariants.
Real-World Validation: On hyperspectral image completion tasks (Indian Pines, Pavia University) at matched training loss:
- GD achieved 43–44% lower held-out error than Adam at low sampling densities, corresponding to a significantly lower effective rank.
- Adam's test error increased as it fit deeper (overfitting the spectral tail), while GD's error decreased.
Algorithm Repair: The author applied their criterion to their previous work, FlowAdam. They identified that coordinate-wise velocity clipping broke equivariance. Replacing this with a global-norm clip restored the bias, improving recovery error by 37% and bringing the method closer to GD's performance.

Significance and Claims
The paper claims that basis choice is not a tuning detail but a fundamental decision about which interpolant the optimizer selects.

Mechanistic Explanation: It provides a testable, constructive criterion (gauge equivariance) to predict whether an optimizer will preserve the low-rank inductive bias of factored models.
Limitations: The author is modest about the scope. The bias is mild and beneficial primarily when the optimal regularization is low and aligned with the low-rank structure. In regimes requiring strong regularization or where per-coordinate adaptation is beneficial (e.g., stiff problems in PINNs), coordinate-wise methods may still outperform equivariant ones.
No Benchmark Claims: The work does not advocate for a new state-of-the-art optimizer but rather offers a mechanistic explanation for observed phenomena, clarifying why certain methods (like Adam) fail to recover low-rank structures in specific settings while others (like GD or Muon) succeed.

In summary, the paper argues that the "loss does not see the basis" (it is gauge-invariant), but "Adam does" (its updates are basis-dependent), and this dependence is the root cause of its failure to select low-rank interpolants in factored models.

Technical Summary: "The Loss Does Not See the Basis, But Adam Does"

More like this