[Paper Review] Fast Convergence of Natural Gradient Descent for Over-Parameterized Neural Networks

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

Stochastic Gradient Optimization Techniques30 citations

TL;DR

This paper establishes the first theoretical convergence analysis of natural gradient descent (NGD) for non-linear neural networks with squared error loss. It proves global convergence under two conditions—full row rank and stability of the Jacobian matrix—showing these hold for over-parameterized two-layer ReLU networks, with extension to K-FAC, an approximate NGD method.

ABSTRACT

Natural gradient descent has proven very effective at mitigating the catastrophic effects of pathological curvature in the objective function, but little is known theoretically about its convergence properties, especially for \emph{non-linear} networks. In this work, we analyze for the first time the speed of convergence to global optimum for natural gradient descent on non-linear neural networks with the squared error loss. We identify two conditions which guarantee the global convergence: (1) the Jacobian matrix (of network's output for all training cases w.r.t the parameters) is full row rank and (2) the Jacobian matrix is stable for small perturbations around the initialization. For two-layer ReLU neural networks (i.e. with one hidden layer), we prove that these two conditions do hold throughout the training under the assumptions that the inputs do not degenerate and the network is over-parameterized. We further extend our analysis to more general loss function with similar convergence property. Lastly, we show that K-FAC, an approximate natural gradient descent method, also converges to global minima under the same assumptions.

Motivation & Objective

• To theoretically analyze the convergence speed of natural gradient descent (NGD) in non-linear neural networks, which remains poorly understood despite NGD's empirical success.
• To identify sufficient conditions under which NGD converges globally to the global optimum for non-linear networks with squared error loss.
• To extend the analysis to over-parameterized two-layer ReLU networks and verify the conditions hold under mild assumptions on input data and network width.
• To investigate whether K-FAC, an approximate NGD method, also achieves global convergence under the same conditions.
• To generalize the convergence results to broader classes of loss functions beyond squared error.

Proposed method

• Theoretical analysis of NGD convergence using the natural gradient update rule, which adjusts parameters based on the Fisher information metric to counteract curvature in the loss landscape.
• Identification of two key conditions: (1) the Jacobian matrix of network outputs w.r.t. parameters must be full row rank, and (2) the Jacobian must be stable under small perturbations around initialization.
• Proof that for two-layer ReLU networks, these conditions are satisfied throughout training when inputs do not degenerate and the network is over-parameterized.
• Use of matrix perturbation theory and properties of ReLU activation to establish Jacobian stability under initialization and weight updates.
• Extension of convergence results to general loss functions by analyzing the Hessian and Fisher information structure under similar assumptions.
• Adaptation of the analysis to K-FAC, showing that its curvature approximation preserves the convergence guarantees under the same conditions.

Experimental results

Research questions

• RQ1 Under what conditions does natural gradient descent converge globally to the global minimum for non-linear neural networks with squared error loss?
• RQ2 Do the Jacobian full row rank and stability conditions hold throughout training in over-parameterized two-layer ReLU networks?
• RQ3 Can the convergence guarantees of NGD be extended to more general loss functions beyond squared error?
• RQ4 Does K-FAC, as an approximate NGD method, also converge to the global minimum under the same theoretical conditions?
• RQ5 How does over-parameterization influence the stability and rank of the Jacobian matrix during training?

Key findings

• Natural gradient descent globally converges to the global minimum for non-linear neural networks when the Jacobian matrix is full row rank and stable under small perturbations.
• For two-layer ReLU networks, the full row rank and stability conditions are satisfied throughout training under mild assumptions on input data and over-parameterization.
• The convergence rate is fast, with the paper establishing global convergence under the specified conditions, though exact rates are not quantified numerically in the provided text.
• The theoretical framework extends to general loss functions, indicating that the convergence properties of NGD are robust beyond squared error loss.
• K-FAC, an approximate natural gradient method, also converges to the global minimum under the same conditions, validating its practical efficacy from a theoretical standpoint.
• The results provide the first theoretical justification for the empirical success of NGD and K-FAC in over-parameterized non-linear networks.