Source-linked AI summary

On Exact Computation with an Infinitely Wide Neural Net

Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang

arXiv:1904.11955v2cs.LGcs.CVcs.NEstat.ML

TL;DR

The paper asks whether fully trained infinitely wide networks can be understood and evaluated through kernel methods, especially for convolutional architectures. It develops an exact, GPU-efficient CNTK algorithm and a rigorous NTK equivalence proof. CNTK reaches strong CIFAR-10 performance, while a remaining gap from finite CNNs defines an open limitation.

  • Problem

    The paper addresses the limited ability to compute and evaluate exact kernels for fully trained infinitely wide convolutional networks, particularly with pooling on CIFAR-10.

  • Method

    The paper develops an exact dynamic-programming algorithm and GPU implementation for ReLU CNTKs, alongside a non-asymptotic NTK equivalence proof.

  • Results

    77% classification accuracy is achieved by an 11-layer CNTK with global average pooling on CIFAR-10, 10% higher than the best reported fixed-kernel Gaussian-process result.

  • Takeaways & Limitations

    Exact CNTK computation provides a practical pure-kernel benchmark for infinitely wide convolutional networks and shows performance close to corresponding finite architectures.

  • Takeaways & Limitations

    CNTKs remain 5%–6% below CNN performance, and the paper leaves infinitely wide networks with Batch Normalization or Residual Layers open.

Abstract

from arXiv · show

How well does a classic deep net architecture like AlexNet or VGG19 classify on a standard dataset such as CIFAR-10 when its width --- namely, number of channels in convolutional layers, and number of nodes in fully-connected internal layers --- is allowed to increase to infinity? Such questions have come to the forefront in the quest to theoretically understand deep learning and its mysteries about optimization and generalization. They also connect deep learning to notions such as Gaussian processes and kernels. A recent paper [Jacot et al., 2018] introduced the Neural Tangent Kernel (NTK) which captures the behavior of fully-connected deep nets in the infinite width limit trained by gradient descent; this object was implicit in some other recent papers. An attraction of such ideas is that a pure kernel-based method is used to capture the power of a fully-trained deep net of infinite width. The current paper gives the first efficient exact algorithm for computing the extension of NTK to convolutional neural nets, which we call Convolutional NTK (CNTK), as well as an efficient GPU implementation of this algorithm. This results in a significant new benchmark for the performance of a pure kernel-based method on CIFAR-10, being $10\%$ higher than the methods reported in [Novak et al., 2019], and only $6\%$ lower than the performance of the corresponding finite deep net architecture (once batch normalization, etc. are turned off). Theoretically, we also give the first non-asymptotic proof showing that a fully-trained sufficiently wide net is indeed equivalent to the kernel regression predictor using NTK.

1 Introduction

The paper asks how fully trained infinitely wide neural networks behave and whether their power can be studied through exact kernel methods. It introduces an efficient exact CNTK computation and reports strong CIFAR-10 performance alongside a rigorous NTK equivalence result.

  • Motivation: Infinite-width analysis asks how classic deep architectures perform on CIFAR-10 as convolutional channels and fully connected layers grow without bound.The motivation is that infinite limits may simplify mathematical analysis and connect deep learning to kernels.
  • Problem setting: Fully trained networks differ from weakly trained networks because all parameters, rather than only the top layer, are optimized by gradient descent.Weakly trained networks can be viewed through random features and kernel regression, whereas this paper studies fully trained nets.
  • Open computational problem: Exact evaluation was previously considered computationally infeasible for convolutional kernels with pooling on large datasets such as CIFAR-10.Earlier work instead used finite Monte Carlo approximations, whose practical accuracy remained uncertain.
  • Contributions: The paper gives an exact, efficient dynamic-programming algorithm for ReLU CNTKs, supplemented by GPU implementation techniques.The method computes ker(x, x′) directly and targets fully trained infinitely wide convolutional architectures.
  • Results: 77% classification accuracy is achieved by an 11-layer convolutional CNTK with global average pooling on CIFAR-10.This is 10% higher than the best reported fixed-kernel Gaussian-process performance, with finite-architecture comparisons requiring batch normalization and data augmentation to be turned off.
  • Theory: The paper also gives a more rigorous non-asymptotic proof that NTK kernel regression captures a sufficiently wide fully trained network under weaker conditions than earlier proofs.The proof establishes the result for sufficiently large width and bounded prediction error with high probability.

2 Related Work

