Adam or Gauss-Newton? A Comparative Study In Terms of Basis Alignment and SGD Noise

TL;DR

The paper investigates how the choice of basis for diagonal preconditioning and gradient noise from mini-batching influence the performance of Adam versus Gauss-Newton (GN) diagonal preconditioners. By separating basis effects (identity vs GN eigenbasis) from diagonal scaling (Adam vs GN$^{-1}$ or GN$^{-1/2}$) and analyzing linear and logistic models, the authors derive conditions under which GN is optimal and where Adam can outperform GN, especially under misaligned bases. They demonstrate that in full-batch settings, Adam can surpass GN and GN$^{-1}$, while in stochastic settings Adam behaves similarly to GN$^{-1/2}$, with equivalence holding across bases for linear regression under Gaussian inputs. These theoretical insights are supported by comprehensive experiments on convex and non-convex tasks, including MLPs and Transformers, confirming the practical relevance of basis alignment and gradient-noise considerations for diagonal preconditioners and suggesting avenues for designing auto-tuning optimizers that retain the favorable properties of Adam across batch regimes.

Abstract

Diagonal preconditioners are computationally feasible approximate to second-order optimizers, which have shown significant promise in accelerating training of deep learning models. Two predominant approaches are based on Adam and Gauss-Newton (GN) methods: the former leverages statistics of current gradients and is the de-factor optimizers for neural networks, and the latter uses the diagonal elements of the Gauss-Newton matrix and underpins some of the recent diagonal optimizers such as Sophia. In this work, we compare these two diagonal preconditioning methods through the lens of two key factors: the choice of basis in the preconditioner, and the impact of gradient noise from mini-batching. To gain insights, we analyze these optimizers on quadratic objectives and logistic regression under all four quadrants. We show that regardless of the basis, there exist instances where Adam outperforms both GN$^{-1}$ and GN$^{-1/2}$ in full-batch settings. Conversely, in the stochastic regime, Adam behaves similarly to GN$^{-1/2}$ for linear regression under a Gaussian data assumption. These theoretical results are supported by empirical studies on both convex and non-convex objectives.

Figures (10)

• Figure 1: Adam converges faster than GN with full batches, (Left) under the identity basis on a linear regression task with block-wise covariance, where GN fails to adapt to the problem curvature; and (Right) under the eigenbasis, for the reparameterized logistic regression task.
• Figure 2: Comparing GN power $p \in \{-\frac{1}{2}, -1\}$. On a regression task where GN$^{-1/2}$ leads to a more favorable condition number, GN$^{-1/2}$ converges faster than GN$^{-1}$ with both small and large batches.
• Figure 3: Learning from a random teacher, comparing Adam, $\text{GN}^{-1}$ and $\text{GN}^{-1/2}$ for the full $2 \times 2$ grid (\ref{['tab:basis-batch-grid']}).
• Figure 4: Sparse parity, comparing Adam, $\text{GN}^{-1}$ and $\text{GN}^{-1/2}$ for the full $2 \times 2$ grid (\ref{['tab:basis-batch-grid']}).
• Figure 5: Staircase, comparing Adam, $\text{GN}^{-1}$ and $\text{GN}^{-1/2}$ for the full $2 \times 2$ grid (\ref{['tab:basis-batch-grid']}). Staircase is a generalization of sparse parity (\ref{['fig:parity-2x2-plots']}).

Theorems & Definitions (15)

• Lemma 1
• Corollary 1
• Lemma 2
• Theorem 2
• Corollary 3
• Lemma : \ref{['lem:adam_gn_half_fisher']}, restated
• proof
• Corollary : \ref{['lem:adam_gn_half_update']}, restated
• proof
• Lemma : \ref{['lem:precond_loss_rate']}, restated