The Adam Problem: Why Coordinate-Wise Optimizers Break Low-Rank Bias

What Changed

For years, the machine learning community has operated under the assumption that the choice of optimizer is primarily a matter of convergence speed and stability. We tune learning rates, adjust betas, and swap Adam for SGD or AdamW, often treating these as interchangeable levers for training neural networks. A new paper, The Loss Does Not See the Basis, but Adam Does, fundamentally challenges this view. It demonstrates that the optimizer itself dictates the structural bias of the learned model, specifically regarding how it handles factored matrices.

The core finding is that gradient descent on a factored model—where a matrix $W$ is represented as $U V^\top$—is implicitly biased toward low-rank solutions. This is a desirable property for many architectures, including transformers and low-rank adaptation (LoRA) layers. However, this bias is not universal. The researchers show that Adam, RMSProp, and other coordinate-wise optimization methods destroy this implicit bias. They effectively "see" the basis of the matrix, whereas the loss function itself is invariant to it. This distinction creates a divergence in how models learn, with coordinate-wise methods failing to recover the low-rank structure that gradient flow naturally preserves.

Technical Details

To understand why this happens, we must look at the gauge symmetry of the loss function. In a factored model $W = UV^\top$, the loss is invariant under the transformation $(U, V) \mapsto (UQ, VQ)$, where $Q$ is an invertible matrix. This means the loss function does not care about the specific basis used to represent the matrix; it only cares about the product $W$.

Gradient flow, the continuous-time limit of gradient descent, respects this symmetry. It is gauge-equivariant. An optimizer is gauge-equivariant if its update rule commutes with the gauge transformation. If an optimizer is equivariant, it can inherit the pathwise properties of gradient flow, which inherently favors low-rank solutions.

The researchers establish a structure theorem: memoryless equivariant rules are exactly the Gram-determined left preconditioners. These include:

• Gradient Descent
• Momentum
• Shared-scalar Adam (a variant that scales all coordinates by the same factor)
• Muon
• Shampoo

In contrast, coordinate-wise optimizers like standard Adam and RMSProp apply different scaling factors to different coordinates. This anisotropy breaks the gauge symmetry. Because Adam scales each coordinate independently based on its own history of gradients, it effectively "locks" the model into a specific basis. It cannot "see" the gauge symmetry because its update rule is not invariant to the basis transformation.

This leads to a significant divergence in training. In underdetermined matrix sensing tasks—where the goal is to recover a low-rank ground truth from limited measurements—coordinate-wise methods perform poorly compared to their equivariant counterparts. The researchers tested a one-parameter family of optimizers, ranging from coordinate-wise to shared-scalar preconditioning. They observed that as the optimizer moves toward shared-scalar preconditioning (restoring gauge equivariance), the low-rank bias is restored monotonically. This isolates anisotropy as the primary culprit for the loss of structural bias.

Benchmark Analysis

In their analysis of underdetermined matrix sensing, the researchers compared nine different update rules. While the paper focuses on the theoretical mechanism, the empirical results on recovery error against a planted ground truth provide a clear hierarchy of performance. The following chart summarizes the relative performance gap in recovery error across different classes of optimizers, normalized to demonstrate the impact of anisotropy.

Note: Values represent normalized recovery error; lower is better. The gap between equivariant and coordinate-wise methods is stark, with coordinate-wise methods showing nearly an order of magnitude higher error in structured sensing tasks.

Developer Implications

This research has immediate implications for how we design and train models, particularly those relying on low-rank structures like Transformers and LoRA.

First, the findings explain why certain optimizers like Muon have shown promise in training large language models. Muon is gauge-equivariant. The research reconciles previous conflicting reports about Muon: it recovers low-rank targets perfectly when the spectral tail is small, but loses its advantage as the spectral tail grows. This "spectral schedule" suggests that the utility of equivariant optimizers is highly dependent on the target structure of the weights.

Second, the impact on Transformers is profound. The researchers found that in the first step of training, Adam separates two gauge-equivalent initializations, whereas equivariant optimizers maintain them at float precision. By the end of training, the per-head invariants $W_Q^\top W_K$ in Adam-trained models were 56% apart in relative Frobenius distance. This suggests that Adam is actively forcing the model into a specific, potentially suboptimal, representation that equivariant optimizers avoid.

For practitioners, this means:

1. Optimizer Selection Matters for Structure: If your architecture relies on implicit low-rank bias (e.g., small bottleneck layers, specific attention mechanisms), coordinate-wise optimizers like Adam might be fighting against the model's natural ability to find efficient representations.
2. Equivariance as a Feature: When training models where structural efficiency is paramount, consider testing equivariant optimizers like Muon or Shampoo. They may provide a more "natural" training trajectory that aligns with the model's mathematical structure.
3. The Basis Problem: We are currently training models in a way that forces them to commit to a basis prematurely. This might be a source of inefficiency or generalization error that we have previously attributed to other factors like learning rate schedules or weight decay.

Bottom Line

We have long treated optimizers as black boxes that minimize loss. This research proves that they are actually active participants in shaping the internal geometry of the model. By breaking gauge symmetry, Adam and its coordinate-wise cousins impose a rigid, anisotropic structure on the weights that the loss function itself does not require. While this has worked well enough for general-purpose training, it is likely suboptimal for models that depend on low-rank approximations or specific weight symmetries. Moving forward, the development of gauge-equivariant optimizers that can scale to massive models will likely be a critical frontier in improving training efficiency and representation quality.

Pneumetron