Spatial-Temporal GP (2)

In this blog post, I will walk through how one could exploit the Kronecker structure of the temporal Gaussian process (GP) regression with one-dimensional space + one-dimensional time inputs and one-dimensional output. This is the second of a series of blog posts on spatial-temporal Gaussian processes.

Separable Kernels

Recall from the last post that we have fitted a temporal GP on an one-dimensional spatial and one-dimensional temporal grid. Since we define the overall kernel as a product of the spatial and temporal component of the kernel, i.e. $k=k_s\times k_t$, we have the Kronecker structure of the Gram matrices $K=K_s\otimes K_t$, visually shown below.

Gram Matrices

Such kernels are known as separable kernels, and in this post we will explore how one could exploit this structure to obtain significant computational speed ups.

Kronecker Facts

Before describing how one could leverage the Kronecker structure, first we state several relevant and helpful facts about matrices with a Kronecker structure.

Consider two matrices $A\in {\mathbb{R}}^{n_1\times n_2}$, $B\in {\mathbb{R}}^{m_1\times m_2}$. The Kronecker product $K=A\otimes B\in {\mathbb{R}}^{n_1{m}_1\times n_2{m}_2}$ is defined by

$$A\otimes B=\left[\begin{array}{ccc}{a}_{11}B& \cdots & a_{1n_2}B\\ ⋮& \ddots & ⋮\\ a_{n_{1}1}B& \cdots & a_{n_1{n}_2}B\end{array}\right].$$

The Kronecker product operator $\otimes$ is bi-linear and associative, so we have

$$\begin{array}{rl}A\otimes (B+C)& =A\otimes B+A\otimes C\\ (B+C)\otimes A& =B\otimes A+C\otimes A\\ (kA)\otimes B& =A\otimes (kB)=k(A\otimes B)\\ A\otimes (B\otimes C)& =(A\otimes B)\otimes C.\end{array}$$ More interesting (and relevant here) properties are the ones related to inverse, Cholesky decomposition, and determinant.

First, we have the inverse property

$$(A\otimes B{)}^{-1}=A^{-1}\otimes B^{-1}$$ for any invertible $A,B$.

Next, we have the mixed-product property

$$(A_1\otimes B_1)(A_2\otimes B_2)=(A_1{A}_2)\otimes (B_1{B}_2).$$ Note that if we have Cholesky decomposition $A=L{L}^{\ast }$ for lower triangular matrix $L$ and its conjugate transpose $L^{\ast }$, we have

$$A\otimes B=(L_A{L}_A^{\ast })\otimes (L_B{L}_B^{\ast })=(L_A\otimes L_B)(L_A^{\ast }\otimes L_B^{\ast }).$$ Similarly, if we have eigendecomposition $A=Q_A{\mathrm{\Lambda }}_A{Q}_A^T$ for diagonal matrix ${\mathrm{\Lambda }}_A$, we have

$$A\otimes B=(Q_A{\mathrm{\Lambda }}_A{Q}_A^T)\otimes (Q_B{\mathrm{\Lambda }}_B{Q}_B^T)=(Q_A\otimes Q_B)({\mathrm{\Lambda }}_A\otimes {\mathrm{\Lambda }}_B)(Q_A^T\otimes Q_B^T).$$

Finally, if we have square matrices $A\in {\mathbb{R}}^{n\times n}$ and $B\in {\mathbb{R}}^{m\times m}$, then

$$|A\otimes B|=|A{|}^m|B{|}^n$$ for matrix determinants.

Kronecker Matrix-Vector Product

We will first show a useful algorithm to compute matrix-vector product of the form $(A\otimes B)z$ where $A,B$ are square matrices of size $N_A\times N_A$ and $N_B\times N_B$ respectively and $z\in {\mathbb{R}}^N$ with $N=N_A{N}_B$. Note that this algorithm can be generalised to matrix-vector product with matrix being a Kronecker product of $D$ square matrices.

We will use the $\mathrm{vec}$ operator where it stacks the columns of a matrix vertically to obtain a single column vector, i.e. for $A=[a_1,a_2,\dots ,a_k]$ with $a_i$ being column vectors, we have

$$\mathrm{vec}(A)=\left[\begin{array}{c}{a}_1\\ ⋮\\ a_k\end{array}\right].$$

