Technical Papers Archive

HyLo: A Hybrid Low-Rank Natural Gradient Descent Method

Authors: Baorun Mu and Saeed Soori (University of Toronto), Bugra Can and Mert Gürbüzbalaban (Rutgers University), and Maryam Mehri Dehnavi (University of Toronto)

Abstract: 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 large-scale. Kronecker factorization methods such as KFAC 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. HyLo uses a computationally-efficient low-rank factorization to achieve superior timing for Fisher inverses. We evaluate HyLo on large models including ResNet-50, U-Net, and ResNet-32 on up to 64 GPUs. HyLo converges 1.4x-2.1x faster than the state-of-the-art distributed implementation of KFAC and reduces the computation and communication time up to 350x and 10.7x on ResNet-50.


Presentation: file


Back to Technical Papers Archive Listing