Richer priors for infinitely wide multi-layer perceptrons

TL;DR

This paper extends the Gaussian process limit of infinitely wide MLPs to more general priors, including non-zero mean and partially exchangeable priors, and empirically demonstrates improved kernel properties and convergence behaviors.

Contribution

It introduces a broader class of priors for deep MLPs, derives their limiting kernels, and empirically shows these priors avoid pathologies of previous models.

Findings

Derived kernels for deep MLPs with richer priors.

Empirical evidence of convergence to the limiting Gaussian process.

Improved kernel properties compared to previous priors.

Abstract

It is well-known that the distribution over functions induced through a zero-mean iid prior distribution over the parameters of a multi-layer perceptron (MLP) converges to a Gaussian process (GP), under mild conditions. We extend this result firstly to independent priors with general zero or non-zero means, and secondly to a family of partially exchangeable priors which generalise iid priors. We discuss how the second prior arises naturally when considering an equivalence class of functions in an MLP and through training processes such as stochastic gradient descent. The model resulting from partially exchangeable priors is a GP, with an additional level of inference in the sense that the prior and posterior predictive distributions require marginalisation over hyperparameters. We derive the kernels of the limiting GP in deep MLPs, and show empirically that these kernels avoid certain pathologies present in previously studied priors. We empirically evaluate our claims of convergence by measuring the maximum mean discrepancy between finite width models and limiting models. We compare the performance of our new limiting model to some previously discussed models on synthetic regression problems. We observe increasing ill-conditioning of the marginal likelihood and hyper-posterior as the depth of the model increases, drawing parallels with finite width networks which require notoriously involved optimisation tricks.