Non-Parametric Variational Inference with Graph Convolutional Networks for Gaussian Processes

Inference for GP models with non-Gaussian noises is computationally expensive when dealing with large datasets. Many recent inference methods approximate the posterior distribution with a simpler distribution defined on a small number of inducing points. The inference is accurate only when data points have strong correlation with these inducing points. In this paper, we consider the inference problem in a different direction: GP function values in the posterior are mostly correlated in short distance. We construct a variational distribution such that the inference for a data point considers only its neighborhood. With this construction, the variational lower bound is highly decomposible, hence we can run stochastic optimization with very small batches. We then train Graph Convolutional Networks as a reusable model to identify variational parameters for each data point. Model reuse greatly reduces the number of parameters and the number of iterations needed in optimization. The proposed method significantly speeds up the inference and often gets more accurate results than previous methods.


Introduction
As a non-parametric Bayesian method, Gaussian Process (GP) (Rasmussen & Williams, 2006) provides flexible and expressive models for various machine learning tasks, such as regression, classification, and Bayesian optimization. The success of a GP model hinges on its efficient inference, which is especially true when the data is large in scale.
The central problem of GP inference is to calculate the posterior distribution of the function values at input points. Though the posterior of GP with non-Gaussian noise is 1 Department of Computer Science, Tufts University, Medford, MA. Correspondence to: Linfeng Liu <linfeng.liu@tufts.edu>, Liping Liu <liping.liu@tufts.edu>.
often not Gaussian, many inference methods nevertheless approximate the posterior with a multivariate Gaussian distribution. Variational inference (Jordan et al., 1999;Wainwright & Jordan, 2008) does so by maximizing the Evidence Lower Bound (ELBO) of the data likelihood. Standard variational inference for GP on large datasets is prehibitively expensive, as it needs to optimize the large covariance matrix of the variational distribution and also to inverse the prior covariance in gradient calculation. It is clear that we need further approximations to speed up the inference procedure.
Variational inference based on inducing points (Quiñonero-Candela & Rasmussen, 2005;Titsias, 2009) uses a lowrank matrix as the covariance of the variational distribution. It first defines a multivariate Gaussian distribution for function values at a small set of M inducing points and then derives distributions of all other data points from the covariance defined by the prior. This method needs space O(N M ) and time O(N M 2 ) in one gradient calculation, with N being the number of data points. With stochastic optimization (Hoffman et al., 2013), GP inference (Hensman et al., 2015;Nguyen & Bonilla, 2014;Sheth et al., 2015;Dezfouli & Bonilla, 2015;Krauth et al., 2016;Cheng & Boots, 2017) can be done in small batches, but each batch still needs to consider all inducing points. In terms of accuracy, inducing-point method has a requirement that all data points need to strongly correlate to a small number of inducing points, which may not be true in many applications. Methods in this category also lose the nonparametric flavor of GP, as the flexibility of the variational distribution is solely decided by the posterior distribution on inducing points.
We consider the approximation in an alternative direction. We propose to 1) localize the inference for a single data point within a realatively small neighborhood around it, and 2) reuse the local inference procedure for all data points. With most popular kernel settings (Chapter 4 of (Rasmussen & Williams, 2006)), the GP function value at a location is only impacted by function values at nearby locations in the prior, so is in the posterior. As a local neighborhood provides most information for the inference at a data point, neglect of weak long-distance correlation brings drastic computational advantage. arXiv:1809.02838v1 [cs.
LG] 8 Sep 2018 In this work, we construct a variational distribution with which the inference for a data point takes the information only from its K surrounding neighbors. The variational distribution supports fast sampling from its marginals, then it works for GP with general noise distributions (Nguyen & Bonilla, 2014;Sheth et al., 2015;Dezfouli & Bonilla, 2015;Krauth et al., 2016) with the reparameterization technique (Kingma & Welling, 2013). This construction makes the ELBO highly decomposible, so it is very cheap to get an unbiased estimation of the ELBO and run stochastic optimization with batch training. The update with a batch of size N b only takes time O(N b K 2 ).
Inspired by recognition networks used in variational inference (Kingma & Welling, 2013;Mnih & Gregor, 2014;Miao et al., 2016), we further train two Graph Convolutional Networks (GCN) (Kipf & Welling, 2017) as the local inference procedure to identify variational parameters for each data point. The two networks take surrounding observations as the input and output variational parameters. The inference procedure reuses the two networks at all data points. This technique greatly reduces the number of variational parameters and optimization iterations in the entire inference procedure.
The proposed inference methods (with and without an inference network) are non-parametric in the sense that the flexibility of the variational distribution grows with the data size. Especially, the method with an inference network does non-parametric inference with a parametric model.
We compare our methods with other state-of-the-art methods on several tasks. The results show that our method achieves better inference performance while using much less computation time.

