Gaussian Processes and Meta-Learning Priors

I’ve been looking into ways of incorporating epistemic uncertainty into learned world models in a way that fits nicely with a control problem. David suggested I look into combining neural networks with Gaussian processes (GPs) like Harrison et al. in Variational Bayesian Last Layers or Meta-Learning Priors for Efficient Online Bayesian Regression (ALPaCA). This work looks super interesting, but once I started digging into it I realized I didn’t really understand how GPs really worked, so in this post I’ll derive GPs, explain what ALPaCA is, and show an implementation of ALPaCA in JAX.

This is meant mostly for me to understand the material better, but maybe you’ll find it useful too.

---

Gaussian Processes (GPs) §

Definition §

Consider a supervised learning problem, where we have a set of $N$ inputs $\mathbf{x}\in \mathbf{X}$ and outputs $\mathbf{y}\in \mathbf{Y}$ and we wish to predict ${\mathbf{y}}^{\ast }\in {\mathbf{Y}}^{\ast }$ at test inputs ${\mathbf{x}}^{\ast }\in {\mathbf{X}}^{\ast }$. In many prediction tasks, we find parameters $\theta$ for a model $f(\mathbf{x}|\theta )$ such that $f(\mathbf{x}|\theta )\approx \mathbf{y},\forall \mathbf{y}\in \mathbf{Y}$. That is, we find parameters that when plugged into the model produce outputs $\mathbf{Y}$ given inputs $\mathbf{X}$. Then, if we assume our model can generalize, we can use the same model to predict what the outputs would look like for unseen inputs ${\mathbf{x}}^{\ast }$ (test inputs). This is what’s done in maximum-likelihood estimation.

Gaussian processes are different because they produce estimates over functions instead of parameters. That is, instead of estimating the most likely $\theta$ given the data by finding $p(\theta |\mathbf{X},\mathbf{Y})$, we estimate the most likely function¹ given the data by finding $p(f|\mathbf{X},\mathbf{Y})$.

To do so, we start by defining a prior distribution over functions that’s conditioned on the training data $\mathbf{X}$. We denote this distribution as $p(\mathbf{f}|\mathbf{X})$ where

$$\mathbf{f}=\left[\begin{array}{c}f({\mathbf{x}}_1)\\ ⋮\\ f({\mathbf{x}}_N)\end{array}\right].$$

Then, we assume this prior is a Gaussian with mean $\mu$ and covariance $\mathbf{K}$ where

$$\mu =\left[\begin{array}{c}m({\mathbf{x}}_1)\\ ⋮\\ m({\mathbf{x}}_N)\end{array}\right]$$

and ${\mathbf{K}}_{ij}=\kappa ({\mathbf{x}}_i,{\mathbf{x}}_j)$. $\kappa$ is known as the kernel function and it measures the similarity between two points ${\mathbf{x}}_i$ and ${\mathbf{x}}_j$, and $m$ is a mean function (usually set to zero). Therefore, $\mathbf{f}\sim \mathcal{N}(\mathbf{f}|\mu ,\mathbf{K})$.

Outputs are jointly Gaussian

Assuming the prior over functions is Gaussian implies the outputs $\mathbf{f}$ are distributed according to a joint Gaussian distribution with mean $\mu$ and covariance $\Sigma$. That is, $p(f(\mathbf{x}),\dots ,f({\mathbf{x}}_N))=\mathcal{N}(\mathbf{f}|\mu ,\mathbf{\Sigma })$.

This is a strong assumption, but it is useful because it allows us to easily compute the posterior distribution over functions given training data (as we will see later) and therefore make predictions.

We can sample the prior to get an idea what the functions look like before seeing the data. To do this, we define a set of $N$ input points (say a grid between -5 and 5), compute $\mathbf{K}$ and $\mu$, and then generate samples from a Gaussian distribution with mean $\mu$ and covariance $\mathbf{K}$. Each sample will be a vector of $N$ outputs corresponding to a realization of $\mathbf{f}$. The shape of each function is implicitly defined by our chosen kernel function $\kappa$. Below is a plot showing samples from a prior with a squared exponential kernel where $\ell =\sigma =1$. Find the code here.

