Adam Optimizer Loses Gradient Descent's Low-Rank Bias

The Invariance Problem in Factored Models

In machine learning, particularly with factored models like W = UV^T, the optimization landscape often exhibits a critical property: invariance to rotations. This means that if you rotate both U and V by the same orthogonal matrix Q, such that W = (UQ)(VQ)^T, the resulting matrix W remains unchanged. Crucially, standard Gradient Descent (GD) respects this property. Its updates naturally preserve this invariance, ensuring that the optimization process does not arbitrarily favor one basis over another for the factors.

However, a new analysis suggests that adaptive optimizers, specifically Adam, fail to maintain this fundamental invariance. The core of the issue lies in Adam's mechanism for tracking per-coordinate second moments. While effective for accelerating convergence in many scenarios, this adaptive scaling mechanism is inherently dependent on the specific basis in which the factors U and V are represented. When Adam updates U and V, its second moment estimates are sensitive to the chosen basis, leading to a loss of the desirable rotational invariance that GD preserves.

This difference is not merely theoretical. The paper posits that this single property—whether an optimizer retains or loses GD's implicit low-rank bias—is a key determinant of its behavior in underdetermined settings. In such scenarios, where the number of parameters exceeds the intrinsic dimensionality of the data, maintaining this bias is crucial for generalization and for finding meaningful solutions.

Empirical Evaluation of Optimizers

To test this hypothesis, the researchers conducted an empirical evaluation across nine different update rules. The experiments focused on underdetermined matrix sensing problems, a classic scenario where low-rank structures are prevalent and important. To ensure a fair comparison, all methods were evaluated at matched training loss levels. This critical step prevents any observed differences from being attributed to one method simply underfitting the data compared to another.

The results revealed a clear dichotomy. Two distinct groups of optimizers emerged based on their performance and their adherence to the low-rank bias property. One group, including standard Gradient Descent, consistently demonstrated the desired behavior, effectively leveraging the underlying low-rank structure of the problem. The other group, notably including Adam, showed a marked departure. These optimizers, by failing to respect the rotational invariance, tended to lose the beneficial low-rank bias inherent in the problem formulation and preserved by GD.

The implication is significant: while Adam might converge faster or achieve a similar loss on a given epoch, its internal mechanism means it is not finding the same types of solutions that GD would. In settings where the low-rank structure is paramount for interpretability, generalization, or avoiding spurious correlations, Adam's approach could be detrimental. It's akin to a carpenter using a power drill that automatically adjusts its torque based on the type of wood, but in doing so, it loses the ability to feel the grain and might inadvertently splinter the wood in ways a hand drill wouldn't. The optimizer is optimizing, but it's optimizing for a different objective, one that is sensitive to the arbitrary choice of coordinate system.

What Does This Mean for Model Training?

The findings challenge the ubiquitous use of Adam as a default optimizer, especially in domains where implicit regularization properties are important. While Adam's adaptive learning rates can be beneficial for navigating complex loss landscapes and achieving faster convergence in terms of raw loss reduction, its disregard for the rotational invariance of factored models means it can destroy the very structure that makes a model learnable and generalizable in underdetermined settings.

For practitioners, this suggests a need for more careful consideration when selecting an optimizer. If your model relies on factored representations (e.g., matrix factorization, low-rank approximations in neural networks) and you're operating in an underdetermined regime, Adam might be actively working against you by corrupting the implicit low-rank bias. In such cases, revisiting simpler optimizers like SGD (Stochastic Gradient Descent) or SGD with momentum might yield better generalization performance, even if raw convergence speed is slightly slower.

The research highlights a subtle but important aspect of optimization: the algorithm's inductive bias. GD, by respecting the rotational invariance, implicitly favors solutions that are robust to the choice of basis, a property often aligned with finding simpler, more generalizable models. Adam, by sacrificing this invariance for adaptive scaling, introduces a bias that can be problematic. The surprise here is not that optimizers differ, but that a widely adopted, seemingly powerful optimizer like Adam fundamentally breaks a desirable geometric property that simpler methods preserve, leading to potentially suboptimal solutions in specific, yet common, machine learning tasks.

What remains an open question is the precise threshold at which Adam's adaptive scaling becomes detrimental. Is there a specific rank deficiency or data sparsity level where the loss of this bias becomes statistically significant and practically harmful? Further research could explore how to modify Adam or develop new adaptive optimizers that can retain this crucial invariance while still benefiting from adaptive learning rates.