Improving Adaptive Moment Optimization via Preconditioner Diagonalization

Son Nguyen, Bo Liu, Lizhang Chen, Qiang Liu
Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, PMLR 300:4312-4320, 2026.

Abstract

Modern deep learning heavily relies on adaptive optimization methods like Adam and its variants, celebrated for their robustness against model scale and ease of 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 \emph{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 $2\times$ speedup in sample efficiency compared to Adam. In addition, our method can also be integrated with memory-efficient optimizers to manage computational overhead.

Cite this Paper

BibTeX
@InProceedings{pmlr-v300-nguyen26e, title = { Improving Adaptive Moment Optimization via Preconditioner Diagonalization }, author = {Nguyen, Son and Liu, Bo and Chen, Lizhang and Liu, Qiang}, booktitle = {Proceedings of The 29th International Conference on Artificial Intelligence and Statistics}, pages = {4312--4320}, year = {2026}, editor = {Khan, Emtiyaz and Li, Yingzhen and Solin, Arno and Ramdas, Aaditya}, volume = {300}, series = {Proceedings of Machine Learning Research}, month = {02--05 May}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v300/main/assets/nguyen26e/nguyen26e.pdf}, url = {https://proceedings.mlr.press/v300/nguyen26e.html}, abstract = { Modern deep learning heavily relies on adaptive optimization methods like Adam and its variants, celebrated for their robustness against model scale and ease of 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 \emph{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 $2\times$ speedup in sample efficiency compared to Adam. In addition, our method can also be integrated with memory-efficient optimizers to manage computational overhead. } }
Endnote
%0 Conference Paper %T Improving Adaptive Moment Optimization via Preconditioner Diagonalization %A Son Nguyen %A Bo Liu %A Lizhang Chen %A Qiang Liu %B Proceedings of The 29th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2026 %E Emtiyaz Khan %E Yingzhen Li %E Arno Solin %E Aaditya Ramdas %F pmlr-v300-nguyen26e %I PMLR %P 4312--4320 %U https://proceedings.mlr.press/v300/nguyen26e.html %V 300 %X Modern deep learning heavily relies on adaptive optimization methods like Adam and its variants, celebrated for their robustness against model scale and ease of 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 \emph{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 $2\times$ speedup in sample efficiency compared to Adam. In addition, our method can also be integrated with memory-efficient optimizers to manage computational overhead.
APA
Nguyen, S., Liu, B., Chen, L. & Liu, Q.. (2026). Improving Adaptive Moment Optimization via Preconditioner Diagonalization . Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 300:4312-4320 Available from https://proceedings.mlr.press/v300/nguyen26e.html.

Related Material