Scalable GAM using sparse variational Gaussian processes

TL;DR

This paper introduces a scalable Bayesian approach to generalized additive models using sparse variational Gaussian processes, enabling efficient and interpretable modeling of complex data.

Contribution

It presents a novel scalable Bayesian GAM framework with sparse GPs and variational inference, improving computational efficiency and model calibration.

Findings

Efficient inference for large datasets using sparse GPs.

Enhanced interpretability of GAM components.

Well-calibrated Bayesian uncertainty estimates.

Abstract

Generalized additive models (GAMs) are a widely used class of models of interest to statisticians as they provide a flexible way to design interpretable models of data beyond linear models. We here propose a scalable and well-calibrated Bayesian treatment of GAMs using Gaussian processes (GPs) and leveraging recent advances in variational inference. We use sparse GPs to represent each component and exploit the additive structure of the model to efficiently represent a Gaussian a posteriori coupling between the components.

Figures

Equations

$$\log p(Y|X)\ge \text{mathds}{E}_{q(\mathcal{F})}[\log p(Y|\mathcal{F},X)]-\mathrm{KL}[q(\mathcal{F})\Vert p(\mathcal{F})]=\mathcal{L}(q).$$

$${\cal L}(q)=\mathds{E}_{q({\cal F})}[\log p(Y\,|\,{\cal F})]+\frac{1}{2}\Big{(}\operatorname{tr}(K_{{\cal F},{\cal F}}^{-1}\Sigma_{{\cal F}})+\mu_{{\cal F}}^{\intercal}K_{{\cal F},{\cal F}}^{-1}\mu_{{\cal F}}-\log|\Sigma_{{\cal F}}|-NC\Big{)}.$$

$$\Sigma^{-1}_{{\cal F}}=K_{{\cal F},{\cal F}}^{-1}+\nabla_{\Sigma_{{\cal F}}}\big{[}\mathds{E}_{q({\cal F})}[\log p(Y\,|\,{\cal F})]\big{]}.$$

$${\Sigma }_{\mathcal{F}}^{-1}=K_{\mathcal{F},\mathcal{F}}^{-1}+(1_C\otimes \Lambda )(1_C\otimes \Lambda {)}^{⊺}.$$

$$\mathcal{L}(q)$$

$${\Sigma }_{\mathbf{U},\mathbf{U}}^{-1}$$

$$s_i(x_i,y_i)=g_i(x_i,y_i)-\frac{{\int }_0^1{g}_i(x_i,s)\mathrm{d}s{\int }_0^1{g}_i(y_i,s)\mathrm{d}s}{{\iint }_0^1{g}_i(s_i,t)\mathrm{d}s\mathrm{d}t}.$$

$${\Sigma }_{\mathcal{F}}^{-1}=K_{\mathcal{F},\mathcal{F}}^{-1}+{\nabla }_{{\Sigma }_{\mathcal{F}}}V(Y,{\mu }_{\mathcal{F}},{\Sigma }_{\mathcal{F}}).$$

$$V(Y,{\mu }_{\mathcal{F}},{\Sigma }_{\mathcal{F}})$$

$$={\sum }_{n}v(y_n,{\mu }_{\rho (x_n)},{\sigma }_{\rho (x_n)}^2),$$

$${\nabla }_{{\Sigma }_{\mathcal{F}}}V(Y,{\mu }_{\mathcal{F}},{\Sigma }_{\mathcal{F}})$$

$$={\sum }_n{\lambda }_n^2{\sum }_{c{c}^{\prime }}{e}_{cn,c^{\prime }n},$$

$${\nabla }_{{\Sigma }_{\mathcal{F}}}V(Y,{\mu }_{\mathcal{F}},{\Sigma }_{\mathcal{F}})$$

$$q(\mathcal{F})=\mathcal{N}(K_{\mathcal{F},\mathcal{F}}\alpha ,{\Sigma }_{\mathcal{F}})$$

$$\Sigma_{{\cal F}}=\big{(}K_{{\cal F},{\cal F}}^{-1}+(1\otimes\Lambda)(1\otimes\Lambda)^{\intercal}\big{)}^{-1},$$

$$\mathcal{L}(q)=\text{mathds}{E}_{q({\sum }_c{f}_c)}[\log p(y|{\sum }_c{f}_c)]-\mathrm{KL}[q(\mathcal{F})\Vert p(\mathcal{F})],$$