The paper builds on Gaussian-process, compositional-kernel, and over-parameterized-network theories, while sharpening the NTK/CNTK connection through exact computation and rigorous analysis.

  • Gaussian-process connections: Infinite-width neural networks were first connected to Gaussian processes, with later work extending the correspondence to deep and convolutional architectures.These GP models typically randomize all parameters and optimize only the top classification layer.
  • Compositional kernels: Compositional-kernel research transformed neural-network architectures into kernels, but the paper distinguishes those kernels from the NTK framework for fully trained networks.Related results also connect sufficiently wide networks and stochastic gradient descent to reproducing kernel Hilbert spaces.
  • Over-parameterized networks: Over-parameterization studies show that large but finite networks can reach zero training error while their weights remain close to initialization.This near-initialization behavior motivates kernel-regime analyses of gradient descent.
  • Comparison scope: The paper restricts its benchmark comparison to fixed kernels defined without using training data, excluding data-tuned kernels and neural feature extractors.This defines the comparison scope for reported kernel performance.
  • Practical significance: The authors report that CNTK performs well on CIFAR-10, addressing concerns that kernel-regime models may be too simple to explain deep-learning success.This claim is presented as evidence about the practical power of CNTK rather than a complete explanation of deep learning.
  • NTK and CNTK: Jacot et al. introduced the NTK, while this paper names and computes its convolutional extension, CNTK, with a more explicit convolutional structure and GPU-friendly implementation.The paper argues that exact kernel evaluation matters because linearization can reduce classification accuracy by 5% on CIFAR-2.

3 Neural Tangent Kernel

This section develops the infinite-width NTK for fully connected networks and proves that sufficiently wide, fully trained networks approximate NTK kernel regression. It also establishes non-asymptotic width conditions and extends the result to generalization on finite test sets.

  • NTK formulation: The NTK arises when training a fully connected network with squared loss by gradient descent in the infinite-width limit.The network-output dynamics are governed by a kernel matrix that becomes constant during training and converges to the deterministic NTK.
  • NTK formulation: In the infinite-width limit, the training kernel matrix H(t) remains equal to its initialization value H*, which is the NTK evaluated on training data.This converts the neural-network dynamics into the corresponding kernel-regression dynamics.
  • NTK formulation: The resulting NTK kernel-regression predictor evaluates as f*(x) = (ker(x, x1), ..., ker(x, xn)) · (H*)^-1y.The expression assumes zero initial output and gives the limiting prediction under kernel-regression gradient flow.
  • Finite-width guarantees: Theorem 3.1 gives a non-asymptotic convergence guarantee when every hidden width is at least Ω(L^6 ε^-4 log(L/δ)) for ReLU networks.Unlike prior results, the theorem only requires the minimum hidden width to be sufficiently large, rather than sequential or equal-rate limits.
  • Finite-width guarantees: Theorem 3.2 proves that a sufficiently wide, fully trained ReLU network predicts within ε of the NTK kernel-regression predictor with probability at least 1 − δ.The required common width depends polynomially on 1/κ, L, 1/λ0, n, and log(1/δ), under the stated unit-norm test-point condition.
  • Finite-width guarantees: The theorem extends from a single test point to exponentially many finite test points via a union bound, yielding the same generalization ability as the corresponding NTK.This conclusion combines the pointwise guarantee with standard hold-out validation analysis.

4 Convolutional Neural Tangent Kernel

This section defines CNTKs for vanilla CNNs and CNNs with global average pooling, then gives a recursive exact-computation procedure. The resulting method scales to CNTK matrices for image datasets despite prior concerns about exact pooled convolutional kernels.

  • Architectures: The paper studies CNTKs for two architectures: vanilla CNNs and CNNs with global average pooling.The corresponding CNTK definitions and formulas are developed for both architectures.
  • CNN setup: CNN inputs are images of width P, height Q, and C(0) channels, with convolution filters of size q × q and unit stride with zero padding.The padding preserves the spatial dimensions across layers.
  • CNTK computation: The CNTK recursion computes covariance and derivative-covariance tensors over pairs of spatial locations and input images.For each layer, the method defines K, ˙K, and Σ tensors indexed by spatial coordinates from both images.
  • CNTK computation: A trace over patches represents the convolution operation before the recursive CNTK computation is completed.The final layer combines ˙K(L) ⊙ Θ(L−1) with K(L), after which the final CNTK value is obtained by a trace.
  • Computational scaling: The exact CNTK matrix for an L-layer convolutional network with global average pooling has time complexity O(n^2P^2Q^2L).This scaling enables exact computation where previous work had treated pooled convolutional kernels as infeasible and used Monte Carlo approximations.

