Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold

AI summary

Overview

Research area: Optimization theory for deep learning — specifically the implicit bias of adaptive gradient methods (AGMs), with Adam as the centerpiece, compared against Stochastic Gradient Descent (SGD).

Technical level: Advanced. The paper works with stochastic differential equations (SDEs), Riemannian-style manifold projections, Hessian spectra, and Itô calculus. The high-level conclusions are accessible, but the derivation machinery is not.

Scope (one sentence): The paper derives a continuous-time "slow SDE" for Adam and a broader class of adaptive optimizers near the manifold of minimizers, proving that Adam implicitly reduces a sharpness measure tr(Diag(H)^{1/2}) rather than the tr(H) measure associated with SGD, and demonstrating when that difference helps or hurts generalization.

What This Paper Is About

Training a neural network can reach zero training error through many different parameter settings, and which one an optimizer lands on affects test performance. Prior theory showed that SGD, once it reaches the manifold of minimizers, drifts around that manifold and is implicitly biased toward flatter solutions by minimizing the trace of the Hessian, tr(H). The goal of this paper is to determine what Adam does instead, since Adam — not SGD — is the default optimizer in modern deep learning, and almost nothing rigorous was known about its implicit bias.

Key Contributions

1. A slow SDE for adaptive gradient methods. The authors generalize the slow SDE framework from SGD to a general class of AGMs (including Adam, RMSProp, Adam-mini, Adalayer, and Shampoo). Their SDE approximates the dynamics of these optimizers near the minimizer manifold Γ and shows they behave like an adaptive semi-gradient descent minimizing a unique, optimizer-specific sharpness measure.

2. A formal SGD-vs-Adam separation. In the well-studied setting of overparameterized models trained with label noise, the paper proves that SGD's implicit regularizer is tr(H), whereas Adam's is tr(Diag(H)^{1/2}), where H is the Hessian.

3. A demonstrated benefit and a demonstrated failure case. In sparse linear regression with diagonal linear networks, the different regularizer leads Adam to better sparsity and generalization than SGD, recovering the sparse ground truth with less data. In deep matrix factorization, Adam does not outperform SGD, likely because minimizing tr(Diag(H)^{1/2}) does not favor low-rank solutions.

4. Technical tools of independent interest. A manifold projection operator tailored to AGMs (a "preconditioner flow projection"), plus a high-probability convergence analysis for AGMs under Polyak-Łojasiewicz conditions that directly bounds L(θ_k) − L* to o(1), rather than bounding only the average of gradient norms as most prior Adam convergence work does.

Main Findings

• Adam minimizes a different sharpness measure than SGD. Under label noise with ℓ₂ loss, the implicit regularizer of Adam is tr(Diag(H)^{1/2}) — the trace of the matrix square root of the Hessian's diagonal — while SGD's is the full tr(H).

• Adam behaves as adaptive semi-gradient descent on the manifold. The drift term of the slow SDE for AGMs acts as a negative semi-gradient of an adaptive loss function, meaning Adam takes semi-gradients on the manifold and treats the preconditioner-dependent noise term as fixed, ignoring its dependence on the parameters.

• The dynamics have two distinct phases. Optimizers take Õ(η^{-1}) steps to converge onto the minimizer manifold, then move along the manifold for O(η^{-2}) steps, during which slower implicit-regularization dynamics dominate. Prior conventional SDE approximations track the first phase and cannot bound error over the second.

• Sharpness reduction does not always mean better generalization. The sparse-linear-regression result shows Adam's regularizer favoring sparsity; the deep-matrix-factorization result shows it does not favor low-rank structure, so Adam's advantage is problem-structure dependent.

• The effect is shared across adaptive optimizers. RMSProp, Adam-mini, Adalayer, Shampoo, Adafactor, and AdaSGD all fit the framework, offering a unified account of how adaptive methods reduce sharpness.

• The result is rigorous where prior Adam analysis was not. Earlier work by Ma et al. (2023) generalized slow SDEs to Adam only for a two-dimensional loss and used a quasistatic approximation lacking full rigor; Cattaneo et al. (2024) analyzed full-batch Adam approximately following an ODE that anti-regularizes sharpness when β₁ < β₂. This paper covers the O(η^{-2}) horizon with gradient noise and reaches the complementary conclusion that Adam regularizes sharpness in the long term.

• A toy illustration. In a two-parameter noisy elliptical loss with label noise δ drawn uniformly from {−0.5, 0.5}, SGD converges to the flattest minima while Adam reduces sharpness too but converges to a different, sparser minimizer closer to the axis.

Methodology in Plain English

The analysis starts from a structural observation: for an overparameterized model, the set of parameter values achieving zero training loss is typically not a scattering of isolated points but a connected surface, or manifold — call it Γ. When the loss is already small, the optimizer's fast convergence dynamics fade out and all that remains is a slow drift along this surface.

Instead of trying to track every individual parameter update, the authors project the optimizer's updates onto the manifold and study only the projected motion, using a stochastic differential equation. This is the "slow SDE" approach: peel off the convergence dynamics, keep the noise-driven motion on the surface, and let the learning rate go to zero. The drift term of the resulting SDE reads as a gradient pointing downhill on a sharpness landscape, but only a partial ("semi-") gradient, because the noise covariance that shapes the landscape is treated as constant.

