PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Networks can Learn Representations with Gradient Descent

Alex Damian, Jason D. Lee, Mahdi Soltanolkotabi

arXiv:2206.15144v1cs.LGcs.ITstat.ML

TL;DR

Kernel or NTK methods can require d^p samples for degree-p polynomials even when targets depend on only a few directions. This paper analyzes gradient descent on two-layer networks beyond the kernel regime, showing that it learns relevant representations, improves sample complexity, and supports dimension-independent target adaptation in transfer learning.

  • Problem

    NTK analyses require n ≍ d^p samples for degree-p polynomials, even when the target depends on only a few relevant directions.

  • Method

    The paper studies gradient descent on two-layer neural networks beyond the NTK/lazy regime for low-dimensional polynomial targets f⋆ = g(Ux), including feature learning and transfer learning.

  • Results

    n ≍ d^2r + dr^p samples suffice for degree-p targets with r relevant dimensions, while retraining only the head achieves target sample complexity N ≍ r^p independent of d.

  • Takeaways & Limitations

    Gradient descent can learn task-relevant representations that separate its sample complexity from kernel methods and enable efficient transfer learning.

  • Takeaways & Limitations

    Without the non-degeneracy assumption, some single-direction polynomials require n ≍ d^(p/2) samples for any gradient-descent-based learner.

Abstract

from arXiv · show

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

1 Introduction

The paper studies how gradient descent learns task-relevant representations beyond the NTK/lazy regime, addressing why practical neural networks outperform associated kernels. For low-dimensional polynomial targets, it establishes improved sample complexity and efficient transfer learning, while identifying a necessary non-degeneracy condition.

  • NTK Limitation: NTK linearization fixes the feature representation, and precise lower bounds show its solutions cannot generalize better than the polynomial kernel.The required hyperparameter choices prevent features from evolving across iterations.
  • Feature Learning: Gradient descent learns features spanning the hidden subspace relevant to f⋆ = g(Ux), enabling vanishing training loss with a small network.These learned features support good generalization performance.
  • Improved Sample Complexity: n ≍ d^2r + dr^p samples suffice for degree-p polynomials depending on r relevant dimensions, versus d^p for random-feature or NTK methods.This breaks the sample-complexity barrier associated with NTK analyses.
  • Transfer Learning: Retraining only the network head learns a target polynomial with N ≍ r^p samples and width m ≍ r^p, independent of ambient dimension d.Learning the target from scratch would require N ≍ d^Ω(p).
  • Lower Bound: Without the non-degeneracy assumption, some single-direction degree-p polynomials require n ≍ d^(p/2) samples for any gradient-descent-based learner.The lower bound shows that the assumption is strictly necessary.

2 Setup

The setup considers noisy polynomial regression under a Gaussian input distribution, with targets depending on a low-dimensional principal subspace. The network uses a symmetrically initialized two-layer ReLU architecture, and the analysis assumes a non-degenerate expected Hessian.

  • Input Distribution and Target Function: The target f⋆: R^d → R is a normalized degree-p polynomial evaluated on x ∼ N(0, I_d).Normalization is E[f⋆(x)^2] = 1, and observations include label noise.
  • Input Distribution and Target Function: The observations satisfy x_i ∼ D and y_i = f⋆(x_i) + ε_i, with ε_i ∼ {−ς, ς}.The parameter ς controls label-noise strength.
  • Latent Structure: Arbitrary degree-p polynomials span a space of dimension Θ(d^p), so learning them generally requires n ≳ d^p samples.The setup therefore imposes low-dimensional latent structure.
  • Latent Structure: The target depends on r relevant dimensions through a function g and linearly independent vectors u_1, …, u_r, whose span defines the principal subspace S⋆.The expected Hessian is assumed non-degenerate on this subspace.
  • Non-degeneracy: The expected Hessian H = E[∇²f⋆(x)] has rank r and span(H) = S⋆ under the non-degeneracy assumption.Its normalized condition number is κ := ∥H†∥√r.
  • Network and Initialization: The model is a two-layer ReLU network with width m and symmetric initialization ensuring fθ0(x) = 0.The symmetry pairs output weights, hidden weights, and biases; small nonsymmetric random initialization is also covered.
  • Notation: The analysis uses empirical and population losses, treats p and ς as O(1), and hides polylogarithmic factors in ˜O and ˜Ω notation.The losses are defined under the Gaussian input and noisy-label model.

