Source-linked AI summary

Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networks

Yu Bai, Jason D. Lee

arXiv:1910.01619v2cs.LGmath.OCstat.ML

TL;DR

NTK theory does not fully explain why trained neural networks can outperform their linearized counterparts, motivating analysis beyond the NTK regime. The paper randomizes over-parametrized networks to couple them with quadratic and higher-order Taylor terms, proving favorable optimization and generalization properties. These networks match NTK sample-complexity bounds and can improve them by a dimension factor under mild distributional assumptions.

  • Problem

    NTK theory has limited expressive power and can leave a performance gap relative to fully trained neural networks, motivating theories beyond linearization.

    Source: §1 · pp. 1–2
  • Method

    The paper randomizes over-parametrized two-layer networks so their training dynamics couple with quadratic or higher-order Taylor terms instead of the linear NTK.

    Source: §1 · pp. 2–3
  • Results

    Randomized networks have favorable optimization landscapes, match NTK sample-complexity bounds, improve polynomial-learning bounds by O(d) under mild distributional assumptions, and extend to higher-order NTKs.

    Source: §1 · pp. 2–3, §7 · p. 13
  • Takeaways & Limitations

    Higher-order approximations provide a principled route to studying neural-network training beyond the NTK regime while retaining comparable expressivity and generalization guarantees.

    Source: §1 · p. 2, §7 · p. 13

    The paper identifies open questions about using quadratic expressivity, explaining full-training success, and designing approximation schemes that remain accurate over larger parameter-space regions.

    Source: §7 · p. 13

Abstract

from arXiv · show

Recent theoretical work has established connections between over-parametrized neural networks and linearized models governed by he Neural Tangent Kernels (NTKs). NTK theory leads to concrete convergence and generalization results, yet the empirical performance of neural networks are observed to exceed their linearized models, suggesting insufficiency of this theory. Towards closing this gap, we investigate the training of over-parametrized neural networks that are beyond the NTK regime yet still governed by the Taylor expansion of the network. We bring forward the idea of \emph{randomizing} the neural networks, which allows them to escape their NTK and couple with quadratic models. We show that the optimization landscape of randomized two-layer networks are nice and amenable to escaping-saddle algorithms. We prove concrete generalization and expressivity results on these randomized networks, which lead to sample complexity bounds (of learning certain simple functions) that match the NTK and can in addition be better by a dimension factor when mild distributional assumptions are present. We demonstrate that our randomization technique can be generalized systematically beyond the quadratic case, by using it to find networks that are coupled with higher-order terms in their Taylor series.

1 Introduction

The paper addresses limitations of NTK linearization by coupling randomized over-parametrized networks with quadratic and higher-order Taylor terms. It establishes optimization, generalization, and expressivity results, including dimension-factor sample-complexity improvements under mild distributional assumptions.

1 Introduction

  • NTK analyses couple nonlinear training to a local linearization, but empirical full training can outperform NTK models and their expressive power may yield poor sample-complexity bounds for simple functions.The gap motivates theories beyond the NTK regime.Source: §1 · pp. 1–2
  • Randomization makes the linear NTK term non-dominant, coupling over-parametrized two-layer networks with quadratic models whose second-order stationary points have low training loss.This landscape supports efficient escape-saddle optimization.Source: §1 · p. 2
  • Randomized networks match NTK sample-complexity bounds without distributional assumptions and improve polynomial-learning bounds by O(d) under relatively uniform base distributions.The result concerns generalization and expressive power of randomized networks.Source: §1 · p. 2
  • The randomization technique extends to networks dominated by k-th order Taylor terms for k > 2, termed higher-order NTKs, with expressive power similar to linear NTKs and potentially better generalization.This extends the approach systematically beyond quadratic approximations.Source: §1 · pp. 2–3

1.1 Prior work

  • Prior work connects wide-network training to kernel methods and NTKs, proves optimization and generalization guarantees, and documents gaps between NTK predictions and fully trained networks.Related approaches also study favorable landscapes, norm-based complexity, mean-field dynamics, and randomization, but their comparisons or guarantees differ from this paper’s focus.Source: §1.1 · pp. 3–4

2 Preliminaries

The preliminaries specify supervised learning with convex smooth losses on bounded d-dimensional features and an over-parametrized two-layer network trained through its first-layer weights. The analysis assumes a second-order smooth activation and introduces norm and stationarity notation.

