MAC: An Efficient Gradient Preconditioning using Mean Activation Approximated Curvature

TL;DR

MAC introduces a computationally efficient second-order optimization method for neural networks by approximating curvature information, outperforming existing methods like KFAC in accuracy and training efficiency.

Contribution

It is the first to apply Kronecker factorization to transformer attention layers and explicitly incorporate attention scores into preconditioning.

Findings

MAC outperforms KFAC in accuracy and training speed.

MAC reduces memory usage compared to existing methods.

MAC converges to global minima under certain conditions.

Abstract

Second-order optimization methods for training neural networks, such as KFAC, exhibit superior convergence by utilizing curvature information of loss landscape. However, it comes at the expense of high computational burden. In this work, we analyze the two components that constitute the layer-wise Fisher information matrix (FIM) used in KFAC: the Kronecker factors related to activations and pre-activation gradients. Based on empirical observations on their eigenspectra, we propose efficient approximations for them, resulting in a computationally efficient optimization method called MAC. To the best of our knowledge, MAC is the first algorithm to apply the Kronecker factorization to the FIM of attention layers used in transformers and explicitly integrate attention scores into the preconditioning. We also study the convergence property of MAC on nonlinear neural networks and provide two conditions under which it converges to global minima. Our extensive evaluations on various network architectures and datasets show that the proposed method outperforms KFAC and other state-of-the-art methods in terms of accuracy, end-to-end training time, and memory usage.