Neural network Gaussian process

An NNGP is derived which is equivalent to a Bayesian neural network with this fully connected architecture.

Consider a fully connected artificial neural network with inputs ${\displaystyle x}$, parameters ${\displaystyle \theta }$ consisting of weights ${\displaystyle W^{l}}$ and biases ${\displaystyle b^{l}}$ for each layer ${\displaystyle l}$ in the network, pre-activations (pre-nonlinearity) ${\displaystyle z^{l}}$, activations (post-nonlinearity) ${\displaystyle y^{l}}$, pointwise nonlinearity ${\displaystyle \phi (\cdot )}$, and layer widths ${\displaystyle n^{l}}$. For simplicity, the width ${\displaystyle n^{L+1}}$ of the readout vector ${\displaystyle z^{L}}$ is taken to be 1. The parameters of this network have a prior distribution ${\displaystyle p(\theta )}$, which consists of an isotropic Gaussian for each weight and bias, with the variance of the weights scaled inversely with layer width. This network is illustrated in the figure to the right, and described by the following set of equations:

${\displaystyle {\begin{aligned}x&\equiv {\text{input}}\\y^{l}(x)&=\left\{{\begin{array}{lcl}x&&l=0\\\phi \left(z^{l-1}(x)\right)&&l>0\end{array}}\right.\\z_{i}^{l}(x)&=\sum _{j}W_{ij}^{l}y_{j}^{l}(x)+b_{i}^{l}\\W_{ij}^{l}&\sim {\mathcal {N}}\left(0,{\frac {\sigma _{w}^{2}}{n^{l}}}\right)\\b_{i}^{l}&\sim {\mathcal {N}}\left(0,\sigma _{b}^{2}\right)\\\phi (\cdot )&\equiv {\text{nonlinearity}}\\y^{l}(x),z^{l-1}(x)&\in \mathbb {R} ^{n^{l}\times 1}\\n^{L+1}&=1\\\theta &=\left\{W^{0},b^{0},\dots ,W^{L},b^{L}\right\}\end{aligned}}}$

We first observe that the pre-activations ${\displaystyle z^{l}}$ are described by a Gaussian process conditioned on the preceding activations ${\displaystyle y^{l}}$. This result holds even at finite width. Each pre-activation ${\displaystyle z_{i}^{l}}$ is a weighted sum of Gaussian random variables, corresponding to the weights ${\displaystyle W_{ij}^{l}}$ and biases ${\displaystyle b_{i}^{l}}$, where the coefficients for each of those Gaussian variables are the preceding activations ${\displaystyle y_{j}^{l}}$. Because they are a weighted sum of zero-mean Gaussians, the ${\displaystyle z_{i}^{l}}$ are themselves zero-mean Gaussians (conditioned on the coefficients ${\displaystyle y_{j}^{l}}$). Since the ${\displaystyle z^{l}}$ are jointly Gaussian for any set of ${\displaystyle y^{l}}$, they are described by a Gaussian process conditioned on the preceding activations ${\displaystyle y^{l}}$. The covariance or kernel of this Gaussian process depends on the weight and bias variances ${\displaystyle \sigma _{w}^{2}}$ and ${\displaystyle \sigma _{b}^{2}}$, as well as the second moment matrix ${\displaystyle K^{l}}$ of the preceding activations ${\displaystyle y^{l}}$,

${\displaystyle {\begin{aligned}z_{i}^{l}\mid y^{l}&\sim {\mathcal {GP}}\left(0,\sigma _{w}^{2}K^{l}+\sigma _{b}^{2}\right)\\K^{l}(x,x')&={\frac {1}{n^{l}}}\sum _{i}y_{i}^{l}(x)y_{i}^{l}(x')\end{aligned}}}$

The effect of the weight scale ${\displaystyle \sigma _{w}^{2}}$ is to rescale the contribution to the covariance matrix from ${\displaystyle K^{l}}$, while the bias is shared for all inputs, and so ${\displaystyle \sigma _{b}^{2}}$ makes the ${\displaystyle z_{i}^{l}}$ for different datapoints more similar and makes the covariance matrix more like a constant matrix.

The pre-activations ${\displaystyle z^{l}}$ only depend on ${\displaystyle y^{l}}$ through its second moment matrix ${\displaystyle K^{l}}$. Because of this, we can say that ${\displaystyle z^{l}}$ is a Gaussian process conditioned on ${\displaystyle K^{l}}$, rather than conditioned on ${\displaystyle y^{l}}$,

${\displaystyle {\begin{aligned}z_{i}^{l}\mid K^{l}&\sim {\mathcal {GP}}\left(0,\sigma _{w}^{2}K^{l}+\sigma _{b}^{2}\right).\end{aligned}}}$

As previously defined, ${\displaystyle K^{l}}$ is the second moment matrix of ${\displaystyle y^{l}}$. Since ${\displaystyle y^{l}}$ is the activation vector after applying the nonlinearity ${\displaystyle \phi }$, it can be replaced by ${\displaystyle \phi \left(z^{l-1}\right)}$, resulting in a modified equation expressing ${\displaystyle K^{l}}$ for ${\displaystyle l>0}$ in terms of ${\displaystyle z^{l-1}}$,

${\displaystyle {\begin{aligned}K^{l}(x,x')&={\frac {1}{n^{l}}}\sum _{i}\phi \left(z_{i}^{l-1}(x)\right)\phi \left(z_{i}^{l-1}(x')\right).\end{aligned}}}$

