Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks

Eshaan Nichani, Alex Damian, Jason D. Lee

TL;DR

This work analyzes how three-layer neural networks can learn nonlinear, hierarchical features that two-layer models cannot efficiently learn. It proves a general theorem that layer-wise gradient descent on a three-layer net learns a low-frequency learned feature approximating $\overline\eta\cdot\mathbb{K}f^*$, enabling hierarchical learning when $f^*=g^*\circ h^*$. The framework is instantiated for single-index models and quadratic-feature targets, yielding improved sample complexities (e.g., $n=\tilde{\Theta}(d)$ for single-index, $n=\tilde{\Theta}(d^4)$ or $\tilde{\Theta}(d^2)$ for quadratic features) relative to two-layer/kernel methods, and a concrete optimization-based depth separation showing 3-layer networks can be learned efficiently where 2-layer networks cannot. These results shed light on the provable benefits of depth in feature learning and open avenues for extending hierarchical feature-learning guarantees to broader function classes and datasets.

Abstract

One of the central questions in the theory of deep learning is to understand how neural networks learn hierarchical features. The ability of deep networks to extract salient features is crucial to both their outstanding generalization ability and the modern deep learning paradigm of pretraining and finetuneing. However, this feature learning process remains poorly understood from a theoretical perspective, with existing analyses largely restricted to two-layer networks. In this work we show that three-layer neural networks have provably richer feature learning capabilities than two-layer networks. We analyze the features learned by a three-layer network trained with layer-wise gradient descent, and present a general purpose theorem which upper bounds the sample complexity and width needed to achieve low test error when the target has specific hierarchical structure. We instantiate our framework in specific statistical learning settings -- single-index models and functions of quadratic features -- and show that in the latter setting three-layer networks obtain a sample complexity improvement over all existing guarantees for two-layer networks. Crucially, this sample complexity improvement relies on the ability of three-layer networks to efficiently learn nonlinear features. We then establish a concrete optimization-based depth separation by constructing a function which is efficiently learnable via gradient descent on a three-layer network, yet cannot be learned efficiently by a two-layer network. Our work makes progress towards understanding the provable benefit of three-layer neural networks over two-layer networks in the feature learning regime.

Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks

TL;DR

This work analyzes how three-layer neural networks can learn nonlinear, hierarchical features that two-layer models cannot efficiently learn. It proves a general theorem that layer-wise gradient descent on a three-layer net learns a low-frequency learned feature approximating , enabling hierarchical learning when . The framework is instantiated for single-index models and quadratic-feature targets, yielding improved sample complexities (e.g., for single-index, or for quadratic features) relative to two-layer/kernel methods, and a concrete optimization-based depth separation showing 3-layer networks can be learned efficiently where 2-layer networks cannot. These results shed light on the provable benefits of depth in feature learning and open avenues for extending hierarchical feature-learning guarantees to broader function classes and datasets.

Abstract

One of the central questions in the theory of deep learning is to understand how neural networks learn hierarchical features. The ability of deep networks to extract salient features is crucial to both their outstanding generalization ability and the modern deep learning paradigm of pretraining and finetuneing. However, this feature learning process remains poorly understood from a theoretical perspective, with existing analyses largely restricted to two-layer networks. In this work we show that three-layer neural networks have provably richer feature learning capabilities than two-layer networks. We analyze the features learned by a three-layer network trained with layer-wise gradient descent, and present a general purpose theorem which upper bounds the sample complexity and width needed to achieve low test error when the target has specific hierarchical structure. We instantiate our framework in specific statistical learning settings -- single-index models and functions of quadratic features -- and show that in the latter setting three-layer networks obtain a sample complexity improvement over all existing guarantees for two-layer networks. Crucially, this sample complexity improvement relies on the ability of three-layer networks to efficiently learn nonlinear features. We then establish a concrete optimization-based depth separation by constructing a function which is efficiently learnable via gradient descent on a three-layer network, yet cannot be learned efficiently by a two-layer network. Our work makes progress towards understanding the provable benefit of three-layer neural networks over two-layer networks in the feature learning regime.
Paper Structure (48 sections, 49 theorems, 304 equations, 2 figures, 1 algorithm)

This paper contains 48 sections, 49 theorems, 304 equations, 2 figures, 1 algorithm.

Table of Contents

  1. Revision (April 2, 2025):
  2. Introduction
  3. Our contributions
  4. Related Work
  5. Neural Networks and Kernel Methods.
  6. Feature Learning.
  7. Depth Separations.
  8. Preliminaries
  9. Problem Setup
  10. Data distribution.
  11. Three-layer neural network.
  12. Training Algorithm.
  13. Technical definitions
  14. Notation
  15. Main Result
  16. ...and 33 more sections

Key Result

Theorem 1

Select $q \in \mathcal{W}^{2, \infty}([-1, 1])$. Let $\eta_1 = \frac{m_1}{m_2}\overline \eta$, and assume $n = \tilde{\Omega}(\norm{\mathbb{K}}_{op}\norm{\mathbb{K} f^*}_{L^2}^{-2})$, $m_2 = \tilde{\Omega}(\norm{\mathbb{K} f^*}_{L^2}^{-1})$ There exist $\overline \eta, \lambda, \eta_2$ such that aft

Figures (2)

  • Figure 1: We ran \ref{['alg:two']} on both the single index and quadratic feature settings described in \ref{['sec:examples']}. Each trial was run with 5 random seeds. The solid lines represent the medians and the shaded areas represent the min and max values. For every trial we recorded both the test loss on a test set of size $2^{15}$ and the linear correlation between the learned feature map $\phi(x)$ and the true intermediate feature $h^\star(x)$ where $h^\star(x) = x \cdot \beta$ for the single index setting and $h^\star(x) = x^T A x$ for the quadratic feature setting. Our results show that the test loss goes to $0$ as the linear correlation between the learned feature map $\phi$ and the true intermediate feature $h^\star$ approaches $1$.
  • Figure 2: Three-layer neural networks learn the target $f^*(x) = \mathop{\mathrm{ReLU}}\nolimits(x^TAx)$ with better sample complexity than two-layer networks.

Theorems & Definitions (100)

  • Definition 1: Sub-Gaussian Vector
  • Definition 2: Kernel objects
  • Definition 3
  • Theorem 1
  • Theorem 2
  • Theorem 3
  • Corollary 1
  • Theorem 4
  • Definition 4: high probability events
  • Example 5
  • ...and 90 more