Bayesian Machine Learning: Function Space

This blog is the fifth in a series in which we consider machine learning from four different viewpoints. We either use gradient descent or a fully Bayesian approach and for each, we can choose to focus on either the model parameters or the output function (figure 1). In part I, part II, and part III, we considered gradient descent in both parameter and function space; this led to the development of the neural tangent kernel.

In part IV of the series we considered the Bayesian approach to model fitting in a simple linear regression model from the perspective of parameter space. Here, we determine a probability distribution over possible parameter values rather than choosing a single “best” estimate. To make a prediction, we take a weighted average the predictions from all possible parameter estimates, where the weights are based on these probabilities. This scheme has the merits of (i) making more moderate predictions and (ii) providing estimates of uncertainty over these predictions.

In this (part V) of this series, we investigate the Bayesian approach from the perspective of function space. This is considerably simpler than the parameter space view and, for linear regression, results in the same predictions. However, as we shall see in the subsequent parts of this series, the implications of the two viewpoints are quite different for neural network models.

This blog is accompanied by a CoLab Python Notebook.

Function Space

In the function space view, we convert the prior $Pr(\boldsymbol\phi)$ over the model parameters $\boldsymbol\phi$ space into a prior $\operatorname{Pr}(\mathbf{f})$ over the function $\textbf{f}[\bullet]$ itself. This allows us to define a joint probability distribution between the function values for the training data and that for the test data. We then compute the conditional distribution of the function value for the test data given the training data.

Linear Model

These ideas are easiest to understand with the concrete example of the linear model:

\begin{align*}
y & =\mathrm{f}[\mathbf{x}, \boldsymbol{\phi}] \\
& =\mathbf{x}^{T} \boldsymbol{\phi}, \tag{1}
\end{align*}

where $\mathbf{x} \in \mathbb{R}^{D}$ is the input data, $y \in \mathbb{R}$ is the model output and $\boldsymbol{\phi} \in \mathbb{R}^{D}$ contains the parameters. As usual, we append the value one to the input data so that $\mathbf{x} \leftarrow$ $\left[\mathbf{x}^{T}, 1\right]^{T}$ to account for the y-offset.

As before, we’ll concatenate all of the $I$ training inputs $\mathbf{x}_{i} \in \mathbb{R}^{D}$ into a single matrix $\mathbf{X} \in \mathbb{R}^{D \times I}$, all of the function outputs $\mathrm{f}[\mathbf{x}, \boldsymbol{\phi}] \in \mathbb{R}$ into a single column vector $\mathbf{f} \in \mathbb{R}^{D}$ and all of the training outputs $y_{i} \in \mathbb{R}$ into a single column vector $\mathbf{y} \in$ $\mathbb{R}^{D}$. We can then write the model in two stages:

\begin{align*}
\mathbf{f} & =\mathbf{X}^{T} \boldsymbol{\phi} \\
\operatorname{Pr}(\mathbf{y} \mid \mathbf{f}) & =\operatorname{Norm}_{\mathbf{y}}\left[\mathbf{f}, \sigma_{n}^{2} \mathbf{I}\right], \tag{2}
\end{align*}

where as before $\sigma_{n}^{2}$ is the noise variance that accounts for the fact that the observed data are not exactly fitted by the model.

Prior distribution

We define a prior distribution over the parameters. As for the parameter-space perspective, we choose a spherical normal distribution with mean zero, and variance $\sigma_{p}^{2}$:

\begin{equation*}
\operatorname{Pr}(\boldsymbol\phi)=\operatorname{Norm}_{\boldsymbol{\phi}}\left[\mathbf{0}, \sigma_{p}^{2} \mathbf{I}\right]. \tag{3}
\end{equation*}

This is illustrated in figure 2a

Gaussian processes

Now we note that the function values $\mathbf{f}=\mathbf{X}^{T} \boldsymbol{\phi}$ are an affine transform of the parameters, and we use the relation (see proof at end of blog) that if $\operatorname{Pr}(\mathbf{z})=$$\operatorname{Norm}_{\mathbf{z}}[\boldsymbol{\mu}, \boldsymbol{\Sigma}]$, then the affine transform $\mathbf{a}+\mathbf{B z}$ is distributed as:

\begin{equation*}
\operatorname{Pr}(\mathbf{a}+\mathbf{B z})=\operatorname{Norm}\left[\mathbf{a}+\mathbf{B} \boldsymbol{\mu}, \mathbf{B} \boldsymbol{\Sigma} \mathbf{B}^{T}\right]. \tag{4}
\end{equation*}

Applying this relation, we can get a prior distribution over the function values $\mathbf{f}$:

\begin{equation*}
\operatorname{Pr}(\mathbf{f} \mid \mathbf{X})=\operatorname{Norm}_{\mathbf{f}}\left[\mathbf{0}, \sigma_{p}^{2} \mathbf{X}^{T} \mathbf{X}\right]. \tag{5}
\end{equation*}

This is technically a Gaussian process, and can equivalently be written as:

\begin{equation*}
\mathbf{f}[\mathbf{X}] \sim \mathcal{G P}\left[\mathbf{0}, \sigma_{p}^{2} \mathbf{X}^{T} \mathbf{X}\right]. \tag{6}
\end{equation*}

The term Gaussian process just refers to a collection of random variables, any finite number of which have a joint Gaussian distribution (see Rasmussen and Williams, 2006). In this case, we can define any input values $\mathbf{X}$ and equation 6 defines a Gaussian distribution between them. For example, we can define a set of equally spaced 1D points and draw samples $\mathbf{f}$ from this prior. Each sample looks like a straight line (figure 2b).

Since the output $\mathbf{y}$ is just a noisy version of the function values, we can similarly write a prior distribution over the output values:

\begin{equation*}
\operatorname{Pr}(\mathbf{y} \mid \mathbf{X})=\operatorname{Norm}_{\mathbf{f}}\left[\mathbf{0}, \sigma_{p}^{2} \mathbf{X}^{T} \mathbf{X}+\sigma_{n}^{2} \mathbf{I}\right]. \tag{7}
\end{equation*}

Figure 2c illustrates samples drawn from this distribution for a set of equally spaced inputs.

Inference

Consider a test set of points $\mathbf{X}^{*}$ which map to function values $\mathbf{f}^{*}$. We can write the joint distribution of the training values $\mathbf{y}$ and the new function value $\mathbf{f}^{*}$:

\[
\operatorname{Pr}\binom{\mathbf{y}}{\mathbf{f}^{*}}=\operatorname{Norm}_{\left[\mathbf{y},\mathbf{f}^{*} \right]^{T}}\left[\left[\begin{array}{l}
0 \tag{8}\\
0
\end{array}\right],\left[\begin{array}{cc}
\sigma_{p}^{2} \mathbf{X}^{* T} \mathbf{X}^{*} & \sigma_{p}^{2} \mathbf{X}^{* T} \mathbf{X} \\
\sigma_{p}^{2} \mathbf{X}^{T} \mathbf{X}^{*} & \sigma_{p}^{2} \mathbf{X}^{T} \mathbf{X}+\sigma_{n}^{2} \mathbf{I}
\end{array}\right]\right].
\]

Now we use the following well-known relation for conditional distributions in Gaussian variables (proof at end of blog): if two sets of variables $\mathbf{z}_{1}$ and $\mathbf{z}_{2}$ are jointly distributed as a multivariate normal:

\[
\operatorname{Pr}\left(\left[\begin{array}{l}
\mathbf{z}_{1} \tag{9}\\
\mathbf{z}_{2}
\end{array}\right]\right)=\operatorname{Norm}_{\mathbf{z}}\left[\left[\begin{array}{l}
\boldsymbol{\mu}_{1} \\
\boldsymbol{\mu}_{2}
\end{array}\right],\left[\begin{array}{ll}
\boldsymbol{\Sigma}_{11} & \boldsymbol{\Sigma}_{12} \\
\boldsymbol{\Sigma}_{21} & \boldsymbol{\Sigma}_{22}
\end{array}\right]\right],
\]

then the conditional distribution of $\mathbf{z}_{2}$ given observed values $\mathbf{z}_{1}$ is:

\begin{equation*}
\operatorname{Pr}\left(\mathbf{z}_{1} \mid \mathbf{z}_{2}\right)=\operatorname{Norm}_{\mathbf{z}_{1}}\left[\boldsymbol{\mu}_{\mathbf{1}}+\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1}\left(\mathbf{z}_{2}-\boldsymbol{\mu}_{2}\right), \boldsymbol{\Sigma}_{11}-\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1} \boldsymbol{\Sigma}_{12}^{T}\right]. \tag{10}
\end{equation*}

Applying this relation to equation 8, we get:

