On the Two Faces of Adam in Separable Linear Classification
TL;DR
This work gives sufficient conditions under which an Adam-trained classifier nearly maximizes the max-norm margin before the updates become gradient-like, and shows that for polynomially decreasing stepsizes with exponent \(a), the updates become approximately proportional to the negative gradient after $\Theta(1/\epsilon)^{1/(1-a)})$ iterations.
Abstract
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 $\epsilon$ is zero, while with a positive $\epsilon$, it is known to approach Euclidean-margin optimality. Our main contribution is the quantitative description of Adam's behavior for small fixed positive $\epsilon$. 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 $\Theta(\log(1/\epsilon)^{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 $\epsilon^{-\Theta(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.