Watch Your Step: Adam's Margin Maximization Bias Depends on the Learning Rate Schedule

Eitan Gronich; Gal Vardi

cs.LGNov 1, 2026

Watch Your Step: Adam's Margin Maximization Bias Depends on the Learning Rate Schedule

Authors: Eitan Gronich, Gal Vardi

Organizations: Weizmann Institute of Science

Abstract

We study the implicit bias of Adam and its similarity to sign gradient descent on homogeneous models in binary classification tasks. Perhaps surprisingly, we uncover that these depend on the learning rate schedule. While previous works studied strongly decaying learning rate schedules $η(t)=\mathrm o \left(t^{1/L-1}\right)$ , where $L$ is the homogeneity degree, and showed $\ell_\infty$ -margin-maximization biases, we focus on the more common complementary regime, including constant and weakly decaying schedules. For a polynomial schedule $t^α$ and homogeneity degree $L$ , we show that when $α> \frac 1L - \frac12$ , Adam has no trajectory-independent margin-maximization bias for deep linear networks (in the KKT sense): first, there is no non-trivial, locally Lipschitz, convex implicit regularizer in this setting, including any squared norm. Second, on width- $1$ networks, the result holds without requiring convexity, implying that any implicit regularizer for wide networks must be trivial when restricted to a width- $1$ network. Nevertheless, we positively characterize the implicit bias in this regime as maximizing a trajectory-dependent weighted- $\ell_2$ margin for any $L$ -homogeneous model. The fundamental mechanism behind our results is convergence of Adam's preconditioner to a trajectory-dependent vector, which also allows showing that the limiting parameter direction depends on initialization. In the intermediate regime $\frac 1L - 1 \leq α\leq \frac1L - \frac12$ we prove an $\ell_\infty$ bias, conditional on gradient directional convergence to a full-support vector. Of independent interest, we establish a loss decay rate for full batch Adam on general homogeneous models, whose leading order depends only on the second-moment parameter.

Figures & tables

Panel of Figure 1 — Figure 1: Initialization-dependent limiting direction forces constancy of $R$ in the fast regime. (a) Varying the initialization produces distinct margin-normalized limits on the hypersurface $\Sigma=\{\bm{\theta}\mid f(\mathbf{x}_{1};\bm{\theta})=1\}$ , due to large early-phase gradients. Since only $\mathbf{x}_{1}$ has an active margin there, KKT forces $\mathbf{k}\perp\Sigma$ for the subgradient $\mathbf{k}\in\partial R(\hat{\bm{\theta}})$ , implying $R$ is constant on $\Sigma$ . (b) Initialization-dependent parameter directions on the unit sphere, reached in experiments running Adam on the width- $1$ model ( 11 ) (with $L=2$ ), each marker representing a different initialization. Blue and orange markers denote the $u_{1}<u_{2}$ and $u_{2}<u_{1}$ sectors, respectively, obtained from two distinct yet symmetric datasets.

Panel of Figure 2 — Figure 2: Adam loss decay rate and margins on a 3-layer network with $\varepsilon=0$ and initialization scale $10^{-4}$ . (a) Late-phase slope of $-\ln\mathcal{L}$ on the $10,000$ last steps of each trajectory, and a reference at $-\ln(\beta_{2})/2$ . (b) The $\ell_{\infty}$ margin. (c)–(d) For $\alpha=0$ and $\alpha=-0.7$ , respectively, the weighted margin $\gamma_{\omega}$ using weights from one reference seed’s endpoint. In (a)–(b), lines are means across 5 seeds and shading denotes one standard error.

Panel of Figure 3 — Figure 3: Adam margins for $\varepsilon=0,10^{-8}$ on a 3-layer network with initialization scale $1$ . (a)–(b) $\ell_{\infty}$ and $\ell_{2}$ margins, respectively. Solid curves and plain bands correspond to $\varepsilon=0$ ; dashed curves and hatched bands correspond to $\varepsilon=10^{-8}$ . Curves show means across five seeds; shading denotes one standard error, with boundaries outlined. (c)–(d) For $\alpha=0$ and $\alpha=-0.7$ , respectively, the weighted margin $\gamma_{\omega}$ using weights from one reference seed’s endpoint.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4:

\ell_{\infty}

margins for

\alpha=-0.3

, for five seeds.

Figure 5: Relative update error for Adam with $\varepsilon=10^{-8}$ . The quantity $\overline{e_{\varepsilon}}(t)$ denotes the trailing 1,000-update moving average of $e_{\varepsilon}(t)$ . Curves show means across five seeds; shading denotes one standard error.

Figure 6: Endpoint-weighted margins for $\varepsilon=0$ in the long-horizon experiment. Rows correspond to learning-rate schedules $\eta_{t}\propto t^{\alpha}$ ; columns identify the reference seed whose endpoint metric is held fixed. Each panel shows all five trajectories under that schedule.

Figure 7: Endpoint-weighted margins for $\varepsilon=10^{-8}$ in the comparison experiment. Rows correspond to learning-rate schedules $\eta_{t}\propto t^{\alpha}$ ; columns identify the reference seed whose endpoint metric is held fixed. Each panel shows all five trajectories under that schedule.