$$\mathrm{KL}[q(\mathcal{F})\Vert p(\mathcal{F})]=\frac{1}{2}[-\log |K_{\mathcal{F},\mathcal{F}}^{-1}{\Sigma }_{\mathcal{F}}|+{\alpha }^{⊺}{K}_{\mathcal{F},\mathcal{F}}\alpha +\mathrm{tr}(K_{\mathcal{F},\mathcal{F}}^{-1}{\Sigma }_{\mathcal{F}})-NC].$$

$${\Sigma }_{\mathcal{F}}$$

$$\displaystyle=K_{{\cal F},{\cal F}}-K_{{\cal F},{\cal F}}(1\otimes\Lambda)\big{(}I+(1\otimes\Lambda)^{\intercal}K_{{\cal F},{\cal F}}(1\otimes\Lambda)\big{)}^{-1}(1\otimes\Lambda)^{\intercal}K_{{\cal F},{\cal F}}$$

$$=K_{\mathcal{F},\mathcal{F}}-K_{\mathcal{F},\mathcal{F}}(1\otimes \Lambda )A^{-1}(1\otimes \Lambda {)}^{⊺}{K}_{\mathcal{F},\mathcal{F}},$$

$$A$$

$$=I+{\sum }_c{\Lambda }^{⊺}{K}_{f_c,f_c}\Lambda .$$

$${\Sigma }_{\text{sum}}$$

$$=(1\otimes I{)}^{⊺}[K_{\mathcal{F},\mathcal{F}}-K_{\mathcal{F},\mathcal{F}}(1\otimes \Lambda )A^{-1}(1\otimes \Lambda {)}^{⊺}{K}_{\mathcal{F},\mathcal{F}}](1\otimes I{)}^{⊺}$$

$$={\sum }_c{K}_{f_c,f_c}-{\sum }_{c,c^{\prime }}{K}_{f_c,f_c}\Lambda A^{-1}\Lambda K_{f_{c^{\prime }},f_{c^{\prime }}}$$

$$={\sum }_c{K}_{f_c,f_c}-{\sum }_{c,c^{\prime }}(L_A^{-⊺}\Lambda K_{f_c,f_c}{)}^{⊺}(L_A^{-⊺}\Lambda K_{f_{c^{\prime }},f_{c^{\prime }}}).$$

$$|K_{\mathcal{F},\mathcal{F}}^{-1}{\Sigma }_{\mathcal{F}}|$$

$$=|K_{\mathcal{F},\mathcal{F}}^{-1}|/|K_{\mathcal{F},\mathcal{F}}^{-1}+(1\otimes \Lambda )(1\otimes \Lambda {)}^{⊺}|$$

$$=|K_{\mathcal{F},\mathcal{F}}^{-1}|/[|I+(1\otimes \Lambda {)}^{⊺}{K}_{\mathcal{F},\mathcal{F}}(1\otimes \Lambda )||I||K_{\mathcal{F},\mathcal{F}}^{-1}|]$$

$$=1/|A|$$

$$\mathrm{tr}(K_{\mathcal{F},\mathcal{F}}^{-1}{\Sigma }_{\mathcal{F}})$$

$$=\mathrm{tr}(I-(1\otimes \Lambda )A^{-1}(1\otimes \Lambda {)}^{⊺}{K}_{\mathcal{F},\mathcal{F}})$$

$$=NC-\mathrm{tr}(A^{-1}(1\otimes \Lambda {)}^{⊺}{K}_{\mathcal{F},\mathcal{F}})$$

$$=NC-{\sum }_c\mathrm{tr}(\Lambda A^{-1}{\Lambda }^{⊺}{K}_{f_c,f_c})$$

$$=NC-\mathrm{tr}(\Lambda A^{-1}{\Lambda }^{⊺}{\sum }_c{K}_{f_c,f_c})$$

$$\mathrm{KL}[q(\mathcal{F})\Vert p(\mathcal{F})]=\frac{1}{2}[\log |A|+{\alpha }^{⊺}{K}_{\mathcal{F},\mathcal{F}}\alpha -\mathrm{tr}(\Lambda A^{-1}{\Lambda }^{⊺}{\sum }_c{K}_{f_c,f_c})].$$

$$A$$

Full text

1 Introduction

