Gaussian Processes Sampling with Sparse Grids
under Additive Schwarz Preconditioner

A GP is a collection of random variables uniquely specified by its mean function $\mu(\cdot)$ and kernel function $K(\cdot,\cdot)$, such that any finite collection of those random variables have a joint Gaussian distribution (Williams and Rasmussen, , 2006). We denote the corresponding GP as $\mathcal{GP}(\mu(\cdot),K(\cdot,\cdot))$.

Let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be an unknown function. Suppose we are given a set of training points $\mathbf{X}=\{\mathbf{x}_{i}\}_{i=1}^{n}$ and the observations $\mathbf{y}=(y_{1},\ldots,y_{n})^{\top}$ where $y_{i}=f(\mathbf{x}_{i})+\epsilon_{i}$ with the i.i.d. noise $\epsilon_{i}\sim\mathcal{N}(0,\sigma_{\epsilon}^{2})$. The GP regression imposes a GP prior over the latent function as $f(\cdot)\sim\mathcal{GP}(\mu(\cdot),K(\cdot,\cdot))$. Then the posterior distribution at $m$ test points $\mathbf{X}^{*}=\{\mathbf{x}_{i}^{*}\}_{i=1}^{m}$ is

with

where $\sigma_{\epsilon}^{2}>0$ is the noise variance, $\mathbf{I}_{n}$ is a $n\times n$ identity matrix, $\mathbf{K}_{*,\mathbf{X}}=K(\mathbf{X}^{*},\mathbf{X})=\big{[}K(\mathbf{x}_{i}^{*},\mathbf{x}_{j})\big{]}_{i=1,j=1}^{i=m,j=n}=K(\mathbf{X},\mathbf{X}^{*})^{\top}=\mathbf{K}_{\mathbf{X},*}^{\top}$, $\mathbf{K}_{\mathbf{X},\mathbf{X}}=\big{[}K(\mathbf{x}_{i},\mathbf{x}_{j})\big{]}_{i,j=1}^{n}$, $\mathbf{K}_{*,*}=\big{[}K(\mathbf{x}_{i}^{*},\mathbf{x}_{j}^{*})\big{]}_{i,j=1}^{m}$, $\bm{\mu}_{\mathbf{X}}=\big{(}\mu(\mathbf{x}_{1}),\cdots,\mu(\mathbf{x}_{n})\big{)}^{\top}$ and $\bm{\mu}_{*}=\big{(}\mu(\mathbf{x}_{1}^{*}),\cdots,\mu(\mathbf{x}_{m}^{*})\big{)}^{\top}$.

The goal is to sample a function $f(\cdot)\sim\mathcal{GP}(\mu(\cdot),K(\cdot,\cdot))$. To represent this in a finite form, we discretize the input space and focus on the function values at a set of grid points $\mathbf{Z}=\{\mathbf{z}_{i}\}_{i=1}^{n_{s}}$. Our goal is to generate samples $\mathbf{f}_{\mathbf{Z}}=\big{(}f(\mathbf{z}_{1}),\cdots,f(\mathbf{z}_{n_{s}})\big{)}^{\top}$ from the multivariate normal distribution $\mathcal{N}(\mu(\mathbf{Z}),K(\mathbf{Z},\mathbf{Z}))=\mathcal{N}(\bm{\mu}_{\mathbf{Z}},\mathbf{K}_{\mathbf{Z},\mathbf{Z}})$. The standard method for sampling from a multivariate normal distribution involves the following steps: (1) generate a vector of samples $\mathbf{f}_{\mathbf{I}}$ whose entries are independent and identically distributed normal, i.e. $\mathbf{f}_{\mathbf{I}_{n_{s}}}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{n_{s}})$; (2) Use the Cholesky decomposition (Golub and Van Loan, , 2013) to factorize the covariance matrix $\mathbf{K}_{\mathbf{Z},\mathbf{Z}}$ as $\mathbf{L}_{\mathbf{Z}}\mathbf{L}_{\mathbf{Z}}^{\top}=\mathbf{K}_{\mathbf{Z},\mathbf{Z}}$; (3) generate the output sample $\mathbf{f}_{\mathbf{Z}}$ as

