A Sequential Marginal Likelihood Approximation Using Stochastic Gradients ^†

Abstract

Existing algorithms like nested sampling and annealed importance sampling are able to produce accurate estimates of the marginal likelihood of a model, but tend to scale poorly to large data sets. This is because these algorithms need to recalculate the log-likelihood for each iteration by summing over the whole data set. Efficient scaling to large data sets requires that algorithms only visit small subsets (mini-batches) of data on each iteration. To this end, we estimate the marginal likelihood via a sequential decomposition into a product of predictive distributions $p\left({\mathbf{y}}_n\right|{\mathbf{y}}_{<n})$. Predictive distributions can be approximated efficiently through Bayesian updating using stochastic gradient Hamiltonian Monte Carlo, which approximates likelihood gradients using mini-batches. Since each data point typically contains little information compared to the whole data set, the convergence to each successive posterior only requires a short burn-in phase. This approach can be viewed as a special case of sequential Monte Carlo (SMC) with a single particle, but differs from typical SMC methods in that it uses stochastic gradients. We illustrate how this approach scales favourably to large data sets with some simple models.

1. Introduction

Marginal likelihood (ML), sometimes call evidence, is a quantitative measure of how well a model can describe a particular data set; it is the probability that the data set occurred within that model. Consider a Bayesian model with parameters $\mathit{\theta }$ for a data set $\mathcal{D}={\left\{{\mathbf{y}}_n\right\}}_{n=1}^N$. The ML is the integral

$$\mathcal{Z}:=p\left(\mathcal{D}\right)=\int p\left(\mathcal{D}\right|\mathit{\theta }\left)p\right(\mathit{\theta })d\mathit{\theta },$$

where $p\left(\mathcal{D}\right|\mathit{\theta })$ is the likelihood and $p\left(\mathit{\theta }\right)$ is the prior. In this paper we consider the case where the data are conditionally independent given the parameters

$$p\left(\mathcal{D}\right|\mathit{\theta })=\prod _{n}p\left({\mathbf{y}}_n\right|\mathit{\theta }),$$

(1)

as is common in many parametric models. The posterior distribution over a set of models is proportional to their ML and so approximations to it are sometimes used for model comparison, and weighted model averaging [1] (chapter 12). This integral is typically analytically intractable for any but the simplest models, so one must resort to numerical approximation methods. Nested sampling (NS) [2] and annealed importance sampling (AIS) [3] are two algorithms able to produce accurate estimates of the ML. NS accomplishes this by sampling from the prior under constraints of increasing likelihood, and AIS by sampling from a temperature annealed distribution $\propto p{\left(\mathcal{D}\right|\mathit{\theta })}^{\beta }p\left(\mathit{\theta }\right)$ and averaging over samples with appropriately calculated importance weights. Although NS and AIS produce accurate estimates of the ML, they tend to scale poorly to large data sets due the fact that they need to repeatedly calculate the likelihood function: for NS this is to ensure staying within the constrained likelihood contour; for AIS the likelihood must be calculated both to sample from the annealed distributions using some Markov chain Monte Carlo (MCMC) method, as well as to calculate the importance weights. Calculation of the likelihood is computationally expensive on large data sets. To combat this problem, various optimization and sampling algorithms rather make use of stochastic approximations of the likelihood by sub-sampling the data set into mini-batches $B\subseteq \mathcal{D}$ [4]. The stochastic log-likelihood approximation is

$$logp\left(\mathcal{D}\right|\mathit{\theta })\approx \frac{\left|\mathcal{D}\right|}{\left|B\right|}\sum _{\mathbf{y}\in B}logp\left(\mathbf{y}\right|\mathit{\theta }),$$