We have already determined that ${\displaystyle z^{l-1}|K^{l-1}}$ is a Gaussian process. This means that the sum defining ${\displaystyle K^{l}}$ is an average over ${\displaystyle n^{l}}$ samples from a Gaussian process which is a function of ${\displaystyle K^{l-1}}$,

${\displaystyle {\begin{aligned}\left\{z_{i}^{l-1}(x),z_{i}^{l-1}(x')\right\}&\sim {\mathcal {GP}}\left(0,\sigma _{w}^{2}K^{l-1}+\sigma _{b}^{2}\right).\end{aligned}}}$

As the layer width ${\displaystyle n^{l}}$ goes to infinity, this average over ${\displaystyle n^{l}}$ samples from the Gaussian process can be replaced with an integral over the Gaussian process:

${\displaystyle {\begin{aligned}\lim _{n^{l}\rightarrow \infty }K^{l}(x,x')&=\int dzdz'\phi (z)\phi (z'){\mathcal {N}}\left(\left[{\begin{array}{c}z\\z'\end{array}}\right];0,\sigma _{w}^{2}\left[{\begin{array}{cc}K^{l-1}(x,x)&K^{l-1}(x,x')\\K^{l-1}(x',x)&K^{l-1}(x',x')\end{array}}\right]+\sigma _{b}^{2}\right)\end{aligned}}}$

So, in the infinite width limit the second moment matrix ${\displaystyle K^{l}}$ for each pair of inputs ${\displaystyle x}$ and ${\displaystyle x'}$ can be expressed as an integral over a 2d Gaussian, of the product of ${\displaystyle \phi (z)}$ and ${\displaystyle \phi (z')}$. There are a number of situations where this has been solved analytically, such as when ${\displaystyle \phi (\cdot )}$ is a ReLU,[18] ELU, GELU,[19] or error function[2] nonlinearity. Even when it can't be solved analytically, since it is a 2d integral it can generally be efficiently computed numerically.[3] This integral is deterministic, so ${\displaystyle K^{l}|K^{l-1}}$ is deterministic.

For shorthand, we define a functional ${\displaystyle F}$, which corresponds to computing this 2d integral for all pairs of inputs, and which maps ${\displaystyle K^{l-1}}$ into ${\displaystyle K^{l}}$,

${\displaystyle {\begin{aligned}\lim _{n^{l}\rightarrow \infty }K^{l}&=F\left(K^{l-1}\right).\end{aligned}}}$

By recursively applying the observation that ${\displaystyle K^{l}\mid K^{l-1}}$ is deterministic as ${\displaystyle n^{l}\rightarrow \infty }$, ${\displaystyle K^{L}}$ can be written as a deterministic function of ${\displaystyle K^{0}}$,

${\displaystyle {\begin{aligned}\lim _{\min \left(n^{1},\dots ,n^{L}\right)\rightarrow \infty }K^{L}&=F\circ F\cdots \left(K^{0}\right)=F^{L}\left(K^{0}\right),\end{aligned}}}$

where ${\displaystyle F^{L}}$ indicates applying the functional ${\displaystyle F}$ sequentially ${\displaystyle L}$ times. By combining this expression with the further observations that the input layer second moment matrix ${\displaystyle K^{0}(x,x')={\frac {1}{n^{0}}}\sum _{i}x_{i}x'_{i}}$ is a deterministic function of the input ${\displaystyle x}$, and that ${\displaystyle z^{L}|K^{L}}$ is a Gaussian process, the output of the neural network can be expressed as a Gaussian process in terms of its input,

${\displaystyle {\begin{aligned}z_{i}^{L}(x)&\sim {\mathcal {GP}}\left(0,\sigma _{w}^{2}F^{L}\left(K^{0}\right)+\sigma _{b}^{2}\right).\end{aligned}}}$