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

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.

• The paper demonstrates that Adam's auto-tuning outperforms GN preconditioners in misaligned (identity) bases, especially in full-batch scenarios.
• It shows that in stochastic regimes, Adam and GN⁻¹/² become nearly equivalent, backed by theoretical bounds linking empirical and true Fisher matrices.
• Empirical analyses on synthetic and real-world tasks highlight scalability challenges and suggest directions for hybrid optimizer designs.

Comparative Analysis of Adam and Gauss-Newton Diagonal Preconditioners: Basis Alignment and SGD Noise

Introduction and Motivation

This paper presents a rigorous comparative study of two classes of diagonal preconditioners for deep learning optimization: those based on Adam (empirical Fisher) and those based on the Gauss-Newton (GN) matrix. The analysis is structured around two principal axes: (1) the choice of basis in which the preconditioner is applied (identity vs. GN eigenbasis), and (2) the impact of gradient noise, as modulated by batch size (full-batch vs. stochastic regime). The work provides both theoretical and empirical evidence that the effectiveness of these preconditioners is highly sensitive to both basis alignment and the stochasticity of the gradient estimates.

Theoretical Framework

The study formalizes preconditioned updates as:

$$\theta_{t+1} = \theta_t - \eta \cdot (U D^{\alpha} U^\top) g_t$$

where $U$ is the basis (identity or GN eigenbasis), $D$ is the diagonal preconditioner (Adam or GN), and $\alpha \in \{-1, -1/2\}$ is the exponent controlling the strength of preconditioning.

Basis Choice

• GN Eigenbasis: GN-based preconditioning ($\GN^{-1}$) achieves optimal convergence for quadratic objectives when applied in the GN eigenbasis, both in full-batch and stochastic regimes.
• Identity Basis: When the basis is misaligned (identity), GN-based preconditioners lose their curvature-adaptive properties, and Adam can outperform both $\GN^{-1}$ and $\GN^{-1/2}$ due to its auto-tuning effect.

Gradient Noise

• Full-Batch: Adam's adaptivity allows it to auto-tune learning rates to local curvature, outperforming GN-based methods when the basis is not aligned.
• Stochastic Regime: For linear regression with Gaussian data, Adam and $\GN^{-1/2}$ become equivalent up to a scaling factor, regardless of basis. This equivalence is formalized via bounds relating the empirical Fisher and true Fisher matrices.

Empirical Results

The theoretical claims are substantiated by extensive experiments on both synthetic and real-world tasks, including linear and logistic regression, MLPs on parity and staircase functions, CIFAR10 classification, and Transformer-based selection tasks.

Full-Batch, Identity Basis

Adam demonstrates rapid convergence on block-structured covariance problems where GN fails to adapt due to basis misalignment.

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.

GN Power Comparison

Empirical results show that $\GN^{-1/2}$ can outperform $\GN^{-1}$ in scenarios where the condition number of the preconditioned Hessian is more favorable for the square-root scaling, even in full-batch settings.

Figure 2: Comparing GN power $\power \in \{-\frac{1}{2}, -1\}$; GN$^{-1/2}$ converges faster than GN$^{-1}$ when the condition number is more favorable.

Non-Convex and Real-World Tasks

Across MLPs trained on random teacher networks, sparse parity, staircase, and CIFAR10, Adam and $\GN^{-1/2}$ closely track each other in the stochastic regime, and Adam is competitive or superior to GN-based methods when the basis is not aligned.

Figure 3: Learning from a random teacher, comparing Adam, $\GN^{-1}$, and $\GN^{-1/2}$ across batch size and basis choices.

Figure 4: Sparse parity, comparing Adam, $\GN^{-1}$, and $\GN^{-1/2}$ across batch size and basis choices.

Figure 5: Staircase, comparing Adam, $\GN^{-1}$, and $\GN^{-1/2}$ across batch size and basis choices.

Figure 6: CIFAR10, comparing Adam, $\GN^{-1}$, and $\GN^{-1/2}$ across batch size and basis choices.