Generalized additive models (GAMs) are a class of interpretable regression models with non-linear yet additive predictors (Hastie, 2017). Their Bayesian treatment requires the specification of priors over functions. Here, we use Gaussian processes (GPs) (Rasmussen and Williams, 2006) and propose an approximate inference algorithm that is scalable with both the number of data points and additive components and that provides accurate posterior uncertainty estimates. We extend the variational pseudo-point GP approximation (Titsias, 2009; Bauer et al., 2016) to posterior dependencies across GPs. This approximation provides state-of-the art performance for GP regression and provides approximations to the posterior distributions in the form of a GP. This approach has been successfully extended to the multiple GP setting using a factorized (mean-field) approximation of the posterior across GPs (Saul et al., 2016; Adam et al., 2016). However, it suffers from the known variance underestimation of mean-field approximations and therefore can lead to poor predictions or can bias learning (Turner and Sahani, 2011). Adam (2017) introduced additional structure to the posterior distribution by allowing some coupling across the inducing variables of the different GPs but this was at the cost of scalability.

2 Background

2.1 Regression with multiple GPs

We consider models with additive predictor and factorizing likelihood $p(Y\,|\,f_{1\dots C},X)=\prod_{n=1}^{N}p(y_{n}\,|\,\textstyle\sum_{c}f_{c}(x_{n}))$, where $f_{1},\dots,f_{C}$ are functions from ${\cal X}_{c}\to\mathds{R}$. The specific form of the likelihood is arbitrary. We denote ${\cal F}=\{f_{1},...,f_{C}\}$ such that $p({\cal F})=\prod_{c}p(f_{c})$ constitutes the joint distribution over the a priori independent processes. ${\cal F}(x)=[f_{1}(x),...,f_{C}(x)]$ is the vector of function evaluations at $x$. To simplify notation, when no argument is given, ${\cal F}={\cal F}(X)\in\mathds{R}^{NC}$. We denote by $K_{{\cal F},{\cal F}}$ the block-diagonal prior covariance matrix over ${\cal F}$. We are interested in computing the joint posterior $p({\cal F}\,|\,X,Y)$.

2.2 Variational Inference

The classical variational lower bound (or ELBO) to the marginal likelihood is given by

[TABLE]

This is the optimization objective in the Variational Free Energy (VFE) approximation. We choose $q({\cal F})$ to be a multivariate normal distribution with mean $\mu_{{\cal F}}$ and covariance $\Sigma_{{\cal F}}$, which is not an approximation in the conjugate likelihood setting. This leads to

[TABLE]

The expectation term in equation (2) is intractable in most cases and needs to be approximated. See Hensman et al. (2015) for deterministic approximations and Salimbeni and Deisenroth (2017) for stochastic approximations.

3 Optimal Gaussian posterior in variational inference

Following Opper and Archambeau (2009), we derive the expression for the optimal $\Sigma_{{\cal F}}$ by noting that at the optimum, $\nabla_{\Sigma_{{\cal F}}}{\cal L}(q)=0$. This implies that

[TABLE]

3.1 Optimality in the additive case

In the additive case considered here, the gradient term in (3) is low-rank and can be parameterized by a vector $\lambda\in\mathds{R}^{N}$ as follows, with $\Lambda=\operatorname{diag}(\lambda)$ and $1_{C}=[\underbrace{1,\dots,1}_{C\text{ times}}]^{\intercal}$:

[TABLE]

This parameterization requires $2N$ values, equal to that of the classical single GP regression setting described in Opper and Archambeau (2009). It also inherits the non-convexity of this objective as highlighted by Khan et al. (2012).

3.2 Optimality in the sparse additive case

Following Adam et al. (2016) we introduce for each GP indexed by $c$ a set of $M$ ‘inducing points’ $Z_{c}=[z_{c}^{(1)},...,z_{c}^{(M)}]\in{\cal X}_{c}^{M}$. The vector of associated function evaluations is given by $\mathbf{U}_{c}=[u_{c}^{(1)},...,u_{c}^{(M)}]=[f_{c}(z_{c}^{(1)}),...,f_{c}(z_{c}^{(M)})]$. We also define the stacked vector $\mathbf{U}=[\mathbf{U}_{1},...,\mathbf{U}_{c}]\in\mathds{R}^{MC}$.

Following Adam (2017), we parameterize $q({\cal F})=q(\mathbf{U})\prod_{c}p(f_{c\neg\mathbf{U}_{c}}\,|\,\mathbf{U}_{c})$. This choice leads to a simplification of the lower bound (2) as

[TABLE]

Saul et al. (2016) considered the mean field case $q(\mathbf{U})=\prod_{c}q(\mathbf{U}_{c})$ with each factor parameterized as a multivariate normal distribution ${\cal N}(\mu_{\mathbf{U}_{c}},\Sigma_{\mathbf{U}_{c}})$. This approach does not capture posterior dependencies across GPs. Adam (2017) parameterized $q(\mathbf{U})$ as a multivariate normal distribution ${\cal N}(\mu_{\mathbf{U}},\Sigma_{\mathbf{U}})$ to include cross-GP coupling through the inducing variables $\mathbf{U}$. We extend this last approach but ask what the optimal $q(\mathbf{U})$ should be. It turns out to be (see Appendix A):

