PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Implicit Bias of Adam

Matias D. Cattaneo, Jason M. Klusowski, Boris Shigida

arXiv:2309.00079v4cs.LGcs.AImath.OCstat.COstat.ML

TL;DR

Backward error analysis has explained finite-step implicit regularization for gradient descent, but its role in Adam and RMSProp was not well understood. The paper derives higher-order ODE approximations and finds that Adam's effect on a perturbed gradient one-norm depends on hyperparameters and training stage, typically hindering its reduction.

  • Problem

    How finite-step backward error analysis characterizes implicit regularization in Adam and RMSProp was missing from prior work, despite these optimizers' widespread use.

    Source: §1 · pp. 1–2

  • Method

    The paper uses backward error analysis to derive a global second-order-in-h ODE approximation for mini-batch Adam and analyze the full-batch dynamics.

    Source: §1 · p. 2, §2 · p. 3, §3 · p. 6

  • Results

    For small ε, Adam penalizes the perturbed one-norm when β > ρ, but typically hinders its decrease when ρ > β; the effect weakens late in training or for non-standard hyperparameters.

    Source: §1 · p. 2, §2 · p. 4

  • Takeaways & Limitations

    The analysis suggests that Adam's anti-regularization effect can disappear near an interpolating solution, consistent with better generalization in that regime.

    Source: §6 · p. 9

    The proposed explanation is incomplete for transformers; more work is needed to connect implicit bias to their training dynamics.

    Source: §6 · p. 9

Abstract

from arXiv · show

In previous literature, backward error analysis was used to find ordinary differential equations (ODEs) approximating the gradient descent trajectory. It was found that finite step sizes implicitly regularize solutions because terms appearing in the ODEs penalize the two-norm of the loss gradients. We prove that the existence of similar implicit regularization in RMSProp and Adam depends on their hyperparameters and the training stage, but with a different "norm" involved: the corresponding ODE terms either penalize the (perturbed) one-norm of the loss gradients or, conversely, impede its reduction (the latter case being typical). We also conduct numerical experiments and discuss how the proven facts can influence generalization.

1. Introduction

The introduction applies backward error analysis to Adam and RMSProp, extending work on finite-step ODE approximations beyond gradient descent. It previews a second-order ODE approximation and a hyperparameter- and training-stage-dependent implicit bias, typically anti-regularization of a perturbed one-norm.

1.1. Related Work

  • Earlier backward-error analyses cover full-batch and mini-batch gradient descent, momentum, and related regularization interpretations.The introduction also surveys research on implicit bias and adaptive-method generalization.Source: §1 · p. 1, §1.1 · p. 3

1. Introduction

  • Prior work identifies backward error analysis for Adam and RMSProp as an open gap, which this paper addresses for mini-batch and full-batch settings.Earlier work had called for applying the approach to adaptive optimizers and noted that their implicit regularization remained poorly understood.Source: §1 · pp. 1–2
  • The paper provides a global second-order-in-h ODE approximation to mini-batch Adam and a linear-in-h correction for its full-batch signGD limit.A similar RMSProp result is placed in the appendix.Source: §1 · p. 2
  • Adam typically impedes reduction of the perturbed one-norm of loss gradients, except under unusual hyperparameters, large ε, or very late training.The paper contrasts this anti-regularization with gradient descent's implicit penalty on the two-norm.Source: §1 · p. 2
  • CIFAR-10 experiments with full-batch Adam find worse generalization where the theory predicts stronger anti-regularization, across ResNets, CNNs, and Vision Transformers.The authors propose this pattern as a possible explanation for reported poor generalization of adaptive gradient algorithms.Source: §1 · p. 2

1.1. Related Work

  • Unlike stochastic differential-equation approximations that take β and ρ toward 1, this analysis is deterministic, keeps them fixed, and goes one order further in h.The paper distinguishes its generic-loss perspective from work focused on sign-based methods and particular model classes.Source: §1.1 · pp. 2–3