Sampling a posterior GP follows a similar procedure. Suppose we have observations $\mathbf{y}=\big{(}y_{1},\cdots,y_{n}\big{)}^{\top}$ on $n$ distinct points $\mathbf{X}=\{\mathbf{x}_{i}\}_{i=1}^{n}$, where $y_{i}=f(\mathbf{x}_{i})+\epsilon_{i}$ with noise $\epsilon_{i}\sim\mathcal{N}(0,\sigma_{\epsilon}^{2})$ and $f(\cdot)\sim\mathcal{GP}(\mu(\cdot),K(\cdot,\cdot))$. The goal is to generate posterior samples $f(\mathbf{X}^{*})\mid\mathbf{y}$ at $m$ test points $\mathbf{X}^{*}=\{\mathbf{x}_{i}^{*}\}_{i=1}^{m}$. Since the posterior samples $f(\mathbf{X}^{*})\mid\mathbf{y}\sim\mathcal{N}(\bm{\mu}_{*|\mathbf{y}},\mathbf{K}_{*,*|\mathbf{y}})$ also follow a multivariate normal distribution, as detailed in eq. 2a and eq. 2b in Section 2.1, we can apply Cholesky decomposition $\mathbf{L}_{*|\mathbf{X}}\mathbf{L}_{*|\mathbf{X}}^{\top}=\mathbf{K}_{*,*|\mathbf{X}}$ to generate GP posterior sample $\mathbf{f}_{*}=f(\mathbf{X}^{*})=\big{(}f(\mathbf{x}_{1}^{*}),\cdots,f(\mathbf{x}_{m}^{*})\big{)}^{\top}$ as

where $\mathbf{f}_{\mathbf{I}_{n}}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{n})$.

Our goal is to sample from GP priors $f(\cdot)\sim\mathcal{GP}(\mu(\cdot),K(\cdot,\cdot))$ over $d$-dimensional grid points $\mathbf{Z}=\{\mathbf{z}_{i}\}_{i=1}^{n_{s}}$, $\mathbf{z}_{i}\in\mathbb{R}^{d}$. According to the conditional distributions of the SoR model in equations eq. 7, the distribution of $\mathbf{f}_{\mathbf{Z}}=f(\mathbf{Z})=(f(\mathbf{z}_{1}),\cdots,f(\mathbf{z}_{n_{s}}))^{\top}$ can be approximated by $\mathbf{K}_{\mathbf{Z},\mathcal{U}}\mathbf{K}_{\mathcal{U},\mathcal{U}}^{-1}\mathbf{f}_{\mathcal{U}}$ with $\mathbf{f}_{\mathcal{U}}\sim\mathcal{N}(\mathbf{0},\mathbf{K}_{\mathcal{U},\mathcal{U}})$, where $\mathbf{f}_{\mathcal{U}}=f(\mathcal{U})$, and $\mathcal{U}=\mathcal{U}(\eta,d)$ is a sparse grid design with level $\eta$ and dimension $d$ defined in eq. 10. Then it reduces to generating samplings from $\mathbf{K}_{\mathcal{U},\mathcal{U}}^{-1}\mathbf{f}_{\mathcal{U}}\sim\mathcal{N}(\mathbf{0},\mathbf{K}_{\mathcal{U},\mathcal{U}}^{-1})$. Apply Smolyak’s algorithm to samplings from GPs (see Theorem A.4 in Appendix A), we can infer that

