Asynchronous Distributed Variational Gaussian Processes for Regression

TL;DR

This paper introduces ADVGP, an asynchronous distributed variational Gaussian process method that significantly improves scalability and efficiency for large-scale regression tasks involving billions of data points.

Contribution

The paper presents the first asynchronous distributed variational GP inference framework, ADVGP, enabling scalable regression on massive datasets with improved efficiency and competitive accuracy.

Findings

Scales GP regression to billions of samples.

Achieves superior prediction accuracy over linear models.

Maintains or improves predictive performance with higher efficiency.

Abstract

Gaussian processes (GPs) are powerful non-parametric function estimators. However, their applications are largely limited by the expensive computational cost of the inference procedures. Existing stochastic or distributed synchronous variational inferences, although have alleviated this issue by scaling up GPs to millions of samples, are still far from satisfactory for real-world large applications, where the data sizes are often orders of magnitudes larger, say, billions. To solve this problem, we propose ADVGP, the first Asynchronous Distributed Variational Gaussian Process inference for regression, on the recent large-scale machine learning platform, PARAMETERSERVER. ADVGP uses a novel, flexible variational framework based on a weight space augmentation, and implements the highly efficient, asynchronous proximal gradient optimization. While maintaining comparable or better predictive performance, ADVGP greatly improves upon the efficiency of the existing variational methods. With ADVGP, we effortlessly scale up GP regression to a real-world application with billions of samples and demonstrate an excellent, superior prediction accuracy to the popular linear models.

Tables

Table 1. Table 1: Root mean square errors (RMSEs) for 700 700 700 K/ 100 100 100 K US Flight data.

Method	$m=50$	$m=100$	$m=200$
Prox GP	$32.9080$	$32.7543$	$32.6143$
GD Dist GP	$32.9411$	$32.8069$	$32.6521$
LBFG Dist GP	$33.0707$	$33.2263$	$32.8729$
SVIGP	$33.1054$	$32.9499$	$32.7802$

Table 2. Table C.1: Negative log evidences for 700 700 700 K/ 100 100 100 K US Flight data.

Method	$m=100$	$m=200$
ADVGP	$\mathrm{𝟗𝟐𝟓𝟐𝟑𝟔}$	$\mathrm{𝟗𝟐𝟐𝟗𝟎𝟕}$
DistGP-GD	$927414$	$924208$
DistGP-LBFGS	$932179$	$927331$

Table 3. Table C.2: Negative log evidences for 2 2 2 M/ 100 100 100 K US Flight data.

Method	$m=100$	$m=200$
ADVGP	$2.58921\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟔}}$	$2.58267\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟔}}$
DistGP-GD	$2.59471\times {10}^6$	$2.58601\times {10}^6$
DistGP-LBFGS	$2.59971\times {10}^6$	$2.59817\times {10}^6$

Table 4. Table D.1: Mean negative log predictive likelihoods for 700 700 700 K/ 100 100 100 K US Flight data.

Method	$m=100$	$m=200$
ADVGP	$1.3106$	$1.3066$
DistGP-GD	$1.3099$	$1.3062$
DistGP-LBFGS	$1.3237$	1.3136
SVIGP	1.3157	1.3096

Table 5. Table D.2: Mean negative log predictive likelihoods for 2 2 2 M/ 100 100 100 K US Flight data.

Method	$m=100$	$m=200$
ADVGP	$1.3301$	$1.3258$
DistGP-GD	$1.3317$	$1.3297$
DistGP-LBFGS	$1.3380$	$1.3355$
SVIGP	1.3335	1.3306

Equations

$$\displaystyle p({\bf y},{\bf f}|{\bf X})=\mathcal{N}\big{(}{\bf f}|{\bf 0},{\bf K}_{nn}\big{)}\mathcal{N}({\bf y}|{\bf f},\beta^{-1}{\bf I}).$$

$$\displaystyle p({\bf y}|{\bf X})=\mathcal{N}\big{(}{\bf y}|{\bf 0},{\bf K}_{nn}+\beta^{-1}{\bf I}\big{)}.$$

$$p(f^{\ast }|{\mathbf{x}}^{\ast },\mathbf{X},\mathbf{y})=\mathcal{N}(f^{\ast }|\alpha ,v),$$

$$\alpha$$

$$v$$

$$p(\mathbf{f}|\mathbf{w})=\mathcal{N}(\mathbf{f}|\mathbf{\Phi }\mathbf{w},{\mathbf{K}}_{nn}-\mathbf{\Phi }{\mathbf{\Phi }}^{\top }),$$

$$p(\mathbf{y},\mathbf{f},\mathbf{w}|\mathbf{X})$$

$$=$$

$$\log p(\mathbf{y}|\mathbf{w})=\log ⟨p(\mathbf{y}|\mathbf{f}){⟩}_{p(\mathbf{f}|\mathbf{w})}\ge ⟨\log p(\mathbf{y}|\mathbf{f}){⟩}_{p(\mathbf{f}|\mathbf{w})}$$

$$=\sum _{i=1}^n\log \mathcal{N}(y_i|{\mathbf{\varphi }}^{\top }({\mathbf{x}}_i)\mathbf{w},{\beta }^{-1})-\frac{\beta }{2}{\tilde{k}}_{ii},$$

$$\log p(\mathbf{y})$$

$$\ge ⟨\log p(\mathbf{y}|\mathbf{w}){⟩}_{q(\mathbf{w})}-\mathrm{KL}(q(\mathbf{w})\Vert p(\mathbf{w})).$$

$$\log p(\mathbf{y})\ge \mathcal{L}=-\mathrm{KL}\left(q(\mathbf{w})\Vert p(\mathbf{w})\right)$$

$$+\sum _{i=1}^n{⟨\log \mathcal{N}(y_i|{\mathbf{\varphi }}^{\top }({\mathbf{x}}_i)\mathbf{w},{\beta }^{-1})⟩}_{q(\mathbf{w})}-\frac{\beta }{2}{\tilde{k}}_{ii}.$$

$$\mathbf{\varphi }(\mathbf{x})={\mathbf{L}}^{\top }{\mathbf{k}}_m(\mathbf{x}),$$

$$L(\mathbf{\theta })={\sum }_{k=1}^r{G}_k(\mathbf{\theta })+h(\mathbf{\theta })$$

