Scalable Bayesian inference for heat kernel Gaussian processes on manifolds

Suppose the real-valued random variable $Y$ has the law of the natural exponential family (NEF), denoted as $Y\sim\text{NEF}(\theta)$. That is, under the Lebesgue measure or the counting measure, we have

where $\theta$ is called the natural parameter, $h(y)$ does not depend on $\theta$, and $\kappa$ is the sufficient statistic of $Y$. The natural exponential family includes varieties of the most common models, such as normal regression and binary classification. It emerges in many fields of statistics and machine learning (Brown et al., 2010; Jonathan R. Bradley and Wikle, 2020; Schweinberger and Stewart, 2020). The expectation of $\kappa$ is given as $\mathbb{E}_{\theta}\{\kappa(Y)\}=J^{\prime}(\theta)$. It is the quantity of interest in the article.

Let $\mathcal{X}$ be a $d$-dimensional Riemannian manifold. Let $\rho$ be the geodesic distance on $\mathcal{X}$. Assume $G$ is a measure on $\mathcal{X}$ with a smooth positive density function $g$ under the volume form $V$. Let $\Delta_{x}$ be the weighted Beltrami-Laplace operator on $(\mathcal{X},G)$. Define the Dirichlet-Laplace operator by $\mathcal{L}=-\Delta_{x}$. Then $\mathcal{L}$ is a non-negative definite self-adjoint operator in $L^{2}(\mathcal{X},G)$. The heat kernel $p_{t}(x,x^{\prime})$ is a fundamental solution of the heat equation on $(\mathcal{X},G)$, satisfying the following conditions,

where $\delta(x,x^{\prime})$ is the Dirac-delta function and $t\in\mathbb{R}^{+}$ is the diffusion time (Grigor’yan, 2009). The heat kernel describes the heat diffusion from $x$ to $x^{\prime}$ on manifolds over time $t$.

Assume that $\mathcal{X}$ is compact. The spectral theorem states that the spectrum of $\mathcal{L}$ consists of an increasing non-negative sequence $\{\lambda_{i}\}_{i=1}^{\infty}$ such that $\lambda_{1}=0$ and $\lim_{i\to\infty}\lambda_{i}=+\infty$. There exists an orthonormal basis $\{e_{i}\}_{i=1}^{\infty}$ in $L^{2}(\mathcal{X},G)$ such that for $i\geq 1$, $e_{i}$ is the eigenfunction of $\mathcal{L}$ with $\lambda_{i}$. Subsequently, from Grigor’yan (2009), for $t\in\mathbb{R}^{+}$ and $x,x^{\prime}\in\mathcal{X}$, $p_{t}(x,x^{\prime})$ can be expanded as

If the eigenpairs of $\mathcal{L}$ are estimated, then we can utilize (2) to approximate the heat kernel. So our goal is to construct a fast numerical estimation for the spectrum of $\mathcal{L}$. We refer readers to Grigor’yan (2006, 2009) for further discussions about heat kernels on manifolds.

According to the desirable property of $p_{t}$, the heat kernel Gaussian process is an excellent alternative for domain-adaptive GPs. However, the heat kernel is analytically intractable except for trivial cases even if $\mathcal{X}$ is known. In the section, we propose a fast graph Laplacian approach to estimate the heat kernel and construct the heat kernel GP in the large-scale data when $\mathcal{X}$ is unknown. Given labeled observations $\{(x_{i},y_{i})\}_{i=1}^{m}$ and unlabeled observations $\{x_{i}\}_{i=m+1}^{n}$, the objective is to approximate a heat kernel covariance matrix over $X^{n}$ in a low time cost. We present the FLGP algorithm in five steps.

Step 1, subsample $s$ induced points from the point cloud $X^{n}$, denoted as $I^{s}=\{u_{i}\}_{i=1}^{s}$. The induced point set is used to simplify $X^{n}$ and represent the domain geometry. In this article, the subsampling method is k-means clustering or random sampling. Step 2, choose a suitable kernel function $k(x,x^{\prime})$, such as the squared exponential kernel, $k(x,x^{\prime})=\exp\left(-||x-x^{\prime}||^{2}/4\epsilon^{2}\right)$ with $\epsilon>0$. Construct the cross kernel matrix $K$ between $X^{n}$ and $I^{s}$ with $K_{ij}=k(x_{i},u_{j})$, for $1\leq i\leq n,1\leq j\leq s$. Then $K$ is an $n\times s$ matrix. Step 3, obtain the cross similarity matrix $Z$ by normalizing $K$. For $1\leq i\leq n,1\leq j\leq s$, in random subsampling, $Z_{ij}$ is given by

and in k-means clustering, $Z_{ij}$ is given by

where $\{n_{i}\}_{i=1}^{s}$ indicates the number of points in each membership. Construct the cross transition matrix $A=D^{-1}Z$, where $D$ is a diagonal matrix with $D_{ii}=\sum_{j=1}^{s}Z_{ij},~{}1\leq i\leq n$. Notably, $A$ can serve as the probability transition matrix of the random walk from $X^{n}$ to $I^{s}$. Step 4, let $\Lambda$ be a diagonal matrix such that $\Lambda_{jj}=\sum_{i=1}^{n}A_{ij},~{}1\leq j\leq s$ and define the subsampling graph Laplacian $L$ as

Do truncated SVD for $A\Lambda^{-1/2}$ to obtain the singular values $\{\sigma_{i,\epsilon}\}_{i=1}^{s}$ and the left singular vectors $\{v_{i,\epsilon}\}_{i=1}^{s}$ with $||v_{i,\epsilon}||_{2}=1,~{}1\leq i\leq s$. Then the eigenpairs of $L$ are given by $\{\lambda_{i,\epsilon}=1-\sigma_{i,\epsilon}^{2},~{}v_{i,\epsilon}\}_{i=1}^{s}$. According to Proposition 1, suppose $\{\lambda_{i,\epsilon}\}_{i=1}^{s}$ are in the ascending order, i.e., $0=\lambda_{1,\epsilon}\leq\cdots\leq\lambda_{s,\epsilon}\leq 1$. Step 5, choose the top $M$ eigenvalues and corresponding eigenvectors of $L$ and estimate the heat kernel based on (2) as

where $t>0$ is the diffusion time.

The whole procedure of FLGP is shown in Figure 2. Finally, we can make inference in GPs with the estimated covariance matrix $C_{\epsilon,M,t}$. Owing to the construction, we have the following proposition for the spectrum of $L$. The proof is given in Appendix A.

The key novelty in FLGP is to construct a reduced-rank approximation for the transition matrix. Then the eigendecomposition of $L$ is replaced by the truncated SVD of $A\Lambda^{-1/2}$. We explain the idea behind the approximation from the random walk viewpoint. According to the graph spectral theory, there is a close relationship between the heat kernel on manifolds and random walks on the point cloud $X^{n}$. The stochastic transition matrix is $n\times n$. Based on the induced points $I^{s}$, we can decompose the direct random walk into two steps. For $x,x^{\prime}\in X^{n}$, the movement from $x$ to $x^{\prime}$ is decomposed into a jump from $x$ to $I^{s}$ and a jump from $I^{s}$ to $x^{\prime}$. Then the transition probability is

If we define the cross transition matrix $A$ from $X^{n}$ to $I^{s}$ as in Step 3, and let $\Lambda^{-1}A^{\top}$ be the transition matrix from $I^{s}$ to $X^{n}$, then the two-step transition probability matrix is given as $A\Lambda^{-1}A^{\top}$. This implies the definition of the graph Laplacian in (5). In contrast, Dunson et al. (2021a) doesn’t utilize the transition decomposition strategy. They construct the random walk with one-step movement. Consequently, they end up with a full-rank transition matrix and have to execute the eigendecomposition on the matrix.

The diagram inside the dashed rectangle in Figure 2 illustrates the transition decomposition graphically. As $I^{s}$ contains only $s$ induced points, the rank of $A\Lambda^{-1}A^{\top}$ is at most $s$. On the contrary, the stochastic matrix’s rank in the direct random walk among $X^{n}$ is $n$, much more than $s$. Thus, we obtain a reduced-rank approximation for the transition probability matrix. Subsequently, the complexity of the truncated SVD is $\mathcal{O}(ns^{2}+s^{3})$ instead of $\mathcal{O}(n^{3})$. This accelerates the computational speed in large-scale data. Sparsity will be introduced into the transition matrix to remove the $ns^{2}$ term further.

Let $r$ be a positive integer. For each $x_{i}\in X^{n}$, let $\{i_{j}\}_{j=1}^{r}$ be the indexes of the closest $r$ induced points to $x_{i}$. The local induced point set of $x_{i}$ is defined as $[x_{i}]\triangleq\{u_{j}\}_{j\in\{i_{1},\ldots,i_{r}\}}$. Denote the sparse kernel matrix by $K^{*}$, which means that for $1\leq i\leq n$, if $j\in\{i_{1},i_{2},\ldots,i_{r}\}$, then $K^{*}_{ij}=K_{ij}$; otherwise $K^{*}_{ij}=0$. Thus $K^{*}$ is an $n\times s$ sparse matrix with $nr$ non-zero elements. Consequently, the cross similarity matrix and the cross transition matrix are sparse matrices with $nr$ non-zero elements. This saves the memory usage. And the time complexity of the truncated SVD decreases to $\mathcal{O}(nr^{2}+s^{3})$. The introduced sparsity accelerates the computation when $r<<s$. In addition, due to the locally Euclidean property (Lee, 2012), the Euclidean distance is almost equivalent to the intrinsic distance in a small neighborhood. Then sparsity will enhance the prediction performance when the global domain geometry is highly non-Euclidean.

It is clear that for the squared exponential kernel, selecting the bandwidth parameter $\epsilon$ involves the computation of the marginal likelihood, which costs redundant float calculation and slows down the algorithm. Moreover, the heat kernel depends on the geometry of $\mathcal{X}$, while the marginal likelihood refers only to the labeled observations. As $X^{n}$ contains more information about $\mathcal{X}$ than $X^{m}$, we want to use the point cloud to construct the kernel. Therefore, it is advantageous to substitute a likelihood-free alternative for the squared exponential kernel. We adopt the local anchor embedding (LAE) kernel of Liu et al. (2010) in the underlying work. The kernel matrix $K^{\mathrm{lae}}$ is calculated by minimizing the square loss of a convex reconstruction problem, i.e., for $1\leq i\leq n$,

The optimal solution is the coefficients of the projection of $x_{i}$ into the convex hull spanned by $[x_{i}]$. Let the remaining elements in $K^{\mathrm{lae}}$ be zero. Then it is also an $n\times s$ sparse matrix with $r$ non-zero elements in each row. Compared to the squared exponential kernel, the LAE method constructs a nonparametric kernel estimation using $X^{n}$ and $I^{s}$, which means it is more flexible and robust.

Given the covariance $\mathcal{C}$, we calculate the marginal likelihood of GPs based on the concrete models as follows,

where $\eta$ is the potential nuisance parameter. The hyperparameters in FLGP consist of predefined parameters including $s,r,M$, and optimized parameters including $\eta,\epsilon$ (if any), and $t$. The optimized parameters are estimated by maximizing the marginal likelihood function. The FLGP algorithm is listed as Algorithm 1. Furthermore, the following proposition demonstrates the time efficiency of FLGP.

As it is usually intractable to calculate the exact heat kernel, we utilize the normalized random walk graph Laplacian to estimate the heat kernel numerically. The approach is also used by Dunson et al. (2021a). The main difference is that we make the approximate inference of exact priors, while they make the exact inference of approximate priors. That is, we treat the prior as the exact heat kernel GPs and prove the oracle contraction rate.

