Source-linked AI summary
Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks
Ziwei Ji, Matus Telgarsky
TL;DR
Existing guarantees for overparameterized networks often require width polynomial in sample size, target accuracy, or failure probability. This paper analyzes two-layer ReLU networks in the NTK regime and shows that gradient descent and SGD achieve arbitrarily small test error with polylogarithmic width, under sample and iteration bounds governed by ε. The separation margin of the limiting kernel supports these guarantees and yields a tight infinite-width sample-complexity result for noisy 2-XOR.
Problem
Prior results require network width polynomial in n, 1/δ, or 1/ε, leaving a gap between theoretical requirements and empirical widths.
Method
The paper analyzes gradient-based training of two-layer ReLU networks in the NTK regime using the separation margin of the limiting kernel and limited parameter movement.
Results
eΩ(1/ε^2) samples and eΘ(1/ε) iterations suffice for gradient descent to achieve ε test error with polylogarithmic width, while SGD achieves ε test error with eΘ(1/ε) samples.
Takeaways & Limitations
The separation margin can distinguish true from random labels and gives a tight infinite-width sample-complexity analysis on noisy 2-XOR data.
Takeaways & Limitations
The analysis is restricted to the NTK regime, binary classification, and two-layer networks; extending it beyond these settings remains open.
Abstract
from arXiv · showhide
Recent theoretical work has guaranteed that overparameterized networks trained by gradient descent achieve arbitrarily low training error, and sometimes even low test error. The required width, however, is always polynomial in at least one of the sample size $n$, the (inverse) target error $1/ε$, and the (inverse) failure probability $1/δ$. This work shows that $\widetildeΘ(1/ε)$ iterations of gradient descent with $\widetildeΩ(1/ε^2)$ training examples on two-layer ReLU networks of any width exceeding $\mathrm{polylog}(n,1/ε,1/δ)$ suffice to achieve a test misclassification error of $ε$. We also prove that stochastic gradient descent can achieve $ε$ test error with polylogarithmic width and $\widetildeΘ(1/ε)$ samples. The analysis relies upon the separation margin of the limiting kernel, which is guaranteed positive, can distinguish between true labels and random labels, and can give a tight sample-complexity analysis in the infinite-width setting
1 Introduction
The paper addresses the gap between empirically modest network widths and theory requiring polynomial width, showing that polylogarithmic-width two-layer ReLU networks can achieve low training and test error under gradient-based training.
- Prior analyses require width polynomial in n, 1/δ, or 1/ε, creating a substantial gap between theoretical guarantees and empirical practice.
- Gradient descent reaches training error ε in eΘ(1/ε) iterations with any polylogarithmic width and constant step size at most 1.
- The generalization analysis combines limited parameter movement with Rademacher complexity bounds to derive the test-error guarantee.
- Polylogarithmic width suffices for gradient descent to reach ε test error using eΩ(1/ε^2) samples and eΘ(1/ε) iterations.
- SGD achieves arbitrarily small test error with polylogarithmic width and eΘ(1/ε) samples.
- The separation margin is positive, can fall to O(1/√n) for random labels, and reaches Ω(1/ln(n)) on noisy 2-XOR data, where the SGD sample complexity is tight.
2 Empirical risk minimization
The section establishes empirical-risk guarantees for two-layer ReLU networks under NTK separability, showing that gradient descent reaches a target risk with polylogarithmic width apart from margin dependence. The proof tracks positive margins from initialization through training and controls the optimization dynamics.
- NTK separability: Theorem 2.2 assumes the training set is separable in the infinite-width NTK feature space with margin γ and a bounded separator.This assumption formalizes the limiting-kernel separability condition used throughout the empirical-risk analysis.
- Main guarantee: T = ⌈2λ^2/(ηǫ)⌉ gradient-descent iterations suffice for empirical risk at most ǫ with probability 1−3δ when η ≤ 1 and m ≥ M.The guarantee holds over random initialization under the theorem’s stated conditions.
- Main guarantee: A two-layer ReLU network with width polylogarithmic in n, 1/δ, and 1/ǫ achieves empirical risk at most ǫ, unlike prior results requiring polynomial width.The width remains polynomial in 1/γ, an adaptive quantity determined by the separation margin.
- Margin control: The analysis proves that the network has a positive margin at initialization and maintains a good margin during training.These properties are supplied by separate high-probability lemmas controlling initialization and the training trajectory.
- Optimization analysis: The proof uses the quantity bQ to control gradient-related terms and avoid an iteration-dependent width factor, enabling polylogarithmic dependence on 1/ǫ.This control is central because the analysis requires Ω(1/ǫ) iterations.
3 Generalization
Under a separability assumption in the first-layer NTK-induced RKHS, the paper derives a test-error bound for gradient descent on shallow networks with sufficiently large width. The result uses a generalization analysis based on the function space explored during training.
- Assumption: The analysis assumes a positive-margin separator in the RKHS induced by the first-layer weights W.This differs from related work that assumes separability in the RKHS induced by the second-layer weights.
- Test error: eΘ(1/ǫ) iterations suffice for gradient descent to find a solution with ǫ test error under the separability assumption.The theorem selects the iterate with minimum empirical risk before a specified iteration bound.
- Proof strategy: The proof controls population risk by bounding the generalization error through the Rademacher complexity of the function space explored by gradient descent.The sigmoid loss derivative is supported on [0, 1] and is 1-Lipschitz, enabling this reduction.
- Caveat: A smoothness-based alternative gives sample complexity eO(1/γ^4ǫ), but introduces complicated bounds and a large constant.The paper identifies a clean smoothness-based analysis as an open question.
4 Stochastic gradient descent
For online stochastic gradient descent, the paper samples one example per update while fixing the second-layer weights. Under the same separability assumption, it establishes a test-error guarantee using a polylogarithmic-width network and a sample bound expressed through the margin.
- Algorithm: SGD samples an example (x_i, y_i) from the data distribution at each step and updates the first-layer weights while fixing the second-layer weights.The formulation uses an online oracle and random initialization.
- Guarantee: Under Assumption 3.1, Theorem 4.1 gives an ǫ test-error guarantee for constant step size and width m = Ω(...).The displayed theorem statement is truncated in the supplied passage, so the complete width expression is not recoverable here.
- Proof strategy: The proof extends the gradient-descent lemmas to SGD and bounds cumulative population risk using martingale concentration.The key observation is that the relevant risk-difference process is a martingale, allowing a martingale Bernstein bound.
5 On separability
This section characterizes separability through the limiting NTK margin, showing when it is positive, how labels affect it, and how it controls sample complexity. It also analyzes the noisy 2-XOR distribution as a case where feature-label structure yields a larger margin.
- NTK separability: The NTK separation margin γ1 has a dual characterization, and when γ1 > 0 it also admits a unit-norm primal classifier with margin γ1.The paper uses this margin to formulate separability assumptions for the limiting kernel.
- NTK separability: When no inputs are parallel, the NTK separability assumption always holds.A sufficient condition based on pairwise input separation gives a positive minimum eigenvalue for K1.
- Label dependence: Replacing true labels with random labels destroys a potentially large margin, consistent with the margin becoming small for random labels.The paper gives a high-probability upper bound on γ1 under random relabeling.
- Label dependence: The required width depends polynomially on 1/γ, which can be polynomial in n for random labels but polylogarithmic in n when features strongly relate to labels.The separation margin is adaptive to the data-label relationship.
- Noisy 2-XOR: A narrow network may be nonseparable on noisy 2-XOR, and widths below the stated margin-dependent scale require analysis beyond the NTK regime.The paper proves a narrow-width obstruction for certain subsets and random initializations.
- Noisy 2-XOR: For noisy 2-XOR, the constructed separator has margin γ = Ω(1/d), so the inverse margin is O(d).The separator depends only on the first two coordinates, while the distribution is supported on 2d points.
- Noisy 2-XOR: The noisy 2-XOR margin yields an infinite-width NTK and SGD sample complexity of eO(1/γ^2ε), which becomes eO(d^2) for constant test accuracy.This matches the d^2 lower bound up to logarithmic factors.
6 Open problems
The paper identifies several directions beyond its current analysis, including learning outside the NTK regime and extending the results to broader architectures and tasks.
- Beyond the NTK: An open question is whether gradient descent learns beyond the NTK after iterates move sufficiently far from their initial weights.The present analysis assumes weights remain close to initialization.
- Architectures: The analysis could be extended to multi-layer, convolutional, and residual networks.The current results concern shallow two-layer networks.
- Tasks: The paper studies binary classification, leaving open whether similar results hold for regression and other tasks.The broader task generality is not established here.
A Omitted proofs from Section 2
The omitted proofs establish concentration, stability, and risk bounds used to prove that gradient descent reaches low training error while remaining near initialization.
- Concentration bounds: Hoeffding’s inequality controls empirical means of bounded random variables appearing in the initialization analysis.These bounds are applied separately across examples and combined by union bounds.
- Concentration bounds: The proofs use Gaussian concentration and union bounds to control random initial features and related quantities simultaneously across training examples.ReLU’s Lipschitz property yields sub-Gaussian controls for the relevant feature norms and inner products.
- Optimization control: The optimization proof decomposes the logistic-loss change into first- and second-order terms, bounding each with convexity, ReLU homogeneity, and Lipschitz arguments.The resulting inequalities are telescoped over iterations.
- Optimization control: The theorem proof shows that the stopping time remains beyond T = ⌈2λ^2/ηε⌉, ensuring the desired training-risk decrease before the argument terminates.A contradiction with the stopping-time definition establishes the lower bound on t1.
B Omitted proofs from Section 3
The omitted proofs derive the test-error guarantee by bounding the Rademacher complexity of the loss-composed function class and combining it with the training result.
- Complexity bound: The generalization proof begins by defining Rademacher complexity for the relevant function class on a sampled dataset.This supplies the complexity measure used in the test-risk bound.
- Complexity bound: A contraction lemma transfers the complexity bound through coordinate-wise Lipschitz transformations, including the logistic-loss derivative.The loss composition is controlled using a 1/4-Lipschitz property.
- Complexity bound: The proof bounds the network class’s Rademacher complexity for a fixed initialization by handling hidden units separately.The argument follows a related standard complexity proof.
- Generalization guarantee: The training theorem is then combined with the complexity and concentration bounds to obtain the stated test-risk guarantee with high probability.The final probability accounts for random initialization and data sampling.
C Omitted proofs from Section 4
The omitted Section 4 proofs establish the SGD extension by combining martingale concentration, convexity, ReLU homogeneity, and telescoping arguments. They yield a high-probability guarantee after n = ⌈2λ^2/(ηε)⌉ steps under the stated sample and width conditions.
- Convexity of ℓ and ReLU homogeneity control the first-order term, while η ≤ 1 and −ℓ′ ≤ ℓ bound the second-order term before telescoping.
- The SGD result extends Theorem 2.2 to the stochastic-gradient setting.
- For any m ≥ M and η ≤ 1, n_0 ≥ ⌈2λ^2/(ηε)⌉ gives the stated guarantee with probability 1 − 3δ.
- The proof samples n_0 examples and feeds one example to SGD at each of the first n_0 steps.
- A martingale Bernstein bound controls the stochastic discrepancy between Q(W_t) and Q_t(W_t).
D Omitted proofs from Section 5
The omitted proofs establish technical Gaussian bounds, apply Fenchel duality to the associated optimization problem, and analyze XOR data through nonseparability. They also bound the relevant quadratic form and parameter γ1 with high probability.
- Optimization analysis: The optimization argument defines a linear map A from the Hilbert space H to R^n and invokes Fenchel duality in both directions.Strong duality yields equality for optimal primal-dual solutions, after which the dual problem is examined.
- Optimization analysis: With the uniform probability vector, Markov’s inequality gives probability 0.9 that the relevant quadratic form is at most 1/(20n), implying an upper bound on γ1.The proof also uses ||w̄||H = γ1 for an optimal solution satisfying w̄ = A* q̄.
- Technical Gaussian bounds: The proofs reduce the Gaussian calculations by symmetry and independence, partitioning cases according to the first two coordinates and bounding the resulting integrals.They introduce Gaussian densities, interval probabilities U(c), and an auxiliary lemma for independent Gaussian variables with unequal variances.
- Technical Gaussian bounds: Lemma D.1 supplies the technical estimate used in the proof of Proposition 5.4.The lemma concerns independent z1 ∼ N(0,1) and z2 ∼ N(0,b^2) with b > 1.
- XOR construction: For XOR data, event As occurs for all 1 ≤ s ≤ m with probability 1/2 when the stated coordinate condition holds, and the data are not linearly separable.This nonseparability establishes Proposition 5.4 in the discussed 0/1-loss setting.