Gaussian Process and Variational Inference
Gaussian Proces: A Gaussian process defines a distribution over the function space. If a function f is from a Gaussian process GP(m(·), κ(·, ·)), then it implies that . . , N ), y = (y i : i = 1, . . . , N ), and we want to model the conditional probability p(y|x) with GP, then we can specify the conditional distribution as follows.
f ∼ GP(0(·), κ(·, ·)), θ = λ(f (x i )), y i ∼ DIST(θ). (1) The link function λ(·) and the distribution DIST can be arbitrary as long as we can calculate log p(y i |f (x i )) and take its gradient with respect to f (x i ). This formulation is general enough for various regression and classification problems.
Denote f = (f i = f (x i ) : i = 1, . . . , N ), then f is a sample of Gaussian distribution with mean 0 and covariance Σ = κ(X, X), the latter of which is referred as kernel matrix. GP inference concerns the calculation of the posterior distribution p(f |x, y). The inference problem is often intractable when the noise distribution is non-Gaussian, then approximate inference becomes inevitable. In this paper we focus on variational inference and approximate p(f |x, y) with a variational distribution q(f ).
Assume now we can calculate the posterior p(f |x, y) or its approximation q(f ), we make prediction for new data points as follows. Let x be a new data point, then we calculate the predictive distribution as follows.
Here q(f |x , X, y) = f p(f |f )q(f )df , which is easy to calculate when q(f ) is Gaussian. The integral in (2) may not have a closed-form solution, but we can easily approximate it with Monte Carlo samples, as the integral variable is in only one dimension.
Variational Inference for GP: Variational inference searches for a good approximation q(f ) by maximizing the ELBO, L(q(f )), in (3) with respect to q(f ).
The bound is tight when q(f ) = p(f |X, y). For computational convenience, q(f ) is often chosen to be a multivariate Gaussian distribution. Denote its mean as µ and variance as V, then the optimization objective L(q(f )) becomes L(µ, V).
Direct calculation of (3) with a general GP model has three issues in real applications. The first term on the right side of (3) needs the inverse of prior covariance Σ. The second term often uses Monte Carlo samples to approximately the expectation E q(fi) [p(y i |f i )], so it needs cheap samples from the marginal distribution q(f i ). Third, the entropy term H [q] needs the determinant of V.
Even one can calculate the ELBO in (3), the optimization problem is still challenging for large problems, as there are a large number of optimization variables in µ and V.

The variational distribution
We construct the variational distribution q(f ) as a multivariate Gaussian distribution with mean µ and covariance V. The covariance V is further parameterized as V = LL , with L being a sparse lower triangular matrix. The off-diagonal non-zero elements in each row The number of free variables in L is N (K + 1) − K(K + 1)/2. The second term K(K + 1)/2 corresponds to the extra number of zeros for the first K rows. The size requirement of α(i)-s allows storing L into a N × (K + 1) matrix, which is convenient for matrix computation in implementation.
With this construction, it is easy to sample from q(f ). Let ∼ N (0, I) be an N -dimensional Gaussian random white noise, then a sample f from q(f ) can be generated by It is fast to draw samples from the marginal q(f i ): we only need to sample a (K + 1)-dimensional white noise ( α(i) , i ) and then calculate This construction of q(f ) corresponds to a directed graphical model defined on all N data points (Rue & Held, 2005), though a directed graphical model does not directly give marginal distributions. In the graph, α(i) is the parent set of data point i. We define α(·) sets from the prior kernel matrix. The principle is that, if two data points i and j, j < i, are strongly correlated in the prior, then j is very likely to be in α(i). Formally, we construct the α(i) for each i in the following three steps. First, we add undirected edges to the graph by neighboring relations: if i or j is one of K nearest neighbors of the other, then we add an undirected edge (i, j) to the graph. Second, we assign each edge a direction from the smaller index to the larger one, e.g. the edge (i, j) points from j to i if j < i. Finally, we make adjustments to make sure that each i has min(K, i − 1) parents. If a data point gets inadequate or exceeding parents from the last step, then we add more neighbors to or remove exceeding parents from its parent set. For the latter case, the parents with large indices would be removed first.
Below we show that q(f ) is able to match any covariance values corresponding to graph edges. This result indicates that the order of data points is not important as long as the graph has edges (in either direction) between near neighbors. In our empirical experience, there are only neglectable performance differences with permutations of data points.
Theorem: Suppose V * is the covariance matrix of the true posterior. By setting appropriate L, the matrix V = LL is able to match V * at entries corresponding to graph edges, that is,

