Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks

Abstract

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 selfconsistent 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.Preprint. Under review.

A Minimum Perturbation Theory of Deep Perceptual Learning

“…Overall, leveraging the large size of the network involved in PL, we have developed a statistical mechanics theory of PL in deep neural networks which provides mechanistic and normative understanding of several important empirical findings of PL. This work complements recent theoretical studies of learning in deep networks [27][28][29][30][31][32][33][34][35], contributing to the understanding of learning and computation in these important architectures.…”

“…In the present work, we directly addressed the issue of PL in a deep network by studying PL of a fine discrimination task in a deep neural network (DNN) model of the sensory hierarchy [24][25][26]. As learning dynamics in DNNs are in general challenging to study [27][28][29][30][31][32][33][34][35], we developed a mean-field theory of information propagation in the model at the limit of large numbers of neurons in every layer and large number of training examples. The theory reveals that during the perceptual task, the DNN effectively behaves like a deep linear neural network.…”

Population codes enable learning from few examples by shaping inductive bias

“…DMFT methods have been used to analyze the test loss dynamics for general linear and spiked tensor models trained with high-dimensional random data (Mannelli et al, 2019;Mignacco et al, 2020;Mignacco & Urbani, 2022) and deep networks dynamics with random initialization (Bordelon & Pehlevan, 2022b;Bordelon et al, 2023). High dimensional limits of SGD have been analyzed with Volterra integral equations in the offline case (Paquette et al, 2021) or with recursive matrix equations in the online case (Varre et al, 2021;Bordelon & Pehlevan, 2022a).…”

“…Further analyses of the after-kernels of feature learning networks are performed in Appendix L. We see that the kernels continue to evolve substantially throughout training. This indicates that a full explanation of the compute optimal scaling exponents will require something resembling a mechanistic theory of kernel evolution (Long, 2021;Fort et al, 2020;Atanasov et al, 2022;Bordelon & Pehlevan, 2022b).…”

“…The simplest way to solve the DMFT equations is to iterate them from a reasonable initial condition (Mignacco et al, 2020;Bordelon & Pehlevan, 2022b). We solve in discrete time for T × T matrices {R 0,2 , R 1 , R 2,4 , R 3 , C 0 , C 1 , C 2 , C 3 } which have entries [R] t,s = R(t, s), [C] t,s = C(t, s), etc.…”