[TABLE]

This form has again $2N$ parameters which becomes an over-parameterization as soon as $N>M^{2}C^{2}/2$. Since we are interested in scalability, it is not of practical interest.

4 A new parameterization for $q(\mathbf{U})$

The second term of the sum in (6) can be expressed as $AA^{\intercal}$ with $A$ of size $MC\times N$. Keeping this structure arising from the additivity of the model, we propose the parameterization

[TABLE]

with $B$ of size $MC\times M$ smaller than $A$. This parameterization preserves the structure of the optimal covariance. It requires storing $M^{2}C$ values, which is less than a direct representation of a Cholesky factor of $\Sigma^{-1}_{\mathbf{U},\mathbf{U}}$ that would require $M^{2}C^{2}$ parameters.

5 Summary of complexities

Time and space complexity of the sparse variational algorithms are summarized in Table LABEL:table:complexity.

6 Related work

Variational inference for the multi-GP setting has so far only used the mean-field (MF) approximation as described in Saul et al. (2016). When posterior dependencies are a quantity of interest, a natural approach is to increase the complexity of the variational posterior to capture these dependencies. This often results in a prohibitive increase in the complexity of the inference. Different solutions have been proposed to tackle this problem. A first approach in Giordano et al. (2015) consists of a two-step scheme where MF inference is assumed to provide accurate posterior mean estimates. A perturbation analysis is then performed around the MF posterior means to provide second order (covariance) estimates. A second approach consists in ‘relaxing’ the MF approximation by extending the variational posterior $q({\cal F})$ with additional multiplicative terms capturing dependencies while keeping the computational complexity of the resulting inference scheme low (Tran et al., 2015; Hoffman and Blei, 2015). Our approach fits in this second family of extensions of the MF parameterization. It is tailored to the VFE approximation to GP models and leverages its sparsity to provide a fast and scalable inference algorithm.

7 Illustration

We consider a simple regression task consisting of approximating the following function: $f(x)=10\sin(\pi x_{1}x_{2})+20(x_{3}-0.5)^{2}+10x_{4}+5x_{5}$ with $x\in[0,1]^{6}$ (note that the last variable has no effect), given 5000 observation points uniformly distributed in the input space and a Gaussian observation noise with unit variance.

We choose a kernel dedicated to sensitivity analysis and tailored to the structure of the function at stake [Durrande et al. (2013)]. Given univariate squared exponential kernels $g_{1},\dots,g_{8}$ we define the kernel as $k(x,y)=\sigma_{0}+\sum_{i=1}^{6}s_{i}(x_{i},y_{i})+s_{7}(x_{1},y_{1})s_{8}(x_{2},y_{2})$ with

[TABLE]

Since the number of observations is relatively large and the kernel has an additive structure (it is the sum of 8 kernels), we choose the sparse additive model described above. We choose 16 regularly spaced one-dimensional inducing points for each kernel $s_{1},\dots,_{6}$ and 16 points distributed as a $4\times 4$ grid for the bi-dimensional kernel $s_{7}s_{8}$. The final model is obtained by maximizing the ELBO with respect to the variational parameters and the hyper-parameters of the $g_{i}$. Given the structure of the model and the fact that inducing inputs are dedicated to model components, it is then possible to decompose the model predictions and to represent separately all the components of the ANOVA representation of the test function. Figure 1 shows that the model accurately approximates the test function and that the proposed framework is helpful to reveal its inner structure.

8 Conclusion

We presented a method that provides a fast, scalable and well-calibrated Bayesian treatment of GAMs. Although motivated by GAMs, our structured variational distribution may be used in models where the predictor is non-additive but where the posterior is well-approximated by a unimodal distribution.

Appendix A Optimal covariance in the additive case

We first define $V(Y,\mu_{{\cal F}},\Sigma_{{\cal F}})=\mathds{E}_{q({\cal F})}\log p(Y\,|\,{\cal F})$. From Opper and Archambeau (2009), we know that the optimal variational precision is structured as

[TABLE]

For factorizing likelihood and additive predictors, and defining $\rho(\cdot)=\textstyle\sum_{c}f_{c}(\cdot)$, we have

[TABLE]

where $q(\rho(x_{n}))$ has variance $\sigma^{2}_{\rho(x_{n})}=1_{C}^{\intercal}\Sigma_{{\cal F}(x_{n})}1_{C}=\textstyle\sum_{c,c^{\prime}}\Sigma_{f_{c}(x_{n}),f_{c^{\prime}}(x_{n})}$.

