Connecting NTK and NNGP: A Unified Theoretical Framework for Wide Neural Network Learning Dynamics

TL;DR

This paper unifies the NTK and NNGP theories for wide neural networks by introducing a new dynamical kernel, revealing two learning phases and explaining neural activity drift in biological systems.

Contribution

It develops an analytical framework connecting NTK and NNGP through a new time-dependent kernel, elucidating the learning dynamics and phases in wide neural networks.

Findings

Identification of two distinct learning phases: gradient-driven and diffusive.

Derivation of a new Neural Dynamical Kernel (NDK) linking NTK and NNGP.

Insights into neural activity drift and its robustness in biological neural circuits.

Abstract

Artificial neural networks have revolutionized machine learning in recent years, but a complete theoretical framework for their learning process is still lacking. Substantial advances were achieved for wide networks, within two disparate theoretical frameworks: the Neural Tangent Kernel (NTK), which assumes linearized gradient descent dynamics, and the Bayesian Neural Network Gaussian Process (NNGP). We unify these two theories using gradient descent learning with an additional noise in an ensemble of wide deep networks. We construct an analytical theory for the network input-output function and introduce a new time-dependent Neural Dynamical Kernel (NDK) from which both NTK and NNGP kernels are derived. We identify two learning phases: a gradient-driven learning phase, dominated by loss minimization, in which the time scale is governed by the initialization variance. It is followed by a slow diffusive learning stage, where the parameters sample the solution space, with a time constant decided by the noise and the Bayesian prior variance. The two variance parameters strongly affect the performance in the two regimes, especially in sigmoidal neurons. In contrast to the exponential convergence of the mean predictor in the initial phase, the convergence to the equilibrium is more complex and may behave nonmonotonically. By characterizing the diffusive phase, our work sheds light on representational drift in the brain, explaining how neural activity changes continuously without degrading performance, either by ongoing gradient signals that synchronize the drifts of different synapses or by architectural biases that generate task-relevant information that is robust against the drift process. This work closes the gap between the NTK and NNGP theories, providing a comprehensive framework for the learning process of deep wide neural networks and for analyzing dynamics in biological circuits.