[Paper Review] On Exact Computation with an Infinitely Wide Neural Net

Sanjeev Arora, Simon S. Du+4 more|arXiv (Cornell University)|Apr 26, 2019

Gaussian Processes and Bayesian InferenceComputer Science36 references61 citations

TL;DR

The paper presents an exact, GPU-friendly algorithm to compute CNTK for CNNs, proves finite-width wide nets converge to CNTK kernel regression, and shows CNTKs achieve strong CIFAR-10 performance.

ABSTRACT

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.

Motivation & Objective

• Motivate understanding of how infinitely wide CNNs perform on standard datasets like CIFAR-10.
• Develop an exact, efficient algorithm to compute the Convolutional Neural Tangent Kernel (CNTK) for CNNs.
• Demonstrate equivalence between fully trained wide nets and kernel regression using CNTK.
• Provide non-asymptotic convergence results and compare CNTK performance to finite nets.
• Offer practical GPU implementations and benchmarks to advance kernel-based understanding of deep learning.

Proposed method

• Define neural network architectures with infinite width limits and describe CNTK as the kernel derived from gradients of outputs with respect to parameters.
• Derive explicit CNTK formulas for vanilla CNNs and CNNs with global average pooling (GAP) including convolutional and pooling steps.
• Prove non-asymptotic convergence: minimum layer widths scale as Omega(L^6/epsilon^4 log(L/delta)) to ensure NTK convergence at initialization (ReLU activation).
• Prove equivalence between fully trained wide nets and NTK-based kernel regression with finite-width perturbation bounds (Theorem 3.2).
• Present an exact dynamic-programming based algorithm to compute CNTK exactly and optimize implementation on GPUs.

Experimental results

Research questions

• RQ1 Can CNTK be computed exactly for convolutional nets with pooling?
• RQ2 Do fully-trained infinitely wide CNNs correspond to kernel regression under NTK?
• RQ3 How close is CNTK-based kernel performance to finite-width CNNs on CIFAR-10?
• RQ4 What are the finite-width requirements for NTK convergence to guarantee kernel-regression-like behavior?
• RQ5 Does depth and global average pooling materially affect CNTK performance on image classification tasks?

Key findings

• CNTKs achieve 77.43% accuracy on CIFAR-10 with an 11-layer CNN-GAP, beating prior GP-based kernels by about 10%.
• An 11-layer CNTK with GAP is within about 5% of the corresponding finite deep-net performance when batch norm and data augmentation are controlled.
• CNTK-based kernels are significantly stronger than earlier fixed-kernel GP methods on CIFAR-10 by up to ~10%.
• Depth and global average pooling significantly affect CNTK performance, with GAP offering substantial gains over vanilla CNNs.
• There remains a 5–6% gap between CNTK (infinite width) and finite CNNs, suggesting finite-width benefits persist.