with each iteration of these algorithms generally using a different mini-batch. Unfortunately NS and AIS cannot trivially use mini-batching to improve scalability. Using stochastic likelihood approximations changes the statistics of the likelihood contours in NS, allowing particles to occasionally move to lower likelihood instead of higher, violating the basic assumptions of the algorithm. AIS could benefit from using stochastic likelihood gradients during the MCMC steps, but it is not obvious how one would calculate the importance weights in this setting. This work presents an approach for large-scale ML approximations using mini-batches. This is done using a sequential decomposition of the ML into predictive distributions, which can each be approximated using stochastic gradient MCMC methods. The particles sampled from each previous posterior distribution can be reused for efficiency since they will typically be close to the next posterior. This can be viewed as a special case of sequential Monte Carlo with Bayesian updating [5].

We illustrate our approach by calculating ML estimates on three simple models. For these models, we obtain roughly an order of magnitude speedup over nested sampling on data sets with one million observations with negligible loss in accuracy.

2. Sequential Marginal Likelihood Estimation

The ML can be decomposed, through the product rule, into a product of predictive distributions of the following form

$$\mathcal{Z}=\prod _{n}p\left({\mathbf{y}}_n\right|{\mathbf{y}}_{<n}),$$

where, from Equation (1),

$$p\left({\mathbf{y}}_n\right|{\mathbf{y}}_{<n})=\int p\left({\mathbf{y}}_n\right|\mathit{\theta })p\left(\mathit{\theta }\right|{\mathbf{y}}_{<n})d\mathit{\theta }.$$

(2)

Assuming one is able to produce accurate estimates $\widehat{p}\left({\mathbf{y}}_n\right|{\mathbf{y}}_{<n})$ of the predictive probabilities, the log-ML can be approximated by

$$log\widehat{\mathcal{Z}}=\sum _{n}log\widehat{p}\left({\mathbf{y}}_n\right|{\mathbf{y}}_{<n}).$$

(3)

In this way the difficult problem of estimating an integral of an extremely peaked function, $p\left(\mathcal{D}\right|\mathit{\theta })$, reduces to the easier problem of estimating many integrals of smoother functions $p\left({\mathbf{y}}_n\right|\mathit{\theta })$. Note that this approach can more generally be applied using any other sequential decomposition of the data set. We only present derivations using the approach above for notational clarity but we decompose the data into varying-sized chunks of observations during our experiments as discussed in Section 4.

Typical sequential Monte Carlo (SMC) methods use a similar approach, and calculate predictive estimates using a combination of importance resampling and MCMC mutation steps [5]. Generic examples of such algorithms are the bootstrap particle filter [6], which is often used for posterior inference in hidden Markov models and other latent variable sequence models [5], and the “left-to-right” algorithm which is used in [7] to evaluate topic models.

The computational efficiency of using this approach depends on the method of approximating Equation (2). One such estimator is

$$\widehat{p}\left({\mathbf{y}}_n\right|{\mathbf{y}}_{<n})=\frac{1}{M}\sum _{i=1}^{M}p\left({\mathbf{y}}_n\right|{\mathit{\theta }}_i),$$

(4)

where each ${\mathit{\theta }}_i$ is drawn from the posterior distribution $p\left(\mathit{\theta }\right|{\mathbf{y}}_{<n})$ using MCMC methods. Again, one might use a chunk of observations for this predictive estimate rather than just one. As described in [8], estimators of this form tend to underestimate the ML, but still converge to the exact ML in the limit $M\to \infty$. Since samples from the previous posterior, $p\left(\mathit{\theta }\right|{\mathbf{y}}_{<n-1})$, would generally be available at each step, we expect only a small number of steps will be needed to accurately sample from the next posterior distribution, $p\left(\mathit{\theta }\right|{\mathbf{y}}_{<n})$. Metropolis-Hastings based MCMC algorithms would have to iterate over the previous $n-1$ data points in order to calculate the acceptance probability for each Markov transition, and so using them to estimate $log\mathcal{Z}$ in this sequential manner would scale at least quadratically in N. The key computational improvement in our approach is instead using stochastic gradient based MCMC algorithms such as stochastic gradient Hamiltonian Monte Carlo [9]. SGHMC utilizes mini-batching, allowing one to efficiently draw samples from the posterior distribution $p\left(\mathit{\theta }\right|{\mathbf{y}}_{<n})$ even when n is large. We now describe how one can use SGHMC to sample from posterior distributions.

