Improving Adaptive Moment Optimization via Preconditioner Diagonalization
Son Nguyen, Bo Liu, Lizhang Chen, Qiang Liu
TL;DR
This work tackles inefficiencies in adaptive moment optimizers by introducing preconditioner diagonalization through an invertible gradient-space transform. By rotating gradients with full-rank projections and applying a diagonal preconditioner in the rotated space, the method improves the estimation of second-order statistics while enabling back-projection to the original parameter space; a periodic SVD-based scheme keeps the extra cost manageable. Convergence guarantees are established via a Hamiltonian descent framework, and empirical results show substantial speedups and accuracy/robustness gains across image classification and language modeling tasks, including 1.8–2x speedups on LLaMA pretraining. The approach also integrates with memory-efficient optimizers (Adafactor, Hfac), offering practical improvements for large-scale models and providing a flexible, theoretically grounded paradigm for adaptive optimization in dynamic subspaces.
Abstract
Modern adaptive optimization methods, such as Adam and its variants, have emerged as the most widely used tools in deep learning over recent years. These algorithms offer automatic mechanisms for dynamically adjusting the update step based on estimates of gradient statistics. Compared to traditional algorithms like Stochastic Gradient Descent, these adaptive methods are typically more robust to model scale and hyperparameter tuning. However, the gradient statistics employed by these methods often do not leverage sufficient gradient covariance information, leading to suboptimal updates in certain directions of the parameter space and potentially slower convergence. In this work, we keep track of such covariance statistics in the form of a structured preconditioner matrix. Unlike other works, our approach does not apply direct approximations to estimate this matrix. We instead implement an invertible transformation that maps the preconditioner matrix into a new space where it becomes approximately diagonal. This enables a diagonal approximation of the preconditioner matrix in the transformed space, offering several computational advantages. Empirical results show that our approach can substantially enhance the convergence speed of modern adaptive optimizers. Notably, for large language models like LLaMA, we can achieve a speedup of 2x compared to the baseline Adam. Additionally, our method can be integrated with memory-efficient optimizers like Adafactor to manage computational overhead.