Simulating the posterior of a Gaussian process

Question

For the first time (excuse imprecission / mistakes) I took a look at Gaussian processes, and more specifically, watched this video by Nando de Freitas. The notes are available online here.

At some point he draws $10$ random samples from a multivariate normal generated by constructing a covariance matrix based on a Gaussian kernel (exponential of squared distances in the $x$ axis). These random samples form the prior smooth plots that become less spread as data become available. Ultimately, the objective is to predict by modifying the covariance matrix, and obtaining the conditional Gaussian distribution at the points of interest.

The entire code is available at an excellent summary by Katherine Bailey here, which in turn credits a code repository by Nando de Freitas here. I have posted it the Python code here for convenience.

It starts with $3$ (instead of $10$ above) prior functions, and introduces a "tuning parameter".

I have translated the code into Python and [R], including the plots:

Here is the first chunk of code in [R] and the resulting plot of three random curves generated via a Gaussian kernel based on proximity in the $x$ values in the test set:

The second chunk of R code is hairier, and starts off by simulating four points of training data, which eventually will help narrow down the spread among the possible (prior) curves around the areas where these training data points lie. The simulation of the $y$ value for these data points is as a $\text{sin}()$ function. We can see the "tightening of the curves around the points":

The third chunk of R code deals with plotting the curve of mean estimated values (the equivalent of the regression curve), corresponding to $50$ ${\bf\mu}$ values (see calculation below), and their confidence intervals:

QUESTION: I want to ask for an explanation of the operations that take place when going from the prior GP to the posterior.

Specifically, I'd like to understand this part of the R code (in the second chunk) to get the means and sd:

# Apply the kernel function to our training points (5 points):

K_train = kernel(Xtrain, Xtrain, param)                          #[5 x 5] matrix

Ch_train = chol(K_train + 0.00005 * diag(length(Xtrain)))        #[5 x 5] matrix

# Compute the mean at our test points:

K_trte = kernel(Xtrain, Xtest, param)                            #[5 x 50] matrix
core = solve(Ch_train) %*% K_trte                                #[5 x 50] matrix
temp = solve(Ch_train) %*% ytrain                                #[5 x 1] matrix
mu = t(core) %*% temp                                            #[50 x 1] matrix

There are two kernels (one of train ($\bf a$) v. train ($\bf a$), K_train, let's call it $\bf \Sigma_{aa}$, with its Cholesky (Ch_train), $\bf \color{orange}{L_{aa}}$, coloring orange all Cholesky's from here on, and the second one of train ($\bf a$) v test ($\bf e$), K_trte, let's call it $\bf \Sigma_{ae}$), and to generate the estimated means $\hat \mu$ for the $50$ points in the testing set the operation is:

\begin{align} {\bf \hat \mu}&={\bf \left [ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right ]^T \, \underset{\color{blue}{[5 \times 5]}}{\color{orange}{L_{aa}}^{-1}} \, \underset{\color{blue}{[5 \times 1]}}{y_{tr}}}\bf\tag{Eq.1}\\ &\text{dimensions}=\color{red}{\left[50 \times 1\right]} \end{align}

# Compute the standard deviation:

tempor = colSums(core^2)                                          #[50 x 1] matrix

# Notice that all.equal(diag(t(core) %*% core), colSums(core^2)) TRUE

s2 = diag(K_test) - tempor                                        #[50 x 1] matrix
stdv = sqrt(s2)                                                   #[50 x 1] matrix

\begin{align} {\bf \hat{\text{var}}}&=\text{diag}\left({\bf \Sigma_{ee}}\right)-\text{diag} \left[\bf \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right]^T \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right] \right] \bf \tag{ Eq.2} \\ &=\text{d}\small{\begin{bmatrix}1&&\dots&\\&1\\&&\ddots\\&&\dots&1&\\ &&&\dots&1\end{bmatrix}}-\bf \text{d} \left[ \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right]^T \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right] \right]\\ &\text{dimensions} = \color{red}{\left[50 \times 1\right]} \end{align}

How does this work?

Also unclear, is the calculation for the color lines (Posterior GP) in the "Three samples from the GP posterior" plot above, where the Cholesky of the testing and training sets seem to come together to generate multivariate normal values, eventually added to $\hat \mu$:

Ch_post_gener = chol(K_test + 1e-6 * diag(n) - (t(core) %*% core))
m_prime = matrix(rnorm(n * 3), ncol = 3)
sam = Ch_post_gener %*% m_prime
f_post = as.vector(mu) + sam 

\begin{align} f_{\text{post}}&=\bf \hat \mu +\small \left[ \underset{\color{blue}{[50 \times 50]}} {\color{orange}{L_{ee}}}\, \, \, - \left[ \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right]^T \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right] \right] \right] \left[\underset{\color{green}{[50 \times 3]}}{\mathscr N(0,1)}\right]\tag{Eq.3} \\ &\text{dimensions}= \color{red}{\left[50 \times 3\right]} \end{align}

Answer 1

When given a test set, $e$, the expected values will be calculated considering a conditional distribution of the value of the function for these new data points, given the data points in the training set, $a$. The idea exposed in the video is that we would have a joint distribution of $a$ and $e$ (in the lecture denoted by an asterisk, $*$) of the form:

$${\bf\begin{bmatrix} a\\ \bf e\end{bmatrix}}\sim \mathscr N\left( \begin{bmatrix}\bf \mu_a\\\mu_e \end{bmatrix}\,,\begin{bmatrix}\bf \Sigma_{aa}&\bf \Sigma_{ae} \\ {\bf \Sigma_{ae}}^T & \bf \Sigma_{ee}\end{bmatrix}\right)$$.

