Learning Hierarchical Polynomials of Multiple Nonlinear Features with Three-Layer Networks

In deep learning theory, a critical question is to understand how neural networks learn hierarchical features. In this work, we study the learning of hierarchical polynomials of textitmultiple nonlinear features using three-layer neural networks. We examine a broad class of functions of the form fstar=gstarcirc bp, where bp:ℝd rightarrow ℝr represents multiple quadratic features with r ll d and gstar:ℝrrightarrow ℝ is a polynomial of degree p. This can be viewed as a nonlinear generalization of the multi-index model citepdamian2022neural, and also an expansion upon previous work that focused only on a single nonlinear feature, i.e. r = 1 citepnichani2023provable,wang2023learning. Our primary contribution shows that a three-layer neural network trained via layerwise gradient descent suffices for beginitemizeitem complete recovery of the space spanned by the nonlinear features item efficient learning of the target function fstar=gstarcirc bp or transfer learning of f=gcirc bp with a different link function enditemize within widetildecO(d4) samples and polynomial time. For such hierarchical targets, our result substantially improves the sample complexity Theta(d2p) of the kernel methods, demonstrating the power of efficient feature learning. It is important to highlight that our results leverage novel techniques and thus manage to go beyond all prior settings such as single-index and multi-index models as well as models depending just on one nonlinear feature, contributing to a more comprehensive understanding of feature learning in deep learning.