Optimization of the ELBO
In this section, we optimize the ELBO with stochastic optimization, which needs a cheap unbiased estimation of the objective and its gradient. The strategy is first decomposing the ELBO in (3), into small factors and then estimating its value based on a random selection of factors. Based on this unbiased estimation, we calculate stochastic gradients with respect to µ and L. For convenience of discussion below, we call the three terms in (3) as the expected like- , and the entropy term L ent = H [q]. Now we decompose the three terms as follows.
Decompose likelihood term Observations y are independent given f , so the likelihood term natually decomposes as in (6). We only need a random batch S of data points to estimate the entire term L ell . We further approximate the expectation term E q(fi) [log p(y i |f i )] for data point i ∈ S with a Monte Carlo samplef i from q(f i ). We use reparameterization technique (Kingma & Welling, 2013): the random samplef i is a function of L i , and the gradient of L ell with respect L i is propagated throughf i . The exact calculation and its estimation is shown in (6).
Decompose the entropy term The entropy term is the determinant of V plus a constant. We have det(V) = N i=1 L 2 i,i by the decomposition of V. Then the exact entropy calculation and its estimation with a random batch S is Decompose the cross entropy term The cross entropy term is Let's neglect the term log det(Σ), which is constant with respect to q(f ), and focus on the trace term. The inverse of Σ is notoriously hard to compute, so people often appeal to approximations.
Nyström method (Rasmussen & Williams, 2006) approximates the kernel matrix with a low-dimensional decomposition. It is not hard to get a low-dimensional decomposition of the precision matrix, and suppose it is Σ −1 ≈ RR , then the approximation is Here j i is a random data point not in α(i), and N i = N − |α(i)| − 1 . The second term in the bracket is to improve the efficiency of the estimation, as L i L j i = 0 for many (i, j i ) pairs when i and j i does not share parents. FITC approximation (Naish-Guzman & Holden, 2008) can be treated likely.
If we use the approximation in (Datta et al., 2016), then the prior distribution approximately decomposes as p(f ) ≈ N i=1 p(f i |f α(i) ) by the graph G we used in the posterior. The conditional p(f i |f α(i) ) is derived from the GP prior and has mean and variance Then the prior is decomposible, and an unbiased estimation of the cross entropy term is We get an unbiased estimation of the ELBO by putting together estimations of the three terms.
L ≈L ell +L cross +L ent .
The approximationL cross can be either (9) or (11). The second one is neater as it uses a prior with same structure as the variational distribution. We use the second one in our experiments. We then run stochastic optimization to maximize the ELBO. This is our first inference method, and we call it NPVI. For a single data point i, the most expensive calculation is (11), which takes time O(K 2 ). For a batch S of N b data points, the calculation in (12) only takes time O(N b K 2 ). The gradient calculation takes a similar amount of time. Prior the run of our algorithm, we need to check nearest neighbors and construct the graph -the time complexity is O(n log n) with advanced data structures such as k-d trees.

