arXiv · 2312.05705
Structured Inverse-Free Natural Gradient: Memory-Efficient & Numerically-Stable KFAC
Abstract
Second-order methods such as KFAC can be useful for neural net training. However, they are often memory-inefficient since their preconditioning Kronecker factors are dense, and numerically unstable in low precision as they require matrix inversion or decomposition. These limitations render such methods unpopular for modern mixed-precision training. We address them by (i) formulating an inverse-free KFAC update and (ii) imposing structures in the Kronecker factors, resulting in structured inverse-free natural gradient descent (SINGD). On modern neural networks, we show that SINGD is memory-efficient and numerically robust, in contrast to KFAC, and often outperforms AdamW even in half precision. Our work closes a gap between first- and second-order methods in modern low-precision training.
Explore related subjects
Keep this discovery
Wu Lin, Felix Dangel, Runa Eschenhagen, Kirill Neklyudov, Agustinus Kristiadi, Richard E. Turner, Alireza Makhzani. 2023-12-09. Structured Inverse-Free Natural Gradient: Memory-Efficient & Numerically-Stable KFAC. https://arxiv.org/abs/2312.05705
Cite the original work for its findings. Save a collection to share your selection of sources.