An Efficient Minibatch Acceptance Test for Metropolis-Hastings

1 Introduction

Markov chain Monte Carlo (MCMC) sampling is a powerful method for computation on intractable distributions. We are interested in large dataset applications, where the goal is to sample a posterior distribution p(θ|x1, . . . , xN) of parameter θ for large N. The Metropolis-Hastings method (M-H) generates sample candidates from a proposal distribution q which is in general different from the target distribution p, and decides whether to accept or reject based on an acceptance test [Metropolis et al., 1953; Hastings, 1970].

Many state-of-the-art machine learning methods are based on minibatch updates (such as SGD) to a model. Such updates produce many improvements to the model for each pass over the dataset and have high sample efficiency. In contrast, conventional M-H requires calculations over the full dataset to produce a new sample. Results from [Korattikara et al., 2014; Bardenet et al., 2014] perform approximate (bounded error) acceptance tests using data subsets, but may still require close to the full data. By contrast, [Maclaurin and Adams, 2014] perform exact tests but require a lower bound on the parameter distribution across its domain, and such bounds are only available for relatively simple distributions.

Here we derive a new test which incorporates the variability in minibatch statistics as a natural part of the test and requires less data per iteration than prior work. We use a Barker test function [Barker, 1965], an idea which was suggested but not explored empirically in Section 6.3 of [Bardenet et al., 2017]. But the asymptotic test statistic CDF and the Barker function are different, which leads to fixed errors for the approach in [Bardenet et al., 2017]. Here, we show that the difference between the distributions can be corrected with an additive random variable. This leads to a test which is fast, and whose error can be made arbitrarily small.

∗This paper is an abridged version of a paper titled "An Efficient Minibatch Acceptance Test for Metropolis-Hastings" that won Honorable Mention for Best Student Paper award at UAI 2017.

We note that this approach is fundamentally different from prior work. It makes no assumptions about the form of, and requires no global bounds on the posterior parameter distribution. It is exact in the limit as batch size increases by the Central Limit Theorem. This is not true of [Korattikara et al., 2014; Bardenet et al., 2014] which use tail bounds and provide only approximate tests even with arbitrarily large batches; we instead use moment estimates from the data to determine how far the minibatch posteriors deviate from a normal distribution. These bounds carry through to the overall test accuracy.

Our test is applicable when the variance over data samples of the log probability ratio between the proposal and the current state is small enough (less than 1), a natural condition for models running M-H sampling with optimal proposals [Roberts and Rosenthal, 2001] on a full dataset. It succinctly captures the condition that the minibatch carries enough information to generate a sample. Though we cannot generally expect to get independent samples from the posterior using only a small subset of the data, we can still exploit minibatches by doing the following:

1. Increase the temperature K of the target distribution. Log likelihoods scale as 1/K, and so the variance of the likelihood ratio will vary as 1/K2 .
2. For continuous distributions, reduce the proposal step size (i.e. generate correlated samples), which is characteristic of Gibbs samplers on large datasets [Dupuy and Bach, 2016]. The variance of the log acceptance probability scales as the square of proposal step size.
3. Utilize Hamiltonian Dynamics [Neal, 2010] for proposals and tests. Here the dynamics itself provide shaping to the posterior distribution, and the M-H test is only needed to correct quantization error.

In sum, the primary contribution of our work is a novel, more efficient (in samples per test) minibatch acceptance test with quantifiable error bounds. The test uses a novel additive correction variable to implement a Barker test based on minibatch mean and variance. We demonstrate several orderof-magnitude improvements in sample efficiency, and that the batch size distribution is short-tailed.

2 Preliminaries

2.1 M-HBackground and RelatedWork

In the Metropolis-Hastings method [Gilks and Spiegelhalter, 1996; Brooks et al., 2011], a difficult-to-compute probability distribution p(θ) is sampled using a Markov chain θ1, . . . , θT . The sample θt+1 at time t+1 is generated using a candidate θ′ from a (simpler) proposal distribution q(θ′ |θt), filtered by an acceptance test. The acceptance test is usually a Metropolis test. The Metropolis test has acceptance probability:

where a ∧ b denotes min(a, b). With probability α(θt, θ′ ), we accept θ′ and set θt+1 = θ′ , otherwise set θt+1 = θt. The test is often implemented with an auxiliary random variable u ∼ U(0, 1) with a comparison u < α(θt, θ′ ); here, U(a, b) denotes the uniform distribution on the interval [a, b]. For simplicity, we drop the subscript t for the current sample θt and denote it as θ. For Bayesian inference, the target distribution is p(θ|x1, . . . , xN) with acceptance probability

where p0(θ) is the prior.

Computing samples this way requires all N data points which is intractable for large datasets. To address this challenge, [Korattikara et al., 2014; Bardenet et al., 2014] perform approximate M-H tests using sequential hypothesis testing; we refer to these as AUSTEREMH and MHSUBLHD. At each iteration, a subset of data is sampled and used to test whether to accept θ′ using an approximation to α(θ, θ′ ). If the approximate test does not yield a decision, the minibatch size increases and the test repeats. Both show useful computational savings but have two drawbacks: (i) they are approximate, and always yield a decision with a finite error, (ii) they require exact, dataset-wide bounds that depend on θ.

There is a separate line of MCMC work based on statistical physics using Langevin Dynamics [Wang and Uhlenbeck, 1945] and Hamiltonian Monte Carlo (HMC) [Neal, 2010; Betancourt, 2017] which can generate distant, high-quality proposals when run on full datasets. Recent work has attempted to combine these methods with SGD-like updates via SGLD [Welling and Teh, 2011; Ahn et al., 2012] and SGHMC [Chen et al., 2014]. Our method can be integrated with these to generate full posterior samples for almost the same cost as SGD.

2.2 Notation

Following [Bardenet et al., 2014], we write the test u < α(θ, θ′ ) equivalently as Λ(θ, θ′ ) > ψ(u, θ, θ′ ), where

and ψ(u, θ, θ′ ) = log (u q(θ′ |θ)p0(θ) q(θ|θ′ )p0(θ′ ) ) . To simplify notation, we assume that temperature K = 1 (saving T to indicate the number of samples to draw). Temperature appears as an exponential on each likelihood, p(xi|θ)1/K , so the effect would be to act as a 1/K factor on Λ(θ, θ′ ).

To reduce computational effort, an unbiased estimate of Λ(θ, θ′ ) based on a minibatch {x∗1 , . . . , x∗b } can be used:

Finally, define Λi(θ, θ′ ) = N log( p(xi|θ′ ) p(xi|θ) ), which are i.i.d. random variables. By the Central Limit Theorem, we expect Λ∗ (θ, θ′ ) to be approximately Normal. The acceptance test then becomes a test of the hypothesis that Λ(θ, θ′ ) > ψ(u, θ, θ′ ) by establishing that Λ∗ (θ, θ′ ) ≫ ψ(u, θ, θ′ ).

3 ANewMHAcceptance Test

For our new M-H test, we denote the exact and approximate log likelihood ratios as ∆ and ∆∗ , respectively. First, ∆ is

where p0, p, and q are the same as in Equation (2). We separate out terms dependent and independent of the data and then write our minibatch estimator ∆∗ for ∆ as

Note that ∆ and ∆∗ are evaluated on the full dataset and a minibatch of size b respectively. The term N/b means ∆∗ (θ, θ′ ) is an unbiased estimator of ∆(θ, θ′ ).

3.1 Barker (Logistic) Acceptance Function

We consider functions other than the classical Metropolis test that satisfy the detailed balance condition needed for accurate posterior estimation, and invoke the following [Barker, 1965]:

Lemma 1. If g(s) is any function such that g(s) = exp(s)g(-s), then the acceptance function α(θ, θ′ ) ≜ g(∆(θ, θ′ )) satisfies detailed balance.

