Neural Networks can Learn Representations with Gradient Descent (2206.15144v1)

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$.

• The paper shows that neural networks use gradient descent to learn task-relevant representations that outperform static kernel methods.
• It establishes that dynamic representation learning significantly reduces sample complexity compared to conventional approaches.
• The work highlights transfer learning advantages and emphasizes the necessity of non-degeneracy in the Hessian for effective learning.

Overview of "Neural Networks can Learn Representations with Gradient Descent"

The paper "Neural Networks can Learn Representations with Gradient Descent" authored by Alex Damian, Jason D. Lee, and Mahdi Soltanolkotabi offers significant insights into understanding the capability of neural networks trained via gradient descent to learn representations. Specifically, it addresses why neural networks often outperform kernel methods, despite theoretical similarities in certain regimes.

Key Contributions

1. Representation Learning Capability: The paper convincingly demonstrates that neural networks can learn task-relevant representations via gradient descent, which enables the learning of function classes that are challenging for kernel methods. This assertion is supported by considering the learning of polynomials dependent on a few relevant directions. By examining the problem of learning functions of the form $f^\star(x) = g(Ux)$ where dimensions $d \gg r$, the authors illustrate that gradient descent can capture intrinsic data geometry and learn with fewer samples than kernel methods, which require $n \asymp d^p$ samples due to their inability to dynamically learn new representations.
2. Improved Sample Complexity: It is shown through rigorous theoretical analysis that gradient descent only needs $n \asymp d^2 r + dr^p$ samples, a marked improvement over kernel methods. This distinction allows neural networks to not only generalize more efficiently but also leverage transfer learning effectively.
3. Transfer Learning Potential: The authors outline how this representation learning process facilitates efficient transfer learning. In scenarios where data distributions between source and target domains share a latent representation $U$, neural networks can perform well in target tasks with sample complexity independent of dimension $d$. This is impossible in the kernel regime, highlighting the neural networks' advantage in flexibility and adaptability.
4. Necessity of Non-degeneracy Assumptions: The paper posits that a non-degeneracy assumption, whereby the expected Hessian possesses full rank corresponding to relevant dimensions, is critical. Without it, learning efficiencies can degrade significantly, requiring sample complexities of $d^{p/2}$ for learning via gradient descent.

Implications and Future Work

• Practical Impact: For real-world applications, such as image and speech recognition, where low-dimensional structures are often latent in high-dimensional data, these theoretical insights imply substantial efficiencies can be achieved. This is vital for systems where computational resources are constrained.
• Theoretical Expansion: While the paper primarily addresses two-layer neural networks, extending the analysis to deeper architectures could provide further understanding of neural networks' hierarchical representation capabilities, possibly closing the gap between empirical success and theoretical underpinning.
• Refinement of Assumptions: Future work might explore reducing reliance on strong assumptions like non-degeneracy, exploring other conditions or constraints that could broaden the applicability of the theoretical results.

Overall, this paper contributes to the ongoing discourse on the theoretical foundation of neural networks, explicitly indicating neural networks' potential beyond the lazy regime and opening avenues for improved learning across varied tasks and settings in AI.