Neural Networks can Learn Representations with Gradient Descent
This paper proves that neural networks trained via gradient descent can efficiently learn high-degree polynomials depending on few features, with sample complexity much lower than kernel methods.
Key Findings
Methodology
The authors analyze a two-layer neural network trained with gradient descent on functions of the form f*(x)=g(Ux), where U maps high-dimensional input to a low-dimensional subspace. Using Hermite polynomial expansions and the non-degeneracy of the average Hessian, they show that gradient descent naturally learns features aligned with the relevant subspace. The analysis combines theoretical bounds on sample complexity, spectral properties of the Hessian, and the dynamics of gradient updates, demonstrating feature extraction beyond the neural tangent kernel regime. The core algorithm involves symmetric initialization, gradient updates, and head regression, with rigorous bounds on generalization and transfer learning performance.
Key Results
- Sample complexity for learning degree p polynomials reduces from dp to d^2r+dr^p, a significant improvement over kernel-based dp requirement.
- In transfer learning, retraining only the output layer on new data sharing the same feature space achieves sample complexity independent of input dimension d, depending only on the latent dimension r.
- Theoretical necessity of the non-degeneracy assumption is validated; violating it leads to a lower bound requiring dp/2 samples, confirming the importance of feature learning in neural networks.
Significance
This work advances the theoretical understanding of neural networks' ability to learn complex functions efficiently by automatic feature extraction. It explains why neural networks outperform kernel methods in high-dimensional settings, especially for functions depending on few relevant directions. The findings provide a foundation for designing more sample-efficient models and transfer learning strategies, impacting both academia and industry by bridging the gap between empirical success and theoretical guarantees.
Technical Contribution
The paper introduces a novel analytical framework based on Hermite expansions and Hessian spectral properties, establishing rigorous bounds on sample complexity for learning low-dimensional polynomial functions. It demonstrates that gradient descent can learn features aligned with the relevant subspace, enabling high-order polynomial learning with significantly fewer samples than traditional kernel methods. Additionally, it proves the effectiveness of simple head retraining for transfer learning, with sample complexity depending only on the latent dimension, not the ambient input dimension.
Novelty
This is the first rigorous proof showing neural networks can learn high-degree polynomials depending on few features with sample complexity much lower than dp, surpassing kernel methods. Unlike prior NTK analyses, this work emphasizes the importance of feature learning driven by the non-degeneracy of the Hessian, providing a new theoretical perspective on deep learning’s generalization and transfer capabilities.
Limitations
- The analysis assumes data is Gaussian and functions satisfy specific non-degeneracy conditions, which may limit real-world applicability.
- Focus is on low-order polynomials; extension to more complex functions remains open.
- The theoretical bounds depend on idealized assumptions about initialization and network width, which may differ in practice.
Future Work
Future research will extend the analysis to non-Gaussian data distributions, deeper networks, and more complex function classes. Investigating the dynamics of feature learning in multi-layer architectures and exploring practical algorithms for efficient transfer learning in real-world scenarios are promising directions.
AI Executive Summary
Deep neural networks have demonstrated remarkable success across various tasks, yet their theoretical foundations remain incomplete. Traditional analyses based on the neural tangent kernel (NTK) suggest that neural networks behave like kernel methods, limiting their ability to learn complex, data-dependent features. However, empirical evidence shows that neural networks often outperform kernels significantly, especially in high-dimensional settings. This paper addresses this gap by analyzing a two-layer neural network trained with gradient descent on functions of the form f*(x)=g(Ux), where U maps the input to a low-dimensional subspace. Using tools such as Hermite polynomial expansions and spectral analysis of the average Hessian, the authors prove that gradient descent naturally learns features aligned with the relevant subspace, enabling efficient learning of high-degree polynomials with sample complexity scaling as d^2r+dr^p, far below the dp bound required by kernel methods. This breakthrough reveals the intrinsic ability of neural networks to perform feature extraction, breaking the limitations of the lazy training regime. Moreover, the study demonstrates that in transfer learning scenarios, simply retraining the output layer on new data sharing the same feature space suffices, with sample complexity depending only on the latent dimension r, independent of the high ambient dimension d. These results provide a rigorous theoretical foundation for the empirical success of deep learning, emphasizing the importance of feature learning beyond linearization. Despite the assumptions of Gaussian data and specific non-degeneracy conditions, the work opens avenues for extending the analysis to more realistic settings, promising to deepen our understanding of neural network generalization and transfer capabilities.
Deep Dive
Abstract
Significant theoretical work has established that in specific regimes, neural networks trained by gradient descent behave like kernel methods. However, in practice, it is known that neural networks strongly outperform their associated kernels. In this work, we explain this gap by demonstrating that there is a large class of functions which cannot be efficiently learned by kernel methods but can be easily learned with gradient descent on a two layer neural network outside the kernel regime by learning representations that are relevant to the target task. We also demonstrate that these representations allow for efficient transfer learning, which is impossible in the kernel regime. Specifically, we consider the problem of learning polynomials which depend on only a few relevant directions, i.e. of the form $f^\star(x) = g(Ux)$ where $U: \R^d \to \R^r$ with $d \gg r$. When the degree of $f^\star$ is $p$, it is known that $n \asymp d^p$ samples are necessary to learn $f^\star$ in the kernel regime. Our primary result is that gradient descent learns a representation of the data which depends only on the directions relevant to $f^\star$. This results in an improved sample complexity of $n\asymp d^2 r + dr^p$. Furthermore, in a transfer learning setup where the data distributions in the source and target domain share the same representation $U$ but have different polynomial heads we show that a popular heuristic for transfer learning has a target sample complexity independent of $d$.