The gradient term in the optimal precision thus can be written as

[TABLE]

where $e_{i,j}$ is the indicator matrix of size $NC\times NC$ with $1$ at location $(i,j)$. With $\Lambda=\operatorname{diag}(\lambda)$, this can be rewritten in matrix form:

[TABLE]

Appendix B ELBO evaluation: additive case

We parameterize the approximate posterior as

[TABLE]

with

[TABLE]

and optimize

[TABLE]

where

[TABLE]

B.1 Computing marginals of $\Sigma_{{\cal F}^{(n)}}$

[TABLE]

where

[TABLE]

To evaluate the ELBO we need, for each data point $(x_{n},y_{n})$, the marginal $q(\sum_{c}f^{n}_{c})$. This corresponds to the diagonal elements of

[TABLE]

B.2 Computing the KL

[TABLE]

[TABLE]

In the end,

[TABLE]

B.3 Summary

$\displaystyle A$ $\displaystyle=I+\textstyle\sum_{c}\Lambda^{\intercal}K_{f_{c},f_{c}}\Lambda$

$\displaystyle\operatorname{KL}[q({\cal F})\,\|\,p({\cal F})]$ $\displaystyle=\frac{1}{2}[\log|A|+\alpha^{\intercal}K_{{\cal F},{\cal F}}\alpha-\operatorname{tr}(\Lambda A^{-1}\Lambda^{\intercal}\textstyle\sum_{c}K_{f_{c},f_{c}})]$

$\displaystyle\mu_{\text{sum}}$ $\displaystyle=\textstyle\sum_{c}K_{f_{c},f_{c}}\alpha_{c}$

$\displaystyle\Sigma_{\text{sum}}$ $\displaystyle=\textstyle\sum_{c}\operatorname{diag}(K_{f_{c},f_{c}})-\textstyle\sum_{c,c^{\prime}}\operatorname{diag}(K_{f_{c},f_{c}}\Lambda A^{-1}\Lambda K_{f_{c^{\prime}},f_{c^{\prime}}})$

Appendix C ELBO evaluation: sparse additive case

We parameterize an approximate posterior over the inducing values as

[TABLE]

with

[TABLE]

where $B=[B_{1},\dots,B_{C}]\in\mathds{R}^{MC\times M}$. We optimize

[TABLE]

with

[TABLE]

C.1 Computing marginals of $\Sigma_{{\cal F}}$

We have

[TABLE]

where $A=I+B^{\intercal}K_{{\cal F},{\cal F}}B$, so

[TABLE]

and

[TABLE]

Therefore

[TABLE]

The Cholesky decomposition of $A=L_{A}L_{A}^{\intercal}$ is of cost $\mathcal{O}(M^{3})$. Solving operations $L^{-1}B_{c}^{\intercal}$ for each additive term costs a total of $\mathcal{O}(CM^{3})$. Computing the marginal predictor variances then costs an extra $\mathcal{O}(NC^{2}M^{2})$. In total, the computational cost of posterior predictions is $\mathcal{O}(CM^{3}+NC^{2}M^{2})$.

C.2 Computing the KL

As in the additive case, we have

[TABLE]

and

[TABLE]

In the end,

[TABLE]

C.3 Summary

$\displaystyle A$ $\displaystyle=I+B^{\intercal}K_{{\cal F},{\cal F}}B$

$\displaystyle\operatorname{KL}[q(\mathbf{U})\,\|\,p(\mathbf{U})]$ $\displaystyle=\frac{1}{2}[\log|A|+\alpha^{\intercal}K_{\mathbf{U},\mathbf{U}}\alpha-\textstyle\sum_{c}\operatorname{tr}(B_{c}A^{-1}B_{c}^{\intercal}K_{\mathbf{U}_{c},\mathbf{U}_{c}})]$

$\displaystyle\mu_{\text{sum}}$ $\displaystyle=\textstyle\sum_{c}K_{f_{c},\mathbf{U}_{c}}\alpha_{c}$

$\displaystyle\Sigma_{\text{sum}}$ $\displaystyle=\textstyle\sum_{c}\operatorname{diag}(K_{f_{c},f_{c}})-\textstyle\sum_{c,c^{\prime}}\operatorname{diag}(K_{f_{c},\mathbf{U}_{c}}B_{c}A^{-1}B_{c^{\prime}}^{\intercal}K_{\mathbf{U}_{c^{\prime}},f_{c^{\prime}}})$