A property about the $\mathrm{vec}$ operator and Kronecker product is the following:

$$(A\otimes B)\mathrm{vec}(Z)=\mathrm{vec}[BZ{A}^T]=\mathrm{vec}[B(A{Z}^T{)}^T]$$ Back to the product of interest $(A\otimes B)z$, we have $z$ as a column vector. To apply the $\mathrm{vec}$ formula above, we need to reshape $z$ to enable sensible matrix products. So, we have, using JAX (and JAX Numpy) notations,

step1 = z.reshape(N_A, N_B)
step2 = A @ step1
step3 = step2.T
step4 = B @ step3
result = step4.T.flatten()

where the .reshape in JAX Numpy is practically transpose then reshape - which is also why we transpose before flatten to get the final result. In terms of computational time, the naive implementation of $(A\otimes B)z$ will be $O((N_A{N}_B{)}^2)$ whereas the Kronecker implementation is only $O(N_A{N}_B)$. The Kronecker implementation will be used whenever it is applicable.

Kronecker Matrix-Matrix Product

One could also easily extend the above Kronecker matrix-vector product to Kronecker matrix-matrix product in the following way. Consider the matrix-matrix product $(A\otimes B)Z$ where $A,B$ are square matrices of size $N_A\times N_A$ and $N_B\times N_B$ respectively and $Z\in {\mathbb{R}}^{N\times M}$. We will break matrix $Z$ down in to $M$ columns and perform Kronecker matrix-vector product to each of the columns. This gives a computational time of $O(N_A{N}_{B}M)$ as opposed to the $O((N_A{N}_B{)}^{2}M)$ of naive implementation. We could also exploit the vectorisation functionalities to further speed up this product using methods such as jax.vmap.

Standard GP Sampling, Training, and Prediction

For a Gaussian process $f\sim \mathcal{GP}(\mu ,k)$ where $\mu$ is the mean function and $k$ is the kernel function, we can draw a sample of $f$ at test locations $X_{\ast }=(x_1^{\ast },x_2^{\ast },\dots ,x_k^{\ast })$

$$f_{\ast }=\mu (X_{\ast })+\sqrt{k(X_{\ast },X_{\ast })}\xi ,\xi \sim N_k(0,I_k)$$ where $k(X_{\ast },X_{\ast })$ is the Gram matrix and the square root denotes the lower Cholesky factor.

Consider we have made $m$ observations of this GP $f$ where the observations are made at locations $X\in {\mathbb{R}}^m$ with values $y\in {\mathbb{R}}^m$ and the observations are noisy with independent additive Gaussian noise of variance ${\sigma }^2$, i.e. $y=f(X)+\xi$ with ${\xi }_i\sim N(0,{\sigma }^2)\mathrm{\forall }i=1,2,\dots ,m$. Denote the existing observations as $\mathcal{D}=\{X,y\}$.

To train the model using data (or conduct MLE¹), we need to optimise the log likelihood

$$\log p(y|X)=-\frac{m}{2}\log (2\pi )-\log |k(X,X)+{\sigma }^2{I}_m|-\frac{1}{2}{y}^T(k(X,X)+{\sigma }^2{I}_m{)}^{-1}y.$$ In addition, we have the (conditional) predictive² distribution

$$\begin{array}{rl}{y}_{\ast }|X_{\ast },\mathcal{D},{\sigma }^2& \sim N_n({\mu }_{y_{\ast }|\mathcal{D}},K_{y_{\ast }|\mathcal{D}}),\\ {\mu }_{y_{\ast }|\mathcal{D}}& =K_{\ast }^T(K+{\sigma }^2{I}_n{)}^{-1}y,\\ K_{y_{\ast }|\mathcal{D}}& =K_{\ast \ast }-K_{\ast }^T(K+{\sigma }^2{I}_n{)}^{-1}{K}_{\ast }.\end{array}$$

which also implies that if we wish to draw a posterior sample we would have

$$f_{\ast }={\mu }_{y_{\ast }|\mathcal{D}}(X_{\ast })+\sqrt{k_{y_{\ast }|\mathcal{D}}(X_{\ast },X_{\ast })}\xi ,\xi \sim N_k(0,I_k).$$

“GP Does Not Scale”