1.2. Notation

  • The notation defines minibatch losses E_k, full-batch loss E, and the two-norm, one-norm, and perturbed one-norm conventions.The perturbed one-norm becomes the ordinary one-norm when ε is zero.Source: §1.2 · p. 3
  • The paper defines Adam's hyperparameters and places ε inside the square root; original outside-root variants require gradients to stay away from zero and are analyzed only for the initial learning period.The main theorem covers the whole period for the inside-root definition.Source: §1.2 · p. 3

2. Implicit Bias of Full-Batch Adam: an Informal Summary

For full-batch Adam, a modified ODE tracks the iterates with error of order h^2 over any fixed time horizon. Its correction can penalize or hinder reduction of a gradient norm, depending on hyperparameters and training stage; typically it hinders reduction of the perturbed one-norm.

2. Implicit Bias of Full-Batch Adam: an Informal Summary

  • Over any fixed time horizon, the full-batch Adam iterates stay within Ch^2 of the modified ODE trajectory under the stated smoothness assumption.The loss is assumed to have bounded partial derivatives through fourth order.Source: §2 · p. 3
  • For small ε, β > ρ penalizes the perturbed one-norm, whereas the typical ρ > β case hinders its reduction.When √ε is small compared with the gradient components, the perturbed one-norm is close to the usual one-norm.Source: §2 · p. 4
  • When ε is large relative to gradient components, the correction’s dependence on ρ cancels and the perturbed one-norm approaches the squared two-norm, up to a constant.This is the other extreme case, associated with late training or non-standard hyperparameters.Source: §2 · p. 4
  • Because ε is usually small, Adam is typically in the first regime during most training, and larger ρ relative to β strengthens the anti-regularization effect.The identified bias is anti-regularization except in unusual cases such as β ≥ ρ, large ε, or very late training.Source: §2 · p. 4
  • The authors suggest this bias may partly explain poorer generalization and sharper minima, while noting that sharpness and generalization are not clearly linked.The theory identifies anti-penalization of an approximation to ℓ∞-sharpness; the reported experiments are consistent with it.Source: §2 · p. 4
  • The overview also covers RMSProp by setting β = 0.The formal result is given in Theorem C.4.Source: §2 · p. 4

3. Main Result: ODE Approximating Mini-Batch Adam

Theorem 3.1 gives a piecewise ODE approximating mini-batch Adam, with a global error of order h^2 over any fixed time horizon. Its derivation uses backward error analysis, and in full-batch training the ODE simplifies at large iteration numbers.

3. Main Result: ODE Approximating Mini-Batch Adam

  • The piecewise ODE solution stays within Ch^2 of Adam's iterates through any fixed horizon T, with C depending on T, ρ, β, and ε.The result assumes the mini-batch losses are four times continuously differentiable with bounded partial derivatives through order four.Source: §3 · p. 5
  • In the full-batch case, the ODE terms rapidly simplify to the forms described by equations (2) and (3) as the iteration number grows.The full-batch setting takes each mini-batch loss to be the same loss E.Source: §3 · p. 5
  • The derivation constructs the update expansions from exponentially weighted gradient and Hessian terms before converting them into a piecewise ODE.The method squares an expansion, weights it by the ρ-based exponential average, and uses the resulting expansions in the backward error analysis.Source: §3 · p. 6
  • Backward error analysis derives the ODE by expanding the update in h and matching terms to obtain its correction to the leading dynamics.The derivation expands the vector field as f(θ) + hf1(θ) + O(h^2), then equates powers of h.Source: §3 · p. 6

4. Illustration: Simple Bilinear Model

In the bilinear model, Adam's trajectory shifts toward or away from regions of smaller gradient one-norm as β, ρ, and the learning rate change. The learning-rate effect depends on the relative sizes of β and ρ.