$${\mathrm{Prox}}_{{\gamma }_t}[\mathbf{\theta }]=\underset{{\mathbf{\theta }}^{\ast }}{\mathrm{argmin}}h({\mathbf{\theta }}^{\ast })+\frac{1}{2{\gamma }_t}\Vert {\mathbf{\theta }}^{\ast }-\mathbf{\theta }{\Vert }_2^2.$$

$$-\mathcal{L}={\sum }_{i=1}^n{g}_i+h$$

$$g_i$$

$$h$$

$$\frac{\partial g_i}{\partial \mathbf{\mu }}$$

$$\frac{\partial g_i}{\partial \mathbf{U}}$$

$${\mu }_i^{(t+1)}$$

$$U_{ij}^{(t+1)}$$

$$U_{ii}^{(t+1)}$$

$${\mu }_i^{\prime (t+1)}$$

$$U_{ij}^{\prime (t+1)}$$

$$\mathbf{\varphi }(\mathbf{x})={\mathrm{diag}(\mathbf{\lambda })}^{-1/2}{\mathbf{Q}}^{\top }{\mathbf{k}}_m(\mathbf{x}),$$

$$\mathbf{\varphi }(\mathbf{x})=\sum _{l=1}^q{q}^{-1/2}{\mathrm{diag}({\mathbf{\lambda }}_l)}^{-1/2}{\mathbf{Q}}_l^{\top }{\mathbf{k}}_{m_l}(\mathbf{x}),$$

$$g_i$$

$$=\frac{1}{2}\ln 2\pi -\frac{1}{2}\ln \beta +\frac{\beta }{2}\left(y_i^2-2y_i{\mathbf{\varphi }}_i^{\mathrm{T}}\mathbf{\mu }+{\mathbf{\mu }}^{\mathrm{T}}{\mathbf{\varphi }}_i{\mathbf{\varphi }}_i^{\mathrm{T}}\mathbf{\mu }+{\mathbf{\varphi }}_i^{\mathrm{T}}\mathbf{\Sigma }{\mathbf{\varphi }}_i+k_{ii}-{\mathbf{\varphi }}_i^{\mathrm{T}}{\mathbf{\varphi }}_i\right),$$

$$h$$

$$=\frac{1}{2}\left(-\ln |\mathbf{\Sigma }|-m+\mathrm{tr}(\mathbf{\Sigma })+{\mathbf{\mu }}^{\mathrm{T}}\mathbf{\mu }\right),$$

$$\displaystyle k(\bm{x}_{i},\bm{x}_{j})=a_{0}^{2}\exp\big{(}-\frac{1}{2}(\bm{x}_{i}-\bm{x}_{j})^{T}\mathrm{diag}(\bm{\eta})(\bm{x}_{i}-\bm{x}_{j})\big{)},$$

$$\frac{\partial g_i}{\partial \ln \sigma }=1-\frac{1}{{\sigma }^2}(y_i^2-2y_i{\mathbf{\varphi }}_i^T\mathbf{\mu }+{\mathbf{\varphi }}_i^T(\mathbf{\Sigma }+\mathbf{\mu }{\mathbf{\mu }}^T){\mathbf{\varphi }}_i+k_{ii}-{\mathbf{\varphi }}_i^T{\mathbf{\varphi }}_i).$$

$$\frac{\partial g_i}{\partial \ln a_0}=\frac{1}{{\sigma }^2}(-y_i{\mathbf{\varphi }}_i^T\mathbf{\mu }+{\mathbf{\varphi }}_i^T(\mathbf{\Sigma }+\mathbf{\mu }{\mathbf{\mu }}^T){\mathbf{\varphi }}_i+k_{ii}-{\mathbf{\varphi }}_i^T{\mathbf{\varphi }}_i).$$

$$\frac{\partial g_i}{\partial \mathbf{Z}}=$$

Full text

1 Introduction

Gaussian processes (GPs) (Rasmussen and Williams, 2006) are powerful non-parametric Bayesian models for function estimation. Without imposing any explicit parametric form, GPs merely induce a smoothness assumption via the definition of covariance function, and hence can flexibly infer various, complicated functions from data. In addition, GPs are robust to noise, resist overfitting and produce uncertainty estimations. However, a crucial bottleneck of GP models is their expensive computational cost: exact GP inference requires $\mathcal{O}(n^{3})$ time complexity and $\mathcal{O}(n^{2})$ space complexity ($n$ is the number of training samples), which limits GPs to very small applications, say, a few hundreds of samples.

To mitigate this limitation, many approximate inference algorithms have been developed (Williams and Seeger, 2001; Seeger et al., 2003; Quiñonero-Candela and Rasmussen, 2005; Snelson and Ghahramani, 2005; Deisenroth and Ng, 2015). Most methods use sparse approximations. Basically, we first introduce a small set of inducing points; and then we develop an approximation that transfers the expensive computations from the entire large training data, such as the covariance and inverse covariance matrix calculations, to the small set of the inducing points. To this end, a typical strategy is to impose some simplified modeling assumption. For example, FITC (Snelson and Ghahramani, 2005) makes a fully conditionally independent assumption. Recently, Titsias (2009) proposed a more principled, variational sparse approximation framework, where the inducing points are also treated as variational parameters. The variational framework is less prone to overfitting and often yields a better inference quality (Titsias, 2009; Bauer et al., 2016). Based on the variational approximation, Hensman et al. (2013) developed a stochastic variational inference (SVI) algorithm, and Gal et al. (2014) used a tight variational lower bound to develop a distributed inference algorithm with the MapReduce framework.

While SVI and the distributed variational inference have successfully scaled up GP models to millions of samples ($\mathcal{O}(10^{6})$), they are still insufficient for real-world large-scale applications, in which the data sizes are often orders of magnitude larger, say, over billions of samples ($\mathcal{O}(10^{9})$). Specifically, SVI (Hensman et al., 2013) sequentially processes data samples and requires too much time to complete even one epoch of training. The distributed variational algorithm in (Gal et al., 2014) uses the MapReduce framework and requires massive synchronizations during training, where a large amount of time is squandered when the Mappers or Reducers are waiting for each other, or the failed nodes are restarted.

