Finite-Width Neural Tangent Kernels from Feynman Diagrams

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$.

Although neural networks (NNs) and their rich learning dynamics are well-studied from an empirical perspective, understanding them theoretically is notoriously difficult and relies on approximations. One successful tool to describe several aspects of NNs and their statistics is via the so-called neural tangent kernel (NTK), which makes an analytic description tractable in the limit of infinitely-wide NNs. However, this limit falls short to capture some effects only present in finite networks such as feature-learning. Recently, progress has been made by deriving corrections to the infinite-width solutions that form a middle ground of capturing effects of realistic NNs while still maintaining some of the analytical tractability of the infinite-width solutions. The derivation of such corrections follows a common pattern but proves tedious due to the vast number of algebraic steps involved, thus hindering progress. In this work, we introduce a diagrammatic framework capable of deriving such corrections for all objects that are necessary to study learning dynamics of NNs. This approach enables the reduction of several pages long algebraic derivations to merely a few lines of diagrammatic steps. We test the framework by computing corrections to several mathematical objects, such as the NTK, describing NNs. Those results are tested extensively by comparing to empirical simulations and show already close agreement to practical neural networks at reasonable widths of $n \gtrsim 20$.