Transformer Experiments

On selection-based regression tasks with Transformers, Adam outperforms GN-based methods even under the GN eigenbasis with full batches, consistent with the theoretical lower bounds for GN in logistic-like settings.

Figure 1: Adam converges faster than GN with full batches, (Right) under the eigenbasis for the reparameterized logistic regression task.

Strong Claims and Contradictory Findings

• Adam Outperforms GN in Full-Batch, Misaligned Basis: The paper provides explicit constructions where Adam's auto-tuning yields faster convergence than GN, contradicting the notion that curvature-based preconditioning is always superior.
• Adam and $\GN^{-1/2}$ Equivalence in Stochastic Regime: Theoretical bounds and empirical results show that Adam and $\GN^{-1/2}$ are nearly indistinguishable in the small-batch regime for linear regression, regardless of basis.
• GN Suboptimality in Logistic Regression: For non-convex objectives (e.g., logistic regression with weight-tied two-layer networks), Adam can converge faster than GN even under the correct eigenbasis, with GN suffering polynomial-in-dimension slowdowns due to global convergence constraints.

Implementation and Practical Considerations

Algorithmic Structure

The preconditioned optimizer is implemented as follows:

for t in range(T): grad = compute_gradient(theta, batch) GN_matrix = estimate_GN_matrix(theta, batch_H) if basis_type == 'identity': U = np.eye(dim) elif basis_type == 'eigenbasis': U = eigenvectors(GN_matrix) grad_rot = U.T @ grad if precond_type == 'adam': D = 1.0 / np.sqrt(running_mean_squared(grad_rot)) elif precond_type == 'gn': D = np.diag(np.diag(GN_matrix)) ** power theta = theta - eta * U @ D @ grad_rot

• Basis Estimation: For large models, Kronecker-factored approximations are used for computational efficiency.
• Diagonal Preconditioner: Adam uses running averages of squared gradients; GN uses diagonal elements of the GN matrix, optionally with exponentiation.
• Batch Size: The batch size for gradient and GN estimation can be decoupled for stability.

Resource and Scaling

• Full Eigenbasis: Computationally expensive for large models; Kronecker approximations are recommended.
• Regularization: Diagonal preconditioners are regularized for numerical stability.
• Hyperparameter Tuning: Learning rate and regularization sweeps are necessary for optimal performance.

Limitations

• Basis Misalignment: GN-based methods lose their curvature-adaptive properties when the basis is not aligned.
• Global Convergence Constraints: GN can suffer from slow convergence in high-dimensional, ill-conditioned logistic regression due to the need for small learning rates.

Implications and Future Directions

The findings have significant implications for optimizer design in deep learning:

• Auto-Tuning vs. Curvature Adaptation: Adam's auto-tuning is robust to basis misalignment, while GN-based methods require accurate basis estimation for optimality.
• Stochastic Regime Equivalence: The equivalence of Adam and $\GN^{-1/2}$ in the stochastic regime suggests that empirical Fisher-based methods are effective proxies for curvature-based preconditioning in practical, large-scale training.
• Algorithmic Design: There is potential for developing new optimizers that combine the auto-tuning benefits of Adam with curvature adaptation, especially for small-batch training.

Future work should investigate the persistence of these phenomena in large-scale, real-world training scenarios, and explore hybrid approaches that leverage both empirical and curvature-based statistics for adaptive optimization.

Conclusion

This paper provides a comprehensive theoretical and empirical analysis of Adam and Gauss-Newton diagonal preconditioners, demonstrating that their relative performance is highly sensitive to basis alignment and gradient noise. Adam's auto-tuning confers robustness in misaligned or ill-conditioned settings, while GN-based methods require precise basis estimation for optimality. In the stochastic regime, Adam and $\GN^{-1/2}$ are nearly equivalent, suggesting a deep connection between empirical Fisher and curvature-based preconditioning. These insights inform the design of future optimizers and highlight the importance of basis and noise considerations in deep learning optimization.