On the Principles Behind Neural Network Optimizers
2026-08-17 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study Adam, a common neural network optimizer, to understand why it sometimes works well and sometimes doesn't. They show that Adam's success depends on choosing the right settings related to batch size, which can prevent it from failing. By analyzing the math behind Transformers, they discover that certain patterns in the training process make Adam's approach particularly effective. Using these insights, they create a new version called Adam-mini that uses less memory without losing performance. Their work also sheds light on broader properties of neural network training.
Adam optimizerneural network trainingbatch sizeHessian matrixpreconditioningblock-diagonal structurerandom matrix theorytransformersnonconvex optimizationmemory efficiency
Authors
Yushun Zhang
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.