It is a consensus / folklore that GP does not scale. This is mostly due to the fact that the training and sampling of GP involves inversions and Cholesky decomposition of $m\times m$ matrices where $m$ is the number of observations. Most commonly used algorithms for matrix inversion and Cholesky decomposition are of $O(m^3)$ time complexity and are serial in natural (so do not enjoy the GPU speed-ups that are prevalent in machine learning) – even a moderately sized data set will induce prohibitive costs.

It is still ongoing research to device tricks and algorithms to make GP more scalable. Some notable approaches includes:

• Inducing Points
• Vecchia Approximations
• SPDE Approach
• Efficiently Posterior Sampling
• Conjugate Gradients.

Here, we will look at one such approach: Kronecker structure exploiting method. Assume we have a 1D space + 1D time temporal GP with $N_s$ spatial grid points and $N_t$ temporal grid points. The naive implementation will have a time complexity of $O((N_s{N}_t{)}^3)$, whereas a Kronecker-aware implementation will only have a time complexity of $O(max\{N_s,N_t{\}}^3)$. Below, we will clarify the precise ways we can leverage the matrix structure to achieve computational speed-ups.

Kronecker Product Gaussian Process

The contents here are largely based on Saatchi (2011).

Sampling from a GP

Naive

$$f_{\ast }=\mu (X_{\ast })+\sqrt{K_s\otimes K_t}\xi ,\xi \sim N_k(0,I_k)$$

Kronecker

$$f_{\ast }=\mu (X_{\ast })+(\sqrt{K_s}\otimes \sqrt{K_t})\xi ,\xi \sim N_k(0,I_k)$$

where we can use the Kronecker matrix-vector product.

GP Likelihood

Naive

$$\log p(y|X)=-\frac{m}{2}\log (2\pi )-\log |K_s\otimes K_t+{\sigma }^2{I}_m|-\frac{1}{2}{y}^T(K_s\otimes K_t+{\sigma }^2{I}_m{)}^{-1}y.$$

Kronecker

There are two places where we need to leverage the Kronecker structure:

• determinant $|K_s\otimes K_t+{\sigma }^2{I}_m|$
• inverse $(K_s\otimes K_t+{\sigma }^2{I}_m{)}^{-1}$.

Consider eigendecompositions $K=Q\mathrm{\Lambda }{Q}^T$, $K_s=Q_s{\mathrm{\Lambda }}_s{Q}_s^T$, and $K_t=Q_t{\mathrm{\Lambda }}_t{Q}_t^T$. We know that $Q{Q}^T=I$ and $|Q|=1$, so since

$$\begin{array}{rl}K+{\sigma }^2{I}_m& =Q\mathrm{\Lambda }{Q}^T+Q({\sigma }^2{I}_m)Q^T=Q(\mathrm{\Lambda }+{\sigma }^2{I}_m)Q^T\\ K_s\otimes K_t+{\sigma }^2{I}_m& =(Q_s\otimes Q_t)({\mathrm{\Lambda }}_s\otimes {\mathrm{\Lambda }}_t+{\sigma }^2{I}_m)(Q_s^T\otimes Q_t^T)\end{array}$$

we have

$$\begin{array}{rl}|K+{\sigma }^2{I}_m|& =|Q|\cdot |\mathrm{\Lambda }+{\sigma }^2{I}_m|\cdot |Q^T|=|{\mathrm{\Lambda }}_s\otimes {\mathrm{\Lambda }}_t+{\sigma }^2{I}_m|\\ (K+{\sigma }^2{I}_m{)}^{-1}& =Q^{-T}(\mathrm{\Lambda }+{\sigma }^2{I}_m{)}^{-1}{Q}^{-1}=Q({\mathrm{\Lambda }}_s\otimes {\mathrm{\Lambda }}_t+{\sigma }^2{I}_m{)}^{-1}{Q}^T\end{array}$$

where the remaining term ${\mathrm{\Lambda }}_s\otimes {\mathrm{\Lambda }}_t+{\sigma }^2{I}_m$ is a diagonal matrix, and we can leverage Kronecker matrix-vector (and matrix-matrix) product whenever necessary.

Footnotes

1. Frequentist MLE is equivalent to a Bayesian MAP with flat priors.

2. since we are using only a summary statistic (MAP with flat prior) for the parameters in the posterior predictive, instead of the full marginal posterior.