where $\text{chol}\Big{(}K(\mathcal{U}_{j,t_{j}},\mathcal{U}_{j,t_{j}})^{-1}\Big{)}$ is the Cholesky decomposition of $K(\mathcal{U}_{j,t_{j}},\mathcal{U}_{j,t_{j}})^{-1}$, $\mathcal{U}=\mathcal{U}(\eta,d)$ is a sparse grid design with level $\eta$ and dimension $d$ defined in eq. 10, $\mathbf{f}_{\mathbf{I}_{n_{\vv{t}}}}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{n_{\vv{t}}})$ is standard multivariate normal distributed, $\mathbf{I}_{n_{\vv{t}}}$ is a $n_{\vv{t}}\times n_{\vv{t}}$ identity matrix, $n_{\vv{t}}$ $=$ $\prod_{j=1}^{d}t_{j}$ is the size of the full grid design $\mathcal{U}_{\vv{t}}=\mathcal{U}_{1,t_{1}}\times\mathcal{U}_{2,t_{2}}\times\cdots\times\mathcal{U}_{d,t_{d}}$ associated with $\vv{t}=(t_{1},\ldots,t_{d})$, $\delta_{x}(\mathcal{U}_{\vv{t}})$ is a Dirac measure on the set $\mathcal{U}$. Therefore, the GP prior $\mathbf{f}_{\mathbf{Z}}$ over grid points $\mathbf{Z}$ can be approximately sampled via

where $\mathbf{f}_{\mathbf{I}_{n_{\vv{t}}}}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{n_{\vv{t}}})$ is standard multivariate normal distributed, $\delta_{x}(\mathcal{U}_{\vv{t}})$ is a Dirac measure on set $\mathcal{U}$ as defined in eq. 14, $\text{chol}\Big{(}K(\mathcal{U}_{j,t_{j}},\mathcal{U}_{j,t_{j}})^{-1}\Big{)}$ is the Cholesky decomposition of $K(\mathcal{U}_{j,t_{j}},\mathcal{U}_{j,t_{j}})^{-1}$.

Since the Kronecker matrix-vector product can be written as the form of vectorization operator (see Proposition 3.7 in (Kolda, , 2006)), the Kronecker matrix-vector product in section 3.1 only costs $\mathcal{O}(\sum_{j=1}^{d}t_{j}\prod_{j=1}^{d}t_{j})$ $=\mathcal{O}(\eta n_{\vv{t}})$ time, hence the time complexity of the prior sampling via section 3.1 is $\mathcal{O}\Big{(}(\eta+n_{s})\sum_{\vv{t}\in\mathbb{P}(\eta,d)}n_{\vv{t}}\Big{)}$ $=\mathcal{O}\Big{(}(\eta+n_{s})n_{sg}\Big{)}$, where $\nu$ is the smoothness parameter of the Matérn kernel, $d$ is the dimension of the sampling points $\mathbf{Z}=\{\mathbf{z}_{i}\}_{i=1}^{n_{s}}$, $n_{s}$ is the size of the sampling points, $n_{sg}$ is the size of inducing points on a sparse grid $\mathcal{U}(\eta,d)$, $n_{\vv{t}}$ $=$ $\prod_{j=1}^{d}t_{j}$ is the size of the full grid design $\mathcal{U}_{\vv{t}}=\times_{j=1}^{d}\mathcal{U}_{j,t_{j}}$, $\vv{t}=(t_{1},\ldots,t_{d})$, $\mathbb{P}(\eta,d)=\{\vv{t}\in\mathbb{N}^{d}|\max(d,\eta-d+1)\leq\sum_{j=1}^{d}t_{j}\leq\eta\}$.

Note that we can obtain the upper bound of the number of points in the Smolyak type sparse grids. According to Lemma 3 in (Rieger and Wendland, , 2017), given that $|\mathcal{U}_{j,t_{j}}|\leq 2^{t_{j}}-1$, when sparse grids are constructed on hyperbolic cross points as defined in eq. 11, the cardinality of a sparse grid $\mathcal{U}(\eta,d)$ defined in eq. 10, is thereby bounded by

Sampling from the posterior GPs can be effectively conducted using Matheron’s rule. The Matheron’s rule was popularized in the geostatistics field by Journel and Huijbregts, (1976), and recently rediscovered by Wilson et al., (2020), who leveraged it to develop a novel approach for GP posterior samplings. Matheron’s rule can be described as follows: Consider $\mathbf{a}$ and $\mathbf{b}$ as jointly Gaussian random vectors. Then the random vector $\mathbf{a}$, conditional on $\mathbf{b}=\bm{\beta}$, is equal in distribution to

where $\mathrm{Cov}(\mathbf{a},\mathbf{b})$ denotes the covariance of $(\mathbf{a},\mathbf{b})$.

According to eq. 17, exact GP posterior samplers can be obtained using two jointly Gaussian random variables. Following this technique, we derive