2 Preliminaries

  • The setup uses i.i.d. labeled data with features in the radius-B_x sphere in R^d and a predictor optimized under empirical and population risk.The loss is convex, 1-Lipschitz, three-times differentiable, and has bounded second and third derivatives, including logistic and soft hinge losses.Source: §2 · p. 4
  • The model is an over-parametrized two-layer network with first-layer matrix W and fixed second-layer vector a, using 1/sqrt(m) NTK-type scaling.Only the first-layer weights are trained.Source: §2 · p. 4
  • The activation is assumed C2 with bounded derivative-growth conditions and Lipschitz second derivative; cubic ReLU is an example, while ReLU is excluded because higher-order smoothness is required.These assumptions support the Taylor-based analysis.Source: §2 · pp. 4–5

2.1 Notation

  • The notation defines mixed matrix norms and corresponding norm balls, while a second-order stationary point satisfies zero gradient and positive-semidefinite Hessian.The Frobenius norm is the special case ∥·∥2,2.Source: §2.1 · p. 5

3 Escaping NTK via randomization

The paper uses sign randomization to suppress the linear NTK term, couple two-layer networks to quadratic Taylor terms, and formulate a trainable randomized risk.

3 Escaping NTK via randomization

  • Ordinary small-weight training keeps the linearized model dominant, and changing only the leading scale or learning rate does not remove that dominance.In the standard regime, higher-order terms are negligible and training follows the linearized network.Source: §3 · p. 5

3.1 Learning randomized neural nets

  • The construction assumes a twice-smooth activation with Lipschitz second derivative; derivative-growth bounds are mainly technical conveniences.The paper states that the Lipschitzness of the second derivative is the restrictive requirement.Source: §3 · p. 5, §3.1 · p. 7
  • The method samples random signs for each neuron and minimizes expected randomized risk after reparameterizing movement from initialization as the trainable variable.An ℓ2,4 regularizer encourages the target weight scale and is intended to improve generalization.Source: §3 · p. 6, §3.1 · p. 7

3 Escaping NTK via randomization

  • At the random weight matrix WΣ, the quadratic component f Q dominates the linear component f L, so the network is coupled with its quadratic part rather than the linear NTK.Source: §3 · p. 6

3.1 Learning randomized neural nets

  • Symmetric initialization pairs positive and negative output weights, making the initial network identically zero while retaining random features for expressivity arguments.Half the output signs are +1 and half are −1, with copied paired initial weights.Source: §3.1 · p. 7

4 Optimization

The randomized objective inherits a favorable quadratic-model landscape: away from near-optimal solutions, descent directions exist, and second-order stationary points achieve competitive loss under sufficient width.

4.1 Nice landscape of clean risk

  • The clean-risk result relies on the connection to matrix sensing and quadratic neural-network landscapes, under smooth-activation and norm-ball conditions.The activation assumptions ensure the required second-order smoothness.Source: §4.1 · p. 7
  • For the clean quadratic risk, every point with loss above the benchmark has either a first-order or second-order descent direction within the relevant norm ball.Thus approximate second-order stationary points are approximate global minima.Source: §4.1 · p. 8

4.2 Nice landscape of randomized neural net risk

  • The randomized neural-network risk transfers the clean quadratic landscape through coupling of the losses, gradients, and Hessians, with regularization controlled separately.The analysis first localizes stationary points using the regularizer and then applies the clean-risk result.Source: §4.2 · p. 8
  • Every second-order stationary point of the regularized randomized loss has low loss comparable to the best quadratic model when width is sufficiently large.The result also bounds the norm of such stationary points and supports saddle-escaping optimization.Source: §4.2 · pp. 8–9
  • Noisy SGD can efficiently seek such stationary points using fresh randomization at each iteration, and learning rates up to O(m^1/4) are allowed within the coupling region.This contrasts with the NTK regime's O(1) learning-rate bound.Source: §4.2 · p. 9

5 Generalization and Expressivity

Quadratic randomized networks obtain operator-norm-based generalization bounds and can match or improve linear-NTK sample complexity for several polynomial and structured-function classes.

