@inproceedings{
lin2023structured,
title={Structured Inverse-Free Natural Gradient: Memory-Efficient \& Numerically-Stable {KFAC} for Large Neural Nets},
author={Wu Lin and Felix Dangel and Runa Eschenhagen and Kirill Neklyudov and Agustinus Kristiadi and Richard E. Turner and Alireza Makhzani},
booktitle={OPT 2023: Optimization for Machine Learning},
year={2023},
url={https://openreview.net/forum?id=JmihRHpCU5}
}
Second-order methods for deep learning---such as KFAC---can be useful for neural network training.
However, they are often memory-inefficient and numerically unstable for low-precision training since their preconditioning Kronecker factors are dense, and require high-precision matrix inversion or decomposition.
Consequently, such methods are not widely used for training large neural networks such as transformer-based models.
We address these two issues by (i) formulating an inverse-free update of KFAC and (ii) imposing structures in each of the Kronecker factors, resulting in a method we term structured inverse-free natural gradient descent (SINGD).
On large modern neural networks, we show that, in contrast to KFAC, SINGD is memory efficient and numerically robust, and often outperforms AdamW even in half precision. Hence, our work closes a gap between first-order and second-order methods in modern low precision training for large neural networks.
Structured Inverse-Free Natural Gradient: Memory-Efficient     Numerically-Stable KFAC for Large Neural Nets
Wu Lin Felix Dangel Runa Eschenhagen Kirill Neklyudov Agustinus Kristiadi Richard E. Turner Alireza Makhzani