where $\mathbf{f}_{\mathbf{X}}=f(\mathbf{X})$ and $\mathbf{f}_{*}=f(\mathbf{X}^{*})$ are jointly Gaussian random variables from the prior distribution, noise variates $\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\sigma_{\epsilon}^{2}\mathbf{I}_{n})$. Clearly, the joint distribution of $(\mathbf{f}_{\mathbf{X}},\mathbf{f}_{*})$ follows the multivariate normal distribution as follows:

We can apply SoR approximation eq. 7 to Matheron’s rule eq. 17 and obtain

where $\bm{\Sigma}_{\mathcal{U}}=\big{[}\mathbf{K}_{\mathcal{U},\mathcal{U}}+\sigma_{\epsilon}^{-2}\mathbf{K}_{\mathcal{U},\mathbf{X}}\mathbf{K}_{\mathbf{X},\mathcal{U}}\big{]}$, $\mathbf{f}_{\mathbf{X}}=f(\mathbf{X})$ and $\mathbf{f}_{*}=f(\mathbf{X}^{*})$ are jointly Gaussian random variables from the prior distribution, noise variates $\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\sigma_{\epsilon}^{2}\mathbf{I}_{n})$. The prior resulting from the SoR approximation can be easily derived from eq. 7, yielding the following expression:

where $\bm{\mu}_{\mathbf{X}+*}=[\bm{\mu}_{\mathbf{X}},\bm{\mu}_{*}]^{\top}$, $\mathbf{K}_{\mathbf{X}+*,\mathcal{U}}=[K(\mathcal{U},\mathbf{X}),K(\mathcal{U},\mathbf{X}^{*})]^{\top}$. The same as above in section 3.1, $(\mathbf{f}_{\mathbf{X}},\mathbf{f}_{*})$ can be sampled from

where $\mathbf{f}_{\mathbf{I}_{n_{\vv{t}}}}$, $\delta_{x}(\mathcal{U}_{\vv{t}})$, and $\text{chol}\Big{(}K(\mathcal{U}_{j,t_{j}},\mathcal{U}_{j,t_{j}})^{-1}\Big{)}$ are the same as that defined in section 3.1.

The only remaining computationally demanding step in eq. 20 is the computation of $\bm{\Sigma}_{\mathcal{U}}^{-1}\mathbf{v}$ with $\mathbf{v}=\mathbf{K}_{\mathcal{U},\mathbf{X}}(\mathbf{y}-\mathbf{f}_{\mathbf{X}}-\bm{\epsilon})$, requiring a time complexity of $\mathcal{O}(n_{sg}^{3})$ where $n_{sg}=|\mathcal{U}(\eta,d)|$ is the size of the sparse grid $\mathcal{U}(\eta,d)$. To reduce this computational intensity, we consider using the conjugate gradient (CG) method (Hestenes et al., , 1952; Golub and Van Loan, , 2013), an iterative algorithm for solving linear systems efficiently via matrix-vector multiplications. Preconditioning (Trefethen and Bau III, , 1997; Demmel, , 1997; Saad, , 2003; Van der Vorst, , 2003) is a well-known tool for accelerating the convergence of the CG method, which introduces a symmetric positive-definite matrix $\mathbf{P}_{\mathcal{U}}\approx\bm{\Sigma}_{\mathcal{U}}$ such that $\mathbf{P}_{\mathcal{U}}^{-1}\bm{\Sigma}_{\mathcal{U}}$ has a smaller condition number than $\bm{\Sigma}_{\mathcal{U}}$. The entire scheme of the preconditioned conjugate gradient (PCG) takes the form in Algorithm 1.