We first estimate the approximate error for the so-called one-step random walk graph Laplacian and then explore the impact of subsampling later. Let $X^{n}$ be the point cloud. Assume that $k$ is the squared exponential kernel with bandwidth $\epsilon$. Construct the kernel matrix $\bar{K}$ on $X^{n}$ with $\bar{K}_{ij}=k(x_{i},x_{j})$. For $1\leq i\leq n$, define $\bar{K}_{i\cdot}=\sum_{j=1}^{n}k(x_{i},x_{j})$. Calculate the similarity matrix by $\bar{Z}_{ij}=k(x_{i},x_{j})/(\bar{K}_{i\cdot}\bar{K}_{j\cdot})$ for $1\leq i,j\leq n$. Let $\bar{D}$ be a diagonal matrix with $\bar{D}_{ii}=\sum_{j=1}^{n}\bar{Z}_{ij},~{}1\leq i\leq n$. Compute the stochastic matrix $\bar{A}$ by $\bar{A}=\bar{D}^{-1}\bar{Z}$. The one-step random walk graph Laplacian is defined as $\bar{L}=I-\bar{A}$.

Furthermore, the kernel normalized graph Laplacian is $\bar{L}/\epsilon^{2}$. Denote the eigenpairs of $\bar{L}/\epsilon^{2}$ by $\{(\bar{\lambda}_{i,\epsilon},\bar{v}_{i,\epsilon})\}_{i=1}^{n}$ with $||\bar{v}_{i,\epsilon}||_{2}=1$. For $t>0$, construct the heat kernel estimation as in (6) by

Let $H_{t}$ be the exact heat kernel matrix over $X^{n}$. Let $\mathcal{L}$ be the Dirichlet-Laplace operator on $\mathcal{X}$. Denote the eigenpairs of $\mathcal{L}$ by $\{(\lambda_{i},e_{i})\}_{i=1}^{\infty}$ with $||e_{i}||_{2}=1$. A technical condition is introduced on the spectrum of $\mathcal{L}$.

The above assumption and the squared exponential kernel are imposed to obtain the spectral convergence of the normalized one-step graph Laplacian in the $l^{\infty}$ norm. Other cases are feasible, as long as the spectral convergence between $\bar{L}$ and $\mathcal{L}$ can be established, such as the compactly supported kernel (García Trillos et al., 2020). Assumption 2 suggests that as $d$ increases and $\epsilon$ decreases, $n$ will increase exponentially. This implies that the current approach is probably inappropriate for processing high-dimensional data. The error bound between $H_{t}$ and $\bar{C}_{\epsilon,M,t}$ will be estimated using the spectral property in the following theorem.

Based on (2), the heat kernel is split into the sum of the first $M$ eigenpairs and the sum of the remaining eigenpairs. They are estimated separately by combining the spectral convergence property of $\bar{L}$ and the fast-decaying spectrum nature of the heat kernel. To limit the maximum distance between $H_{t}$ and $\bar{C}_{\epsilon,M,t}$, we first choose $M$ large enough to bound the latter term and subsequently choose sufficiently small $\epsilon$ to bound the former term in the right-hand side of (7). Furthermore, as $n$ increases, the point cloud becomes denser in the domain. Thus according to Theorem 1 and Remark 2, the approximation error of the heat kernel will decrease. Therefore, unlabeled observations contribute to an accurate estimation of the heat kernel by reflecting the geometry of $\mathcal{X}$.

In FLGP, we construct a two-step random walk instead of the standard one-step random walk. The errors from subsampling and transition decomposition will be estimated by the perturbation analysis. A similar technique is used in the fast spectral clustering algorithm of Yan et al. (2009). Assume that the subsampling method is k-means clustering, the kernel is the squared exponential kernel and the kernel matrix is not sparse. Suppose for $1\leq i\leq s$, the number of points in the membership of $u_{i}$ is $n_{i}$. Define the extended induced point set $\tilde{I}^{n}$ as $(\overbrace{u_{1},\ldots,u_{1}}^{n_{1}},\overbrace{u_{2},\ldots,u_{2}}^{n_{2}},\ldots,\overbrace{u_{s},\ldots,u_{s}}^{n_{s}})$. Denote the extended cross kernel matrix by $\tilde{K}$. Then $\tilde{K}$ is an $n\times n$ matrix. Let $N_{0}=0$ and $N_{i}=\sum_{j=1}^{i}n_{j},~{}1\leq i\leq s$. Construct the extended cross similarity matrix $\tilde{Z}$ by

