Learning Over-Parametrized Two-Layer ReLU Neural Networks beyond NTK

TL;DR

This paper demonstrates that over-parameterized two-layer ReLU neural networks trained with gradient descent can efficiently learn certain functions beyond the capabilities of kernel methods like NTK, with provable guarantees.

Contribution

It provides the first theoretical proof that over-parameterized neural networks can learn beyond the NTK regime in polynomial time with polynomial samples.

Findings

Neural networks achieve population loss at most o(1/d).

Kernel methods have population loss at least Ω(1/d).

Gradient descent can learn the target function efficiently.

Abstract

We consider the dynamic of gradient descent for learning a two-layer neural network. We assume the input $x\in\mathbb{R}^d$ is drawn from a Gaussian distribution and the label of $x$ satisfies $f^{\star}(x) = a^{\top}|W^{\star}x|$, where $a\in\mathbb{R}^d$ is a nonnegative vector and $W^{\star} \in\mathbb{R}^{d\times d}$ is an orthonormal matrix. We show that an over-parametrized two-layer neural network with ReLU activation, trained by gradient descent from random initialization, can provably learn the ground truth network with population loss at most $o(1/d)$ in polynomial time with polynomial samples. On the other hand, we prove that any kernel method, including Neural Tangent Kernel, with a polynomial number of samples in $d$, has population loss at least $\Omega(1 / d)$.