Source-linked AI summary
On the Principles Behind Neural Network Optimizers
Yushun Zhang
TL;DR
The thesis addresses the fragile theoretical foundation and poorly understood behavior of Adam in neural-network training. It combines convergence analysis, Hessian-structure investigation, and random matrix theory to ground optimizer behavior and motivate new designs. The results establish problem-dependent safety conditions, explain structural advantages in matrix-based networks, and motivate Adam-mini.
Problem
Adam’s convergence and practical advantage over SGD remain incompletely understood despite its widespread use in neural-network training.
Method
The thesis combines convergence analysis, Hessian-structure analysis, and random matrix theory to study Adam and derive optimizer-design principles.
Results
The thesis establishes a theoretical grounding for neural-network optimizers, including convergence insights for Adam and structural explanations connecting Hessians to optimizer behavior.
Takeaways & Limitations
The findings support safer, more reliable, more efficient, and better-understood neural-network training, while informing designs such as Adam-mini and analysis of Muon.
Takeaways & Limitations
Constant-step-size Adam with β2 < 1 can exhibit limit-cycle non-convergence on quadratic objectives; the reported convergence results use diminishing step sizes.
Abstract
from arXiv · showhide
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.
1. Introduction
The thesis develops a principled foundation for neural-network optimizers by studying Adam’s convergence, its performance relative to SGD, the Hessian structures underlying these behaviors, and optimizer design.
- Motivation: Neural-network training is formulated as empirical risk minimization, requiring stable and computationally efficient optimization at scale.The objective averages mini-batch training losses over trainable parameters.
- Motivation: Adam is widely used in modern neural networks but has faced theoretical criticism because it can diverge even on simple examples.This creates concerns about reliable deployment and predictable behavior in large-scale AI training.
- Contributions: The work aims to establish a theoretical grounding for current neural-network optimizers and provide insights relevant to Adam, Hessian structure, and Muon.Its stated scope extends beyond Adam to local structures in matrix-based nonconvex problems and recent optimizers.
- Research questions: The thesis asks whether Adam converges, why it outperforms SGD on Transformers, when neural-network structure emerges, and how that structure can guide optimizer design.These questions span safety guarantees, optimizer behavior, problem structure, and principled algorithm development.
- Approach: The thesis combines new concentration tools for Adam’s convergence analysis with random-matrix-theoretic decoupling methods for Hessian-structure analysis.The thesis proceeds from generic nonconvex convergence guarantees toward neural-network-specific structure and optimizer design.
2. Divergence or Convergence? A Story of Adam
Adam exhibits a problem-dependent divergence–convergence phase transition: properly chosen, batch-size-dependent hyperparameters yield convergence guarantees, while small β2 can cause divergence. The results provide tuning guidance and clarify scope boundaries of the theory.
- Convergence guarantee: Properly chosen problem-dependent hyperparameters can make unmodified Adam converge without changing its update rules.The hyperparameters depend particularly on the number of mini-batches, equivalently the batch size, without requiring per-coordinate tuning.
- Divergence result: Small-β2 Adam can diverge, and under the stated counterexample conditions its iterates, gradients, and function values all diverge to infinity.The divergent region expands toward the whole [0, 1]^2 region as n grows to infinity.
- Main results: Adam has a problem-dependent phase transition in the (β1, β2) plane, with a safe region yielding convergence and a danger region allowing divergence to infinity.The phase transition combines positive and negative theoretical results for unmodified Adam.
- Convergence scope: Convergence means convergence to a neighborhood of critical points when D0 > 0, but to exact critical points when D0 = 0.The neighborhood result is also observed for simple convex quadratic functions with D0 > 0.
- Hyperparameter guidance: Larger β2 is required when the number of mini-batches is larger, or equivalently when the batch size is smaller.Experiments on MNIST align with this batch-size-dependent theoretical guidance.
- Limitations and caveats: The sufficient threshold for β2 is not claimed tight, and larger β2 is not necessarily better after it exceeds the batch-size-dependent threshold.The text reports worse performance for β2 = 0.999 than β2 = 0.95 in an empirical LLM trend.
2.3. Key Lemmas for the Convergence Result in Theorem 2.2
The analysis explains how Adam’s preconditioner behaves across β2 regimes and why large β2 supports convergence. It addresses the stochastic denominator challenge using concentration and decoupling arguments.
- When β2 is large, 1/√vk becomes sufficiently predictable for Adam’s normalized stochastic gradient to approach the population gradient direction.The proof exploits Adam’s geometric-sum structure in vk and uniform sampling of stochastic-gradient indices.
- When β2 is small, Adam’s update direction can deviate toward the opposite direction of the gradient, causing divergence.Figure 2.5 illustrates the change in update direction as β2 varies.
- For large β2, the normalized stochastic gradient becomes a descent direction because slower changes in vk make its behavior more predictable.The momentum term mk is handled using a new potential function.
- Adam’s main analytical challenge is that the random preconditioner vk appears in the denominator and may reach zero.This creates a stochastic nonlinear dynamical system that is difficult to analyze.
- Lemma 2.1 formalizes concentration of Adam’s preconditioner, with bounds whose constants approach 1 as β2 approaches 1.The result holds without a boundedness condition on stochastic gradients.
- The proof maps the dynamics of 1/√vk to bounded i.i.d. Bernoulli proxies, then uses decoupling steps to recover the original dynamics.This tracks potentially unbounded random variables through bounded surrogates.
3. Why Transformers Need Adam: An Investigation into Hessian
Adam’s advantage depends on Hessian structure: its diagonal preconditioner is effective when curvature becomes near-block-diagonal and heterogeneous, as observed in Transformers. This structure emerges during training from large matrix products, while training suppresses additional Hessian components.
- Diagonal preconditioning: Adam’s diagonal preconditioner typically performs poorly on dense Hessians because it cannot capture dominant off-diagonal curvature couplings.This explains Adam’s poor performance on dense random Gaussian quadratic problems.
- Diagonal preconditioning: As the Hessian becomes closer to diagonal, DAdam consistently improves across Gaussian quadratic problems with varied sizes and condition numbers.The experiments vary Hessian structure, size, and condition number.
- Hessian structure: Neural-network Hessians evolve toward near-block-diagonal structure, with blocks corresponding mainly to weight-matrix rows or Transformer attention heads.The block-diagonal pattern persists along training, while initialization can also exhibit block-circulant and dense components.
- Origins and dynamics: Consecutive matrix multiplication is identified as the major source of the special Hessian structure, rather than cross-entropy loss alone.Large output dimension increases block-diagonality, while training reduces intra-layer components in proportion to the optimality gap.
- Origins and dynamics: Random-matrix analysis handles the nonlinear dependent product through asymptotic independence, Lindeberg-based decoupling, and the generalized Marcenko-Pastur law.The dependence between the data matrix and diagonal loss/activation factors diminishes as input dimension grows.
- Block heterogeneity: Transformer Hessians also show strong block-heterogeneity, giving Adam an advantage because coordinate-wise learning rates can address differences among blocks.SGD can match Adam without block heterogeneity but performs worse when heterogeneity exists.
4. Adam-mini: A New Optimizer Motivated by the Hessian Structure
Adam-mini uses Hessian-informed parameter blocks to retain Adam-like performance with substantially less optimizer-state memory. Its design replaces per-parameter learning rates with block-level rates while preserving block heterogeneity.
- Motivation and overview: 50% less optimizer-state memory accompanies Adam-mini’s on-par or better performance than Adam.The method makes more than 99.9% of Adam’s v negligible scalars.
- Design principle: The method uses customized block learning rates because dense Hessian blocks make additional per-parameter learning rates unnecessary.The design targets block heterogeneity without retaining Adam’s full learning-rate resource.
- Method: Adam-mini partitions gradient vectors according to dense Hessian sub-blocks and computes one mean second-moment statistic per block.Each block receives η/√vmean,b as its learning rate.
- Partition principle: In Transformers, Query and Key matrices are partitioned by heads, while other weight matrices are partitioned by rows or output neurons.These partitions are compatible with distributed setups without communication overhead for computing block means.
- Empirical results: Adam-mini closely tracks AdamW across compute and model scales from 39M to 1B parameters while using half the optimizer-state memory.The experiments pre-train Llama 2 architectures on C4 using Chinchilla-style token budgets.
- Extensions: The Adam-mini idea extends to Muon through neuron-wise learning rates, producing NorMuon without additional memory or computational overhead.NorMuon addresses Muon’s inability to handle neuron-wise heterogeneity.
5. Conclusion
The thesis connects Adam’s theory, Hessian structure, and optimizer design into a systematic account of neural-network optimization. It concludes with practical principles and open questions about the boundary of the theory and broader architectures.
- Contributions: The thesis links convergence guarantees, Hessian-based explanations of Adam’s advantages, structural origins, and Adam-mini’s optimizer design.Its stated goal is to narrow the gap between optimization theory and large-scale practice.
- Future directions: The precise shape and uniqueness of Adam’s divergence-convergence critical boundary remain uncharacterized.The thesis establishes existence of the boundary but not its complete characterization.
- Future directions: Extending rigorous Hessian analysis to attention and SwiGLU-based mixture-of-experts architectures remains an open direction.The proposed computational-graph analysis predicts structure by checking connected paths between distant links.
- Implications: Hessian structures may provide design principles for more effective second-order optimizers.
Major References
The listed major references include the thesis’s foundational Adam convergence papers and their journal extension. These references establish the work’s connection to prior convergence analysis.
- Adam convergence: The references include “Adam Can Converge Without Any Modification On Update Rules,” published at NeurIPS 2022 as a Spotlight.
- Adam convergence: They also include “Adam Converges Without Any Modification On Update Rules,” identified as the journal extended version of the first paper.
A Hessian Perspective.
The references cover work on Hessian structure in neural networks and the Adam-mini optimizer. Together, they identify the structural-analysis and optimizer-design lines represented in this thesis.
- Hessian structure: “Towards Quantifying the Hessian Structure of Neural Networks” is listed as a NeurIPS 2024 reference.
- Optimizer design: “Adam-mini: Use Fewer Learning Rates To Gain More” is listed among the references as an ICLR 2025 work.