Prediction §

Given training data consisting of inputs $\mathbf{X}$ and noiseless outputs $\mathbf{f}$, where ${\mathbf{f}}_i=f({\mathbf{x}}_i),\forall i\in \{1,..,N\}$ we would like to predict outputs ${\mathbf{y}}^{\ast }$ at test inputs ${\mathbf{x}}^{\ast }\in {\mathbf{X}}^{\ast }$. To do so using a GP, we must find a distribution over functions conditioned on the test inputs ${\mathbf{X}}^{\ast }$, training inputs $\mathbf{X}$, and training outputs $\mathbf{f}$. That is, we seek a distribution $p({\mathbf{f}}^{\ast }|{\mathbf{X}}^{\ast },\mathbf{X},\mathbf{f})$ where ${\mathbf{f}}^{\ast }$ are the outputs at the test inputs.

Recall we assumed that function outputs are jointly Gaussian. Therefore, we can write the joint distribution of the training and test outputs as

$$\begin{array}{rl}\left[\begin{array}{c}\mathbf{f}\\ {\mathbf{f}}^{\ast }\end{array}\right]& \sim \mathcal{N}\left(\left[\begin{array}{c}\mu \\ {\mu }^{\ast }\end{array}\right],\left[\begin{array}{cc}\mathbf{K}& {\mathbf{K}}^{\ast }\\ {\mathbf{K}}^{\ast T}& {\mathbf{K}}^{\ast \ast }\end{array}\right]\right),\end{array}$$

where $\mathbf{K}=\kappa (\mathbf{X},\mathbf{X})$, ${\mathbf{K}}^{\ast }=\kappa (\mathbf{X},{\mathbf{X}}^{\ast })$, and ${\mathbf{K}}^{\ast \ast }=\kappa ({\mathbf{X}}^{\ast },{\mathbf{X}}^{\ast })$. Then, we can write the conditional distribution analytically using known results as follows:

$$\begin{array}{rl}p({\mathbf{f}}^{\ast }|{\mathbf{X}}^{\ast },\mathbf{X},\mathbf{f})& =\mathcal{N}({\mathbf{f}}^{\ast }|{\mu }^{\mathbf{c}},{\mathbf{\Sigma }}^{\mathbf{c}})\\ {\mu }^{\mathbf{c}}& ={\mu }^{\ast }+{\mathbf{K}}^{\ast T}{\mathbf{K}}^{-1}(\mathbf{f}-\mu )\\ {\mathbf{\Sigma }}^{\mathbf{c}}& ={\mathbf{K}}^{\ast \ast }-{\mathbf{K}}^{\ast T}{\mathbf{K}}^{-1}{\mathbf{K}}^{\ast },\end{array}$$

Below is a plot of sampled functions from a posterior distribution given a set using a squared exponential kernel with $\ell =\sigma =1$. Find the code here.

Posterior samples given 5 noiseless data pairs (black crosses). Shaded region denotes the 95% confidence interval (2 stds from mean).

In practice, observations are noisy. I will not cover that case here, but the derivation conditional distribution is very similar and can be found in Ch. 15 of Murphy’s Probabilistic ML.

Meta-Learning Priors (ALPaCA) §

Notation

I don’t have infinite time so in this section I’ll mostly follow the notation Harrison et al. used in Meta-Learning Priors for Efficient Online Bayesian Regression (ALPaCA). Read carefully, as there may be differences in notation and definitions from the previous sections.

Now that we (hopefully) understand GPs, let’s see what Harrison et al. did with them in ALPaCA. First, I’ll work through some preliminaries to understand the problem formulation and Bayesian regression. Then, I’ll talk about ALPaCA’s algorithm, where the meta-learning is happening, and share my JAX implementation. Finally, I’ll discuss how ALPaCA can be interpreted as an approximation of a GP.

Formulation §

Consider a function $f$ with unknown latent parameters $\theta$. Let’s assume we can observe samples $y$ of $f$ corrupted by additive Gaussian noise. That is, we observe a sequence of samples (x, y) where $y=f(x;\theta )+ϵ$ where $ϵ\sim \mathcal{N}(0,{\Sigma }_{ϵ})$. Therefore, the likelihood of the data is given by

