Title: Neural Networks Learn Multi-Index Models Near the Information-Theoretic Limit
Primary area: learning theory
Keywords: Representation Learning, Multi-Index Models, Two-Layer Network, Gradient Descent, Sample Complexity

Abstract: In deep learning, a central issue is to understand how neural networks efficiently learn high-dimensional features. To this end, we explore the gradient descent learning of a general Gaussian Multi-index model $f(\boldsymbol{x})=g(\boldsymbol{U}\boldsymbol{x})$ with hidden subspace $\boldsymbol{U}\in \mathbb{R}^{r\times d}$, which is the canonical setup to study representation learning. Our results show that under mild generic assumptions on the link function, a standard two-layer neural network trained via layer-wise gradient descent can agnostically learn the target with $o_d(1)$ test error using $\widetilde{\mathcal{O}}(d)$ samples and $\widetilde{\mathcal{O}}(d^2)$ time. The sample and time complexity both align with the information-theoretic limit up to leading order and are therefore optimal. During the first stage of training, the proof proceeds via showing that the inner weights can perform a self-correcting process that gradually reduces finite-sample noise and recovers the hidden subspace. It surprisingly indicates that optimal results can only be achieved if the first layer is trained for more than $\mathcal{O}(1)$ steps. This work demonstrates the ability of neural networks to effectively learn hierarchical functions with respect to both sample and time efficiency.

TL;DR: We prove that neural networks trained via GD can efficiently learn Multi-Index models near the information-theoretic limit.