Finite-Width Neural Tangent Kernels from Feynman Diagrams

TL;DR

This paper introduces a Feynman diagram-based method to compute finite-width corrections to neural tangent kernels, enabling better understanding of neural network training dynamics beyond the infinite-width approximation.

Contribution

The authors develop a novel Feynman diagram framework for calculating finite-width corrections to NTK statistics, extending analysis to layer-wise recursion relations and higher-order tensors.

Findings

Feynman diagrams simplify finite-width correction calculations.

Finite-width effects are negligible for scale-invariant nonlinearities like ReLU.

Numerical results match sampled neural network statistics for widths greater than 20.

Abstract

Neural tangent kernels (NTKs) are a powerful tool for analyzing deep, non-linear neural networks. In the infinite-width limit, NTKs can easily be computed for most common architectures, yielding full analytic control over the training dynamics. However, at infinite width, important properties of training such as NTK evolution or feature learning are absent. Nevertheless, finite width effects can be included by computing corrections to the Gaussian statistics at infinite width. We introduce Feynman diagrams for computing finite-width corrections to NTK statistics. These dramatically simplify the necessary algebraic manipulations and enable the computation of layer-wise recursion relations for arbitrary statistics involving preactivations, NTKs and certain higher-derivative tensors (dNTK and ddNTK) required to predict the training dynamics at leading order. We demonstrate the feasibility of our framework by extending stability results for deep networks from preactivations to NTKs and proving the absence of finite-width corrections for scale-invariant nonlinearities such as ReLU on the diagonal of the Gram matrix of the NTK. We numerically implement the complete set of equations necessary to compute the first-order corrections for arbitrary inputs and demonstrate that the results follow the statistics of sampled neural networks for widths $n\gtrsim 20$.