To tackle this problem, we propose Asynchronous Distributed Variational Gaussian Process inference (ADVGP), which enables GP regression on applications with (at least) billions of samples. To the best of our knowledge, this is the first variational inference that scales up GPs to this level. The contributions of our work are summarized as follows: first, we propose a novel, general variational GP framework using a weight space augmentation (Section 3). The framework allows flexible constructions of feature mappings to incorporate various low-rank structures and to fulfill different variational model evidence lower bounds (ELBOs). Furthermore, due to the simple standard normal prior of the random weights, the framework enables highly efficient, asynchronous proximal gradient-based optimization, with convergence guarantees as well as fast, element-wise and parallelizable variational posterior updates. Second, based on the new framework, we develop a highly efficient, asynchronous variational inference algorithm in the recent distributed machine learning platform, ParameterServer (Li et al., 2014b) (Section 4). The asynchronous algorithm eliminates an enormous amount of waiting time caused by the synchronous coordination, and fully exploits the computational power and network bandwidth; as a result, our new inference, ADVGP, greatly improves on both the scalability and efficiency of the prior variational algorithms, while still maintaining a similar or better inference quality. Finally, in a real-world application with billions of samples, we effortlessly train a GP regression model with ADVGP and achieve an excellent prediction accuracy, with $17\%$ improvement over the popular linear regression implemented in Vowpal Wabbit (Agarwal et al., 2014), the state-of-the-art large-scale machine learning software widely used in industry.

2 Gaussian Processes Review

In this paper, we focus on Gaussian process (GP) regression. Suppose we aim to infer an underlying function $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ from an observed dataset $\mathcal{D}=\{{\bf X},{\bf y}\}$, where ${\bf X}=[{\bf x}_{1}^{\top},\ldots,{\bf x}_{n}^{\top}]^{\top}$ is the input matrix and ${\bf y}$ is the output vector. Each row of ${\bf X}$, namely ${\bf x}_{i}$ ($1\leq i\leq n$), is a $d$-dimensional input vector. Correspondingly, each element of ${\bf y}$, namely $y_{i}$, is an observed function value corrupted by some random noise. Note that the function $f$ can be highly nonlinear. To estimate $f$ from $\mathcal{D}$, we place a GP prior over $f$. Specifically, we treat the collection of all the function values as one realization of the Gaussian process. Therefore, the finite projection of $f$ over the inputs ${\bf X}$, i.e., ${\bf f}=[f({\bf x}_{1}),\ldots,f({\bf x}_{n})]$ follows a multivariate Gaussian distribution: {\bf f}\sim\mathcal{N}\big{(}{\bf f}|\bar{{\bf f}},{\bf K}_{nn}\big{)}, where $\bar{{\bf f}}=[\bar{f}({\bf x}_{1}),\ldots,\bar{f}({\bf x}_{n})]$ are the mean function values and ${\bf K}_{nn}$ is the $n\times n$ covariance matrix. Each element of ${\bf K}_{nn}$ is a covariance function $k(\cdot,\cdot)$ of two input vectors, i.e., $[{\bf K}_{nn}]_{i,j}=k({\bf x}_{i},{\bf x}_{j})$. We can choose any symmetric positive semidefinite kernel as the covariance function, e.g., the ARD kernel: $k({\bf x}_{i},{\bf x}_{j})=a_{0}^{2}\exp\left(-\frac{1}{2}({\bf x}_{i}-{\bf x}_{j})^{\top}{\rm diag}(\bm{\eta})({\bf x}_{i}-{\bf x}_{j})\right)$, where $\bm{\eta}=[1/a_{1}^{2},...,1/a_{d}^{2}]$. For simplicity, we usually use the zero mean function, namely $\bar{f}(\cdot)=0$.

Given ${\bf f}$, we use an isotropic Gaussian model to sample the observed noisy output ${\bf y}$: $p({\bf y}|{\bf f})=\mathcal{N}({\bf y}|{\bf f},\beta^{-1}{\bf I})$. The joint probability of GP regression is

[TABLE]

Further, we can obtain the marginal distribution of ${\bf y}$, namely the model evidence, by marginalizing out ${\bf f}$:

[TABLE]

The inference of GP regression aims to estimate the appropriate kernel parameters and noise variance from the training data $\mathcal{D}=\{{\bf X},{\bf y}\}$, such as $\{a_{0},\bm{\eta}\}$ in ARD kernel and $\beta^{-1}$. To this end, we can maximize the model evidence in (2) with respect to those parameters. However, to maximize (2), we need to calculate the inverse and the determinant of the $n\times n$ matrix ${\bf K}_{nn}+\beta^{-1}{\bf I}$ to evaluate the multivariate Gaussian term. This will take $\mathcal{O}(n^{3})$ time complexity and $\mathcal{O}(n^{2})$ space complexity and hence is infeasible for a large number of samples, i.e., large $n$.

For prediction, given a test input ${\bf x}^{*}$, since the test output $f^{*}$ and training output ${\bf f}$ can be treated as another GP projection on ${\bf X}$ and ${\bf x}^{*}$, the joint distribution of $f^{*}$ and ${\bf f}$ is also a multivariate Gaussian distribution. Then by marginalizing out ${\bf f}$, we can obtain the posterior distribution of $f^{*}$:

[TABLE]

where

[TABLE]

and ${\bf k}_{n*}=[k({\bf x}^{*},{\bf x}_{1}),\ldots,k({\bf x}^{*},{\bf x}_{n})]^{\top}$. Note that the calculation also requires the inverse of $K_{nn}+\beta^{-1}{\bf I}$ and hence takes $\mathcal{O}(n^{3})$ time complexity and $\mathcal{O}(n^{2})$ space complexity.

3 Variational Framework Using Weight Space Augmentation

Although GPs allow flexible function inference, they have a severe computational bottleneck. The training and prediction both require $\mathcal{O}(n^{3})$ time complexity and $\mathcal{O}(n^{2})$ space complexity (see (2), (4) and (5)), making GPs unrealistic for real-world, large-scale applications, where the number of samples (i.e., $n$) are often billions or even larger. To address this problem, we propose ADVGP that performs highly efficient, asynchronous distributed variational inference and enables the training of GP regression on extremely large data. ADVGP is based on a novel variational GP framework using a weight space augmentation, which is discussed below.