5.1 Generalization

  • Operator-norm control of quadratic feature maps yields generalization bounds that are never worse than Frobenius-norm bounds and can improve with additional structure in x.The bound uses the empirical covariance operator norm, with M_x,op ≤ 1 always and M_x,op ≤ κ/d under stated sub-Gaussian isotropic conditions.Source: §5.1 · pp. 9–10
  • The randomized network's generalization bound scales at most logarithmically with width for m ≳ n^4, supporting small samples under extreme over-parameterization.Its main term adapts to feature-distribution structure and can improve the naive bound by O(1/d).Source: §5.1 · p. 10

5.2 Expressivity and Sample Complexity through Quadratic Models

  • The quadratic model expresses sums of one-directional polynomials through a reduction from degree-p representation to degree-(p−2) random-feature representation.With sufficiently large width, the theorem guarantees a corresponding quadratic-model parameter achieving low training loss.Source: §5.2 · p. 10
  • For polynomial learning under relatively uniform distributions, quadratic sample complexity is eO(p/d) times the linear-NTK bound, yielding a dimension-factor reduction unless p ≍ d.This comparison concerns upper bounds, since the linear-NTK lower bound is generally unknown.Source: §5.2 · p. 11
  • For noisy 2-XOR, the randomized quadratic network is provably eO(d) better in sample complexity than the linear-NTK lower bound of Ω(d^2).The comparison is rigorous because the linear-NTK lower bound is known for this task.Source: §5.2 · p. 11
  • For low-rank matrix sensing, the quadratic representation supports sample complexity scaling through the rank-dependent polynomial structure and compares favorably with linear NTK.The construction represents the target using eigenvalue-weighted one-directional quadratic terms.Source: §5.2 · p. 11

6 Higher-order NTKs

Randomization can make higher-order Taylor terms dominate neural-network behavior, extending the quadratic construction to systematic k-th order couplings. These higher-order models retain expressive and generalization properties while potentially improving sample complexity under distributional assumptions.

6 Higher-order NTKs

  • Randomization constructs over-parametrized networks whose k-th order Taylor term dominates the expansion, producing higher-order NTKs.The construction assumes an analytic activation and uses symmetric initialization; the standard linear NTK is the first-order term.Source: §6 · p. 12, §6.1 · p. 12

6.1 Extracting the k-th order term

  • Taking ∥w⋆,r∥2 = O(m−1/2k) yields, with high probability, lower-order terms f (1), . . . , f (k−1) and the remainder controlled, while f (k) can express an O(1) function.The remainder term f −P has order O(m−1/2k).Source: §6.1 · p. 12
  • The k-th order model extends the quadratic model’s generalization and expressivity results to learning polynomials of degree at least k.The supplied passage introduces these results but does not provide the corresponding numerical bound.Source: §6.1 · p. 13
  • For large n, higher-order coupling can improve sample complexity over the linear NTK by a factor of d^(k−1) under mild approximate-isotropy assumptions.The stated condition is that the kth moment tensor has a constant condition number.Source: §6.1 · p. 13
  • The paper leaves the existence of a deterministic W approximating the k-th order NTK with vanishing error as future work.The approximation error is stated as ε_m → 0 as m → ∞.Source: §6.1 · p. 13

7 Conclusion

The conclusion frames randomized higher-order coupling as a route beyond linearized neural-network analysis, combining quadratic optimization guarantees with improved distribution-dependent generalization and systematic higher-order constructions. It also identifies unresolved questions about using quadratic expressivity and explaining fully trained networks.

7 Conclusion

  • Randomized two-layer networks coupled to quadratic models have low-loss second-order stationary points and can be minimized with escape-saddle algorithms.The conclusion describes the optimization landscape as suitable for efficient minimization.Source: §7 · p. 13
  • These randomized networks match linearized models in expressivity and generalization guarantees, while improving generalization by a dimension factor under distributional assumptions.The conclusion states the comparison at the level of guarantees rather than a specific task or metric.Source: §7 · p. 13
  • The paper extends randomization beyond quadratic coupling by establishing networks whose Taylor series is dominated by the k-th order term.This is the paper’s higher-order generalization of the randomized construction.Source: §7 · p. 13
  • Open questions concern better exploiting quadratic expressivity, explaining full training through higher-order expansions, and designing approximations valid over larger parameter regions.The conclusion presents Taylor expansion as only one possible approximation route.Source: §7 · p. 13