3. Stochastic Gradient Hamiltonian Monte Carlo

SGHMC [9] simulates a Brownian particle in a potential, by numerically integrating the Langevin equation

$$\begin{array}{cc}d\mathit{\theta }\hfill & =\mathbf{v}dt\hfill \\ d\mathbf{v}\hfill & =-\nabla U\left(\mathit{\theta }\right)dt-\gamma \mathbf{v}dt+\sqrt{2\gamma }dW,\hfill \end{array}$$

where $U\left(\mathit{\theta }\right)$ is the potential energy, $\gamma$ is the friction coefficient and W is the standard Wiener process [10]. That is, each increment $\Delta W$ is independently normally distributed with mean zero and variance $\Delta t$. It can be shown through the use of a Fokker-Planck equation [11], that the above dynamics converge to the stationary distribution

$$p(\mathit{\theta },\mathbf{v})\propto exp\left(-U\left(\mathit{\theta }\right)-\frac{{\mathbf{v}}^2}{2}\right).$$

We can use this to sample from the full data posterior by using a potential energy equal to the negative log-joint $U\left(\mathit{\theta }\right):=-logp(\mathcal{D},\mathit{\theta })$. The numeric integration is typically discretized [9,12] as follows

$$\begin{array}{cc}\Delta \mathit{\theta }\hfill & =\mathbf{v}\hfill \\ \Delta \mathbf{v}\hfill & =-\eta \nabla \widehat{U}\left(\mathit{\theta }\right)-\alpha \mathbf{v}+\mathbf{ϵ}\sqrt{2(\alpha -\widehat{\beta })\eta },\hfill \end{array}$$

(5)

where $\eta$ is called the learning rate, $1-\alpha$ is the momentum decay, $\mathbf{ϵ}$ is a standard Gaussian random vector, $\widehat{\beta }$ is an optional parameter to offset the variance of the stochastic gradient term and

$$\widehat{U}\left(\mathit{\theta }\right)=-\frac{\left|\mathcal{D}\right|}{\left|B\right|}\sum _{\mathbf{y}\in B}logp\left(\mathbf{y}\right|\mathit{\theta })-logp\left(\mathit{\theta }\right)$$

is an unbiased estimate of $U\left(\mathit{\theta }\right)$. A new mini-batch B is sampled for each iteration of Equation (5). Since $U\left(\mathit{\theta }\right)$ grows with the size of the data set, a small learning rate $\eta \sim \mathcal{O}\left(\frac{1}{\left|\mathcal{D}\right|}\right)$ is required to minimize the discretization error. The variance of the stochastic gradient term is proportional to ${\eta }^2$ while the variance of the injected noise is proportional to $\eta$, so in the limit $\eta \to 0$, stochasticity in the gradient estimates becomes negligible and the correct continuum dynamics are recovered, even if one ignores the errors from the stochastic gradient noise and $\widehat{\beta }=0$ is used. We refer the reader to [9,13,14] for an in-depth analysis of the algorithm parameters.

For the purposes of Bayesian updating, we use SGHMC to sample from the n^th posterior distribution, $p\left(\mathit{\theta }\right|{\mathbf{y}}_{\le n})$, using the stochastic potential energy

$${\widehat{U}}_n\left(\mathit{\theta }\right)=-\frac{n-1}{\left|B\right|}\sum _{\mathbf{y}\in B}logp\left(\mathbf{y}\right|\mathit{\theta })-logp\left({\mathbf{y}}_n\right|\mathit{\theta })-logp\left(\mathit{\theta }\right),$$

where the mini-batch is drawn i.i.d. with replacement from the set of all previous data points i.e., $B\subset \left\{{\mathbf{y}}_k\right|k<n\}$. The extra term here is to ensure that the previously unseen data point is always taken into account. If the extra term was not included, there would be some chance that the new data point does not get taken into account during the SGHMC steps. With this potential energy SGHMC still converges to the required posterior distribution, since it is still an unbiased estimator of the negative log-joint.

4. Experiments