First, we construct an equivalent augmented model by introducing an $m\times 1$ auxiliary random weight vector ${\bf w}$ ($m\ll n$). We assume ${\bf w}$ is sampled from the standard normal prior distribution: $p({\bf w})=\mathcal{N}({\bf w}|{\bf 0},{\bf I})$. Given ${\bf w}$, we sample an $n\times 1$ latent function values ${\bf f}$ from

[TABLE]

where $\bm{\Phi}$ is an $n\times m$ matrix: $\bm{\Phi}=[\bm{\phi}({\bf x}_{1}),\ldots,\bm{\phi}({\bf x}_{n})]^{\top}$. Here $\bm{\phi}(\cdot)$ represents a feature mapping that maps the original $d$-dimensional input into an $m$-dimensional feature space. Note that we need to choose an appropriate $\bm{\phi}(\cdot)$ to ensure the covariance matrix in (6) is symmetric positive semidefinite. Flexible constructions of $\bm{\phi}(\cdot)$ enable us to fulfill different variational model evidence lower bounds (ELBO) for large-scale inference, which we will discuss more in Section 5.

Given ${\bf f}$, we sample the observed output ${\bf y}$ from the isotropic Gaussian model $p({\bf y}|{\bf f})=\mathcal{N}({\bf y}|{{\bf f}},\beta^{-1}{\bf I})$. The joint distribution of our augmented model is then given by

[TABLE]

This model is equivalent to the original GP regression—when we marginalize out ${\bf w}$, we recover the joint distribution in (1); we can further marginalize out ${\bf f}$ to recover the model evidence in (2). Note that our model is distinct from the traditional weight space view of GP regression (Rasmussen and Williams, 2006): the feature mapping $\bm{\phi}(\cdot)$ is not equivalent to the underlying (nonlinear) feature mapping induced by the covariance function (see more discussions in Section 5). Instead, $\bm{\phi}(\cdot)$ is defined for computational purpose only—that is, to construct a tractable variational evidence lower bound (ELBO), shown as follows.

Now, we derive the tractable ELBO based on the weight space augmented model in (7). The derivation is similar to (Titsias, 2009; Hensman et al., 2013). Specifically, we first consider the conditional distribution $p({\bf y}|{\bf w})$. Because $\log p({\bf y}|{\bf w})=\log\int p({\bf y}|{\bf f})p({\bf f}|{\bf w}){\rm d}{\bf f}=\log\langle p({\bf y}|{\bf f})\rangle_{p({\bf f}|{\bf w})}$, where $\langle\cdot\rangle_{p(\bm{\theta})}$ denotes the expectation under the distribution $p(\bm{\theta})$, we can use Jensen’s inequality to obtain a lower bound:

[TABLE]

where $\tilde{k}_{ii}$ is the $i$th diagonal element of ${\bf K}_{nn}-\bm{\Phi}\bm{\Phi}^{\top}$.

Next, we introduce a variational posterior $q({\bf w})$ to construct the variational lower bound of the log model evidence,

[TABLE]

where $\mathrm{KL(\cdot\|\cdot)}$ is the Kullback–Leibler divergence. Replacing $\log p({\bf y}|{\bf w})$ in (9) by the right side of (8), we obtain the following lower bound,

[TABLE]

Note that this is a variational lower bound: the equality is obtained when $\bm{\Phi}\bm{\Phi}^{T}={\bf K}_{nn}$ and $q({\bf w})=p({\bf w}|{\bf y})$. To achieve equality, we need to set $m=n$ and have $\bm{\phi}(\cdot)$ map the $d$-dimensional input into an $n$-dimensional feature space. In order to reduce the computational cost, however, we can restrict $m$ to be very small and choose any family of mappings $\bm{\phi}(\cdot)$ that satisfy ${\bf K}_{nn}-\bm{\Phi}\bm{\Phi}^{\top}\succeq{\bf 0}$. The flexible choices of $\bm{\phi}(\cdot)$ allows us to explore different approximations in a unified variational framework. For example, in our practice, we introduce an $m\times d$ inducing matrix ${\bf Z}=[{\bf z}_{1},\ldots,{\bf z}_{m}]^{\top}$ and define

[TABLE]

where ${\bf k}_{m}({\bf x})=[k({\bf x},{\bf z}_{1}),\ldots,k({\bf x},{\bf z}_{m})]^{\top}$ and ${\bf L}$ is the lower triangular Cholesky factorization of the inverse kernel matrix over ${\bf Z}$, i.e., $[{\bf K}_{mm}]_{i,j}=k({\bf z}_{i},{\bf z}_{j})$ and ${\bf K}_{mm}^{-1}={\bf L}{\bf L}^{\top}$. It can be easily verified that $\bm{\Phi}\bm{\Phi}^{\top}={\bf K}_{nm}{\bf K}_{mm}^{-1}{\bf K}_{nm}^{\top}$, where ${\bf K}_{nm}$ is the cross kernel matrix between ${\bf X}$ and ${\bf Z}$, i.e., $[{\bf K}_{nm}]_{ij}=k({\bf x}_{i},{\bf z}_{j})$. Therefore ${\bf K}_{nn}-\bm{\Phi}\bm{\Phi}$ is always positive semidefinite, because it can be viewed as a Schur complement of ${\bf K}_{nn}$ in the block matrix $\left[\begin{smallmatrix}{\bf K}_{mm}\quad{\bf K}_{nm}^{\top}\\ {\bf K}_{nm}\quad{\bf K}_{nn}\end{smallmatrix}\right]$. We discuss other choices of $\bm{\phi}(\cdot)$ in Section 5.

4 Delayed Proximal Gradient Optimization for ADVGP

A major advantage of our variational GP framework is the capacity of using the asynchronous, delayed proximal gradient optimization supported by ParameterServer (Li et al., 2014a), with convergence guarantees and scalability to huge data. ParameterServer is a well-known, general platform for asynchronous machine learning algorithms for extremely large applications. It has a bipartite architecture where the computing nodes are partitioned into two classes: server nodes store the model parameters and worker nodes the data. ParameterServer assumes the learning procedure minimizes a non-convex loss function with the following composite form:

[TABLE]

