Understanding Approximate Fisher Information for Fast Convergence of Natural Gradient Descent in Wide Neural Networks

TL;DR

This paper provides a theoretical foundation explaining why various approximate Fisher information methods enable fast convergence of Natural Gradient Descent in wide neural networks, by analyzing their behavior in function space via neural tangent kernels.

Contribution

It reveals that under certain conditions, approximate Fisher information in NGD achieves the same convergence speed as exact NGD in the infinite-width limit, supported by a unified theoretical analysis.

Findings

Approximate Fisher information methods converge as fast as exact NGD in wide neural networks.

Layer-wise and unit-wise approximations maintain fast convergence under specific assumptions.

Isotropic gradient in function space is key to the success of these approximations.

Abstract

Natural Gradient Descent (NGD) helps to accelerate the convergence of gradient descent dynamics, but it requires approximations in large-scale deep neural networks because of its high computational cost. Empirical studies have confirmed that some NGD methods with approximate Fisher information converge sufficiently fast in practice. Nevertheless, it remains unclear from the theoretical perspective why and under what conditions such heuristic approximations work well. In this work, we reveal that, under specific conditions, NGD with approximate Fisher information achieves the same fast convergence to global minima as exact NGD. We consider deep neural networks in the infinite-width limit, and analyze the asymptotic training dynamics of NGD in function space via the neural tangent kernel. In the function space, the training dynamics with the approximate Fisher information are identical to those with the exact Fisher information, and they converge quickly. The fast convergence holds in layer-wise approximations; for instance, in block diagonal approximation where each block corresponds to a layer as well as in block tri-diagonal and K-FAC approximations. We also find that a unit-wise approximation achieves the same fast convergence under some assumptions. All of these different approximations have an isotropic gradient in the function space, and this plays a fundamental role in achieving the same convergence properties in training. Thus, the current study gives a novel and unified theoretical foundation with which to understand NGD methods in deep learning.