Let $\tilde{D}$ be a diagonal matrix with $\tilde{D}_{ii}=\sum_{j=1}^{n}\tilde{Z}_{ij},~{}1\leq i\leq n$. Define the extended stochastic matrix from $X^{n}$ to $\tilde{I}^{n}$ by $\tilde{A}=\tilde{D}^{-1}\tilde{Z}$. Let $\tilde{\Lambda}$ be a diagonal matrix with $\tilde{\Lambda}_{ii}=\sum_{j=1}^{n}\tilde{A}_{ji}$. The extended two-step graph Laplacian is defined as $\tilde{L}=I-\tilde{A}(\tilde{\Lambda})^{-1}\tilde{A}^{\top}$.

Without loss of generality, assume that for $1\leq i\leq s$, $x_{N_{i-1}+1},\ldots,x_{N_{i}}$ belong to the same membership with the clustering center $u_{i}$ in coincidence. Consider the one-step stochastic matrix $\bar{A}$ on $X^{n}$. Let $\bar{\Lambda}$ be a diagonal matrix with $\bar{\Lambda}_{ii}=\sum_{j=1}^{n}\bar{A}_{ji},~{}1\leq i\leq n$. Subsequently, the exact two-step graph Laplacian is defined as $\hat{L}=I-\bar{A}\bar{\Lambda}^{-1}\bar{A}^{\top}$.

The assumption is used to constrain the spectral norm of stochastic matrices. It imposes a prior on the distribution of the point cloud. The prior reveals that for each $x_{i}\in X^{n}$, there are sufficiently many points in its neighborhood. Now we can estimate the spectral norm error between $L$ and $\hat{L}$.

The theorem comes from the measurement error analysis of graph Laplacian. More relevant results refer to Karoui and Wu (2016). According to Theorem 2, the distortion error bound in the k-means clustering derives the spectral norm bound of the graph Laplacian in FLGP. This spectral result further implies the error estimation of the covariance matrix. Theorem 2 suggests that for $x$ located far away from $I^{s}$, the corresponding prediction by FLGP may be terrible. This agrees with the intuition.

Let $X^{m}=\{x_{i}\}_{i=1}^{m}$ and $Y^{m}=\{y_{i}\}_{i=1}^{m}$. For any test point $x^{*}$, let $f^{*}=f(x^{*})$ and $\varphi^{*}=\varphi(x^{*})$. Then the posterior expectation of $\varphi^{*}$ is given by $\mathbb{E}\{\varphi^{*}|Y^{m}\}=\mathbb{E}\left[\mathbb{E}\{l(f^{*})|f^{m}\}|Y^{m}\right]$. Let $\Sigma^{1}$ and $\Sigma^{2}$ be two covariance matrices of $f$ over $\{x_{1},\ldots,x_{m},x^{*}\}$. For $i=1,2$, rewrite $\Sigma^{i}$ as $\begin{pmatrix}\Sigma^{i}_{f^{m}f^{m}}&\Sigma^{i}_{f^{m}f^{*}}\\ \Sigma^{i}_{f^{*}f^{m}}&\Sigma^{i}_{f^{*}f^{*}}\\ \end{pmatrix}$, where $\Sigma^{i}_{f^{m}f^{m}}\in\mathbb{R}^{m\times m}$, then $\pi_{\Sigma^{i}}(f^{*}|f^{m})\sim N(\mu_{i},\sigma_{i}^{2})$ with

Let $H_{t}$ be the exact heat kernel matrix on $X^{n}$ and let $C_{\epsilon,M,t}$ be the approximate covariance matrix on $X^{n}$.