where $\bm{\theta}$ are the model parameters. Here $G_{k}(\bm{\theta})$ is a (possibly non-convex) function associated with the data in worker $k$ and therefore can be calculated by worker $k$ independently; $h(\bm{\theta})$ is a convex function with respect to $\bm{\theta}$.

To efficiently minimize the loss function in (12), ParameterServer uses a delayed proximal gradient updating method to perform asynchronous optimization. To illustrate it, let us first review the standard proximal gradient descent. Specifically, for each iteration $t$, we first take a gradient descent step according to $\sum\nolimits_{k}G_{k}(\bm{\theta})$ and then perform a proximal operation to project $\bm{\theta}$ toward the minimum of $h(\cdot)$, i.e., $\bm{\theta}^{(t+1)}=\mathrm{Prox}_{\gamma_{t}}[\bm{\theta}^{(t)}-\gamma_{t}\sum\nolimits_{k}\nabla G_{k}(\bm{\theta}^{(t)})]$, where $\gamma_{t}$ is the step size. The proximal operation is defined as

[TABLE]

The standard proximal gradient descent guarantees to find a local minimum solution. However, the computation is inefficient, even in parallel: in each iteration, the server nodes wait until the worker nodes finish calculating each $\nabla G_{k}(\bm{\theta}^{(t)})$; then the workers wait for the servers to finish the proximal operation. This synchronization wastes much time and computational resources. To address this issue, ParameterServer uses a delayed proximal gradient updating approach to implement asynchronous computation.

Specifically, we set a delay limit $\tau\geq 0$. At any iteration $t$, the servers do not enforce all the workers to finish iteration $t$; instead, as long as each worker has finished an iteration no earlier than $t-\tau$, the servers will proceed to perform the proximal updates, i.e., $\bm{\theta}^{(t+1)}=\mathrm{Prox}_{\gamma_{t}}[\bm{\theta}^{(t)}-\gamma_{t}\sum\nolimits_{k}\nabla G_{k}(\bm{\theta}^{(t_{k})})]$ ($t-\tau\leq t_{k}\leq t$), and notify all the workers with the new parameters $\bm{\theta}^{(t+1)}$. Once received the updated parameters, the workers compute and push the local gradient to the servers immediately. Obviously, this delay mechanism can effectively reduce the wait between the server and worker nodes. By setting different $\tau$, we can adjust the degree of the asynchronous computation: when $\tau=0$, we have no asynchronization and return to the standard, synchronous proximal gradient descent; when $\tau=\infty$, we are totally asynchronous and there is no wait at all.

A highlight is that given the composite form of the non-convex loss function in (12), the above asynchronous delayed proximal gradient descent guarantees to converge according to Theorem 4.1.

Theorem 4.1**.**

*(*Li et al., 2013) Assume the gradient of the function $G_{k}$ is Lipschitz continuous, that is, there is a constant $C_{k}$ such that $\|\nabla G_{k}(\bm{\theta})-\nabla G_{k}(\bm{\theta}^{\prime})\|\leq C_{k}||\bm{\theta}-\bm{\theta}^{\prime}||$ for any $\bm{\theta}$, $\bm{\theta}^{\prime}$, and $k=1,...,r$. Define $C=\sum_{k=1}^{r}C_{k}$. Also, assume we allow a maximum delay for the updates by $\tau$ and a significantly-modified filter on pulling the parameters with threshold $\mathcal{O}(t^{-1})$. For any $\epsilon>0$, the delayed proximal gradient descent converges to a stationary point if the learning rate $\gamma_{t}$ satisfies $\gamma_{t}\leq((1+\tau)C+\epsilon)^{-1}$.

Now, let us return to our variational GP framework. A major benefit of our framework is that the negative variational evidence lower bound (ELBO) for GP regression has the same composite form as (12). Thereby we can apply the asynchronous proximal gradient descent for GP inference on ParameterServer. Specifically, we explicitly assume $q({\bf w})=\mathcal{N}({\bf w}|\bm{\mu},\bm{\Sigma})$ and obtain the negative variational ELBO (see (10))

[TABLE]

where

[TABLE]

Instead of directly updating $\bm{\Sigma}$, we consider ${\bf U}$, the upper triangular Cholesky factor of $\bm{\Sigma}$, i.e., $\bm{\Sigma}={\bf U}^{\top}{\bf U}$. This not only simplifies the proximal operation but also ensures the positive definiteness of $\bm{\Sigma}$ during computation. The partial derivatives of $g_{i}$ with respect to $\bm{\mu}$ and ${\bf U}$ are

[TABLE]

where $\mathrm{triu}[\cdot]$ denotes the operator that keeps the upper triangular part of a matrix but leaves any other element zero.It can be verified that the partial derivatives of $g_{i}$ with respect to $\bm{\mu}$ and ${\bf U}$ are Lipschitz continuous and $h$ is also convex with respect to $\bm{\mu}$ and ${\bf U}$. According to Theorem 4.1, minimizing $\mathcal{-L}$ (i.e., maximizing $\mathcal{L}$) with respect to the variational parameters, $\bm{\mu}$ and ${\bf U}$, using the asynchronous proximal gradient method can guarantee convergence. For other parameters, such as kernel parameters and inducing points, $h$ is simply a constant. As a result, the delayed proximal updates for these parameters reduce to the delayed gradient descent optimization such as in (Agarwal and Duchi, 2011).

We now present the details of ADVGP implementation on ParameterServer. We first partition the data for $r$ workers and allocate the model parameters (such as the kernel parameters, the parameters of $q({\bf w})$ and the inducing points ${\bf Z}$) to server nodes. At any iteration $t$, the server nodes aggregate all the local gradients and perform the proximal operation in (13), as long as each worker $k$ has computed and pushed the local gradient on its own data subset $D_{k}$ for some prior iteration $t_{k}$ ($t-\tau\leq t_{k}\leq t$), $\nabla G_{k}^{(t_{k})}=\sum_{i\in D_{k}}\nabla g_{i}^{(t_{k})}$. Note that the proximal operation is only performed for the parameters of $q({\bf w})$, namely $\bm{\mu}$ and ${\bf U}$; since $h$ is constant for the other model parameters, such as the kernel parameters and the inducing points, their gradient descent updates remain unchanged. Minimizing (13) by setting the derivatives to zero, we obtain the proximal updates for each element in $\bm{\mu}$ and ${\bf U}$:

[TABLE]

where

[TABLE]

The proximal operation comprises of element-wise, closed-form computations, therefore making the updates of the variational posterior $q({\bf w})$ highly parallelizable and efficient. The gradient calculation for the other parameters, including the kernel parameters and inducing points, although quite complicated, is pretty standard and we give the details in the supplementary material (Appendix A). Finally, ADVGP is summarized in Algorithm 1.

5 Discussion and Related Work

Exact GP inference requires computing the full covariance matrix (and its inverse), and therefore is infeasible for large data. To reduce the computational cost, many sparse GP inference methods use a low-rank structure to approximate the full covariance. For example, Williams and Seeger (2001); Peng and Qi (2015) used the Nyström approximation; Bishop and Tipping (2000) used relevance vectors, constructed from covariance functions evaluated on a small subset of the training data. A popular family of sparse GPs introduced a small set of inducing inputs and targets, viewed as statistical summary of the data, and define an approximate model by imposing some conditional independence between latent functions given the inducing targets; the inference of the inexact model is thereby much easier. Quiñonero-Candela and Rasmussen (2005) provided a unified view of those methods, such as SoR (Smola and Bartlett, 2001), DTC (Seeger et al., 2003), PITC (Schwaighofer and Tresp, 2003) and FITC (Snelson and Ghahramani, 2005).

Despite the success of those methods, their inference procedures often exhibit undesirable behaviors, such as underestimation of the noise and clumped inducing inputs (Bauer et al., 2016). To obtain a more favorable approximation, Titsias (2009) proposed a variational sparse GP framework, where the approximate posteriors and the inducing inputs are both treated as variational parameters and estimated by maximizing a variational lower bound of the true model evidence. The variational framework is less prone to overfitting and often yields a better inference quality (Titsias, 2009; Bauer et al., 2016). Based on Titsias (2009)’s work, Hensman et al. (2013) developed a stochastic variational inference for GP (SVIGP) by parameterizing the variational distributions explicitly. Gal et al. (2014) reparameterized the bound of Titsias (2009) and developed a distributed optimization algorithm with MapReduce framework. Further, Dai et al. (2014) developed a GPU acceleration using the similar formulation, and Matthews et al. (2017) developed GPflow library, a TensorFlow implementation that exploit GPU hardwares.

To further enable GPs on real-world, extremely large applications, we proposed a new variational GP framework using a weight space augmentation. The proposed augmented model, introducing an extra random weight vector ${\bf w}$ with standard normal prior, is distinct from the traditional GP weight space view (Rasmussen and Williams, 2006) and the recentering tricks used in GP MCMC inferences (Murray and Adams, 2010; Filippone et al., 2013; Hensman et al., 2015). In the conventional GP weight space view, the weight vector is used to combine the nonlinear feature mapping induced by the covariance function and therefore can be infinite dimensional; in the recentering tricks, the weight vector is used to reparameterize the latent function values, to dispose of the dependencies on the hyper-parameters, and to improve the mixing rate. In our framework, however, the weight vector ${\bf w}$ has a fixed, much smaller dimension than the number of samples ($m\ll n$), and is used to introduce an extra feature mapping $\bm{\phi}(\cdot)$ — $\bm{\phi}(\cdot)$ plays the key role to construct a tractable variational model evidence lower bound (ELBO) for large scale GP inference.

The advantages of our framework are mainly twofold. First, by using the feature mapping $\bm{\phi}(\cdot)$, we are flexible to incorporate various low rank structures, and meanwhile still cast them into a principled variational inference framework. For example, in addition to (11), we can define

[TABLE]

where ${\bf Q}$ are $\bm{\lambda}$ are eigenvectors and eigenvalues of ${\bf K}_{mm}$. Then $\bm{\phi}(\cdot)$ is actually a scaled Nyström approximation for eigenfunctions of the kernel used in GP regression. This actually fulfills a variational version of the EigenGP approximation (Peng and Qi, 2015). Further, we can extend (21) by combining $q$ Nyström approximations. Suppose we have $q$ groups of inducing inputs $\{{\bf Z}_{1},\ldots,{\bf Z}_{q}\}$, where each ${\bf Z}_{l}$ consists of $m_{l}$ inducing inputs. Then the feature mapping can be defined by

[TABLE]

where $\bm{\lambda}_{l}$ and ${\bf Q}_{l}$ are the eigenvalues and eigenvectors of the covariance matrix for ${\bf Z}_{l}$. This leads to a variational sparse GP based on the ensemble Nyström approximation (Kumar et al., 2009). It can be trivially verified that both (21) and (22) satisfied ${\bf K}_{nn}-\bm{\Phi}\bm{\Phi}^{\top}\succeq{\bf 0}$ in (6).

In addition, we can also relate ADVGP to GP models with pre-defined feature mappings, for instance, Relevance Vector Machines (RVMs) (Bishop and Tipping, 2000), by setting $\bm{\phi}({\bf x})={\rm diag}(\bm{\alpha}^{1/2}){\bf k}_{m}({\bf x})$, where $\bm{\alpha}$ is an $m\times 1$ vector. Note that to ensure ${\bf K}_{nn}-\bm{\Phi}\bm{\Phi}^{\top}\succeq{\bf 0}$, we have to add some constraint over the range of each $\alpha_{i}$ in $\bm{\alpha}$.

The second major advantage of ADVGP is that our variational ELBO is consistent with the composite non-convex loss form favored by ParameterServer, therefore we can utilize the highly efficient, distributed asynchronous proximal gradient descent in ParameterServer to scale up GPs to extremely large applications (see Section $6.3$). Furthermore, the simple element-wise and closed-form proximal operation enables exceedingly efficient and parallelizable variational posterior update on the server side.

6 Experiments

6.1 Predictive Performance

First, we evaluated the inference quality of ADVGP in terms of predictive performance. To this end, we used the US Flight data111http://stat-computing.org/dataexpo/2009/ (Hensman et al., 2013), which recorded the arrival and departure time of the USA commercial flights between January and April in 2008. We performed two groups of tests: in the first group, we randomly chose $700$K samples for training; in the second group, we randomly selected $2$M training samples. Both groups used $100$K samples for testing. We ensured that the training and testing data are non-overlapping.