We use mini-batches of size 500, with the following SGHMC parameters: $\eta =0.1/n$, $\alpha =0.2$ and $\widehat{\beta }=0$. Predictive distributions are approximated using $M=10$ samples and 20 burn-in steps for each new posterior. As mentioned in Section 2, rather than Bayesian updating by adding a single observation at a time, we add chunks of data at a time. In the following experiments we use chunks of 20 data points when $n\le 80$, chunks of size $⌊\frac{n}{4}⌋$ when $80<n<2000$, and chunks of size 500 thereafter. Our motivation for this is because we expect that smaller chunk sizes will give a lower variance in the estimator Equation (4) when n is small and so the posterior is less peaked, but for large n using larger chunk sizes is more efficient.

We use NS as our reference standard of accuracy. We implement NS with 20 SGHMC steps to sample from the constrained prior. For SGHMC used with NS we used parameters $\eta ={10}^{-3}$, $\alpha =0.1$ and $\widehat{\beta }=0$ because there is no gradient noise when sampling from the prior. Results reported are for two particles; more behave similarly but are slower. We allow NS to run until the additive terms are less than 1% of the current $\widehat{\mathcal{Z}}$ estimate. This is a popular stopping criterion and is also used in [8]. For more information on the constrained sampler and NS see Appendix A. Our experiments were run on a laptop with an Intel i7 CPU. For fair comparison, all code is single threaded. Multithreading gives a considerable speedup when calculating the likelihood on large data sets but can introduce subtle complexities that are difficult to control for in tests of runtime performance.

We evaluate our approach on the following three models using simulated data sets:

4.1. Linear Regression

The data set consists of pairs $(\mathbf{x},y)$ related by

$$y={\mathbf{w}}^T\mathbf{x}+b+ϵ,$$

where $ϵ$ is zero mean Gaussian distributed with known variance ${\sigma }^2$. We do not assume any distribution over $\mathbf{x}$ as it always appears on the right hand side of the conditional.

$$p\left(y\right|\mathbf{x},\mathbf{w},b)=\frac{1}{\sqrt{2\pi {\sigma }^2}}exp\left(-\frac{{(y-{\mathbf{w}}^T\mathbf{x}-b)}^2}{2{\sigma }^2}\right)$$

Parameters are $\mathbf{w}$ and b, with standard Gaussian priors. ML can be calculated analytically for this model. We used 5 dimensional vectors $\mathbf{x}$; this model has 6 parameters.

4.2. Logistic Regression

The data set consists of pairs $(\mathbf{x},y)$, where $\mathbf{x}$ is an observation vector which is assigned a class label $y\in \{1,\dots ,K\}$. The labels have a discrete distribution with probabilities given by

$$p\left(y\right|\mathbf{x},\mathit{\theta })=\frac{exp({\mathbf{w}}_y^T\mathbf{x}+b_y)}{{\sum }_{k}exp({\mathbf{w}}_k^T\mathbf{x}+b_k)}.$$

Parameters are $\mathit{\theta }=({\mathbf{w}}_{1:K},b_{1:K})$, with standard Gaussian priors. Again we do not assume any distribution over $\mathbf{x}$ as it always appears on the right hand side of the conditional. We used 10 dimensional vectors $\mathbf{x}$ with 4 classes; this model has 44 parameters.

4.3. Gaussian Mixture Model

The data are modeled by a mixture of multivariate Gaussian distributions with diagonal covariance matrices. Mixture weights, means and variances are treated as parameters. This type of model is often treated as a latent variable model, where the mixture component assignments of each data point are the latent variables. Here we marginalize out the latent variables to obtain the following conditional distribution:

$$p\left(\mathbf{y}\right|\mathit{\theta })=\sum _{k=1}^K{\beta }_k\prod _{j=1}^d\frac{1}{\sqrt{2\pi {\sigma }_{k,j}^2}}exp\left(-\frac{{(y_j-{\mu }_{k,j})}^2}{2{\sigma }_{k,j}^2}\right),$$