The choice of the precondition $\mathbf{P}_{\mathcal{U}}$ has a trade-off between the smallest time complexity of applying $\mathbf{P}_{\mathcal{U}}^{-1}$ operation and the optimal condition number of $\mathbf{P}_{\mathcal{U}}^{-1}\bm{\Sigma}_{\mathcal{U}}$. Cutajar et al., (2016) first investigated several groups of preconditioners for radial basis function kernel, then Gardner et al., (2018) proposed a specialized precondition for general kernels based on pivoted Cholesky decomposition. Wenger et al., 2022a analyzed and summarized convergence rates for arbitrary multi-dimensional kernels and multiple preconditioners. Due to the hierarchical structure of the sparse grid $\mathcal{U}$, instead of using common classes of preconditioners, we consider two-level additive Schwarz (TAS) preconidtioner (Toselli and Widlund, , 2004) for the system matrix $\bm{\Sigma}_{\mathcal{U}}$, which is formed as an additive Schwarz (AS) preconditioner (Dryja and Widlund, , 1994; Cai and Sarkis, , 1999) coupled with an additive coarse space correction. Since the sparse grid $\mathcal{U}(\eta,d)$ defined in eq. 10 is covered by a set of overlapping domains $\{\mathcal{U}_{\vv{t}}\}_{\vv{t}\in\mathbb{G}(\eta)}$ where $\mathbb{G}(\eta)=\{\vv{t}\in\mathbb{N}^{d}|\sum_{j=1}^{d}t_{j}=\eta\}$, the TAS precondition for the system matrix $\bm{\Sigma}_{\mathcal{U}}$ is then defined as

On the right hand side of eq. 23, the first term is the one-level AS preconditioner, and the second term is the additive coarse space correction. $\mathbf{S}_{\vv{t}}\in\mathbb{R}^{n_{\vv{t}}\times n_{sg}}$ is the selection matrix that projects $\mathcal{U}_{\vv{t}}$ to $\mathcal{U}$, $\mathbf{P}_{\vv{t}}=\mathbf{S}_{\vv{t}}\bm{\Sigma}_{\mathcal{U}}\mathbf{S}_{\vv{t}}^{\top}\in\mathbb{R}^{n_{\vv{t}}\times n_{\vv{t}}}$ is a symmetric positive definite sub-matrix defined on $\mathcal{U}_{\vv{t}}$. $\mathbf{S}_{c}\in\mathbb{R}^{n_{c}\times n_{sg}}$ is the selection matrix whose columns spanning the coarse space, and $\mathbf{P}_{c}=\mathbf{S}_{c}\bm{\Sigma}_{\mathcal{U}}\mathbf{S}_{c}^{\top}\in\mathbb{R}^{n_{c}\times n_{c}}$ is the coarse space matrix. $n_{c}$ is the dimension of the coarse space, $n_{sg}:=|\mathcal{U}(\eta,d)|$ is the size of the sparse grid $\mathcal{U}(\eta,d)$, $n_{\vv{t}}:=|\mathcal{U}_{\vv{t}}|=\prod_{j=1}^{d}t_{j}$ is the size of the full grid design $\mathcal{U}_{\vv{t}}=\times_{j=1}^{d}\mathcal{U}_{j,t_{j}}$, $\vv{t}=(t_{1},\ldots,t_{d})$. Figure 2 illustrates an example of the selection matrices $\{\mathbf{S}_{\vv{t}}\}_{\vv{t}\in\mathbb{G}(\eta)}$ for the sparse grid $\mathcal{U}(\eta=3,d=2)$.

Some recent work (Toselli and Widlund, , 2004; Spillane et al., , 2014; Dolean et al., , 2015; Al Daas et al., , 2023) provide robust convergence estimates of TAS preconditioner, which states that the condition number of $\mathbf{P}_{\text{TAS}}^{-1}\bm{\Sigma}_{\mathcal{U}}$ is bounded by the number of distinct colors required to color the graph of $\bm{\Sigma}_{\mathcal{U}}$ so that $span\{\mathbf{S}_{\vv{t}}\}_{\vv{t}\in\mathbb{G}(\eta)}$ of the same color are mutually $\bm{\Sigma}_{\mathcal{U}}$-orthogonal, the maximum number of subdomains that share an unknown, and a user defined tolerance $\tau$. Specifically, given a threshold $\tau>0$, a GenEO coarse space (see Section 3.2 in (Al Daas et al., , 2023)) can be constructed, and the following inequality holds for the preconditioned matrix $\mathbf{P}_{\text{TAS}}^{-1}\bm{\Sigma}_{\mathcal{U}}$: