On the Principles Behind Neural Network Optimizers

Yushun Zhang

cs.LGAug 17, 2026

On the Principles Behind Neural Network Optimizers

Authors: Yushun Zhang

Organizations: School of Data Science, The Chinese University of Hong Kong, Shenzhen, China

Abstract

Reliable optimization is central to neural network (NN) training, yet Adam, the default optimizer for modern LLMs, rests on a fragile foundation. This thesis develops a principled grounding for Adam and motivates new designs. First, we revisit Adam's divergence--convergence debate and show the existence of a problem-dependent phase transition: with properly chosen, batch-size-dependent hyperparameters, Adam converges, whereas under small- $β_2$ regimes it can diverge. Second, we investigate why Adam substantially outperforms SGD on Transformers through Hessian structure. We find that the Hessian evolves toward a near-block-diagonal form along training, accompanied by strong block heterogeneity. We prove that this structure makes Adam's diagonal preconditioner effective. We further show that this special Hessian structure originates from consecutive multiplications of large matrix variables, and we provide a rigorous analysis based on random matrix theory. Finally, these insights motivate Adam-mini, a new optimizer that reduces Adam's memory footprint by 50% while preserving its performance. Our results also have broader implications beyond Adam: they reveal new local structures in matrix-based nonconvex problems, and also help understand and improve recent NN optimizers, such as Muon.

Figures & tables

Fig. 1.1.: Overview of the thesis. We begin with general nonconvex objectives and then narrow down to the special structures found in NNs.

Panel of Figure 2 — Fig. 2.1.: (a) : The divergent region of Adam claimed by [ Reddi et al., 2018 ] . They fix $(\beta_{1},\beta_{2})$ first and then pick a problem to construct the divergence example. (b) : An illustration of our contribution in $(\beta_{1},\beta_{2})$ phase diagram. We fix the problem before picking $(\beta_{1},\beta_{2})$ . Note that this is a different setting from (a) , so there is no contradiction. Both boundaries of the red and blue regions depend on batch size (shown later). The shape of the region follows the solution to our analytic conditions. The dotted curve satisfies $\beta_{1}=\sqrt{\beta_{2}}$ . (c), (d) : The training loss on MNIST and CIFAR-10 over an extensive sweep of $\beta_{1}$ and $\beta_{2}$ . The performance of Adam reconciles with our theoretical characterization in (b) .

Panel of Figure 3 — Fig. 2.2.: (a, b): We restate Figure 1 from [ Reddi et al., 2018 ] . The divergence of Adam happens under both cyclic and random update orders, so randomization cannot prevent divergence. Since they consider constrained problems, the term “divergence” here means getting stuck at the sub-optimal solution $\theta=1$ . In the figures, AMSGrad is a different method, and it is not our focus. (c): Diminishing stepsize $\eta_{k}=\frac{1}{\sqrt{k}}$ does not prevent divergence.

Panel of Figure 4 — Fig. 2.3.: (a): The training loss of Adam on MNIST under different batch size and $\beta_{2}$ . The trend aligns with our theory: we need a larger $\beta_{2}$ when batch size is small. Here, we used the default $\beta_{1}=0.9$ . (b) (c): large- $\beta_{2}$ Adam converges to a neighborhood of critical points when $D_{0}>0$ and converges to exact critical points when $D_{0}=0$ . We use diminishing stepsize $\eta_{k}=0.1/\sqrt{k}$ as in our theory. The labels in (c) stand for $[\beta_{1},\beta_{2}]$ .

Panel of Figure 5 — Fig. 2.4.: (a): On our counterexample with $n=20$ , Adam diverges to infinity in the colored region. The region is plotted by solving our analytic divergence conditions in NumPy . The blue curve satisfies $\beta_{1}=\sqrt{\beta_{2}}$ . (b): When $\beta_{2}$ is small, Adam diverges to infinity. The labels in (b) stand for $[\beta_{1},\beta_{2}]$ . (c, d) The optimality gap $|\theta-\theta^{*}|$ after running 50k iterations of Adam on our counterexample under $n=10$ and $50$ .

Fig. 2.5.: An illustration of the changes of Adam’s update direction when $\beta_{2}$ changes.

Panel of Figure 7 — Fig. 3.1.: Adam’s practical advantage is not universal. (a) Adam performs worse than GD on simple random Gaussian quadratic problems. (b) Adam performs substantially better than SGD on a Transformer with 200M parameters. All methods were well-tuned through an extensive hyperparameter search.

Panel of Figure 8 — Fig. 3.2.: The effectiveness of Adam’s preconditioner $D_{\text{Adam}}$ on different structures of a Hessian. The $y$ -axis, denoted by $\tau\in[0,1]$ , serves as a measure of Hessian diagonality; the smaller the value of $\tau$ , the closer the Hessian is to a diagonal matrix. In the heat map, yellow corresponds to cases where $D_{\text{Adam}}$ effectively reduces the Hessian’s condition number, whereas blue denotes cases where it does not. Across different sizes and condition numbers of the Hessian, we find that $D_{\text{Adam}}$ is less effective when Hessian is dense.

Panel of Figure 9 — Fig. 3.3.: (a-c): The Hessian of a 1-hidden-layer network on Gaussian synthetic data under CE loss. At initialization, we observe the “block-circulant” pattern in $H_{wv}$ , and the near-block-diagonal structure in $H_{ww}$ and $H_{vv}$ (with $m+C=508$ blocks in total). We refer to it as “block-circulant-block-diagonal matrix” . We notice that the “block-circulant” pattern in $H_{wv}$ vanishes along training, while the near-block-diagonal patterns in $H_{ww}$ and $H_{vv}$ are preserved. (d-f): The Hessian of a 4-layer network. We observe additional dense blocks in distant inter‑layer Hessian. Other structures are similar to those in 1-hidden-layer networks.