3 Main Results

The paper analyzes gradient-based training beyond the NTK regime for low-dimensional polynomial representations. It shows that learned features improve sample complexity and enable transfer learning, while a non-degeneracy assumption is necessary for avoiding degree-dependent hardness.

  • Training procedure: The training procedure preprocesses data, updates the first-layer weights, reinitializes biases, and then jointly updates network parameters with weight decay.The returned prediction function combines an affine term with a two-layer nonlinear network.
  • Sample complexity: n ≳ d^2r + dr^p samples suffice to learn degree-p polynomials depending on r relevant dimensions, using a small network with m ≳ r^p.The dependence on r is near optimal because the minimax complexity remains Θ(r^p) even when the relevant subspace is known.
  • Sample complexity: Resampling between stages reduces the sample complexity to d^2r + r^p by removing a factor of d from the second term.The additional factor arises from dependence between data used in the two stages, and the authors suggest sharper analysis may remove it.
  • Assumptions and lower bounds: Assumption 2 is necessary: violating it yields function classes requiring d^Ω(p) samples, whereas incorporating it permits efficient learning with d^2r + dr^p samples.The non-degeneracy condition allows the network to extract useful features for robust learning of high-degree polynomials.
  • Transfer learning: Pretraining learns the relevant subspace, after which retraining only the final layer learns any target polynomial sharing those dimensions with sample complexity independent of d.The transfer-learning guarantee requires n ≳ d^2r pretraining samples and N, m ≳ r^p target-side quantities.

4 Related work

Related work connects neural-network training to kernel, mean-field, and non-lazy feature-learning analyses. The paper positions its setting against methods that lack feature learning, require restrictive widths or rates, or rely on special data distributions.

  • NTK and kernel methods: NTK linearization treats gradient descent as optimization with a fixed feature representation, enabling convergence and generalization analyses in the kernel regime.The approach commonly uses large initialization, small learning rates, or wide networks.
  • NTK and kernel methods: Because NTK methods cannot learn new features, degree-p polynomials in dimension d require Θ(d^p) samples, matching polynomial-kernel scaling.Prior empirical work also reports a test-accuracy gap between kernel predictors and deep networks trained with SGD.
  • Beyond the lazy regime: Non-lazy studies analyze feature learning through small initialization, spectral bias, low-rank structure, orthogonal weights, or higher-order training-dynamics approximations.These results span low-rank reconstruction, tensor decomposition, and neural networks with orthogonal weights.
  • Related feature-learning settings: For isotropic Gaussian data, relevant features must be extracted from higher-order correlations between data and labels rather than from distributional imbalance.This distinguishes the paper from parity-learning results where an unbalanced distribution directs gradients toward the planted set.
  • Mean-field analyses: Mean-field analyses couple wide-network dynamics to probability distributions described by PDEs, but generally lack explicit convergence rates and require unrealistically large widths.A concurrent Boolean-cube study identifies the merged staircase property under zero hidden-weight initialization, which can hinder feature learning.

5 Proof Sketches