4. Illustration: Simple Bilinear Model

  • Increasing β moves Adam toward smaller gradient one-norm, while increasing ρ moves it toward larger one-norm when ε is sufficiently small.The loss is minimized on θ1θ2 = 3/4.Source: §4 · p. 6, Figure 2 · p. 7
  • Increasing the learning rate moves Adam toward smaller gradient one-norm when β exceeds ρ, but toward larger one-norm when ρ exceeds β.The figure caption specifies the smaller-one-norm direction when β is significantly larger than ρ.Source: §4 · p. 6, Figure 2 · p. 7
  • Figure 2 marks the one-norm minimizers on loss level sets and shows trajectories starting from (2.8, 3.5).The minimizers are dashed; their limit point is marked by a cross on the level sets 4θ1θ2 − 3 = c.Source: Figure 2 · p. 7

5. Numerical Experiments

The experiments report close agreement between Adam and its ODE approximations, then examine how hyperparameters relate to the perturbed gradient one-norm and test accuracy. The reported norm and accuracy trends align with the theory's predicted anti-regularization pattern.

5. Numerical Experiments

  • Over 100 epochs on truncated MNIST, the Adam trajectory remains close in infinity norm to the piecewise ODE solution.The experiment uses full-batch Adam and a roughly 10^5-parameter fully connected network.Source: §5 · p. 7, Figure 3 · p. 7
  • In the stable oscillation regime, increasing ρ appears to increase the perturbed one-norm and decrease test accuracy.The experiments focus on hyperparameters that are not too far apart, excluding the divergence regime and reducing large spikes in the spike regime.Source: §5 · p. 8
  • As ρ increases, ResNet-50's perturbed norm rises and test accuracy falls; as β increases, the norm falls and accuracy rises.These full-batch CIFAR-10 results align with the theoretical prediction; similar patterns are reported for CNNs and Vision Transformers.Source: Figure 4 · p. 8, Figure 5 · p. 8, §5 · p. 8
  • During most of training, larger β corresponds to a smaller perturbed one-norm; norm peaks are higher with smaller β and larger ρ.The norm decreases, rises, then decreases again before flattening; the reported pattern matches the correction term's predicted impediment to norm reduction.Source: §5 · p. 8, Figure 6 · p. 8

6. Limitations and Future Directions

The paper flags assumptions that may be uncertain near optimality, leaves mini-batch first-order behavior for future work, and notes limits on interpreting Adam’s generalization and small-ε approximations.

6. Limitations and Future Directions

  • Near optimality, edge-of-stability behavior can make it unclear whether the required higher-order derivative bounds are safe to assume.The authors note that large-batch methods may operate near or at the edge of stability, where the Hessian’s largest eigenvalue can be large.Source: §6 · p. 8
  • The paper leaves qualitative analysis of Theorem 3.1’s first-order terms in the mini-batch setting for future work.Source: §6 · p. 9
  • Adam does not always generalize worse than SGD; the authors say their explanation remains incomplete for transformer training dynamics.They note that Adam often outperforms SGD for transformers and call for more work connecting implicit bias to transformer training dynamics.Source: §6 · p. 9
  • As ε approaches zero, the theorem’s constant diverges, and later-stage trajectory closeness may require small, suboptimal learning rates.The authors distinguish this possibility from the initial learning period, where they say it is not a problem.Source: §6 · p. 9

A. Overview

  • The appendix collects omitted proofs and details, including additional numerical-experiment details.Source: A. Overview · p. 12

A.1. Notation

  • The notation uses E_k(θ) for the kth minibatch loss and E(θ) for the full-batch loss, with indexed nabla denoting partial derivatives.Source: §A.1 · p. 12

B. RMSProp with ε Outside the Square Root

B.1. RMSProp with ε outside: full-batch

  • For full-batch RMSProp with ε outside the square root, the ODE terms simplify at large iteration numbers when ε is small.Source: §B.1 · p. 15

C. RMSProp with ε Inside the Square Root

