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

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Moniri, Behrad, Lee, Donghwan, Hassani, Hamed, Dobriban, Edgar
Formato: Preprint
Publicado: 2023
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866908310252290048
author Moniri, Behrad
Lee, Donghwan
Hassani, Hamed
Dobriban, Edgar
author_facet Moniri, Behrad
Lee, Donghwan
Hassani, Hamed
Dobriban, Edgar
contents 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 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 training and test errors, we demonstrate that these non-linear features can enhance learning.
format Preprint
id arxiv_https___arxiv_org_abs_2310_07891
institution arXiv
publishDate 2023
record_format arxiv
spellingShingle A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks
Moniri, Behrad
Lee, Donghwan
Hassani, Hamed
Dobriban, Edgar
Machine Learning
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 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 training and test errors, we demonstrate that these non-linear features can enhance learning.
title A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks
topic Machine Learning
url https://arxiv.org/abs/2310.07891