Search papers, labs, and topics across Lattice.
This thesis rigorously analyzes the Adam optimizer, revealing a problem-dependent phase transition that dictates its convergence behavior based on batch-size-dependent hyperparameters. It also uncovers the evolving Hessian structure during training, demonstrating that its near-block-diagonal form enhances Adam's performance over SGD in Transformers. The introduction of Adam-mini, a memory-efficient variant of Adam that retains performance, highlights the practical implications of these findings for optimizing neural networks.
Adam's performance can be dramatically improved with a new memory-efficient variant that retains its advantages while cutting memory usage by 50%.
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.