$$p(y|x,\theta )=\mathcal{N}(y|f(x;\theta ),{\Sigma }_{ϵ}).$$

Given a prior over the latent parameters, i.e., $p(\theta )$ the posterior predictive density of $\tau$ data points ${\mathcal{D}}_{\tau }^{\ast }=\{(x_t,y_t){\}}_{t=1}^{\tau }$ generated from $f(x;{\theta }^{\ast })$ is given by

$$p(y|x,{\mathcal{D}}_{\tau }^{\ast })=\int p(y|x,\theta )p(\theta |{\mathcal{D}}_{\tau }^{\ast })d\theta$$

Unfortunately, this integral is intractable because we don’t have analytic expressions for $f(x;\theta )$ and $p(\theta )$, and even if we did, computing $p(\theta |{\mathcal{D}}_{\tau }^{\ast })$ over all possible $\theta$ is likely too computationally expensive.

Instead, let’s use a surrogate model $q_{\xi }(y|x,D_{\tau }^{\ast })$, parameterized by $\xi$, to approximate the true posterior predictive density $p(y|x,D^{\ast }\tau )$, and then let’s optimize this model so that it’s as close to the true posterior predictive density for all likely ${\theta }^{\ast }$. The bolded part is where the meta-learning is happening: we’re learning a model that will work well for all possible ${\theta }^{\ast }$.

We consider a scenario in which the data comes in as a stream: at each timestep, the agent is given a new input $x_t$, and after estimating the output ${\hat{y}}_t$, the true output $y_t$ is revealed. An example of this is a Markovian dynamical system, where the agent wishes to predict the distributions of the next state given the current state.

In this setting, the problem of learning the surrogate model can be formulated as

$$\underset{\xi }{\min }{D}_{KL}(p(y_{t+1}|x_{t+1},{\mathcal{D}}_t^{\ast })||q_{\xi }(y_{t+1}|xt+1,{\mathcal{D}}_t^{\ast })).$$

Note that we don’t know $t$ (dataset size), $x_{t+1}$, or $D_t^{\ast }$ ahead of time, so the best we can do is minimize the objective in expectation (I’ll elaborate on this later). This implies that whatever $\xi$ we choose will be optimal for all possible datasets. Very cool.

Unfortunately, this expected objective is intractable because need access to $p(\theta )$ and $p(D_t^{\ast }|{\theta }^{\ast })$ which are unknown. Instead, we assume we have access to various datasets ${\mathcal{D}}_t^j$ generated from iid samples of ${\theta }^j\sim p(\theta )$, $x_t\sim p(x)$, and $y_t\sim p(y|x,{\theta }^j)$. Each dataset can be thought of as trajectories of the system generated by different latent parameters ${\theta }^j$. The full dataset is defined as $\mathcal{D}:=\{{\mathcal{D}}_t^j{\}}_{j=1}^M$.

Bayesian Regression §

ALPaCA uses Bayesian linear regression to compute $q_{\xi }(y|x,{\mathcal{D}}_t^{\ast })$. If we consider a set of basis functions $\varphi$, the regression problem can be written as finding $K$ such that

$$y_t^{\top }={\varphi }^{\top }(x_t)K+{ϵ}_t,$$

where $K$ is a coefficient matrix and ${ϵ}_t\sim \mathcal{N}(0,{\Sigma }_{ϵ})$. Let $Y^{\top }=[y_1,\dots ,y_{\tau }],{\Phi }^{\top }=[\varphi (x_1),\dots ,\varphi (x_{\tau })]$, and E = $[{ϵ}_1,\dots ,{ϵ}_{\tau }]$, then we can re-write the regression problem as

$$Y=\Phi K+E.$$

Therefore, the likelihood of the data is given by

$$\begin{array}{c}p(Y|\Phi ,K,{\Sigma }_{ϵ})=\mathcal{N}(Y|\Phi K,{\Sigma }_{ϵ}).\end{array}$$