The technical obstacle is that Adam's per-coordinate scaling changes what counts as a tangent direction. The authors resolve this by defining a preconditioner flow projection — a projection built for the specific preconditioner matrix that the optimizer uses — so the SDE can describe how Adam moves on Γ. They then specialize the general AGM result to the label-noise setting, where the Hessian replaces the general noise covariance, and compute the resulting regularizer explicitly. The theoretical prediction is then checked empirically in sparse linear regression with diagonal linear networks, using an "AdamE-λ" variant introduced in the paper as a tool to tune Adam's implicit bias.

Why This Matters

Impact on research. The paper answers a long-open question — what Adam is implicitly biased toward — with a rigorous, quantitative answer rather than a heuristic. It reframes optimizer choice as a choice of which sharpness measure gets minimized, and by extending the slow SDE beyond SGD, it gives the theory community a template for analyzing any optimizer that fits the general AGM form. It also clarifies an apparent contradiction with prior work: conventional, short-horizon SDE analysis and long-horizon slow SDE analysis can give different signs for what Adam does to sharpness.

Real-world applications (as studied or implied in the paper):

• Sparse linear regression with diagonal linear networks — the setting where Adam's tr(Diag(H)^{1/2}) bias buys better sparsity and generalization, recovering sparse ground truth with less data.
• Language model training — the paper cites BERT, Transformer, and GPT, all of which use β₁ = 0.9, as within its assumptions (β₁ ≤ 0.9).
• Computer vision training — pix2pix (β₁ = 0.5), U-Net and ViT (β₁ = 0.9) are cited as standard practice consistent with the paper's assumptions.
• Deep matrix factorization — a cautionary application where the same implicit bias fails to help, showing that adaptive optimization is not uniformly better.

Industry relevance. Adam and its descendants dominate practical training at scale, and the paper's framework covers memory- and compute-efficient variants such as Adam-mini, Adalayer, and Shampoo, all of which are used to cut training costs. A theory that predicts which solution an adaptive optimizer will select — and when that selection is favorable — is directly relevant to choosing optimizers, tuning β₁ and β₂, and deciding whether extra sparsity or low-rank structure is desirable for a given workload.

Future Directions

• Designing optimizers from the sharpness measure backward. Since the framework characterizes which sharpness measure each AGM reduces, a natural next step is constructing optimizers that target a chosen regularizer (for example, one that favors low-rank solutions where tr(Diag(H)^{1/2}) does not).

• Understanding the failure of Adam in deep matrix factorization. The paper attributes the lack of improvement to tr(Diag(H)^{1/2}) not favoring low-rank solutions, but a more complete characterization of when each measure helps — and how it relates to the parameterization — is left open.

• Extending the slow SDE to optimizers outside the general AGM class. Weight decay, bias correction, and update clipping are explicitly excluded or simplified in the formulation; incorporating them would bring the analysis closer to deployed AdamW and related variants.

• Closing the gap with short-horizon analyses. The paper presents its long-horizon, sharpness-regularizing conclusion as complementary to Cattaneo et al. (2024)'s short-horizon anti-regularization result when β₁ < β₂; a unified account of both regimes remains to be built.

• Empirical validation at scale. The empirical verification reported is in sparse linear regression with diagonal linear networks; broader tests on the large models cited in the assumptions are an obvious open direction.

Target Audience

Readers who benefit most are optimization theorists and machine learning researchers studying implicit bias, sharpness, and generalization; graduate students comfortable with SDEs, manifolds, and Hessian analysis; and practitioners or engineers working on optimizer design and large-scale training who want a principled account of why Adam and its variants reach different solutions than SGD. The paper assumes familiarity with the slow SDE literature (Li et al. 2021b; Gu et al. 2023a) and with standard smoothness assumptions, so it is not an entry-level read; the abstract and the noisy-ellipse example are the most accessible entry points.

Authors’ abstract

Despite the popularity of the Adam optimizer in practice, most theoretical analyses study Stochastic Gradient Descent (SGD) as a proxy for Adam, and little is known about how the solutions found by Adam differ. In this paper, we show that Adam implicitly reduces a unique form of sharpness measure shaped by its adaptive updates, leading to qualitatively different solutions from SGD. More specifically, when the training loss is small, Adam wanders around the manifold of minimizers and takes semi-gradients to minimize this sharpness measure in an adaptive manner, a behavior we rigorously characterize through a continuous-time approximation using stochastic differential equations. We further demonstrate how this behavior differs from that of SGD in a well-studied setting: when training overparameterized models with label noise, SGD has been shown to minimize the trace of the Hessian matrix, $\tr(\mH)$, whereas we prove that Adam minimizes $\tr(\Diag(\mH)^{1/2})$ instead. In solving sparse linear regression with diagonal linear networks, this distinction enables Adam to achieve better sparsity and generalization than SGD. Finally, our analysis framework extends beyond Adam to a broad class of adaptive gradient methods, including RMSProp, Adam-mini, Adalayer and Shampoo, and provides a unified perspective on how these adaptive optimizers reduce sharpness, which we hope will offer insights for future optimizer design.