Inference Network to Reduce Model Parameters
The number of parameters in the inference method above is in the same scale as the data size. Though each training batch is fast, it needs many iterations to converge when the data size is large. In this subsection, we train an inference network to directly recognize variational parameters µ i and L i from observations and the prior.
We define an inference network M such that, (µ i , L i ) = M(y (i,α(i)) , Σ (i,α(i)),(i,α(i)) ). Ideally, all observations y and the entire covariances Σ should be in the input, as they all affect µ i and L i for data point i. For the sake of feasibility, here we only use a small part of related data as the input and hope that the network can still reach good values for the two parameters without complete information. Note that the network only needs to output non-zero elements of L i .
We train two separate Graph Convolutional Networks (GCN) (Kipf & Welling, 2017) for µ i and L i , as the inference network M. A GCN computes hidden layers with the adjacency matrix of a graph and takes input from graph node values. Here we treat data points i ∪ α(i) as a small graph with Σ (i,α(i)),(i,α(i)) as the adjacency matrix and y (i,α(i)) as values at graph nodes. In this graph, i is unique because we are computing parameters for i. To break the symmetry relation between i and any other data point in as the adjacency matrix for the GCN. The input to the GCN is y (i,α(i)) . In the GCN, the hidden layer H (l+1) is computed from a previous layer H (l) as follows, Here D is a diagonal matrix, whose diagonal is the row sum of A. W (l) is the trainable weights for the l th layer. H (0) = y (i,α(i)) . The activation function σ(·) is set to be identity for the last layer and ReLU for other layers. The variational mean µ i is obtained by taking the mean of the last layer.
Reuse of this network for different data points greatly re-duces the number of optimization parameters. After learning from several iterations of inferences, it will directly identify good values for variational parameters in later iterations. Therefore, it dramatically reduces the number of iterations for convergence. The GCN learning process is especially suitable for the inference task here: GCN uses the adjacency matrix to consider correlations with nearby observations, which is exactly needed in the inference procedure.
We call this inference method as NPVI-NN. The complexity of each batch is similar to NPVI. The only extra calculation is the neural network, which takes time O(K 2 ) if we treat the size of the neural network as constant.