For our test we use the logistic function g(s) = (1 + exp(-s))-1 , which satisfies Lemma 1. To understand this choice, a test function f(x) for M-H must satisfy Lemma 1. In addition, it must be monotone, bounded by [0, 1] and be such that limx→-∞f(x) = 0 and limx→∞f(x) = 1. The Logistic is the unique function in this class which is antisymmetric about 0.5, representing the CDF of a symmetric random variable. Our method requires approximating this with the sum of a Normal random variable (also symmetric) and a correction (Section 3.3). The Logistic CDF L and Normal CDF Φ are extremely close even without correction; the CDF error from the closest Normal CDF satisfies supx |L(x) - Φ(x/1.7)| < 0.01. With our correction we can make this error orders of magnitude smaller.

Assume we have current and candidate samples θ and θ′ , and have V ∼ U(0, 1). We accept θ′ if g(∆(θ, θ′ )) > V , and reject otherwise. Since g(s) is monotonically increasing, its inverse g-1 (s) is well-defined and unique. So equivalently, we accept θ′ if and only if ∆(θ, θ′ ) > X = g-1 (V ) where X is a random variable with the logistic function as its CDF. The density of X is symmetric, so we equivalently test

for a logistic random variable X.

3.2 AMinibatch Acceptance Test

We now describe acceptance testing using the minibatch estimator ∆∗ (θ, θ′ ). From Equation (6), ∆∗ (θ, θ′ ) can be represented as a constant term plus the mean of b IID terms Λi(θ, θ′ ) of the formN log p(x∗i |θ′ ) p(x∗i |θ) . As b increases, ∆∗ (θ, θ′ ) therefore has a distribution which approaches Normal by the CLT, with mean ∆. In the limit, we can write

where ¯N(0, σ2 (∆∗ )) denotes a distribution which is approximately normal with variance σ2 (∆∗ ). But to perform the test in Equation (7) we want ∆+X for a logistic random variable X (call it Xlog from now on). In [Bardenet et al., 2017] it was proposed to use ∆∗ in a Barker test, and tolerate the fixed error between the logistic and normal distributions. Our approach is to instead decompose Xlog as

where we assume Xnorm ∼ N(0, σ2 ) and that Xcorr is a zero-mean "correction" variable with density Cσ(X). The two variables are added (i.e., their distributions convolve) to formXlog. Using Xcorr ∼ Cσ(X), the acceptance test is now

Therefore, assuming the variance of ∆∗ is small enough, if we have an estimate of ∆∗ from the current data minibatch, we test acceptance by adding a random variable Xcorr and then accept θ′ if the result is positive (and reject otherwise).

If ¯N(0, σ2 (∆∗ )) is exactly N(0, σ2 (∆∗ )), the above test is exact, and as we show in Section 4, if there is a maximum error ϵ between the CDF of ¯N(0, σ2 (∆∗ )) and the CDF of N(0, σ2 (∆∗ )), then our test has an error of at most ϵ relative to the full batch version.

3.3 The Correction Distribution

Our test in Equation (10) requires knowing the distribution of Xcorr. In Section 4, we show that the test accuracy depends on the absolute error between the CDFs of Xnorm + Xcorr and Xlog. Consequently, we need to minimize this in our construction of Xcorr; we reduce the problem to a least squares minimization (see [Seita et al., 2017], Section 4).

4 Analysis

We now derive error bounds for our M-H test and the target distribution it generates. See [Seita et al., 2017] for proofs.

4.1 Bounding the Error Of ∆∗ FromNormal

We use the following from Corollary 2 in [Novak, 2005]:

Lemma 2. Let X1, . . . ,Xn be a set of zero-mean, independent, identically-distributed random variables with sample mean ¯X and sample variance s2X where:

Then the t-statistic t = ¯X/sX has a distribution which is approximately normal, with error bounded by:

Lemma 2 demonstrates that if we know E[|X|] and E[|X|3 ], we can bound the error of the normal approximation, which decays as O(n- 1 2 ). Making the change of variables y = xsX, Equation (12) becomes

showing that the distribution of ¯X approaches the normal distribution N(0, sX) whose standard deviation is sX, as measured from the sample. To apply this to our test, let Xi = Λi(θ, θ′ )-Λ(θ, θ′ ), so that the Xi are zero-mean, i.i.d. variables. If we only extract a subset of b samples from our minibatch, then ¯X = ∆∗ (θ, θ′ )-∆(θ, θ′ ), so that sX = s∆∗ . We can now substitute into Equation (13):