5 Experiments

The experiments evaluate CNNs and corresponding CNTKs on CIFAR-10, finding strong kernel performance, architectural effects, and a remaining gap between infinite- and finite-width models.

  • Benchmark results: 77.43% classification accuracy is achieved by the best kernel, an 11-layer CNTK with global average pooling, establishing a new CIFAR-10 pure-kernel benchmark.This is 10% higher than methods reported in Novak et al. (2019).
  • Architectural effects: Depth affects classification accuracy for both CNNs and CNTKs, showing that depth also influences kernel performance.
  • Architectural effects: 8% - 10% higher classification accuracy is obtained with global average pooling for both CNNs and CNTKs.The authors suggest that some neural-network performance techniques may also benefit kernel methods.
  • Infinite versus finite width: A 5% - 6% performance gap remains between CNTKs and CNNs, indicating that finite width has benefits.The authors identify characterizing this gap as an open problem for theories operating in the NTK regime.
  • Small-data evaluation: CNTK performance on small training subsets is highly correlated with corresponding full-dataset CNN and CNTK performance, supporting a possible neural architecture search application.CNTK-GAP-2Ks outperform CNTK-V-2Ks by about 8% - 9%.

6 Conclusion

The paper makes infinitely wide neural networks experimentally tractable through exact CNTK computation and relates sufficiently wide networks to kernel regression.

  • The first practical exact CNTK algorithm enables investigation of infinitely wide, infinitely over-parameterized deep nets, whose behavior is not much worse than finite counterparts.
  • A rigorous proof shows that a sufficiently wide net is approximately equivalent to the kernel regression predictor.The paper leaves infinitely wide nets with Batch Normalization or Residual Layers as an open problem.

A Experiment Details

The experiments use controlled CNN and CNTK settings without common training tricks, and compute predictions through CNTK kernel regression.

  • Setup: No data augmentation, batch normalization, dropout, or weight decay is used, while convolution filters have size 3, stride 1, and zero padding.The experiments use 2, 3, 5, 10, or 20 convolution layers and quadratic loss.
  • Kernel regression: CNTK training solves (H*)^-1Y using encoded ten-class labels, and testing computes kernel similarities against training inputs.
  • Kernel regression: The predicted class is the one with the largest value in the kernel-regression output.
  • CNN training: Finite CNNs use 1024 channels and SGD with fixed learning rates selected from {0.1, 1, 10}, with reported performance averaged over 10 epochs after full training accuracy.The best average performance over 3 trials is reported.
  • Implementation: Neural-network training uses PyTorch and Tesla V100 GPUs, while kernel values are computed with CuPy and native CUDA code.Experiments are performed on Amazon Web Services.

B Additional Experiments on Random Features

Additional experiments compare exact CNTKs with random-feature kernels on CIFAR-2 and find a large accuracy advantage for exact computation.

  • Experimental design: Random-feature kernels and exact CNTKs are compared on CIFAR-2 using CNNs with 128 channels, with each random kernel tested 10 times and summarized by its median.
  • Results: Table 2 reports classification accuracies for random kernels generated from random features and exact CNTKs on CIFAR-2.
  • Results: Exact kernels achieve much higher accuracies than random features on CIFAR-2, demonstrating the importance of exact kernel computation.

C Proof of Lemma 3.1

The proof derives the fully connected network’s infinite-width NTK by tracking parameter and output evolution, Gaussian-process limits, and layerwise gradient quantities. Gaussian initialization enables fresh-weight replacements whose validity is established for ReLU networks.

  • Dynamics: The parameters evolve continuously, and the network outputs on training inputs are collected into a vector whose dynamics depend on desired outputs.The proof introduces continuous time t, output vector u(t), and target vector y.
  • NTK derivation: The fully connected NTK is derived from layerwise derivatives of the prediction with respect to weight matrices.The proof writes these derivatives in compact form and combines them inductively into the final NTK expression.
  • Concentration ingredients: The proof uses Gaussian-process covariance operators and ReLU-specific continuity properties to control the layerwise kernel quantities.It defines covariance operators and normalized expectations, then establishes uniform continuity for arcsin and the ReLU-induced map.
  • Gaussian replacement: Gaussian initialization permits replacing a dependent weight matrix with a fresh independent sample without changing its infinite-width limit.The text states that this approximation is made rigorous for ReLU activation in Theorem 3.1 and applied inductively.
  • Concentration ingredients: The argument organizes concentration events for covariance, derivative, and diagonal quantities across input pairs and layers.These events are combined through conditional-probability bounds, projection arguments, and Gaussian conditioning.