$$\mathit{\theta }=({\beta }_{1:K},{\mu }_{1:K,1:d},{\sigma }_{1:K,1:d}^2).$$

Mixture weights ${\beta }_{1:K}$ are modeled by a Dirichlet prior with $\mathbf{\alpha }=1$; means ${\mu }_{k,j}$ are modeled by Gaussian priors, centered around zero and with variance $4{\sigma }_{k,j}^2$; variances ${\sigma }_{k,j}^2$ are modeled by inverse gamma priors with shape and scale parameters equal to 1. We used 5 Gaussian components and observations were 2 dimensional; this model has 25 parameters with 24 degrees of freedom.

5. Results and Discussion

The log-ML typically grows linearly in the number of data points. For this reason, it is natural to measure errors in $\frac{log\mathcal{Z}}{N}$ rather than $log\mathcal{Z}$. In [8], the authors suggest that errors in $\frac{log\mathcal{Z}}{N}$ of $0.1$ are acceptable. For each model, we measured the runtime performance of nested sampling and our sequential estimator for various data set sizes up to one million data points. Due to computational constraints we run NS only for data set sizes at logarithmically increasing intervals, while the sequential sampler naturally produces many more intermediate results in a single run. In some of the figures below, the sequential sampler initially underestimates the log-ML. We believe this is due to higher variance of Equation (4) when n is small, and so we also give a hybrid result which replaces the initial terms in Equation (3) with a NS estimate of the ML for the first 100 data points.

5.1. Linear Regression

For the linear regression model, the exact ML is available analytically and is shown in Figure 1a for comparison. Both algorithms are able to produce accurate results for this model for all data set sizes. The final error of the sequential sampler on one million data points is only about ${10}^{-4}N$ (roughly 0.01%). For this model, our method was faster than NS by about a factor of 3 on one million observations.

5.2. Logistic Regression

For the largest data set, NS and our sequential sampler produced estimates which differed by $3\times {10}^{-4}N$ (roughly 0.7%), which is negligible. Our sequential sampler was almost a factor 17 faster than the nested sampler on one million observations for this model.

5.3. Gaussian Mixture Model

The posterior distribution for this model is multimodal. Some modes are due to permutation symmetries; these modes do not have to be explored since each one contains the same information. There are also some local modes which do not necessarily capture meaningful information about the data; for example, fitting a single Gaussian to the whole data set may be a poor local optimum of the likelihood function. If an MCMC walker finds one of these modes it can get trapped. However, we find that by Bayesian updating, the MCMC walkers tend to leave the poor local modes early on, before they become extremely peaked. This is similar to how annealing can help prevent MCMC and optimization algorithms from getting trapped in poor local optima. The estimates produced by NS and our sequential sampler differed on the largest data set by $2\times {10}^{-3}N$ (roughly 0.06%). For this model our sequential sampler was about a factor 11 faster than the nested sampler on one million observations.

In all the experiments our sequential sampler seems to converge to the same result as NS within a negligible error for large n. The initial disagreement between NS and our sequential sampler on the first few thousand data points, seen in Figure 2a and Figure 3a, seems as if it can be attributed to the initial terms in Equation (3), since the proposed hybrid approach, replacing early terms in the sequential sampler by estimates based on NS, matches NS closely for all data set sizes.

6. Materials and Methods

Code for this work was implemented using pytorch [15]. The code for our experiments is available at https://gitlab.com/pleased/sequential-evidence.

7. Conclusions

We found that our sequential sampler using stochastic gradients was able to produce accurate ML estimates with a speedup over NS. Furthermore, since the marginal cost of updating the ML estimates when new data arrives does not depend on the number of previous data points, this approach may be effective for weighted model averaging in a setting where one periodically gets access to new data, such as in streaming applications. Potential future work may involve a more in-depth analysis of the algorithm parameters, for example automatic tuning of the number of samples, M, used for approximating predictive probabilities. One can imagine further exploring the effects of stochastic gradients for ML calculation in the full SMC setting, including latent variable models and resampling steps, with more general stochastic gradient MCMC algorithms such as those in [13].