[Paper Review] A Local Convergence Theory for Mildly Over-Parameterized Two-Layer Neural Network

Mo Zhou, Rong Ge+1 more|arXiv (Cornell University)|Feb 4, 2021

Neural Networks and Applications15 references7 citations

TL;DR

This paper establishes a local convergence theory for mildly over-parameterized two-layer neural networks, showing that gradient descent converges to global optima where all student neurons align with teacher neurons when the loss is below a threshold polynomial in separation and width parameters. The key contribution is a novel characterization of the optimization landscape via a special Lojasiewicz-type gradient property, enabling convergence independent of network size and valid beyond the Neural Tangent Kernel regime.

ABSTRACT

While over-parameterization is widely believed to be crucial for the success of optimization for the neural networks, most existing theories on over-parameterization do not fully explain the reason -- they either work in the Neural Tangent Kernel regime where neurons don't move much, or require an enormous number of neurons. In practice, when the data is generated using a teacher neural network, even mildly over-parameterized neural networks can achieve 0 loss and recover the directions of teacher neurons. In this paper we develop a local convergence theory for mildly over-parameterized two-layer neural net. We show that as long as the loss is already lower than a threshold (polynomial in relevant parameters), all student neurons in an over-parameterized two-layer neural network will converge to one of teacher neurons, and the loss will go to 0. Our result holds for any number of student neurons as long as it is at least as large as the number of teacher neurons, and our convergence rate is independent of the number of student neurons. A key component of our analysis is the new characterization of local optimization landscape -- we show the gradient satisfies a special case of Lojasiewicz property which is different from local strong convexity or PL conditions used in previous work.

Motivation & Objective

• To explain why mildly over-parameterized two-layer neural networks converge to global optima in practice, despite limitations of existing theories.
• To address the gap in understanding local convergence when the number of neurons is only slightly larger than the teacher network.
• To develop a theory that allows neurons to move significantly from initialization, unlike the Neural Tangent Kernel regime.
• To characterize the local optimization landscape using a new gradient property distinct from strong convexity or PL conditions.
• To prove convergence to global optima with a rate independent of the number of student neurons.

Proposed method

• Introduce a new local optimization landscape characterization based on a special case of the Lojasiewicz property, which governs gradient behavior near global optima.
• Define a threshold τ = poly(Δ/r) based on the separation Δ and number of teacher neurons r, below which convergence is guaranteed.
• Use gradient descent with a step size η ≤ O(min{r^{-1/2}w_{max}^{-1/2}κ, r^{-1}w_{max}^{-1}}) to ensure descent and convergence.
• Establish concentration bounds for stochastic gradients using finite-sample estimates, ensuring robustness to data sampling error.
• Prove that the loss decreases at a rate O(1/t) when below threshold, with convergence to zero loss.
• Leverage the Lojasiewicz-type condition to derive a linear convergence rate in terms of loss squared, independent of network width.

Experimental results

Research questions

• RQ1 Can gradient descent converge to global optima in mildly over-parameterized two-layer networks when the initial loss is below a certain threshold?
• RQ2 Does the convergence behavior remain robust regardless of the number of student neurons, as long as it is at least as large as the number of teacher neurons?
• RQ3 What is the nature of the local optimization landscape that enables convergence without relying on strong convexity or PL conditions?
• RQ4 How does the convergence rate depend on network width and initialization, and can it be independent of the number of neurons?
• RQ5 Can the theory explain the empirical observation that student neurons align with teacher neuron directions in practical training?

Key findings

• The paper proves that if the training loss is below τ = poly(Δ/r), gradient descent converges to a global optimum where all student neurons match the directions of teacher neurons.
• Convergence occurs with a rate O(1/t), independent of the number of student neurons, even when the network is only mildly over-parameterized.
• The key mechanism is a novel Lojasiewicz-type gradient property that holds locally around global optima, distinct from PL or strong convexity.
• The convergence threshold τ depends only on the separation Δ and number of teacher neurons r, not on the size of the student network.
• The analysis holds for any number of student neurons ≥ r, and the convergence rate is robust to finite-sample stochastic gradient estimates.
• The theory explains the empirical phenomenon of neuron matching in over-parameterized settings, even when neurons move significantly from initialization.