A.1 A matrix operator norm concentration bound

  • The appendix develops matrix concentration and random-feature arguments used to support polynomial representation results.The supplied passages include a matrix operator-norm lemma and feature-map equivalence arguments for polynomial functions.Source: §A.1 · pp. 17–18, §A.2 · p. 18

A.2 Expressing polynomials with random features

  • The random-feature construction represents functions such as α(β^T x)^p through an equivalent feature map and associated RKHS.The passages relate the kernel feature map to tensor-product features and polynomial representation.Source: §A.2 · pp. 18–19

A.3 Proof of Equation (4)

The supplied appendix passages assemble concentration, coupling, and landscape-closeness arguments for the quadratic approximation. They connect the neural-network loss and its derivatives to the clean quadratic risk through bounds on functions, gradients, and Hessians.

A.3 Proof of Equation (4)

  • A covering-net argument bounds the relevant operator quantities with failure probability controlled by d, log m, and the target confidence parameter.The supplied proof passages use a 1/2-covering with size at most 5^d and a union-bound tail estimate.Source: §A.3 · p. 20

B.2 Coupling lemmas

  • The appendix bounds the quadratic model, couples the neural-network output to it, and establishes closeness of their risks.These are presented as separate lemmas for f^Q, the coupling between f and f^Q, and L^Q versus L.Source: §B.2 · pp. 21–22, §B.3 · p. 22

B.3 Closeness of landscapes

  • Directional gradients and Hessian quadratic forms are controlled using Taylor expansions, smoothness, Cauchy–Schwarz, and the coupling bounds.The proof explicitly applies the bounds for the model and coupling lemmas to both gradient and Hessian terms.Source: §B.3 · pp. 23–24

B.4 Proof of Theorem 2

  • Theorem 2 follows by combining risk, gradient, and Hessian closeness and choosing the width m large enough to satisfy the resulting error requirements.The proof collects the conditions from the preceding lemmas and merges them into a sufficient width requirement.Source: §B.4 · p. 26

B.5 Proof of Corollary 3

  • The remaining proof steps bound the corollary’s terms using Cauchy–Schwarz, AM–GM, and a tuned choice α = 5/14.The supplied passages identify the inequalities and parameter choice but do not state the corollary itself.Source: §B.5 · p. 27

B.6 Proof of Theorem 4

The proof of Theorem 4 localizes second-order stationary points, controls the neural-network and quadratic risks within a bounded region, and combines these estimates to obtain the stated loss guarantee. A separate concentration argument bounds the relevant network complexity.

B.6 Proof of Theorem 4

  • The proof chooses regularization and accuracy parameters, then uses stationarity and Hessian nonnegativity to derive the loss comparison.The argument introduces λ_0 and ε before evaluating the second-order stationary point condition.Source: §B.6 · p. 27, Bounding loss and norm Choosing · p. 28
  • Every second-order stationary point is first shown to lie in a controlled norm region, using convexity and the bounded zero-loss assumption.The localization argument relies on the clean risk and the assumption ℓ(y, 0) ≤ 1.Source: §B.6 · pp. 27–28
  • Within that region, the coupling error is bounded by ε after choosing the width sufficiently large relative to the problem parameters.The proof states a uniform coupling bound over the bounded norm ball.Source: §B.6 · p. 28

C.1 Proof of Lemma 5

  • The auxiliary complexity bound uses the 1-Lipschitz property of the loss and the Rademacher contraction theorem.The final supplied passage attributes the remaining norm control to the power-mean or Cauchy–Schwarz inequality.Source: §C.1 · p. 29

C.2 Proof of Theorem 6

The proof relates the generalization of the randomized network to its quadratic model, then establishes that infinitely wide quadratic models can express specified one-directional polynomials under the stated activation and degree conditions.

C.2 Proof of Theorem 6

  • The generalization analysis decomposes the network error into coupling and quadratic-model terms, with the latter bounded through symmetrization and matrix concentration.The argument applies uniformly over the parameter constraint set and combines the resulting bounds with the coupling error.Source: §C.2 · p. 29, Generalization of quadratic model By… · p. 30
  • Under uniform-sphere or hypercube feature assumptions, covariance concentration supplies the distribution-dependent bound used in the proof.The corresponding condition is invoked after a distribution-independent argument that holds for all x in the sphere.Source: §C.2 · p. 29, Generalization of quadratic model By… · p. 30

