A Gaussian Process View on Observation Noise and Initialization in Wide Neural Networks

Sergio Calvo Ordoñez, Jonathan Plenk, Richard Bergna, Alvaro Cartea, José Miguel Hernández-Lobato, Konstantina Palla, Kamil Ciosek

TL;DR: We formally propose a framework for wide neural networks that shows equivalence to adding aleatoric noise in an NTK-GP posterior mean, preserving linearization and enabling initialization with arbitrary prior means.

Abstract: Performing gradient descent in a wide neural network is equivalent to computing the posterior mean of a Gaussian Process with the Neural Tangent Kernel (NTK-GP), for a specific choice of prior mean and with zero observation noise. However, existing formulations of this result have two limitations: i) the resultant NTK-GP assumes no noise in the observed target variables, which can result in suboptimal predictions with noisy data; ii) it is unclear how to extend the equivalence to an arbitrary prior mean, a crucial aspect of formulating a well-specified model. To address the first limitation, we introduce a regularizer into the neural network's training objective, formally showing its correspondence to incorporating observation noise into the NTK-GP model. To address the second, we introduce a \textit{shifted network} that enables arbitrary prior mean functions. This approach further allows us to obtain the posterior mean with gradient descent on a single neural network, without expensive ensembling or kernel matrix inversion. Our theoretical insights are validated empirically, with experiments exploring different values of observation noise, datasets, and network architectures. These results remove key obstacles that have limited the practical use of NTK-GP equivalence in applied Gaussian process modeling.

Submission Number: 1683