Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks
We analyze feature learning in infinite width neural networks trained with
gradient flow through a self-consistent dynamical field theory. We construct a
collection of deterministic dynamical order parameters which are inner-product
kernels for hidden unit activations and gradients in each layer at pairs of
time points, providing a reduced description of network activity through
training. These kernel order parameters collectively define the hidden layer
activation distribution, the evolution of the neural tangent kernel, and
consequently output predictions. For deep linear networks, these kernels
satisfy a set of algebraic matrix equations. For nonlinear networks, we provide
an alternating sampling procedure to self-consistently solve for the kernel
order parameters. We provide comparisons of the self-consistent solution to
various approximation schemes including the static NTK approximation, gradient
independence assumption, and leading order perturbation theory, showing that
each of these approximations can break down in regimes where general
self-consistent solutions still provide an accurate description. Lastly, we
provide experiments in more realistic settings which demonstrate that the loss
and kernel dynamics of CNNs at fixed feature learning strength is preserved
across different widths on a CIFAR classification task.