\begin{align*}
\operatorname{Pr}\left(\mathbf{f}^{*} \mid \mathbf{X}^{*}, \mathbf{X}, \mathbf{y}\right)&= \operatorname{Norm}_{\mathbf{f}^{*}}\left[\sigma_{p}^{2} \mathbf{X}^{* T} \mathbf{X}\left(\sigma_{p}^{2} \mathbf{X}^{T} \mathbf{X}+\sigma_{n}^{2} \mathbf{I}\right)^{-1} \mathbf{y}\right., \tag{11}\\
&\hspace{3cm}\left.\sigma_{p}^{2} \mathbf{X}^{* T} \mathbf{X}^{*}-\sigma_{p}^{2} \mathbf{X}^{* T} \mathbf{X}\left(\sigma_{p}^{2} \mathbf{X}^{T} \mathbf{X}+\sigma_{n}^{2} \mathbf{I}\right)^{-1} \sigma_{p}^{2} \mathbf{X}^{T} \mathbf{X}^{*}\right] \\
&= \operatorname{Norm}_{\mathbf{f}^{*}}\Biggl[\mathbf{X}^{* T} \mathbf{X}\left(\mathbf{X}^{T} \mathbf{X}+\frac{\sigma_{n}^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{y}, \\
&\hspace{3cm}\sigma_{p}^{2} \mathbf{X}^{* T} \mathbf{X}^{*}-\sigma_{p}^{2} \mathbf{X}^{* T} \mathbf{X}\left(\mathbf{X}^{T} \mathbf{X}+\frac{\sigma_{n}^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{X}^{T} \mathbf{X}^{*}\Biggr]
\end{align*}

where we have moved the terms $\sigma_{p}^{2}$ into the inverse in the second line. If we draw samples from this distribution for equally spaced test points in $\mathbf{X}^{*}$, each sample looks like a line that is broadly compatible with the training data $\mathbf{X}$ (figure 3a).

Finally, we note that the original data points are drawn from distribution over the noisy output $\mathbf{y}^{*}$, which has the same distribution plus an the extra noise term on the covariance.

\begin{align*}
& \operatorname{Pr}\left(\mathbf{y}^{*} \mid \mathbf{X}^{*}, \mathbf{X}, \mathbf{y}\right)=\operatorname{Norm}_{\mathbf{y}^{*}} {\Biggl[\mathbf{X}^{* T} \mathbf{X}\left(\mathbf{X}^{T} \mathbf{X}+\frac{\sigma_{n}^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{y}} \\
&\hspace{3cm}\sigma_{p}^{2} \mathbf{X}^{* T} \mathbf{X}^{*}-\sigma_{p}^{2} \mathbf{X}^{* T} \mathbf{X}\left(\mathbf{X}^{T} \mathbf{X}+\frac{\sigma_{n}^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{X}^{T} \mathbf{X}^{*}+\sigma_{n}^{2} \mathbf{I}\Biggl]. \tag{12}
\end{align*}

Samples from this distribution look like noisy lines that are broadly compatible with the data points (figure 3b).

Figure 4 visualizes the solution in a different way. Here, we have plotted the mean of the posterior (the MAP solution) and the marginal uncertainty at each of a set of equally spaced positions (represented by a $\pm 2$ standard deviation region). In this way, we can visualize the function and its uncertainty, or the equivalent result in output space.

Parameter space vs. function space

Recall from part IV of this series of blogs that the predictive distribution for the parameter space (Bayesian) view of the linear model was:

\begin{align*}
\operatorname{Pr}\left(y^{*} \mid \mathbf{x}^{*}, \mathbf{X}, \mathbf{y}\right)=\operatorname{Norm}_{y^{*}}\Biggl[ & \mathbf{x}^{* T} \mathbf{X}\left(\mathbf{X}^{T} \mathbf{X}+\frac{\sigma_{n}^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{y} \tag{13}\\
& \sigma_{p}^{2} \mathbf{x}^{* T} \mathbf{x}^{*}-\sigma_{p}^{2} \mathbf{x}^{* T} \mathbf{X}\left(\mathbf{X}^{T} \mathbf{X}+\frac{\sigma^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{X}^{T} \mathbf{x}^{*}+\sigma_{n}^{2}\Biggr].
\end{align*}

It’s clear that for this model, the parameter space view and the function space view lead to the same result (although we developed the function space approach to make multiple predictions simultaneously). However, this is NOT true in general. In future articles in this series, we will see that these two perspectives have quite different implications for neural networks.

Non-linear regression

For both the parameter space and function space views, we only yield neat closed-form solution for the predictions because the model is linear; although the same principles apply for non-linear models, there is not generally a closed-form solution. What then, should we do when faced with a training dataset for which a linear model is clearly inadequate? One tractable approach is to transform the initial data by a fixed non-linear transformation $\mathbf{z}=\mathbf{g}[\mathbf{x}]$ and then apply a linear model $\mathrm{f}[\mathbf{z}, \boldsymbol{\phi}]$ to the transformed data.

Example: polynomial regression

A simple example of this type of model is polynomial regression for 1D data $x$ in which we might use the non-linear transformation (figure 2a):

\[
\mathbf{z}=\left[\begin{array}{c}
1 \tag{2}\\
x \\
x^{2} \\
x^{3} \\
x^{4}
\end{array}\right].
\]

If we draw samples from the prior from this model, each looks like a polynomial of order five (figure 5).

When we condition the output function on a training dataset, we samples from the posterior look like order-5 polynomial curves that broadly agree with the data (figure 6a), or we can alternatively visualize the MAP solution and the associated marginal uncertainty (figure 6b).

In this way, we can make closed-form predictions based on a polynomial model and yield estimates of uncertainty, without ever explicitly ‘fitting’ the model and committing to a single set of parameters.

Example: ReLU model

A second example is a transformation based on a collection of offset ReLU functions:

\[
\mathbf{z}=\left[\begin{array}{c}
1 \tag{15}\\
\operatorname{ReLU}\left[\beta_{1}+\omega_{1} x\right] \\
\operatorname{ReLU}\left[\beta_{2}+\omega_{2} x\right] \\
\operatorname{ReLU}\left[\beta_{3}+\omega_{3} x\right] \\
\operatorname{ReLU}\left[\beta_{4}+\omega_{4} x\right] \\
\operatorname{ReLU}\left[\beta_{5}+\omega_{5} x\right]
\end{array}\right].
\]

Note that this is not exactly the same as a neural network because in this case the biases $\beta_{0}$ and slopes $\omega_{0}$ are chosen in advance and fixed; they are not optimized (or marginalized over) as part of the training process.

The prior and posterior functions are visualized in figure 7 and 8, respectively. In each case the output function is a piecewise linear, with the position of the joints determined by the joints in the component ReLU functions (which are themselves determined by the predefined $\beta_{\bullet}$ and $\omega_{\bullet}$ parameters.

Disadvantages of non-linear regression

Although, the approach of transforming the input data using a fixed non-linearity is simple and effective, these methods do not necessarily scale well to high dimensions. For example, if we have 1000 dimensional input, then the set of third degree polynomials contains $\left(D_{z}=1+1000+1000^{2}+1000^{3}\right)$ terms. However, the complexity of the fitted function should ultimately depend on the amount of training data. If this is significantly less than the number of transformed dimensions, we should be able to make efficiency gains.

Kernels and the kernel trick

To make progress, we note that the final equation for the predictive distribution:

\begin{align*}
\operatorname{Pr}\left(\mathbf{y}^{*} \mid \mathbf{Z}^{*}, \mathbf{Z}, \mathbf{y}\right)=\operatorname{Norm}_{\mathbf{y}^{*}} & {\Biggl[\mathbf{Z}^{* T} \mathbf{Z}\left(\mathbf{Z}^{T} \mathbf{Z}+\frac{\sigma_{n}^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{y}}, \\
& \sigma_{p}^{2} \mathbf{Z}^{* T} \mathbf{Z}^{*}-\sigma_{p}^{2} \mathbf{Z}^{* T} \mathbf{Z}\left(\mathbf{Z}^{T} \mathbf{Z}+\frac{\sigma_{n}^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{Z}^{T} \mathbf{Z}^{*}+\sigma_{n}^{2} \mathbf{I}\Biggr] \tag{16}
\end{align*}

depends only on matrix products of the form $\mathbf{Z}^{T} \mathbf{Z}$ where $\mathbf{Z} \in \mathbb{R}^{D_{z} \times I}$ contains the transformed data vectors in its columns.

When the dimension $D_{z}$ of the transformed vectors $\mathbf{z}=\mathbf{g}[\mathbf{x}]$ is high, computing the non-linear transformations and computing the dot products that make the matrix product can be time-consuming. A viable alternative is kernel substitution in which we directly define a single kernel function $\mathrm{k}\left[\mathbf{x}_{i}, \mathbf{x}_{j}\right]$ as a replacement for the operations $\mathbf{g}\left[\mathbf{x}_{i}\right]^{T} \mathbf{g}\left[\mathbf{x}_{j}\right]$. For many transformations $\mathbf{g}[\bullet]$ it is more efficient to evaluate the kernel function directly than to transform the variables separately and then compute the dot product.

Taking this idea one step further, it is possible to choose a kernel function $\mathrm{k}\left[\mathbf{x}_{i}, \mathbf{x}_{j}\right]$ without knowledge of what transformation $\mathbf{g}[\bullet]$ it corresponds to. Clearly, the function must be carefully chosen so that it does in fact correspond to computing some function $\mathbf{z}=\mathbf{g}[\mathbf{x}]$ of each data vector and taking the inner product of the results: for example, since $\mathbf{z}_{i}^{T} \mathbf{z}_{j}=\mathbf{z}_{j}^{T} \mathbf{z}_{i}$ it must treat its arguments symmetrically so that $\mathrm{k}\left[\mathbf{x}_{i}, \mathbf{x}_{j}\right]=\mathrm{k}\left[\mathbf{x}_{j}, \mathbf{x}_{i}\right]$.

When we use kernel functions, we no longer explicitly compute the transformed vector $\mathbf{z}$. One advantage of this is we can define kernel functions that correspond to projecting the data into very high dimensional or even infinite spaces.

We now return to the question of which constraints must be put on the kernel function to ensure that there exists an underlying function $\mathbf{f}\left[\mathbf{x}_{i}\right]=\mathbf{z}_{i}$ so that $\mathrm{k}\left[\mathbf{x}_{i}, \mathbf{x}_{j}\right]=\mathbf{z}_{i}^{T} \mathbf{z}_{j}$. Mercer’s theorem states that this is true when the kernel’s arguments are in a measurable space, and the kernel is positive semi-definite so that:

\begin{equation*}
\sum_{i j} \mathrm{k}\left[\mathbf{x}_{i}, \mathbf{x}_{j}\right] a_{i} a_{j} \geq 0, \tag{17}
\end{equation*}

for any finite subset $\left\{\mathbf{x}_{n}\right\}_{n=1}^{N}$ of vectors in the space and any real numbers $\left\{a_{n}\right\}_{n=1}^{N}$. Examples of valid kernels include:

The last of these is particularly interesting. It can be shown that this kernel function corresponds to computing infinite length vectors $\mathbf{z}$ and taking their dot product. The entries of $\mathbf{z}$ correspond to evaluating a radial basis function at every possible point in the space of $\mathbf{x}$.

It is also possible to create new kernels by combining two or more existing kernels. For example, sums and products of valid kernels are guaranteed to be positive semi-definite and so are also valid kernels.

Using kernel notation, we can now write the predictive equation as:

\begin{align*}
& \operatorname{Pr}\left(\mathbf{y}^{*} \mid \mathbf{X}^{*}, \mathbf{X}, \mathbf{y}\right)=\operatorname{Norm}_{\mathbf{y}^{*}}\Biggl[\mathbf{K}[\mathbf{X}, \mathbf{X}]\left(\mathbf{K}[\mathbf{X}, \mathbf{X}]+\frac{\sigma_{n}^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{y}, \\
& \hspace{2cm}\quad \sigma_{p}^{2} \mathbf{K}\left[\mathbf{X}^{*}, \mathbf{X}^{*}\right]-\sigma_{p}^{2} \mathbf{K}\left[\mathbf{X}^{*}, \mathbf{X}\right]\left(\mathbf{K}[\mathbf{X}, \mathbf{X}]+\frac{\sigma_{n}^{2}}{\sigma_{p}^{2}} \mathbf{I}\right)^{-1} \mathbf{K}\left[\mathbf{X}, \mathbf{X}^{*}\right]+\sigma_{n}^{2} \mathbf{I}\Biggr], \tag{21}
\end{align*}

where the $(i, j)^{\text {th }}$ element of $\mathbf{K}\left[\mathbf{X}^{*}, \mathbf{X}\right]$ would comprise the element $\mathrm{k}\left[\mathbf{x}_{i}^{*}, \mathbf{x}_{j}\right]$ and so on. Figure 2.9 shows the prior and posterior distribution for kernel regression using the RBF kernel; it results in a set of smooth functions that pass close to the data points.

Conclusion

In this blog we described a method to define a prior over the space of output functions. We can use this prior find a joint distribution over the training and test data points; we then find the conditional distribution over the test points, given the training points. For linear models, this gives the same results as the parameter-space perspective on Bayesian modelling. We saw how this approach can be extended to model more complex data by applying a non-linear transformation to the input data before applying the linear model.

The next part of this series of blogs applies the Bayesian approach to deep neural networks. We will see that the implications of parameter-space and function-space perspectives are quite different here.

Affine transformations of normal distributions

The multivariate normal distribution is defined as:

$$
\operatorname{Pr}(\mathbf{x})=\frac{1}{(2 \pi)^{D}|\boldsymbol{\Sigma}|^{1 / 2}} \exp \left[-0.5(\mathbf{x}-\boldsymbol{\mu}) \boldsymbol{\Sigma}^{-1}(\mathbf{x}-\boldsymbol{\mu})\right].
$$

To transform the distribution as $\mathbf{y}=\mathbf{A} \mathbf{x}+\mathbf{b}$, we substitute in $\mathbf{x}=\mathbf{A}^{-1}(\mathbf{y}-\mathbf{b})$ and divide by the determinant of the Jacobian of the transformation, which for this case is just $|\mathbf{A}|$ giving:

$$
\operatorname{Pr}(\mathbf{x})=\frac{1}{(2 \pi)^{D}|\boldsymbol{\Sigma}|^{1 / 2}|\mathbf{A}|} \exp \left[-0.5\left(\mathbf{A}^{-1}(\mathbf{y}-\mathbf{b})-\boldsymbol{\mu}\right) \boldsymbol{\Sigma}^{-1}\left(\mathbf{A}^{-1}(\mathbf{y}-\mathbf{b})-\boldsymbol{\mu}\right)\right].
$$

We’ll first work with the constant $\kappa$ from outside the exponential to show that it is the constant for a distribution with variance $\mathbf{A} \boldsymbol{\Sigma} \mathbf{A}^{T}$. We have:

\begin{aligned}
\kappa & =\frac{1}{(2 \pi)^{D}|\mathbf{\Sigma}|^{1 / 2}|\mathbf{A}|} \\
& =\frac{1}{(2 \pi)^{D}|\mathbf{\Sigma}|^{1 / 2}|\mathbf{A}|^{1 / 2}|\mathbf{A}|^{1 / 2}} \\
& =\frac{1}{(2 \pi)^{D}|\mathbf{\Sigma}|^{1 / 2}|\mathbf{A}|^{1 / 2}\left|\mathbf{A}^{T}\right|^{1 / 2}} \\
& =\frac{1}{(2 \pi)^{D}|\mathbf{A}|^{1 / 2}|\mathbf{\Sigma}|^{1 / 2}\left|\mathbf{A}^{T}\right|^{1 / 2}} \\
& =\frac{1}{(2 \pi)^{D}\left|\mathbf{A} \mathbf{\Sigma} \mathbf{A}^{T}\right|^{1 / 2}},
\end{aligned}

which is exactly the required constant.

Now we’ll work with the quadratic in the exponential term to show that it corresponds to a normal distribution in $\mathbf{y}$ with variance $\mathbf{A} \boldsymbol{\Sigma} \mathbf{A}^{T}$ and mean $\mathbf{A} \boldsymbol{\mu}+\mathbf{b}$.

\begin{aligned}
\left(\mathbf{A}^{-1}(\mathbf{y}-\mathbf{b})-\boldsymbol{\mu}\right) \boldsymbol{\Sigma}^{-1}\left(\mathbf{A}^{-1}(\mathbf{y}-\mathbf{b})-\boldsymbol{\mu}\right) &\\
&\hspace{-4cm}=\mathbf{y}^{T} \mathbf{A}^{-T} \boldsymbol{\Sigma}^{-1} \mathbf{A}^{-1} \mathbf{y}-2\left(\mathbf{b}^{T} \mathbf{A}^{-T}+\boldsymbol{\mu}^{T}\right) \boldsymbol{\Sigma}^{-1} \mathbf{A}^{-1} \mathbf{y} \\
&+\left(\mathbf{A}^{-1} \mathbf{b}+\boldsymbol{\mu}\right)^{T} \boldsymbol{\Sigma}^{-1}\left(\mathbf{A}^{-1} \mathbf{b}+\boldsymbol{\mu}\right) \\
&\hspace{-4cm}=\mathbf{y}^{T} \mathbf{A}^{-T} \boldsymbol{\Sigma}^{-1} \mathbf{A}^{-1} \mathbf{y}-2\left(\mathbf{b}^{T}+\boldsymbol{\mu}^{T} \mathbf{A}^{T}\right) \mathbf{A}^{-T} \boldsymbol{\Sigma}^{-1} \mathbf{A}^{-1} \mathbf{y} \\
&+\left(\mathbf{A}^{-1} \mathbf{b}+\mathbf{A} \boldsymbol{\mu}\right)^{T} \mathbf{A}^{-T} \boldsymbol{\Sigma}^{-1} \mathbf{A}^{-1}(\mathbf{b}+\mathbf{A} \boldsymbol{\mu}) \\
&\hspace{-4cm}=(\mathbf{y}-(\mathbf{A} \boldsymbol{\mu}+\mathbf{b}))^{T}\left(\mathbf{A} \boldsymbol{\Sigma} \mathbf{A}^{T}\right)^{-1}(\mathbf{y}-(\mathbf{A} \boldsymbol{\mu}+\mathbf{b})).
\end{aligned}

This is clearly the quadratic term from a normal distribution in $\mathbf{y}$ with variance $\mathbf{A} \boldsymbol{\Sigma} \mathbf{A}^{T}$ and mean $\mathbf{A} \boldsymbol{\mu}+\mathbf{b}$ as required.

Conditional distribution of a Gaussian

Consider a joint Gaussian distribution with variables $\mathbf{z}=\left[\mathbf{z}_{1}, \mathbf{z}_{2}\right]$:

$$\operatorname{Pr}(\mathbf{z})=\operatorname{Pr}\left(\left[\begin{array}{l}
\mathbf{z}_{1} \\
\mathbf{z}_{2}
\end{array}\right]\right)=\operatorname{Norm}_{\mathbf{x}}\left[\left[\begin{array}{l}
\boldsymbol{\mu}_{1} \\
\boldsymbol{\mu}_{2}
\end{array}\right],\left[\begin{array}{ll}
\boldsymbol{\Sigma}_{11} & \boldsymbol{\Sigma}_{12} \\
\boldsymbol{\Sigma}_{21} & \boldsymbol{\Sigma}_{22}
\end{array}\right]\right].$$

The conditional distribution of the varaibles $\mathbf{z}_{1}$ given observed variables $\mathbf{z}_{2}$ is:

\begin{equation*}
\operatorname{Pr}\left(\mathbf{z}_{1} \mid \mathbf{z}_{2}\right)=\operatorname{Norm}_{\mathbf{z}_{1}}\left[\boldsymbol{\mu}_{\mathbf{1}}+\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1}\left(\mathbf{z}_{2}-\boldsymbol{\mu}_{2}\right), \boldsymbol{\Sigma}_{11}-\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1} \boldsymbol{\Sigma}_{12}^{T}\right]
\end{equation*}

Proof: We can write out the joint distribution as:

\begin{aligned}
\operatorname{Pr}\left(\left[\begin{array}{l}
\mathbf{z}_{1} \\
\mathbf{z}_{2}
\end{array}\right]\right) & =\kappa_{1} \exp \left[-0.5\left(\left[\begin{array}{l}
\mathbf{z}_{1} \\
\mathbf{z}_{2}
\end{array}\right]-\left[\begin{array}{l}
\boldsymbol{\mu}_{1} \\
\boldsymbol{\mu}_{2}
\end{array}\right]\right)^{T}\left[\begin{array}{ll}
\boldsymbol{\Sigma}_{11} & \boldsymbol{\Sigma}_{12} \\
\boldsymbol{\Sigma}_{12}^{T} & \mathbf{\Sigma}_{22}
\end{array}\right]^{-1}\left(\left[\begin{array}{l}
\mathbf{z}_{1} \\
\mathbf{z}_{2}
\end{array}\right]-\left[\begin{array}{l}
\boldsymbol{\mu}_{1} \\
\boldsymbol{\mu}_{2}
\end{array}\right]\right)\right] \\
& =\kappa_{1} \exp \left[-0.5\left(\left[\begin{array}{l}
\mathbf{z}_{1} \\
\mathbf{z}_{2}
\end{array}\right]-\left[\begin{array}{l}
\boldsymbol{\mu}_{1} \\
\boldsymbol{\mu}_{2}
\end{array}\right]\right)^{T}\left[\begin{array}{ll}
\mathbf{A} & \mathbf{B} \\
\mathbf{C} & \mathbf{D}
\end{array}\right]\left(\left[\begin{array}{l}
\mathbf{z}_{1} \\
\mathbf{z}_{2}
\end{array}\right]-\left[\begin{array}{l}
\boldsymbol{\mu}_{1} \\
\boldsymbol{\mu}_{2}
\end{array}\right]\right)\right],
\end{aligned}

where $\kappa_{1}$ is the standard constant associated with the normal distribution and the terms $\mathbf{A}, \mathbf{B}, \mathbf{C}$ and $\mathbf{D}$ are the four components of the inverse covariance matrix. We can find these using Schur’s complement relation to write:

\begin{aligned}
& \mathbf{A}=\left(\boldsymbol{\Sigma}_{11}-\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1} \boldsymbol{\Sigma}_{12}^{T}\right)^{-1} \\
& \mathbf{B}=-\left(\boldsymbol{\Sigma}_{11}-\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1} \boldsymbol{\Sigma}_{12}^{T}\right)^{-1} \boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1} \\
& \mathbf{C}=-\boldsymbol{\Sigma}_{22}^{-1} \boldsymbol{\Sigma}_{12}^{T}\left(\boldsymbol{\Sigma}_{11}-\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1} \boldsymbol{\Sigma}_{12}^{T}\right)^{-1} \\
& \mathbf{D}=\boldsymbol{\Sigma}_{22}^{-1}+\boldsymbol{\Sigma}_{22}^{-1} \boldsymbol{\Sigma}_{12}^{T}\left(\boldsymbol{\Sigma}_{11}-\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1} \boldsymbol{\Sigma}_{12}^{T}\right)^{-1} \boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1}.
\end{aligned}

This can easily be confirmed by multiplying the resulting inverse matrix with the original covariance matrix, which yields the identity.

Now we exploit the conditional probability relation $\operatorname{Pr}\left(\mathbf{z}_{1} \mid \mathbf{z}_{2}\right) \operatorname{Pr}\left(\mathbf{z}_{2}\right)=\operatorname{Pr}\left(\mathbf{z}_{1}, \mathbf{z}_{2}\right)$ and note that $\operatorname{Pr}\left(\mathbf{z}_{1} \mid \mathbf{z}_{2}\right) \propto \operatorname{Pr}\left(\mathbf{z}_{1}, \mathbf{z}_{2}\right)$. We hence take the approach of reformulating the joint distribution in terms of a normal distribution in $\mathbf{z}_{1}$ alone. Multiplying out the matrices in the exponent, we see that:

$$\operatorname{Pr}\left(\mathbf{z}_{1} \mid \mathbf{z}_{2}\right)=\kappa_{2} \cdot \exp \left[-0.5\left(\left(\mathbf{z}_{1}-\boldsymbol{\mu}_{1}\right)^{T} \mathbf{A}\left(\mathbf{z}_{1}-\boldsymbol{\mu}_{1}\right)+2\left(\mathbf{z}_{1}-\boldsymbol{\mu}_{1}\right)^{T} \mathbf{B}\left(\mathbf{z}_{2}-\boldsymbol{\mu}_{2}\right)\right)\right],$$

where we have subsumed the terms that do not depend on $\mathbf{z}_{1}$ to create a new constant $\kappa_{2}$. Now we both multiply and divide by a new term to complete the square:

\begin{array}{r}
\operatorname{Pr}\left(\mathbf{z}_{1} \mid \mathbf{z}_{2}\right)=\kappa_{3} \cdot \exp \left[-0.5\left(\left(\mathbf{z}_{1}-\boldsymbol{\mu}_{1}\right)^{T} \mathbf{A}\left(\mathbf{z}_{1}-\boldsymbol{\mu}_{1}\right)+2\left(\mathbf{z}_{1}-\boldsymbol{\mu}_{1}\right)^{T} \mathbf{B}\left(\mathbf{z}_{2}-\boldsymbol{\mu}_{2}\right)\right.\right. \\
\left.\left.+\left(\mathbf{z}_{2}-\boldsymbol{\mu}_{2}\right)^{T} \mathbf{B}^{T} \mathbf{A}^{-1} \mathbf{B}\left(\mathbf{z}_{2}-\boldsymbol{\mu}_{2}\right)\right)\right],
\end{array}

where the constant $\kappa_{3}$ has changed by a factor that balances the extra term in the exponential. Writing this out as a normal distribution, we have:

\begin{equation*}
\operatorname{Pr}\left(\mathbf{z}_{1} \mid \mathbf{z}_{2}\right)=\kappa_{3} \cdot \operatorname{Norm}_{\mathbf{z}_{1}}\left[\boldsymbol{\mu}_{1}+\mathbf{A}^{-1} \mathbf{B}\left(\mathbf{z}_{2}-\boldsymbol{\mu}_{2}\right), \mathbf{A}^{-1}\right].
\end{equation*}

The constant of proportionality must be $1 /\left(\kappa_{3}|\mathbf{D}|^{1 / 2} 2 \pi^{D / 2}\right)$ so that it is a proper normal distribution. Hence the conditional distribution is given by:

\begin{align*}
\operatorname{Pr}\left(\mathbf{z}_{1} \mid \mathbf{z}_{2}\right) & =\operatorname{Norm}_{\mathbf{z}_{1}}\left[\boldsymbol{\mu}_{1}-\mathbf{A}^{-1} \mathbf{B}\left(\mathbf{z}_{2}-\boldsymbol{\mu}_{2}\right), \mathbf{A}^{-1}\right]\\
& =\operatorname{Norm}_{\mathbf{z}_{2}}\left[\boldsymbol{\mu}_{1}+\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1}\left(\mathbf{z}_{2}-\boldsymbol{\mu}_{2}\right), \boldsymbol{\Sigma}_{11}-\boldsymbol{\Sigma}_{12} \boldsymbol{\Sigma}_{22}^{-1} \boldsymbol{\Sigma}_{12}^{T}\right],
\end{align*}

as required.