The proof shows that initialization gradients reveal the target’s relevant subspace, after which fixed-feature regression exploits the reduced dimension. The analysis also establishes statistical-query lower bounds through nearly orthogonal function families.

  • Gradient expansion: The gradient expansion is an asymptotic series in d^-1/2: higher-order terms reveal more information about f⋆ but require at least d^k samples to extract.The analysis focuses on the first informative term, involving the average Hessian H, which requires n ≥ d^2 samples to isolate.
  • Feature learning: The population gradient at initialization points hidden features toward the principal subspace S⋆, and the gradients span that subspace under Assumption 2.This follows because the gradient depends only on S⋆ and the initialization is symmetric.
  • Feature learning: n ≥ d^2 samples are required to extract the relevant subspace information from empirical gradients, forming the dominant term in the sample complexity.The population signal is order d^-1/2, while only empirical gradients are observed.
  • Feature learning: After one gradient step, the network resembles a random-feature model whose features lie in the hidden subspace S⋆, reducing the polynomial-learning burden from d dimensions to r dimensions.The resulting reduced-dimensional regime suggests n ≳ r^p samples for learning degree-p polynomials.
  • Generalization: The remaining algorithm fixes the hidden features and applies ridge regression to the output head, with generalization controlled using Rademacher complexity.The proof constructs a suitable head, relates ridge regression to norm-constrained regression, and applies a standard neural-network bound.
  • Statistical-query lower bound: A correlational statistical-query lower bound is obtained by constructing many normalized functions with small pairwise correlations, limiting what tolerance-τ queries can learn.Approximately orthogonal vectors supply the construction, while Hermite features transfer their geometric separation into small correlations.

6 Experiments

The experiments compare Algorithm 1 with random features and the neural tangent kernel on high-degree, low-dimensional examples. They report sharp sample-efficiency advantages for representation learning and transfer after sufficient pretraining.

  • 6.1 Sample Complexity: Algorithm 1 isolates the single relevant direction u and then fits a one-dimensional random-feature model to the target function.The full process generalizes with n ≍ d^2 samples.
  • 6.1 Sample Complexity: n ≍ d^p samples are necessary for NTK and random-feature methods, whereas Algorithm 1 learns the relevant degree-p structure with substantially fewer samples.For the example, kernel methods plateau at the optimal quadratic predictor while Algorithm 1 reaches near-zero excess risk.
  • Experimental setup: The experiments use d = 10, p = 4, label-noise strength σ^2 = 1, widths m = 100 and m = 1000, and averages over 10 random seeds.Weight decay is tuned on a holdout set, with performance evaluated on a separate test set.
  • 6.1 Sample Complexity: At width m = 100, random features and the NTK fit only the quadratic term 1/2 He2(u · x), matching theoretical predictions.Algorithm 1 instead exhibits a sudden drop from trivial risk to vanishing excess risk without the intermediate plateau.
  • 6.2 Transfer Learning: The transfer experiment pretrains on f⋆ and then trains the output layer for a target polynomial sharing the same relevant direction.The target head is a different degree-3 Hermite polynomial, while the representation is reused.
  • 6.2 Transfer Learning: n ≥ d^2 pretraining samples enable fine-tuning to achieve nontrivial target loss with N = O(1) target samples.With n = d^0 or d^1, target fine-tuning remains at trivial risk until N ≳ d^3.

7 Discussion and Future Work

The discussion separates gradient-based training from kernel methods by emphasizing learned representations as the shared mechanism behind sample-efficiency and transfer results. It identifies tighter dimension dependence and joint-layer training as key future directions.

  • Discussion: Gradient descent efficiently learns a broad family of degree-p polynomials with n ≍ d^2 samples, contrasting with the d^p lower bound for random-feature and NTK analyses.The separation is attributed to representation learning beyond the kernel regime.
  • Discussion: Useful learned representations are the central idea connecting the sample-complexity result and the transfer-learning result.The same representation-learning mechanism underlies Theorems 1 and 3.
  • Future work: The upper bound’s dimension dependence remains potentially loose: its d^2 term comes from learning degree-2 Hermite structure, while the lower bound rules out only n ≤ d.Sparse degree-2 polynomials are known to be efficiently learnable with n ≍ d.
  • Future work: Jointly training the hidden and output layers remains an open direction because dependencies between the layers are difficult to control.The authors suggest that analyzing this setting could clarify learning order and inductive bias.

