Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions

TL;DR

This paper demonstrates that data repetition in training neural networks significantly enhances their ability to learn high-dimensional multi-index functions efficiently, surpassing previous theoretical limitations.

Contribution

It introduces a simple modification to gradient descent with data repetition, enabling neural networks to learn relevant features in high-dimensional settings more effectively.

Findings

Data repetition improves learning efficiency in high dimensions.

Networks can learn almost all directions within O(d log d) steps.

Hierarchical mechanisms enable learning of complex coupled functions.

Abstract

Neural networks can identify low-dimensional relevant structures within high-dimensional noisy data, yet our mathematical understanding of how they do so remains scarce. Here, we investigate the training dynamics of two-layer shallow neural networks trained with gradient-based algorithms, and discuss how they learn pertinent features in multi-index models, that is target functions with low-dimensional relevant directions. In the high-dimensional regime, where the input dimension $d$ diverges, we show that a simple modification of the idealized single-pass gradient descent training scenario, where data can now be repeated or iterated upon twice, drastically improves its computational efficiency. In particular, it surpasses the limitations previously believed to be dictated by the Information and Leap exponents associated with the target function to be learned. Our results highlight the ability of networks to learn relevant structures from data alone without any pre-processing. More precisely, we show that (almost) all directions are learned with at most $O(d \log d)$ steps. Among the exceptions is a set of hard functions that includes sparse parities. In the presence of coupling between directions, however, these can be learned sequentially through a hierarchical mechanism that generalizes the notion of staircase functions. Our results are proven by a rigorous study of the evolution of the relevant statistics for high-dimensional dynamics.