We compared ADVGP with two existing scalable variational inference algorithms: SVIGP (Hensman et al., 2013) and DistGP (Gal et al., 2014). SVIGP employs an online training, and DistGP performs a distributed synchronous variational inference. We ran all the methods on a computer node with $16$ CPU cores and $64$ GB memory. While SVIGP uses a single CPU core, DistGP and ADVGP use all the CPU cores to perform parallel inference. We used ARD kernel for all the methods, with the same initialization of the kernel parameters. For SVIGP, we set the mini-batch size to $5000$, consistent with (Hensman et al., 2013). For DistGP, we tested two optimization frameworks: local gradient descent (DistGP-GD) and L-BFGS (DistGP-LBFGS). For ADVGP, we initialized $\bm{\mu}={\bf 0},{\bf U}={\bf I}$, and used ADADELTA (Zeiler, 2012) to adjust the step size for the gradient descent before the proximal operation. To choose an appropriate delay $\tau$, we sampled another set of training and test data, based on which we tuned $\tau$ from $\{0,8,16,24,32,40\}$. These tunning datasets do not overlap the test data in the evaluation. Note that when $\tau=0$, the computation is totally synchronous; larger $\tau$ results in more asynchronous computation. We chose $\tau=32$ as it produced the best performance on the tunning datasets.

Table 2 and Table 2 report the root mean square errors (RMSEs) of all the methods using different numbers of inducing points, i.e., $m\in\{50,100,200\}$. As we can see, ADVGP exhibits better or comparable prediction accuracy in all the cases. Therefore, while using asynchronous computation, ADVGP maintains the same robustness and quality for inference. Furthermore, we examined the prediction accuracy of each method along with the training time, under the settings $m\in\{100,200\}$. Figure 1 shows that during the same time span, ADVGP achieves the highest performance boost (i.e., RMSE is reduced faster than the competing methods), which demonstrates the efficiency of ADVGP. It is interesting to see that in a short period of time since the beginning, SVIGP reduces RMSE as fast as ADVGP; however, after that, RMSE of SVIGP is constantly larger than ADVGP, exhibiting an inferior performance. In addition, DistGP-LBFGS converges earlier than both ADVGP and SVIGP. However, RMSE of DistGP-LBFGS is larger than both ADVGP and SVIGP at convergence. This implies that the L-BFGS optimization converged to a suboptimal solution.

We also studied how the delay limit $\tau$ affects the performance of ADVGP. Practically, when many machines are used, some worker may always be slower than the others due to environmental factors, e.g., unbalanced workloads. To simulate this scenario, we intentionally introduced a latency by assigning each worker a random sleep time of [math], $10$ or $20$ seconds at initialization; hence a worker would pause for its given sleep time before each iteration. In our experiment, the average per-iteration running time was only $0.176$ seconds; so the fastest worker could be hundreds of iterations ahead of the slowest one in the asynchronous setting. We examined $\tau=0,5,10,20,40,80,160$ and plotted RMSEs as a function of time in Figure 2. Since RMSE of the synchronous case ($\tau=0$) is much larger than the others, we do not show it in the figure. When $\tau$ is larger, ADVGP’s performance is more fluctuating. Increasing $\tau$, we first improved the prediction accuracy due to more efficient CPU usage; however, later we observed a decline caused by the excessive asynchronization that impaired the optimization. Therefore, to use ADVGP for workers at various paces, we need to carefully choose the appropriate delay limit $\tau$.

6.2 Scalability

Next, we examined the scalability of our asynchronous inference method, ADVGP. To this end, we used the $700$K/$100$K dataset and compared with the synchronous inference algorithm DistGP (Gal et al., 2014). For a fair comparison, we used the local gradient descent version of DistGP, i.e., DistGP-GD. We conducted two experiments on 4 c4.8xlarge instances of Amazon EC2 cloud, where we set the number of inducing points $m=100$. In the first experiment, we fixed the size of the training data, and increased the number of CPU cores from $4$ to $128$. We examined the per-iteration running time of both ADVGP and DistGP-GD. Figure 3(A) shows that while both decreasing with more CPU cores, the per-iteration running time of ADVGP is much less than that of DistGP-GD. This demonstrates the advantage of ADVGP in computational efficiency. In addition, the per-iteration running time of ADVGP decays much more quickly than that of DistGP-GD as the number of cores approaches $128$. This implies that even the communication cost becomes dominant, the asynchronous mechanism of ADVGP still effectively reduces the latency and maintains a high usage of the computational power. In the second experiment, we simultaneously increased the number of cores and the size of training data. We started from $87.5$K samples and $16$ cores and gradually increased them to $700$K samples and $128$ cores. As shown in Figure 3(B), the average per-iteration time of DistGP-GD grows linearly; in contrast, the average per-iteration time of ADVGP stays almost constant. We speculate that without synchronous coordination, ADVGP can fully utilize the network bandwidth so that the increased amount of messages, along with the growth of the data size, affect little the network communication efficiency. This demonstrates the advantage of asynchronous inference from another perspective.

6.3 NYC Taxi Traveling Time Prediction

Finally, we applied ADVGP for an extremely large problem: the prediction of the taxi traveling time in New York city. We used the New York city yellow taxi trip dataset 222http://www.nyc.gov/html/tlc/html/about/trip_record_data.shtml, which consist of $1.21$ billions of trip records from January 2009 to December 2015. We excluded the trips that are outside the NYC area or more than 5 hours. The average traveling time is 764 seconds and the standard derivation is 576 seconds. To predict the traveling time, we used the following 9 features: time of the day, day of the week, day of the month, month, pick-up latitude, pick-up longitude, drop-off latitude, drop-off longitude, and travel distance. We used Amazon EC2 cloud, and ran ADVGP on multiple Amazon c4.8xlarge instances, each with 36 vCPUs and 60 GB memory. We compared with the linear regression model implemented in Vowpal Wabbit (Agarwal et al., 2014). Vowpal Wabbit is a state-of-the-art large scale machine learning software package and has been used in many industrial-scale applications, such as click-through-rate prediction (Chapelle et al., 2014).