The conditional of a multivariate Gaussian distribution has a mean $E({\bf x}_1 | {\bf x}_2)= {\boldsymbol \mu}_1 + \Sigma_{12} \Sigma^{-1}_{22} ({\bf x}_2- {\boldsymbol \mu}_2)$. Now, considering that the first row of the block matrix of covariances above is $[50 \times 50]$ for $\bf \Sigma_{aa}$, but only $[50 \times 5]$ for $\bf \Sigma_{ae}$, a transposed will be necessary to make the matrices congruous in:

$$E ({\bf e\vert a}) = {\bf \mu_e} + {\bf \Sigma_{ae}}^T {\bf \Sigma_{aa}}^{-1}\,\left ({\bf y}-{\bf \mu_{a}}\right)$$ Because the model is planned with ${\bf \mu_{a}} = {\bf \mu_{e}}=0$, the formula simplifies nicely into:

$$E ({\bf e\vert a}) = {\bf \Sigma_{ae}}^T {\bf \Sigma_{aa}}^{-1}\,{\bf y_{tr}}$$

Enter the Cholesky decomposition (which again I will code in orange as in OP):

\begin{align*} E ({\bf e\vert a}) &= {\bf \Sigma_{ae}}^T\,\, \,\underset{\color{gray}{<--\alpha-->}}{{\bf \Sigma_{aa}}^{-1}\,{\bf y_{tr}}}\\ &={\bf \Sigma_{ae}}^T {\bf \color{orange}{(L_{aa}L_{aa}^T)}}^{-1}\,{\bf y_{tr}}\\ &= {\bf \Sigma_{ae}}^T {\bf \color{orange}{L_{aa}^{-T}L_{aa}^{-1}}}\,{\bf y_{tr}}\\ &= {\bf \Sigma_{ae}}^T {\bf \color{orange}{L_{aa}^{-T}}\,\,\,\,\,\, \underset {\color{gray}{ <-m->}}{\color{orange}{L_{aa}^{-1}}{\bf y_{tr}}}} \tag {*} \end{align*}

If $\bf m =\color{orange}{{\bf L_{aa}}^{-1}}\,{\bf y_{tr}}$, then $\color{orange}{\bf L_{aa}} \bf m= {\bf y_{tr}}$, and we end up with a linear system that we can solve, obtaining $\bf m$. Here is the key slide in the original presentation:

---

---

Since $\bf B^T A^T = (A\,B)^T$, Eq. (*) is equivalent to Eq.(1) equation in the OP:

\begin{align} {\bf \hat \mu}&={\bf \left [ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right ]^T \, \underset{\color{blue}{[5 \times 5]}}{\color{orange}{L_{aa}}^{-1}} \, \underset{\color{blue}{[5 \times 1]}}{y_{tr}}}\\ &=\bf \left( \Sigma_{ae}^T \color{orange}{ L_{aa}^{-T}} \right) \left(\color{orange}{ L_{aa}^{-1}}\, y_{tr} \right)\\ &\text{dimensions} = \color{red}{\left[50 \times 1\right]} \end{align}

given that

$$\bf \left [ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right ]^T = \underset{\color{blue}{[50 \times 5]}}{\Sigma_{ae}}^T \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1T}}\, \, \, $$

Similar reasoning would be applied to the variance, starting with the formula for the conditional variance in a multivariate Gaussian:

$${\rm var}({\bf x}_1|{\bf x}_2)= \Sigma_{11} -\Sigma_{12}\Sigma^{-1}_{22}\Sigma_{21}$$

which in our case would be:

\begin{align*} \bf \text{var}_{\hat\mu_{\bf e}} &= \bf \Sigma_{ee} - \Sigma_{ae}^T\Sigma_{aa}^{-1}\Sigma_{ae}\\ &= \bf \Sigma_{ee} - \Sigma_{ae}^T \left[ L_{aa}L_{aa}^T\right]^{-1}\Sigma_{ae}\\ &= \bf \Sigma_{ee} - \Sigma_{ae}^T \left[ L_{aa}^{-1}\right]^TL_{aa}^{-1}\Sigma_{ae}\\ &= \bf \Sigma_{ee} - \left[ L_{aa}^{-1} \Sigma_{ae}\right]^T L_{aa}^{-1}\Sigma_{ae} \end{align*}

and arriving at Eq.(2):

\begin{align} \text{var}_{\hat\mu_{\bf e}}&=\text{d}\left[ \bf K_{ee} - \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right]^T \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right] \right]\\ &\text{dimensions}=\color{red}{\left[50 \times 1\right]} \end{align}

We can see that Eq.(3) in the OP is a way of generating posterior random curves conditional on the data (training set), and utilizing a Cholesky form to generate three multivariate normal random draws:

\begin{align} f_{\text{post}} &= {\bf \hat \mu} + \left[ \text{var}_{\hat\mu_{\bf e}}\right][\text{rnorm}\sim (0,1)]\\ &=\bf \hat \mu + \left[ \underset{\color{blue}{[50 \times 50]}} {\color{orange}{L_{ee}}}\, \, \, - \left[ \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right]^T \left[ \underset{\color{blue}{[5 \times 5]}} {\color{orange}{L_{aa}}^{-1}}\, \, \, \underset{\color{blue}{[5 \times 50]}}{\Sigma_{ae}} \right] \right] \right] \left[\underset{\color{green}{[50 \times 3]}}{\text{rand.norm's}}\right]\\ &\text{dimensions}= \color{red}{\left[50 \times 3\right]} \end{align}