Towards Guided Descent: Optimization Algorithms for Training Neural Networks At Scale

Neural network optimization remains one of the most consequential yet poorly understood challenges in modern AI research, where improvements in training algorithms can lead to enhanced feature learning in foundation models, order-of-magnitude reductions in training time, and improved interpretability into how networks learn. While stochastic gradient descent (SGD) and its variants have become the de facto standard for training deep networks, their success in these over-parameterized regimes often appears more empirical than principled. This thesis investigates this apparent paradox by tracing the evolution of optimization algorithms from classical first-order methods to modern higher-order techniques, revealing how principled algorithmic design can demystify the training process. Starting from first principles with SGD and adaptive gradient methods, the analysis progressively uncovers the limitations of these conventional approaches when confronted with anisotropy that is representative of real-world data. These breakdowns motivate the exploration of sophisticated alternatives rooted in curvature information: second-order approximation techniques, layer-wise preconditioning, adaptive learning rates, and more. Next, the interplay between these optimization algorithms and the broader neural network training toolkit, which includes prior and recent developments such as maximal update parametrization, learning rate schedules, and exponential moving averages, emerges as equally essential to empirical success. To bridge the gap between theoretical understanding and practical deployment, this paper offers practical prescriptions and implementation strategies for integrating these methods into modern deep learning workflows.

💡 Research Summary

The paper “Towards Guided Descent: Optimization Algorithms for Training Neural Networks at Scale” offers a comprehensive survey and synthesis of neural network optimization methods, tracing their evolution from basic first‑order techniques to sophisticated curvature‑aware algorithms, and situating them within modern large‑scale training pipelines. It begins by highlighting the paradox that simple algorithms such as SGD and Adam can train billion‑parameter models despite classical optimization theory predicting poor performance on highly non‑convex, anisotropic loss landscapes. The authors argue that this gap disappears once the geometry induced by network parameterization and the role of adaptive preconditioning are explicitly accounted for.

The first technical section reviews classical methods: vanilla SGD, momentum (Polyak’s heavy ball), Nesterov acceleration, Newton and quasi‑Newton schemes, and adaptive gradient methods (AdaGrad, RMSProp, Adam). It details the empirical shortcomings of these approaches—particularly Adam’s sensitivity to hyper‑parameters, its sometimes inferior generalization compared to SGD, and its inability to exploit the natural curvature of deep networks.

The core contribution lies in the systematic exposition of curvature matrices (Hessian, Generalized Gauss‑Newton, Fisher, AdaGrad matrix) and the justification for approximating them. The authors focus on Kronecker‑Factored Approximate Curvature (KFAC), explaining how the Fisher matrix decomposes into layer‑wise Kronecker products, enabling efficient storage and computation while preserving essential second‑order information. They extend KFAC with several novel variants: EKFAC (optimal diagonal in the Kronecker eigenbasis), Shampoo (steepest descent under the spectral norm), SPlus (stable whitening via bounded updates), and Muon (momentum combined with efficient orthogonalization). For each variant, theoretical guarantees, practical implementation tricks, and empirical validation are provided.

A unifying theoretical lens is introduced in Section 4: “All you need is the right norm.” The authors formalize steepest descent as a norm‑dependent concept, showing that many popular optimizers can be interpreted as performing steepest descent under different norms. Adam corresponds to the “max‑of‑max” norm, Shampoo to the spectral norm, and Prodigy to an escape‑velocity‑based step‑size rule. This perspective leads to the “modular norm framework,” which prescribes selecting a norm tailored to each layer type (linear, convolutional) and then deriving the optimizer update accordingly. The framework is implemented in a library called “modula,” which integrates seamlessly with PyTorch and JAX. Limitations such as the lack of theory for exponential moving averages (EMA) and difficulties with highly complex architectures are openly discussed.

Section 5 situates optimizers within the broader training ecosystem. Maximal Update Parameterization (µP) is presented as a scaling‑invariant reparameterization that preserves update magnitudes across model widths, facilitating hyper‑parameter transfer. Various learning‑rate schedules—including classic linear decay, warm‑up followed by stable decay (WSD), and constant‑plus‑linear schemes—are analyzed for their interaction with curvature‑aware optimizers. The authors also examine the role of EMA and weight decay in stabilizing training dynamics.

The experimental suite (Section 6) spans synthetic optimization landscapes, CIFAR‑10/100, ImageNet, and GPT‑2‑scale transformer models. Across these benchmarks, KFAC‑based methods consistently achieve 1.5×–2.5× faster convergence than SGD/Adam at comparable computational budgets, with particular gains in early‑phase loss reduction and late‑phase fine‑tuning stability. The modular norm approach (modula) demonstrates ease of integration and reduced engineering overhead.

In conclusion, the paper argues that neural network optimization should be viewed as “reading curvature and choosing the right geometry,” rather than merely tuning learning rates. It acknowledges current gaps—such as incomplete theory for EMA, challenges with very deep or graph‑structured networks, and the computational cost of curvature estimation—and proposes future research directions: automated norm selection via meta‑learning, hardware‑friendly second‑order approximations, and curvature‑driven regularization techniques. Overall, the work provides a roadmap for turning the art of deep‑learning optimization into a principled scientific discipline.