NKFAC: A Fast and Stable KFAC Optimizer for Deep Neural Networks

In recent advances in second-order optimizers, computing the inverse of second-order statistics matrices has become critical. One such example is the Kronecker-factorized approximate curvature (KFAC) algorithm, where the inverse computation of the two second-order statistics to approximate the Fisher information matrix (FIM) is essential. However, the time-consuming nature of this inversion process often limits the extensive application of KFAC. What's more, improper choice of the inversion method or hyper-parameters can lead to instability and fail the entire optimization process. To address these issues, this paper proposes the Newton-Kronecker factorized approximate curvature (NKFAC) algorithm, which incorporates Newton's iteration method for inverting second-order statistics. As the FIM between adjacent iterations changes little, Newton's iteration can be initialized by the inverse obtained from the previous step, producing accurate results within a few iterations thanks to its fast local convergence. This approach reduces computation time and inherits the property of second-order optimizers, enabling practical applications. The proposed algorithm is further enhanced with several useful implementations, resulting in state-of-the-art generalization performance without the need for extensive parameter tuning. The efficacy of NKFAC is demonstrated through experiments on various computer vision tasks. The code is publicly available at https://github.com/myingysun/NKFAC.