These appendices state local and global error bounds for RMSProp and Adam under their respective assumptions, and describe a large-iteration simplification for full-batch RMSProp with ε inside the square root.

C. RMSProp with ε Inside the Square Root

C.1. RMSProp with ε Inside: Full-Batch

  • For full-batch RMSProp with ε inside the square root, the terms simplify rapidly at large iteration numbers.Source: §C.1 · p. 17

D. Adam with ε Outside the Square Root

E. Adam with ε Inside the Square Root

F. Bounding the Derivatives of the ODE Solution

Under Assumption B.2, the appendix establishes uniform bounds on the first, second, and third derivatives of the piecewise ODE solution’s coordinates.

F. Bounding the Derivatives of the ODE Solution

G. Proof of Theorem B.3

The appendix proves the RMSProp local error bound through preparatory lemmas and mean-value arguments, and describes the models used in experiments and trajectory plots.

G. Proof of Theorem B.3

H.1. Models

  • The CIFAR ResNet models modify the first convolution and remove max pooling, then use global average pooling and a 10- or 100-way output layer.Source: §H.1 · p. 38
  • The MLP used to show trajectory closeness has two 32-unit GeLU layers followed by a 10-unit fully connected layer.Source: §H.1 · p. 38

H.2. Data Augmentation

The experiments describe CIFAR-10 data augmentation and training details, including how training was stopped and how batches were handled when the full dataset did not fit in GPU memory.

H.2. Data Augmentation

  • Each training pass pads 32 × 32 images to 40 × 40, randomly crops them back to 32 × 32, and applies a random horizontal flip with probability 0.5.Images are also centered per pixel and divided by their standard deviation.Source: §H.2 · p. 38

H.3. Experiment Details

  • Training stops when train accuracy is near-perfect and test accuracy no longer significantly improves; reported maximal test accuracy is the final value reached.The maximal perturbed one-norm is taken at a post-initial-fall peak of its curve.Source: §H.3 · p. 38
  • Because the full dataset does not fit in GPU memory, gradients from 100 ghost batches are accumulated before each optimization step.This uses ghost batch normalization rather than full-dataset batch normalization.Source: §H.3 · p. 38

H.4. Additional Evidence

The additional CIFAR-10 experiments report the same hyperparameter-linked norm and test-accuracy patterns for a simple CNN and a Vision Transformer as for the main-paper results.

H.4. Additional Evidence

  • The simple CNN robustness check uses three convolutional and max-pooling blocks followed by flattening and dense layers.The listed convolutional block channel widths are 32, 64, and 128.Source: §H.4 · pp. 38–39
  • The same conclusions are reported for a Vision Transformer using SimpleViT with 4 × 4 patches, six blocks, 16 heads, embedding size 512, and MLP dimension 1024.The architecture follows Andriushchenko et al. (2023).Source: §H.4 · p. 39
  • For a simple CNN, increasing ρ raises the perturbed one-norm and lowers test accuracy, while increasing β lowers the perturbed one-norm and raises test accuracy.Figure 7 averages five runs across initialization seeds; Figure 8 averages three.Source: Figure 7 · p. 39, Figure 8 · p. 40

I. Adam with ε Inside the Square Root: Informal Derivation

The informal derivation uses Taylor expansions, coefficient matching, and summation steps to obtain the stated modified equation for each time interval.

  • The derivation replaces t_n−1 with t_n in h-terms when that replacement affects only higher-order terms, then sums the resulting equations.Source: §I · p. 41
  • It disregards unbounded n−k terms on the stated basis of exponential averaging, then sums over k.Source: §I · pp. 41–42
  • Result I.2 states that the modified equation for t_n ≤ t < t_n+1 is (32).Source: §I · p. 43
  • Taylor expansion and matching coefficients of powers of h in (64) and (65) provide the next derivation step.Source: §I · p. 44
Loading 2309.00079v4…