E.2 Proof Sketch

This proof sketch establishes convergence to the NTK at initialization and then compares neural-network and NTK gradient-flow trajectories. Kernel perturbation control and linear convergence yield equivalence in the sufficiently wide regime.

  • Initialization: Theorem 3.1 states that sufficiently wide ReLU networks converge to the NTK at initialization with high probability.The theorem assumes a lower bound on the minimum hidden-layer width and provides a probability-controlled convergence statement.
  • Initialization: The initialization proof inductively controls Gaussian-process covariance and derivative quantities across layers and selected input pairs.The argument uses concentration, union bounds, Gaussian conditioning, projection matrices, and a fresh-weight approximation.
  • Training dynamics: Kernel-value perturbation bounds reduce prediction perturbation to controlling the training kernel and test-kernel deviations.The reduction assumes small initialization error and bounds on H(t) and test-point kernel discrepancies.
  • Training dynamics: The NTK predictor is represented by gradient flow, allowing its output dynamics to be compared directly with the neural-network predictor.Both predictors have time derivatives of the same form, and their difference is analyzed in integral form.
  • Training dynamics: Under these bounds, both neural-network and NTK predictors converge linearly toward the target outputs.The proof uses positive-semidefinite kernel structure and a lower bound on the minimum eigenvalue during training.

F.1 Kernel Perturbation During Training

The training analysis bounds how much finite-width networks deviate from their initialization. It connects small weight movement to small gradient and kernel perturbations, preserving convergence behavior when width is sufficiently large.

  • Kernel stability: Lemma F.2 states that sufficiently wide networks maintain small kernel perturbations throughout training with high probability.The bound applies for all training times and pairs drawn from training and test inputs.
  • Gradient and weight control: Kernel perturbation is reduced to perturbations in forward and backward network quantities, which are controlled through weight-matrix perturbations.The proof separately tracks Δg(h) and Δb(h), then invokes adapted lemmas to bound them.
  • Kernel stability: Small perturbations in the gradient imply small perturbations in kernel values during training.The proof combines gradient-Lipschitz control with the initialization result to obtain the kernel bound.
  • Gradient and weight control: If the training residual decreases linearly while other matrices move little, each individual weight matrix also moves little.Lemma F.7 establishes this layerwise implication under sufficiently small perturbation parameters.
  • Convergence: Small weight movement preserves a positive kernel eigenvalue and therefore yields linear convergence of the training dynamics.The argument uses a lower bound on λmin(H(t)) to obtain the convergence result.
  • Convergence: The proof closes by showing that weight matrices remain close to initialization, completing the chain from weight control to gradient and kernel stability.The final argument uses the preceding lemmas and a contradiction based on the first time a movement bound could fail.

G.2 Derivation

The derivation expresses CNTK as layerwise gradient interactions induced by convolution, then adapts the final computation for global average pooling. The resulting entrywise formulas support exact GPU-oriented computation.

  • Vanilla CNN: The derivation approximates layerwise gradient terms using the initialized weight matrix and propagates the resulting recursion through the network.The final fully connected layer contributes G(L)(x, x′) ≈ I, enabling the induction that derives the CNTK formula.
  • Global average pooling: Global average pooling changes the final kernel computation by replacing diagonal summation with an average over all spatial entries.Cross-variances between every pair of patches contribute because of global average pooling.
  • Global average pooling: For the global-average-pooling architecture, the kernel recursion starts with Θ(0)(x, x′) = 0 because the first layer is not trained.The final layer is also fixed, so its K(L) contribution is ignored.
  • Efficient computation: The exact CNTK computation is organized around entrywise matrix and tensor operations, enabling efficient GPU implementations.Memory can be reduced by storing diagonal covariance entries and the cross-covariance matrix rather than the full A(h) matrix.
  • Efficient computation: For ReLU with normalized diagonal entries, closed-form expectation formulas accelerate computation of the covariance and derivative kernels.The method applies these formulas to the relevant matrices after diagonal normalization.
Loading 1904.11955v2…