A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks

Behrad Moniri, Donghwan Lee, Hamed Hassani, Edgar Dobriban

Primary Area: learning theory

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: feature learning, deep learning theory, random matrix theory, high dimensional asymptotics

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Abstract:

Feature learning is thought to be one of the fundamental reasons for the success of deep neural networks.
It is rigorously known that in two-layer fully-connected neural networks under certain conditions, one step of gradient descent on the first layer followed by ridge regression on the second layer can lead to feature learning; characterized by the appearance of a separated rank-one component---spike---in the spectrum of the feature matrix.
However, with a constant gradient descent step size, this spike only carries information from the linear component of the target function and therefore learning non-linear components is impossible.
We show that with a learning rate that grows with the sample size,
such training in fact introduces
multiple rank-one components,
each corresponding to a specific polynomial feature.
We further prove that the limiting large-dimensional and large sample training and test errors of the updated neural networks are fully characterized by these spikes.
By precisely analyzing the improvement in the loss, we demonstrate that these non-linear features can enhance learning.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 8250