We first randomly selected $100$M training samples and $500$K test samples. We set $m=50$ and initialized the inducing points as the the K-means cluster centers from a subset of $2$M training samples. We trained a GP regression model with ADVGP, using $5$ Amazon instances with $200$ processes. The delay limit $\tau$ was selected as $20$. We used Vowpal Wabbit to train a linear regression model, with default settings. We also took the average traveling time over the training data to obtain a simple mean prediction. In Figure 4(A), we report RMSEs of the linear regression and the mean prediction, as well as the GP regression along with running time. As we can see, ADVGP greatly outperforms the competing methods. Only after $6$ minutes, ADVGP has improved RMSEs of the linear regression and the mean prediction by $9\%$ and $41\%$, respectively; the improvements continued for about $30$ minutes. Finally, ADVGP reduced the RMSEs of the linear regression and the mean prediction by $22\%$ and $49\%$, respectively. The RMSEs are {ADVGP: $333.4$, linear regression: $424.8$, mean-prediction: $657.7$}.

To further verify the advantage of GP regression in extremely large applications, we used $1$B training and $1$M testing samples. We used $50$ inducing points, initialized by the K-means cluster centers from a $1$M training subset. We ran ADVGP using $28$ Amazon instances with $1000$ processes and chose $\tau=100$. As shown in Figure 4(B), the RMSE of GP regression outperforms the linear models by a large margin. After $12$ minutes, ADVGP has improved the RMSEs of the linear regression and the mean prediction by $9\%$ and $66\%$, respectively; the improvement kept growing for about $1.5$ hours. At the convergence, ADVGP outperforms the linear regression and the mean prediction by $17\%$ and $80\%$, respectively. The RMSEs are {ADVGP: $309.7$, linear regression: $362.8$, mean-prediction: $556.3$}. In addition, the average per-iteration time of ADVGP is only $0.21$ seconds. These results confirm the power of the nonlinear regression in extremely large real-world scenarios, comparing with linear models, while the latter are much easier to be scaled up and hence more popular.

7 Conclusion

We have presented ADVGP, an asynchronous, distributed variational inference algorithm for GP regression, which enables real-world extremely large applications. ADVGP is based on a novel variational GP framework, which allows flexible construction of low rank approximations and can relate to many sparse GP models.

Appendices

Appendix A Derivatives

A.1 Objective

As described in the paper, the objective function to be minimized is $-\mathcal{L}=\sum_{i=1}^{n}g_{i}+h,$ where

[TABLE]

and we define $\beta=\sigma^{-2}$ and $\bm{\phi}_{i}=\bm{\phi}(\bm{x}_{i})$.

A.2 Kernel

A common choice for the kernel is the anisotropic squared exponential covariance function:

[TABLE]

in which the hyperparameters are the signal variance $a_{0}$ and the lengthscales $\bm{\eta}=\{1/a_{k}^{2}\}_{k=1}^{d}$, controlling how fast the covariance decays with the distance between inputs. Using this covariance function, we can prune input dimensions by shrinking the corresponding lengthscales based on the data (when $\eta_{d}=0$, the $d$-th dimension becomes totally irrelevant to the covariance function value). This pruning is known as Automatic Relevance Determination (ARD) and therefore this covariance is also called the ARD squared exponential.

Derivative over $\ln\sigma$

The derivative of $g_{i}$ over $\ln\sigma$ is

[TABLE]

**Derivative over $\ln a_{0}$ **

The derivative of $g_{i}$ over $\ln a_{0}$ is

[TABLE]

**Derivative over $\bm{Z}$ **

By defining $\bm{L}$ the lower triangular Cholesky factor of $\bm{K}_{mm}^{-1}$, the derivative of $g_{i}$ over $\bm{Z}$ is

[TABLE]

where

[TABLE]

The symbol $\circ$ denotes the Hadamard product, and $\bm{\Psi}$ is an upper triangular matrix with diagonal elements all equal to $0.5$ and strictly upper triangular elements all equal to $1$, as follows:

[TABLE]

**Derivative over $\ln\bm{\eta}$ **

The derivative of $g_{i}$ over $\ln\bm{\eta}$ is

[TABLE]

Appendix B Properties of the ELBO of ADVGP

By defining $\bm{U}$ as the upper triangular Cholesky factor of $\bm{\Sigma}$, i.e., $\bm{\Sigma}=\bm{U}^{\mathrm{T}}\bm{U}$, we have

Lemma B.1**.**

The gradient of $g_{i}$ in Equation 23, $\nabla g_{i}$, is Lipschitz continuous with respect to each element in $\bm{\mu}$ and $\bm{U}$.

We can prove this by showing the first derivative of $\nabla g_{i}$ with respect to each element of $\bm{\mu}$ and $\bm{U}$ is bounded, which is constant in our case. As shown in our paper, the gradients of $g_{i}$ with respect to $\bm{\mu}$ and $\bm{U}$ are:

[TABLE]

which are affine functions for $\bm{\mu}$ and $\bm{\Sigma}$ respectively. Therefore, the first derivative of $\nabla g_{i}$ is constant.

Lemma B.2**.**

$h$* in Equation 24 is a convex function with respect to $\bm{\mu}$ and $\bm{U}$.*

This can be proved by verifying that the Hessian matrices of $h$ with respect to $\bm{\mu}$ and $\mathrm{vec(\bm{U})}$ are both positive semidefinite, where we denote $\mathrm{vec(\cdot)}$ as the operator that stacks the columns of a matrix as a vector. To show this, we first compute the partial derivatives of $h$ with respect to $\bm{\mu}$ and $\bm{U}$ as

[TABLE]

The Hessian matrix of $h$ with respect to $\bm{\mu}$ is

[TABLE]

The Hessian matrix of $h$ with respect to $\mathrm{vec}(\bm{U})$ is

[TABLE]

where $\bm{h}=[\frac{\partial h}{\partial U_{11}^{2}},\ldots,\frac{\partial h}{\partial U_{1m}^{2}},\ldots,\frac{\partial h}{\partial U_{m1}^{2}},\ldots,\frac{\partial h}{\partial U_{mm}^{2}}]$, and $\frac{\partial h}{\partial U_{ij}^{2}}=1+\delta(i,j)\frac{1}{U^{2}_{i,i}}$.

Appendix C Negative Log Evidences on US Flight Data

Appendix D Mean Negative Log Predictive Likelihoods (MNLPs) on US Flight Data