[Paper Review] Understanding Approximate Fisher Information for Fast Convergence of Natural Gradient Descent in Wide Neural Networks

Ryo Karakida, Kazuki Osawa|arXiv (Cornell University)|Oct 2, 2020

Stochastic Gradient Optimization TechniquesComputer Science30 references15 citations

TL;DR

This paper establishes a theoretical foundation showing that Natural Gradient Descent (NGD) with various approximate Fisher information matrices—such as block diagonal, block tridiagonal, K-FAC, and unit-wise approximations—achieves the same fast convergence to global minima as exact NGD in wide, fully-connected neural networks. The key insight is that these approximations yield isotropic gradients in function space, which enables rapid convergence independent of the Neural Tangent Kernel (NTK), even under layer-wise or unit-wise approximations when specific conditions on the network width and learning rate are met.

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.

Motivation & Objective

• To resolve the theoretical gap in understanding why heuristic approximations of the Fisher information matrix (FIM) in Natural Gradient Descent (NGD) work well in practice despite high computational cost.
• To analyze the asymptotic training dynamics of NGD with approximate FIMs in the infinite-width limit of deep neural networks.
• To identify the precise conditions under which different FIM approximations (e.g., block diagonal, K-FAC, unit-wise) preserve the fast convergence of exact NGD.
• To clarify the role of isotropy in the function space gradient as a unifying principle for fast convergence across diverse NGD approximations.

Proposed method

• The study uses the Neural Tangent Kernel (NTK) framework to analyze training dynamics in function space for wide, fully-connected neural networks in the infinite-width limit.
• It derives the asymptotic dynamics of NGD with approximate FIMs and shows equivalence to exact NGD in function space under specific conditions.
• The analysis focuses on layer-wise approximations (block diagonal, block tridiagonal, K-FAC) and unit-wise approximations, proving that they yield identical convergence behavior to exact NGD when the learning rate is appropriately scaled.
• The paper introduces a damping parameter ρ > 0 to stabilize the inverse of the FIM and derives bounds on the deviation between linearized and true dynamics, showing convergence in the limit M → ∞.
• It establishes that the key mechanism enabling fast convergence is the isotropy of the gradient in function space, which emerges from the structure of the approximated FIM.
• Numerical experiments validate the theoretical predictions, particularly for unit-wise NGD, by comparing theoretical convergence rates with empirical training dynamics.

Experimental results

Research questions

• RQ1 Under what conditions do approximate NGD methods with block-wise or unit-wise FIM approximations achieve the same fast convergence as exact NGD in wide neural networks?
• RQ2 Why do practical FIM approximations—such as K-FAC or block diagonal—work well in practice despite theoretical uncertainty?
• RQ3 What structural property of the FIM approximation ensures fast convergence in function space, and how does it relate to the NTK?
• RQ4 How does the isotropy of the gradient in function space emerge from different FIM approximations, and why is it crucial for fast convergence?

Key findings

• Approximate NGD with block diagonal, block tridiagonal, K-FAC, and unit-wise FIM approximations achieves the same fast convergence to global minima as exact NGD in the infinite-width limit of deep neural networks.
• The convergence dynamics in function space are identical for all these approximations and exact NGD, provided the learning rate is scaled appropriately with network width or sample size.
• The key mechanism enabling fast convergence is the isotropy of the gradient in function space, which is induced by the structure of the approximated FIM and is independent of the NTK.
• Unit-wise NGD achieves fast convergence under the condition that the damping parameter ρ > 0 is small but non-zero, and the convergence rate is bounded by A³ρ⁻⁶/√M, which vanishes as M → ∞.
• Numerical experiments confirm that the isotropy condition holds for layer-wise and unit-wise approximations but fails for entry-wise diagonal approximations, explaining their poor performance.
• The training dynamics in parameter space differ across approximations, leading to different global minima and test predictions, even though function space dynamics are identical.