A Proofs

The appendix defines high-probability events using an exponentially small failure term and conditions the remaining proofs on such an event. The definition is designed to remain stable under polynomially many union bounds.

  • High-probability events: The analysis sets ι = Cι log(nmd) and assumes ι ≤ cd for a sufficiently small absolute constant c.This parameter tracks failure probabilities throughout the appendix.
  • High-probability events: High probability means an event occurs with probability at least 1 − poly(n, m, d)e^-ι.The polynomial factor does not depend on the sufficiently large constant Cι.
  • High-probability events: High-probability events remain valid under union bounds over sets of size poly(n, m, d), and subsequent proofs are conditioned on the event.This closure supports applying multiple lemmas simultaneously.

A.1 Hermite Expansions

This section develops Hermite-based identities and concentration bounds for ReLU features, including the features learned after one gradient-descent step.

  • ReLU is defined as σ(x) = max(0, x), and its Hermite expansion is introduced.
  • Scale invariance makes Algorithm 1 independent of the hidden-weight norms, allowing the analysis to assume unit-norm weights uniformly distributed on the sphere.
  • The residual function is defined by subtracting the constant and linear components from f⋆.
  • The functions g(w) and gn(w) capture features that can be learned after one gradient-descent step.
  • The analysis bounds dependent and independent random terms involving gn(w) and the data using concentration lemmas.

A.2.1 Univariate Random Feature Approximation

This section shows that reinitialized biases enable random features to transform ReLU into degree-p polynomial features, supporting polynomial approximation.

  • Reinitialized-bias random features transform σ(x) = ReLU(x) into σ(x) = x^p, which is more natural for learning polynomials.
  • For random signs a and bounded random variables b, the construction provides functions vk(a, b) representing degree-k behavior on |x| ≤ 1.
  • The construction also works when b is Gaussian, yielding the corresponding univariate random-feature representation on |x| ≤ 1.

A.2.2 Multivariable Random Feature Approximation

This section establishes that multivariable random features can efficiently represent degree-p polynomials restricted to a low-dimensional principal subspace, under stated sample and dimension conditions.

  • Additional concentration arguments establish the required high-probability bounds for the constructed random features and their moments.
  • Random features gn(w) are sufficiently expressive to represent any polynomial of degree p restricted to the principal subspace S⋆.
  • The tensor projection ΠSymk(S⋆) identifies the symmetric k-tensor components restricted to the principal subspace.
  • Under n ≥ C d^2 r κ^2 ι^(p+1) and d ≥ Cκr^(3/2), tensor-supported representations exist for every k ≤ p.
  • The constructed representation has bounds involving (rdκ^2)^k, the tensor Frobenius norm, and the feature norm.

A.3 Proof of Theorem 1

This section assembles the approximation and concentration lemmas to prove the theorem for two-layer networks, using empirical losses, tensor representations, and norm-constrained regression.

  • The proof treats empirical L2 losses with noisy labels yi = f⋆(xi) + ϵi, where ϵi takes values in {−σ, σ}.
  • The theorem assumes n ≥ C d^2 r κ^2 ι^(p+1) and m ≥ r^p κ^(2p) ι^(6p+1), otherwise it returns the zero predictor.
  • The proof constructs a parameter vector θ⋆ using a coefficient vector a⋆ and the first-step weights and biases.
  • Gradient-descent iterates can approximate the constructed predictor while preserving a loss bound and controlling the output-layer norm.
  • Concentration tools include sphere nets, Gaussian and chi-squared bounds, Hoeffding’s inequality, and Bernstein’s inequality.
  • The tensor argument uses symmetric tensors, contractions, Hermite-to-monomial basis changes, and projections onto symmetric tensor spaces.
Loading 2206.15144v1…