Corollary 1.

where the upper bound can be expressed as ϵ(θ, θ′ , b). Corollary 1 shows that the distribution of ∆∗ approximates a Normal distribution with mean ∆ and variance s2∆∗. Furthermore, it bounds the error with estimable quantities: both E[|X|] and E[|X|3 ] can be estimated as means of |Λi - Λ| and |Λi -Λ|3 , respectively, on each minibatch.

4.2 Adding RandomVariables

By applying Lemma 4 from [Seita et al., 2017], we have:

where S(x) is the standard logistic function, and Xnc and Xcorr are generated as per Algorithm 1.

Corollary 2 shows that the bounds from Section 4.1 are preserved after adding random variables, so our test remains accurate. Furthermore, bounds on the error of an M-H test imply bounds on the stationary distribution of the Markov chain under appropriate conditions [Korattikara et al., 2014].

5 Experiments

We compare with the most similar prior works [Korattikara et al., 2014] and [Bardenet et al., 2014] on Mixtures of Gaussians and Logistic Regression. For full experiment details, see [Seita et al., 2017]; here we provide a brief summary.

5.1 Mixture Of Gaussians

This model is adapted from [Welling and Teh, 2011] by increasing the number of samples to 1 million. The parameters are θ = 〈θ1, θ2〉, and the generation process is

We set σ2 1 = 10, σ2 2 = 1, σ2x = 2, and θ = 〈0, 1〉. We are interested in large-scale problems and sampled 1,000,000 points to form posterior ∝ p(θ) ∏1,000,000 i=1 p(xi|θ)1/K with the same prior from Equation (15). This produces a much sharper posterior with two very narrow peaks. To reproduce the original posterior, we set temperature K = 10, 000.

We benchmark with AUSTEREMH(C) and MHSUBLHD. All methods collect 3000 samples using a random walk proposer with covariance matrix diag(0.15, 0.15), which means the M-H test is responsible for shaping the sample distribution. Section 6.1 in [Seita et al., 2017] indicates that there are no obvious differences in the posterior samples. Table 1 shows that MHMINIBATCH dominates in terms of efficiency with about a 100x improvement over alternative methods.

5.2 Logistic Regression

We test logistic regression for binary classification of 1s vs. 7s on MNIST [LeCun and Cortes, 1998] and a subset of infinite MNIST [Loosli et al., 2007]. For the former, extracting all 1s and 7s resulted in 13,000 training samples ("MNIST13k"), and for the latter, we used 87,000 additional (augmented) 1s and 7s to get 100,000 total ("MNIST-100k").

Both datasets use the same test set, with 2,163 samples. For all methods, we impose a uniform prior and again use a random walk proposer with covariance 0.05I for MNIST-13k and 0.01I for MNIST-100k. The temperature is K = 100.

The first two subplots of Figure 1 display the prediction accuracy on both datasets for all methods as a function of the cumulative training points processed (note that the methods consume different amounts of data). To generate the curves, for each sample θt, t ∈ {1, . . . , T}, we use θt as the logistic regression parameter. The results suggest that our test obtains convergence more than 10x faster than the others.

Table 2 contains the average of the average minibatch sizes (± one standard deviation) across all trials; see [Seita et al., 2017] for detailed histograms and the right subplot of Figure 1 for a representative result. MHMINIBATCH, with sizes of 125.4 and 216.5 for MNIST-13k and MNIST-100k, respectively, consumes more than 7x and 4x fewer data points than the next-best method, AUSTEREMH(NC). We note that both AUSTEREMH(NC) and MHSUBLHD require computing log p(xi|θ) and log p(xi|θ′ ) for all xi each iteration.

6 Conclusions

We have derived a minibatch M-H which approximates full data tests and presented promising theoretical and empirical results. A priority is to integrate the test with HMC and variants [Homan and Gelman, 2014; Tripuraneni et al., 2017]. Another idea is to merge our test with [Korattikara et al., 2014] by applying both each iteration and utilizing variance reduction techniques [Chen and Ghahramani, 2016].

Acknowledgments

We thank the reviewers who provided helpful comments. Daniel Seita is supported by an NPSC fellowship.