[Paper Review] Convergence of gradient descent for deep neural networks

Sourav Chatterjee|arXiv (Cornell University)|Mar 30, 2022

Stochastic Gradient Optimization Techniques4 citations

TL;DR

This paper establishes a novel convergence criterion for gradient descent in nonconvex optimization, proving that with proper initialization, gradient descent converges to a global minimum in deep neural networks with smooth, strictly increasing activations when the input dimension is at least the number of data points. The key contribution is that network width need not grow with data size, enabling convergence for fixed-width networks under mild conditions.

ABSTRACT

We give a simple local Polyak-Lojasiewicz (PL) criterion that guarantees linear (exponential) convergence of gradient flow and gradient descent to a zero-loss solution of a nonnegative objective. We then verify this criterion for the squared training loss of a feedforward neural network with smooth, strictly increasing activation functions, in a regime that is complementary to the usual over-parameterized analyses: the network width and depth are fixed, while the input data vectors are assumed to be linearly independent (in particular, the ambient input dimension is at least the number of data points). A notable feature of the verification is that it is constructive: it leads to a simple "positive" initialization (zero first-layer weights, strictly positive hidden-layer weights, and sufficiently large output-layer weights) under which gradient descent provably converges to an interpolating global minimizer of the training loss. We also discuss a probabilistic corollary for random initializations, clarify its dependence on the probability of the required initialization event, and provide numerical experiments showing that this theory-guided initialization can substantially accelerate optimization relative to standard random initializations at the same width.

Motivation & Objective

• To establish a general convergence criterion for gradient descent in nonconvex optimization that guarantees convergence to a global minimum.
• To address the gap in understanding why gradient descent often finds global minima in deep learning despite nonconvexity.
• To show that convergence is possible even when network width is fixed, rather than growing with data size.
• To provide conditions under which gradient descent with small step sizes converges exponentially fast to a global minimum in a local ball around initialization.
• To extend the analysis to deep neural networks with smooth, strictly increasing activation functions and arbitrary depth.

Proposed method

• Introduces a new convergence criterion based on the ratio of gradient norm squared to function value, defined as $ \alpha(x_0,r) = \inf_{x \in B(x_0,r), f(x) \neq 0} \frac{|\nabla f(x)|^2}{f(x)} $, with $ \alpha = \infty $ if $ f \equiv 0 $ on the ball.
• Establishes that if $ 4f(x_0) < r^2 \alpha(x_0,r) $, then gradient flow converges exponentially fast to a global minimum in $ B(x_0,r) $, with bounds $ |\phi(t) - x^*| \leq r e^{-\alpha t/2} $ and $ f(\phi(t)) \leq e^{-\alpha t} f(x_0) $.
• Proves an analogous result for gradient descent with small step size $ \eta $, showing convergence to a global minimum with linear convergence rate: $ |x_k - x^*| \leq (1 - \delta)^{k/2} r $, $ f(x_k) \leq (1 - \delta)^k f(x_0) $, where $ \delta = \min\{1, (1 - \epsilon)\alpha\eta\} $.
• Applies the criterion to deep neural networks by analyzing the training loss $ S(w) $, showing that under proper initialization (e.g., i.i.d. Gaussian weights), the criterion $ 4S(w) < \frac{\delta^2}{4} \inf_{w' \in B(w,\delta/2)} \frac{|\nabla S(w')|^2}{S(w')} $ holds with positive probability.
• Uses probabilistic arguments to show that with high probability, the initialization satisfies the convergence criterion when input dimension $ d \geq \text{some constant} $, and extends to all $ d $ via continuity and positive probability of satisfying the condition.
• Leverages bounds on derivatives of $ S(w) $ and activation functions (uniformly bounded below and above) to control the gradient-to-loss ratio and ensure the criterion is met.

Experimental results

Research questions

• RQ1 Under what conditions does gradient descent converge to a global minimum in nonconvex deep learning problems, even without growing network width?
• RQ2 Can the convergence of gradient descent to a global minimum be guaranteed for fixed-width deep neural networks under mild assumptions on data and activation functions?
• RQ3 How does the proposed convergence criterion, based on the ratio $ |\nabla f|^2 / f $, differ from or strengthen existing tools like the Kurdyka–Łojasiewicz inequality?
• RQ4 What role does initialization play in ensuring that gradient descent avoids local minima and converges globally?
• RQ5 Can the convergence criterion be satisfied with positive probability for deep networks when the input dimension is at least the number of data points?

Key findings

• Gradient descent with small step size converges to a global minimum of the training loss in deep neural networks with smooth, strictly increasing activation functions, provided the input dimension is at least the number of data points.
• The convergence is exponential in continuous time (gradient flow) and linear in discrete time (gradient descent), with explicit bounds on the rate of convergence.
• The network width can be fixed and independent of the number of data points, which contrasts with prior results requiring width to grow with data size.
• With high probability, a random initialization using i.i.d. $ \mathcal{N}(0, c/d) $ weights satisfies the convergence criterion when the input dimension $ d $ is sufficiently large.
• For all input dimensions $ d $, there exists a positive probability that the initialization satisfies the convergence criterion, ensuring almost sure convergence of gradient descent to a global minimum.
• The analysis shows that the training loss $ S(w) $ satisfies the key inequality $ 4S(w) < \frac{\delta^2}{4} \inf_{w' \in B(w,\delta/2)} \frac{|\nabla S(w')|^2}{S(w')} $ with positive probability under the specified initialization, enabling convergence.