Fig. 3.4.: (a-f): The Hessian of different modules in a Transformer.

Fig. 3.5.: The Hessian structure of NNs can be understood through their computation graphs: any pair of connected links yields a multiplicative relation in the loss, which further gives a non-zero Hessian entry; while the non-connected pair gives a vanishing or exactly zero Hessian entry.

Panel of Figure 12 — Fig. 3.6.: The hidden-layer Hessian $H_{ww}$ of a 1-hidden-layer network. The network has 8 hidden neurons. We observe that the block-diagonal Hessian structure in $H_{ww}$ becomes clearer as $C$ increases.

Panel of Figure 13 — Fig. 3.7.: (a-c): The Jensen-Shannon (JS) distance among blockwise Hessian spectra at initialization. We find that the JS distance of blockwise spectra in CNNs is significantly smaller than that in Transformers.

Fig. 4.1.: An illustration of Adam-mini. Adam-mini assigns learning rates (lrs) by Hessian structure. It uses more lrs than SGD but fewer than Adam.

Neural Network Optimization

Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.

See all 321 papers

Mar 2, 2026

Adam Converges Without Any Modification On Update Rules

Adam is the default algorithm for training neural networks, including large language models (LLMs). However, \citet{reddi2019convergence} provided an example that Adam diverges, raising concerns for its deployment in AI model training. We identify a key mismatch between the divergence example and practice: \citet{reddi2019convergence} pick the problem after picking the hyperparameters of Adam, i.e.,

(β_1,β_2)

; while practical applications often fix the problem first and then tune

(β_1,β_2)

. In this work, we prove that Adam converges with proper problem-dependent hyperparameters. First, we prove that Adam converges when

β_2

is large and

β_1 < \sqrt{β_2}

. Second, when

β_2

is small, we point out a region of

(β_1,β_2)

combinations where Adam can diverge to infinity. Our results indicate a phase transition for Adam from divergence to convergence when changing the

(β_1, β_2)

combination. To our knowledge, this is the first phase transition in

(β_1,β_2)

2D-plane reported in the literature, providing rigorous theoretical guarantees for Adam optimizer. We further point out that the critical boundary

(β_1^*, β_2^*)

is problem-dependent, and particularly, dependent on batch size. This provides suggestions on how to tune

β_1

and

β_2

: when Adam does not work well, we suggest tuning up

β_2

inversely with batch size to surpass the threshold

β_2^*

, and then trying

β_1< \sqrt{β_2}

. Our suggestions are supported by reports from several empirical studies, which observe improved LLM training performance when applying them.

Yushun Zhang, Bingran Li, Congliang Chen +2

The Chinese University of Hong Kong, Shenzhen, China · Shenzhen Research Institute of Big Data

Optimization Convergence Analysis Adaptive Gradient Methods Hyperparameter Optimization

Ask PDF

Aug 28, 2026

Blog: Survey of Optimizers

Neural-network optimization in 2025-2026 is no longer well described as a succession of new Adam variants. The design space has expanded from coordinates to matrices and layers, from fixed training horizons to policies over time, and from mathematical update rules to state representations that must survive sharding and low-precision computation. This survey organizes recent optimizers and training optimization methods along four largely independent axes: temporal estimation, update geometry, horizon management, and representation and systems. It connects the spectral normalization of Muon, the historical matrix statistics of Shampoo and SOAP, adaptive and hybrid matrix methods, memory-efficient optimizers, schedule-free training, small-batch corrections, and quantized optimizer states. The central empirical conclusion is deliberately non-triumphal: matrix-aware methods represent a genuine advance, but there is no context-independent replacement for AdamW. Rankings change with model scale, data-to-parameter ratio, batch size, schedule, parameter partition, tuning budget, and whether the target metric is tokens, FLOPs, wall-clock time, or memory. The practical consequence is a compositional view of optimizer design and a stricter protocol for evaluating optimizer claims.

Ruoran Xu

Neural Network Optimization Memory-Efficient Optimization

Ask PDF

Jun 29, 2026

Characterizing Optimizer-Dependent Training Dynamics Through Hessian Eigenvector Displacement and Localization

Hessian spectral properties are a standard tool in analysing neural-network training, with eigenvalues linked to sharpness, generalization, and optimization dynamics. Eigenvalues quantify curvature magnitude, while eigenvectors identify which parameters generate that curvature. In this work, we study how the leading Hessian eigenvectors evolve during training and how they affect the learning trajectories. We track the training dynamics of multilayer perceptrons on a classification problem and measure eigenvector dynamics through two complementary statistics: (i) displacement over time, inspired by analyses of glassy systems, and (ii) localization via the inverse participation ratio. The metrics are compared against a random null model of the Hessian induced by the architecture. Our results reveal clear optimizer-dependent behaviour. SGD leads to progressively more stable leading curvature directions, while Adam exhibits substantially stronger reorganization of eigenvectors throughout training. We also observe a localization phenomenon under Adam, where a small subset of parameters contributes disproportionately to the leading curvature directions. These results suggest that Hessian eigenvector dynamics capture key differences in optimizer behaviour and the resulting training trajectories.

Marcelina Marjankowska, Valerio Modugno, Paolo Barucca

Financial Computing and Analytics Group, University College London · Robot Perception and Learning Lab, University College London

Spectral Methods Adaptive Gradient Methods Neural Network Training Dynamics Stochastic Gradient Descent

Ask PDF

Abstract

Neural Network Optimization

Explore similar work

Adam Converges Without Any Modification On Update Rules

Blog: Survey of Optimizers

Characterizing Optimizer-Dependent Training Dynamics Through Hessian Eigenvector Displacement and Localization