C.3 Expressive power of infinitely wide quadratic models

  • For f⋆(x) = α(β⊤x)^p with p ≥2 and p−2 in {1} ∪ {2ℓ}, an infinitely wide quadratic model can represent f⋆ under σ′′(t) = relu(t).The construction reduces representation of the degree-p polynomial to a random-feature representation of degree p−2 and also provides a norm bound.Source: §C.3 · p. 31
  • The quadratic construction chooses a random scalar depending on initialization so that its induced random-feature expectation equals the required lower-degree polynomial.This choice yields the desired expressivity while satisfying the stated expected fourth-power norm bound.Source: §C.3 · p. 31

C.4 Proof of Theorem 7

The paper extends quadratic expressivity results from one-directional polynomials to the main theorem, then randomizes network weights so higher-order Taylor terms dominate while lower-order terms and remainders vanish with width.

C.4 Proof of Theorem 7

  • With sufficiently large width, symmetric initialization yields a finite quadratic network that expresses f⋆(x) = α(β⊤x)^p with probability at least 1−δ.The proof first establishes the one-directional case and then applies it repeatedly with disjoint initial weights to obtain the main theorem.Source: §C.4.1 · p. 32, §C.4.2 · p. 34

C.4.1 Expressing a single “one-directional” polynomial

  • The one-directional construction transfers the infinitely wide representation to finitely many neurons through concentration and controls the resulting parameter norm.The argument chooses the radius and width to satisfy the concentration requirements.Source: §C.4.1 · pp. 32–34

C.4.2 Proof of main theorem

  • The main theorem combines k applications of the one-directional result and a union bound over disjoint initial-weight sets.This produces a single finite network satisfying the required representation and norm guarantee.Source: §C.4.2 · p. 34

D.1 Coupling f and f (k) via randomziation

  • Randomization chooses paired weights so the network is approximately the k-th order NTK plus an error that vanishes as m →∞, thereby escaping the linear NTK regime.The scaling ∥w⋆,r∥2 = O(m−1/2k) makes lower-order Taylor terms and the remainder have order O(m−1/2k), while the k-th order term can express an O(1) function.Source: §D.1 · p. 35

D.2 Generalization and expressivity of f (k)

The higher-order NTK analysis bounds generalization through tensor operator norms and proves expressivity for polynomial targets, with uniform feature distributions yielding a dimension-dependent advantage over distribution-free bounds.

D.2.1 Generalization

  • For uniform sphere features, generalization is bounded through the operator norm of a k-th-order tensor feature and improves over the distribution-free bound by min{d^(k−1), n}.The result assumes bounded k and constrains the parameter set through the corresponding 2k norm.Source: §D.2 · p. 36, §D.2.1 · p. 38

D.2 Generalization and expressivity of f (k)

  • When n ≥ d^(k−1), higher-order NTKs achieve sample complexity lower by a factor of d^(k−1) than the linear-NTK comparison.The leading multiplicative factor is the same for all k, including k = 1, while the dimension-dependent improvement appears under the stated distributional setting.Source: §D.2 · p. 36

D.2.1 Generalization

  • The tensor-operator-norm framework converts a Rademacher complexity bound into the stated higher-order NTK generalization bound.The argument applies to nonnegative losses that are 1-Lipschitz in the prediction and bounded at zero.Source: §D.2.1 · pp. 37–38

D.2.2 Expressivity

  • For f⋆(x) = α(β⊤x)^p with p ≥ k and p−k in {1} ∪ {2ℓ}, sufficiently wide higher-order NTKs express f⋆ with high probability.The activation is chosen so that its k-th derivative equals relu, and the construction supplies a norm bound for the representing weights.Source: §D.2.2 · p. 38, Expressivity with infinitely many… · p. 40

D.3 Proof of Lemma 18

  • The proof of the tensor norm bound uses truncation, covering arguments, concentration inequalities, and the uniform-sphere feature assumption.The expressivity proof likewise passes from an infinite-neuron construction to finite width through concentration.Source: §D.3 · pp. 38–39, §D.4 · p. 40
Loading 1910.01619v2…