On the Principles Behind Neural Network Optimizers

This research develops a theoretical foundation for neural-network optimizers, focusing on Adam, the method used widely to train models such as ChatGPT, Llama, and DeepSeek. It resolves part of Adam’s divergence debate by showing a problem- and batch-size-dependent phase transition: with sufficiently large beta2, unmodified Adam converges toward critical points, while small-beta2 settings can make iterates, gradients, and losses diverge. The explanation for Adam’s advantage over SGD on Transformers is geometric: training drives their Hessians toward near-block-diagonal structure with strongly heterogeneous block spectra, allowing Adam’s diagonal preconditioner to assign effectively different learning rates across blocks. Random-matrix analysis attributes this structure primarily to consecutive multiplications of large weight matrices, rather than specifically to cross-entropy loss. These findings also clarify why Adam is less advantageous on CNNs and dense quadratic problems, and motivate head-wise Muon variants for Transformer attention. The practical contribution is Adam-mini, which groups parameters by Hessian-aligned rows or attention heads and replaces most coordinate-level second-moment statistics with block means; it reduces optimizer-state memory by 50% while matching or sometimes exceeding AdamW. In Llama 2 pretraining on C4, experiments covered models from 39M to 1B parameters. The thesis also estimates that Adam could be 5 times faster than gradient descent on heterogeneous block-diagonal quadratics, while noting that Adam’s states alone require about 104 GB for a 13B model.