Now let’s select the prior for $K$ as $p(K)=\mathcal{M}\mathcal{N}(K|{\bar{K}}_0,{\Lambda }_0^{-1},{\Sigma }_{ϵ})$, where $\mathcal{M}\mathcal{N}$ denotes the matrix normal distribution, and ${\Lambda }_0$ is a precision matrix. Then, the posterior distribution of $K$, conditioned on $Y$ and $\Phi$, is given by

$$\begin{array}{rl}p(K|Y,\Phi )& =\mathcal{M}\mathcal{N}(K|{\bar{K}}_{\tau },{\Lambda }_{\tau }^{-1},{\Sigma }_{ϵ})\\ {\bar{\Lambda }}_{\tau }& ={\Phi }^{\top }\Phi +{\Lambda }_0\\ {\bar{K}}_{\tau }& ={\bar{\Lambda }}_{\tau }^{-1}({\Phi }^{\top }Y+{\Lambda }_0{\bar{K}}_0).\end{array}$$

The posterior distribution is then given by

$$p(y_{\tau +1}|\varphi (x_{\tau +1}),\Phi ,Y)=\mathcal{N}({\bar{K}}_{\tau }^{\top }\varphi (x_{\tau +1}),{\Sigma }_{\tau +1})$$

where

$${\Sigma }_{\tau +1}=(1+{\varphi }^{\top }(x_{\tau +1}){\Lambda }_{\tau }^{-1}\varphi (x_{\tau +1})){\Sigma }_{ϵ}.$$

Now we have a posterior over $y_{\tau +1}$ given the value of the basis function at $x_{\tau +1}$, i.e., $\varphi (x_{\tau +1})$, the value of the basis functions for all previous data points, i.e., $\Phi$, and the observed outputs, i.e., $Y$. The paper goes over details on how this was computed.

ALPaCA in JAX §

In ALPaCA, the basis functions are outputs of a neural network and we do a Bayesian regression on a linear transformation applied to the final output of the network. Then, we have two phases:

• Phase 1 (offline): learn the basis functions (the neural network weights $w$) and the prior parameters $K_0$ and ${\Lambda }_0$ using a sample of datasets.
• Phase 2 (online): Update the posterior parameters ${\bar{K}}_{\tau }$ and ${\Lambda }_{\tau }$ as new data comes in.

Phase 1 is where the “meta-learning” happens: we learn a prior that works well for all expected datasets. Since our prior is optimized in Phase 1, we can easily adapt to new data without having to retrain the neural network (Phase 2). Additionally, we also get live, calibrated, uncertainty estimates for every predictions. Below are the algorithms for the two phases, taken directly from the paper (note equation numbers are different from the ones I used above).

I wrote an implementation of ALPaCA in JAX and you can find it here. Below is my reproduction of Figure 2 in the paper. The plot shows predictive performance for a sinusoidal function as a function of the number of training samples. Rows correspond to different number of training samples and columns correspond to different methods. As you can see, ALPaCA (first column) is able to make some pretty good predictions even with few samples.

Blue line is the ground truth. Orange line is the prediction (wth uncertainty estimates). Black crosses are training samples.

Connection to GP’s §

Bayesian linear regression can be thought of as a GP with a kernel that’s a function of the prior $p(K)$ (see paper for more details). Therefore, Phase 1 in ALPaCA can be interpreted as learning a GP kernel that works well for all expected datasets—essentially shaping the inductive bias of the GP around the data. And then, Phase 2 updates the posterior in a more efficient way than a regular GP.

Resources §

• Murphy’s Probabilistic ML
• Rasmussen’s Gaussian Processes for Machine Learning
• ALPaCA paper: Meta-Learning Priors for Efficient Online Bayesian Regression
• My JAX implementation of ALPaCA: ALPaCA

Footnotes §

1. For the development of GPs, we will refer to functions as a finite-dimensional vector of function values at a set of input points. Technically, functions are infinitely-dimensional objects that require an infinite number of (input,output) pairs to be fully described (unless we have an explicit functional form for a function, e.g., $f(x)=x^2$). However, when working with GPs we are able to describe functions with a finite number of points because we assume that function values at different points are correlated. This correlation is defined according to a selected kernel function. However, the underlying function is still infinite-dimensional. Disclaimer: I’m still trying to wrap my head around these details. But I think this is the gist of it. ↩