Figure 8: Endpoint-weighted margins for $\varepsilon=0$ in the comparison experiment. Rows correspond to learning-rate schedules $\eta_{t}\propto t^{\alpha}$ ; columns identify the reference seed whose endpoint metric is held fixed. Each panel shows all five trajectories under that schedule.

Implicit Bias

Implicit bias studies how training algorithms and model structure favor particular solutions, representations, and generalization behaviors without explicit penalties.

See all 34 papers

Feb 18, 2026

The Implicit Bias of Adam and Muon on Smooth Homogeneous Neural Networks

We study the implicit bias of momentum-based optimizers on smooth homogeneous models. We show that momentum steepest descent algorithms like Muon (spectral norm), MomentumGD (

\ell_2

norm), and Signum (

\ell_\infty

norm) are approximate steepest descent trajectories under a decaying learning rate schedule, proving that these algorithms have a bias towards KKT points of the corresponding margin maximization problem. We extend the analysis to Adam (without the stability constant), which maximizes the

\ell_\infty

margin, and to Muon-Signum and Muon-Adam, which maximize a hybrid norm. Our experiments corroborate the theory and show that the identity of the margin maximized depends on the choice of optimizer. Overall, our results extend earlier lines of work on steepest descent in homogeneous models and momentum-based optimizers in linear models.

Eitan Gronich, Gal Vardi

Department of Computer Science and Applied Mathematics, Weizmann Institute of Science, Rehovot, Israel

Gradient Descent Adaptive Gradient Methods Algorithmic Bias

Ask PDF

Sep 27, 2026

On the Two Faces of Adam in Separable Linear Classification

We consider the behavior of deterministic, full-batch, bias-corrected Adam in separable linear classification with softmax parametrization under log-loss. In this setting, under a wide range of conditions Adam is known to approach max-norm-margin optimality when its stability constant

ε

is zero, while with a positive

ε

, it is known to approach Euclidean-margin optimality. Our main contribution is the quantitative description of Adam's behavior for small fixed positive

ε

. We give sufficient conditions under which an Adam-trained classifier nearly maximizes the max-norm margin before the updates become gradient-like. We also show that the classifier reaches a fixed target Euclidean margin only much later. Specifically, we show that for polynomially decreasing stepsizes with exponent

a

, where

1/3<a<1

, the updates become approximately proportional to the negative gradient after

Θ(\log(1/ε)^{1/(1-a)})

iterations. At that time, the classifier still nearly maximizes the max-norm margin. Reaching a fixed target Euclidean margin above that of every max-norm-optimal classifier, but below the optimum, is shown to require

ε^{-Θ(1)/(1-a)}

iterations. Under inverse-linear stepsize decay (

a=1

), the update transition takes polynomially many iterations, whereas reaching the target margin takes exponentially many. Experiments support these predictions. The later change in the classifier can improve or worsen generalization after training error reaches zero, connecting the analysis to grokking and its reverse.

Chen Fan, Csaba Szepesvári

University of Alberta

Implicit Bias Classification Adaptive Gradient Methods Gradient Descent Dynamics Linear Classification

Ask PDF

Mar 9, 2026

Minor First, Major Last: A Depth-Induced Implicit Bias of Sharpness-Aware Minimization

We study the implicit bias of Sharpness-Aware Minimization (SAM) when training

L

-layer linear diagonal networks on linearly separable binary classification. For linear models (

L=1

), both

\ell_\infty

- and

\ell_2

-SAM recover the

\ell_2

max-margin classifier, matching gradient descent (GD). However, for depth

L = 2

, the behavior changes drastically -- even on a single-example dataset. For

\ell_\infty

-SAM, the limit direction depends critically on initialization and can converge to

\mathbf{0}

or to any standard basis vector, in stark contrast to GD, whose limit aligns with the basis vector of the dominant data coordinate. For

\ell_2

-SAM, we show that although its limit direction matches the

\ell_1

max-margin solution as in the case of GD, its finite-time dynamics exhibit a phenomenon we call "sequential feature amplification", in which the predictor initially relies on minor coordinates and gradually shifts to larger ones as training proceeds or initialization increases. Our theoretical analysis attributes this phenomenon to

\ell_2

-SAM's gradient normalization factor applied in its perturbation, which amplifies minor coordinates early and allows major ones to dominate later, giving a concrete example where infinite-time implicit-bias analyses are insufficient. Synthetic and real-data experiments corroborate our findings.

Chaewon Moon, Dongkuk Si, Chulhee Yun

Kim Jaechul Graduate School of AI, KAIST · Mobilint, Inc.

Sharpness-Aware Minimization Implicit Bias Deep Linear Networks

Ask PDF

Abstract

Implicit Bias

Explore similar work

The Implicit Bias of Adam and Muon on Smooth Homogeneous Neural Networks

On the Two Faces of Adam in Separable Linear Classification

Minor First, Major Last: A Depth-Induced Implicit Bias of Sharpness-Aware Minimization