Experiment
We evaluate our two methods on three tasks, bird abundency estimation, raster data modeling, and flight delay regression. We put the details of the three tasks in their respective subsections.
We compared our methods with three state-of-the-art methods, SVGP (Hensman et al., 2015), SAVIGP (Dezfouli & Bonilla, 2015), and DGP (Cheng & Boots, 2017), for general GP inference. All three methods follow the scheme of inducing points. We use the implementation of SVGP by GPFlow (Matthews et al., 2017), the implementations of SVAIGP by its authors, and the implementation of DGP by Faust (2018). We implement noise distribution for these methods when necessary.
We use the RBF kernel for all experiments. The length scale is a hyper-parameter, which is selected from the candidate set {0.05, 0.1, 0.5, 1.0, 1.5, 2.0, 2.5} by checking performance on a validation set. We randomly split each dataset into three subsets for training (70%), validation (10%), and testing (20%). We report negative loglikelihood on the test dataset as the performance measure for the five algorithms in comparison. The maximum training time for each algorithm on the training dataset is 13.89 hours (50ks).
SVGP, DGP, NPVI, and NPVI-NN are all optimized by stochastic optimization with mini-batches. The optimizer for these methods is Tensorflow AdaGrad. The learning rates of NPVI is set to 0.2 while the learning rates of SVGP, DGP, and NPVI-NN are set to 0.1, based on their respective convergence speeds on the validation set. The size of each mini-batch is 50 for the four methods. SAVIGP uses L-BFGS-B optimizer with default settings from the package. We test SVGP and SAVIGP with M = 200, 1000, 2000 random inducing points. With any number over 2000, SVGP and SAVIGP will take very long time to converge. DGP uses separate inducing points to estimate the variational mean and covariance. The paper suggests to use less number of inducing points for covariance estimation. We use 200 inducing points for covariance estimation and vary the number of inducing points (M = 200, 1000, 2000) for mean estimation. We also have tried M = 4000 for DGP and found only insignificant performance improvement. We collect the result when the algorithm converges on the validation set or reaches the time limit. For NPVI and NPVI-NN, we vary K as K = 10, 20, 40. We have tested different GCN structures for NPVI-NN and found typical configurations of the network can give reasonably good performances. We finally decide a network structure with three hidden layers with dimensions 20, 10, and 1.
The eBird data The eBird project (Sullivan et al., 2009) hosts a repository for the bird-watching community to report bird observations at different locations around the world. In this experiment, we consider the population of the bird species Savannah Sparrow in February from 2012 to 2016 across the contiguous United States. In our data preprocessing step, we retain all positive observations and a random selection from exceedingly many zero observations -the dataset after preprocessing has 15,085 observations with 6,477 positive data points. We also rescale the input location values such that an Euclidean distance of 0.1 corresponds to a geographic distance of 20 miles. The link function in (1) is defined as the softplus function, λ(x) = log(1+exp(x)). The distribution of an observation y i at location i is Poisson with mean λ(f i ). Table 1 shows the negative predictive log-likelihood on the test set. Smaller values mean better predictive performances. We observe that NPVI and NPVI-NN significantly outperform the other three approaches. The standard deviations obtained from NPVI and NPVI-NN are also the smaller than the SVGP, SAVIGP, and DGP. This result also indicates that NPVI and NPVI-NN is able to achieve good inference performance with a small number K.
We also investigate how the inference method impacts the choice of hyper-parameters. Figure 1(a) shows the length scale chosen by different methods through cross validation.
The results indicate that SVGP, SAVIGP, and DGP prefer a smooth prior. As we have analyzed before, a small length scale weakens the correlation between data points and inducing points, then the three methods with a too small number of inducing points will get poor inference results.     (b), we hypothesize that the length scale 0.1 chosen by NPVI is ideal for the application. This choice limit the influence of a data point within the distance of 50 miles, which is also reasonable in ecology.
In Figure 2, we plot the data (a) and the surface predicted by the NPVI (b), NPVI-NN (c), and SVGP (d). For visualization purpose, we put use large pixels to indicate data points. The surfaces inferred by NPVI and NPVI-NN keep more detailed information, which is good for prediction considering its performance in Table 1.
The NightLight raster data In this task, we fit spatial raster data, the NightLight image, with GP models. The data is a satellite image of the night time light over US contiguous states by Earth at Night project ( Figure 3). We use GP models to fit the pixel values over the image. In this experiment, we preprocess the image by first converting the RGB color to greyscale and then downscaling the image. Finally, we get 56,722 pixels with their values rescaled to the range [0, 1]. We use a log-normal distribution with variance parameter σ 2 = 0.01 as the noise distribution DIST(the actual standard deviation of the distribution is between 0.1 and 0.27). The mean parameter of the noise distribution DIST is the Gaussian output f . Table 2 shows the negative predictive log-likelihood of all methods. NPVI and NPVI-NN both outperform the other three methods. On this dataset, NPVI-NN performs even better than NPVI. It is because NPVI does not fully converge within the time limit. NPVI-NN often converges in the first training epoch.
We compare the running speed of all methods. Figure 4 shows the negative log-likelihood values in valida-tion against training time for all five methods. There is not an equivalent setting for our methods and a setting for inducing-point methods, so we just choose typical settings for both type of methods. NPVI outperforms SVGP and SAVIGP within reasonable time. NPVI-NN converges much faster than all other methods.
Flight Delay dataset This data set contains 5.9 million records of flight arrivals and departures within the US in 2008 1 . Following Hensman et al. (2013), we select 8 dimensions as the attributes and randomly choose 800,000 instances from the original dataset as the dataset for this experiment. The task is to use the 8 attributes to fit the delay time. In the data pre-processing step, we standardize input features, remove extreme delay time values, and then translate and scale of all delay time to [0, 1]. We set the noise distribution DIST to be the log-normal distribution with σ 2 = 0.01 and µ being the GP output.
The predictive performances and training time of the four methods are shown in Table 3. We can see that NPVI-NN has the best predictive performance and also runs much faster than the other three methods. We observe that NPVI-NN converges after training with about 100k data points (about one sixth epoch). NPVI cannot run on this dataset due to too many parameters.

Conclusion
In this work, we propose two new variational inference methods, NPVI and NPVI-NN, for GP models with general noises. The two methods has a highly decomposible ELBO , so they can run stochastic optimization efficiently    to learn the variational distribution. NPVI and NPVI-NN retains the flavor of non-parametric learning by inferring function values directly at data points instead of through inducing points. NPVI-NN reduces the number of variational parameters by training an inference network to recognize variational parameters from nearby data points. It only uses a relatively small number of iterations to converge and greatly speeds up the inference procedure.