[Paper Review] Fast Convergence of Natural Gradient Descent for Overparameterized Neural Networks

Guodong Zhang, James Martens+1 more|arXiv (Cornell University)|May 27, 2019

Stochastic Gradient Optimization TechniquesComputer Science59 references41 citations

TL;DR

The paper proves global convergence and linear rates for natural gradient descent (NGD) on nonlinear overparameterized neural networks, under two conditions on the Jacobian, and extends results to K-FAC and general losses while maintaining favorable generalization.

ABSTRACT

Natural gradient descent has proven effective at mitigating the effects of pathological curvature in neural network optimization, but little is known theoretically about its convergence properties, especially for \emph{nonlinear} networks. In this work, we analyze for the first time the speed of convergence of natural gradient descent on nonlinear neural networks with squared-error loss. We identify two conditions which guarantee efficient convergence from random initializations: (1) the Jacobian matrix (of network's output for all training cases with respect to the parameters) has full row rank, and (2) the Jacobian matrix is stable for small perturbations around the initialization. For two-layer ReLU neural networks, we prove that these two conditions do in fact hold throughout the training, under the assumptions of nondegenerate inputs and overparameterization. We further extend our analysis to more general loss functions. Lastly, we show that K-FAC, an approximate natural gradient descent method, also converges to global minima under the same assumptions, and we give a bound on the rate of this convergence.

Motivation & Objective

• Motivate the use of natural gradient descent to address pathological curvature in neural network optimization.
• Identify simple, generic conditions on the network Jacobian that guarantee efficient convergence from random initializations.
• Show convergence results for two-layer ReLU networks under overparameterization and mild input assumptions.
• Extend analysis to general loss functions and to approximate NGD methods like K-FAC.
• Demonstrate that NGD can achieve faster convergence without sacrificing generalization.

Proposed method

• Define NGD update using the Fisher/ Gauss-Newton matrix and its generalized inverse when F is singular.
• Introduce two conditions on the Jacobian: (i) full row rank at initialization, (ii) stability of the Jacobian under small parameter perturbations.
• Prove linear convergence of NGD under these conditions with a step size bound.
• Apply the abstract analysis to a specific overparameterized two-layer ReLU network with random initialization and normalized inputs.
• Show NGD improves convergence rate by O(lambda_min(G∞)/n) over GD and that K-FAC also achieves linear convergence under similar assumptions.

Experimental results

Research questions

• RQ1 Under what conditions does natural gradient descent converge to a global minimizer in nonlinear, overparameterized neural networks?
• RQ2 How does NGD compare to gradient descent in convergence rate and learning rate tolerance for two-layer ReLU networks?
• RQ3 Can NGD and K-FAC provide provable global convergence for nonlinear networks with general loss functions beyond squared error?
• RQ4 What are the implications of NGD on generalization relative to standard gradient descent?

Key findings

• NGD achieves linear convergence to a global minimum when the Jacobian has full row rank at initialization and remains stable nearby.
• For two-layer ReLU networks with overparameterization, NGD converges with a constant step size up to O(1) and can converge in O(1) iterations in the infinite-width limit.
• NGD provides an O(lambda_min(G∞)/n) improvement in convergence rate over gradient descent for the specified two-layer network setup.
• K-FAC, under the same assumptions and sufficient overparameterization, also converges linearly to a global minimum, with its rate tied to the data Gram matrix.
• Generalization bounds for NGD match those proven for gradient descent in the two-layer ReLU setting, indicating no loss in generalization despite faster convergence.