HyLo: A Hybrid Low-Rank Natural Gradient Descent Method

This work presents a Hybrid Low-Rank Natural Gradient Descent method, called HyLo, that accelerates the training time of deep neural networks. Natural gradient descent (NGD) requires computing the inverse of the Fisher information matrix (FIM), which is typically expensive at largescale. Kronecker factorization methods such as K F A C attempt to improve NGD's running time by approximating the FIM with Kronecker factors. However, the size of Kronecker factors increases quadratically as the model size grows. Instead, in HyLo, we use the Sherman-Morrison-Woodbury variant of NGD (SNGD) and propose a reformulation of SNGD to resolve its scalability issues. HyL o uses a computationally-efficient low-rank factorization to achieve superior timing for Fisher inverses. We evaluate HyL o on large models including ResNet-50, U-Net, and ResNet-32 on up to 64 GPUs. H yL o converges 1.4×-2.1× faster than the state-of-the-art distributed implementation of K FA C and reduces the computation and communication time up to 350× and 10.7× on ResNet-50.