Quasi-Monte Carlo Quasi-Newton for variational Bayes

Many practical problems take the form

and $\boldsymbol{z}$ is random vector with a known distribution $p$. Classic problems of simulation-optimization (Andradóttir,, 1998) take this form and recently it has become very important in machine learning with variational Bayes (VB) (Blei et al.,, 2017) and Bayesian optimization (Frazier,, 2018) being prominent examples.

First order optimizers (Beck,, 2017) use a sequence of steps $\theta_{k+1}\leftarrow\theta_{k}-\alpha_{k}\nabla F(\theta_{k})$ for step sizes $\alpha_{k}>0$ and an operator $\nabla$ that we always take to be the gradient with respect to $\theta$. Very commonly, neither $F$ nor its gradient is available at reasonable computational cost and a Monte Carlo (MC) approximation is used instead. The update

for $\boldsymbol{z}_{i}\,{\buildrel\text{iid}\over{\sim}\,}p$ is a simple form of stochastic gradient descent (SGD) (Duchi,, 2018). There are many versions of SGD, notably AdaGrad (Duchi et al.,, 2011) and Adam (Kingma and Ba,, 2014) which are prominent in optimization of neural networks among other learning algorithms. Very often SGD involves sampling a mini-batch of observations. In this paper we consider SGD that samples instead some random quantities from a continuous distribution.

While a simple SGD is often very useful, there are settings where it can be improved. SGD is known to have slow convergence when the Hessian of $F$ is ill-conditioned (Bottou et al.,, 2018). When $\theta$ is not very high dimensional, the second order methods such as Newton iteration using exact or approximate Hessians can perform better. Quasi-Newton methods such as BFGS and L-BFGS (Nocedal and Wright,, 2006) that we describe in more details below can handle much higher dimensional parameters than Newton methods can while still improving upon first order methods. We note that quasi second-order methods cannot be proved to have better convergence rate than SGD. See Agarwal et al., (2012). They can however have a better implied constant.

A second difficulty with SGD is that MC sampling to estimate the gradient can be error prone or inefficient. For a survey of MC methods to estimate a gradient see Mohamed et al., (2020). Improvements based on variance reduction methods have been adopted to improve SGD. For instance Paisley et al., (2012) and Miller et al., (2017) both employ control variates in VB. Recently, randomized quasi-Monte Carlo (RQMC) methods, that we describe below, have been used in place of MC to improve upon SGD. Notable examples are Balandat et al., (2020) for Bayesian optimization and Buchholz et al., (2018) for VB. RQMC can greatly improve the accuracy with which integrals are estimated. The theory in Buchholz et al., (2018) shows how the improved integration accuracy from RQMC translates into faster optimization for SGD. The primary benefit is that RQMC can use a much smaller value of $n$. This is also seen empirically in Balandat et al., (2020) for Bayesian optimization.

Our contribution is to combine RQMC with a second order limited memory method known as L-BFGS. We show theoretically that improved integration leads to improved optimization. We show empirically for some VB examples that the optimization is improved. We find that RQMC regularly allows one to use fewer samples per iteration and in a crossed random effects example it found a better solution than we got with MC.

At step $k$ of our stochastic optimization we will use some number $n$ of sample values, $\boldsymbol{z}_{k,1},\dots,\boldsymbol{z}_{k,n}$. It is notationally convenient to group these all together into ${\mathcal{Z}}_{k}=(\boldsymbol{z}_{k,1},\dots,\boldsymbol{z}_{k,n})$. We also write

for gradient estimation at step $k$.

The closest works to ours are Balandat et al., (2020) and Buchholz et al., (2018). Like Balandat et al., (2020) we incorporate RQMC into L-BFGS. Our algorithm differs in that we take fresh RQMC samples at each iteration where they had a fixed sample of size $n$ that they used in a ‘common random numbers’ approach. They prove consistency as $n\to\infty$ for both MC and RQMC sampling. Their proof for RQMC required the recent strong law of large numbers for RQMC from Owen and Rudolf, (2021). Our analysis incorporates the geometric decay of estimation error as the number $K$ of iterations increases, similar to that used by Buchholz et al., (2018) for SGD. Our error bounds include sampling variances through which RQMC brings an advantage.

An outline of this paper is as follows. Section 2 reviews the optimization methods we need. Section 3 gives basic properties of scrambled net sampling, a form of RQMC. Section 4 presents our main theoretical findings. The optimality gap after $K$ steps of stochastic quasi-Newton is $F(\theta_{K})-F(\theta^{*})$ where $\theta^{*}$ is the optimal value. Theorem 4.1 bounds the expected optimality gap by a term that decays exponentially in $K$ plus a second term that is linear in a measure of sampling variance. RQMC greatly reduces that second non-exponentially decaying term which will often dominate. A similar bound holds for tail probabilities of the optimality gap. Theorem 4.2 obtains a comparable bound for ${\mathbb{E}}(\|\theta_{K}-\theta^{*}\|^{2})$. Section 5 gives numerical examples on some VB problems. In a linear regression problem where the optimal parameters are known, we verify that RQMC converges to them at an improved rate in $n$. In logistic regression, crossed random effects and variational autoencoder examples we see the second order methods greatly outperform SGD in terms of wall clock time to improve $F$. Since the true parameters are unknown we cannot compare accuracy of MC and RQMC sampling algorithms for those examples. In some examples RQMC finds a better ELBO than MC does when both use SGD, but the BFGS algorithms find yet better ELBOs. Section 6 gives some conclusions. The proofs of our main results are in an appendix.

We write $\nabla^{2}F(\theta)$ for the Hessian matrix of $F(\theta)$. The classic Newton update is

Under ideal circumstances it converges quadratically to the optimal parameter value $\theta^{*}$. That is $\|\theta_{k+1}-\theta^{*}\|=O(\|\theta_{k}-\theta^{*}\|^{2})$. Newton’s method is unsuitable for the problems we consider here because forming $\nabla^{2}F\in^{d\times d}$ may take too much space and solving the equation in (4) can cost $O(d^{3})$ which is prohibitively expensive. The quasi-Newton methods we consider do not form explicit Hessians. We also need to consider stochastic versions of them.

Scrambled nets are a form of RQMC sampling. We begin by briefly describing plain quasi-Monte Carlo (QMC) sampling. QMC is most easily described for computing expectations of $f(\boldsymbol{u})$ for $\boldsymbol{u}\sim\mathbf{U}[0,1]^{s}$. In our present context $f(\cdot)$ will be a component of $g(\cdot;\theta)$ and not the same as the $f$ in equation (1). In practice we must ordinarily transform the random variables $\boldsymbol{u}$ to some other distribution such as a Gaussian via $\boldsymbol{x}=\psi(\boldsymbol{u})$. We suppose that $\psi(\boldsymbol{u})\sim p$ for some transformation $\psi(\cdot)$. The text by Devroye, (1986) has many examples of such transformations. Here we subsume any such transformation $\psi(\cdot)$ into the definition of $f$.

In QMC sampling, we estimate $\mu=\int_{[0,1]^{s}}f(\boldsymbol{u})\mathrm{\,d}\boldsymbol{u}$ by $\hat{\mu}=(1/n)\sum_{i=1}^{n}f(\boldsymbol{u}_{i})$, just like in MC sampling except that distinct points $\boldsymbol{u}_{i}$ are chosen so that the discrete uniform distribution $\mathbf{U}\{\boldsymbol{u}_{1},\dots,\boldsymbol{u}_{n}\}$ is made very close to the continuous $\mathbf{U}[0,1]^{s}$ distribution. The difference between these two distributions can be quantified in many ways, called discrepancies (Chen et al.,, 2014). For a comprehensive treatment of QMC see Dick and Pillichshammer, (2010) or Niederreiter, (1992) or Dick et al., (2013).

When $f$ is of bounded variation in the sense of Hardy and Krause (BVHK) (see Owen, (2005)) then QMC attains the asymptotic error rate $|\hat{\mu}-\mu|=O(n^{-1}\log(n)^{s-1})=O(n^{-1+\epsilon})$ for any $\epsilon>0$. QMC is deterministic and to get practical error estimates RQMC methods were introduced. In RQMC, each individual $\boldsymbol{u}_{i}\sim\mathbf{U}[0,1]^{d}$ while collectively $\boldsymbol{u}_{1},\dots,\boldsymbol{u}_{n}$ still retain the low discrepancy property. Uniformity of $\boldsymbol{u}_{i}$ makes RQMC unbiased: ${\mathbb{E}}(\hat{\mu})=\mu$. Then if $f\in\mathrm{BVHK}$ we get an RMSE of $O(n^{-1+\epsilon})$. The whole RQMC process can then be replicated independently to quantify uncertainty. See Cranley and Patterson, (1976) and Owen, (1995) for methods and a survey in L’Ecuyer and Lemieux, (2002).

Scrambled net sampling (Owen,, 1995) is a form of RQMC that operates by randomly permuting the bits (more generally digits) of QMC methods called digital nets. The best known are those of Sobol’, (1969) and Faure, (1982). In addition to error estimation, scrambled nets give the user some control over the powers of $\log(n)$ in the QMC rate and also extend the domain of QMC from Riemann integrable functions of which BVHK is a subset to much more general functions including some with integrable singularities. For any integrand $f\in L^{2}[0,1]^{s}$, MC has RMSE $O(n^{-1/2})$ while scrambled nets have RMSE $o(n^{-1/2})$ without requiring $f\in\mathrm{BVHK}$ or even that $f$ is Riemann integrable (Owen, 1997a, ). For fixed $n$, each construction of scrambled nets has a ‘gain constant’ $\Gamma<\infty$ so that the RMSE is below $\sqrt{\Gamma}n^{-1/2}$ for any $f\in L^{2}[0,1]^{s}$. This effectively counters the powers of $\log(n)$. For smooth enough $f$, an error cancellation phenomenon for scrambled nets yields an RMSE of $O(n^{-3/2}\log(n)^{(s-1)/2})=O(n^{-3/2+\epsilon})$ (Owen, 1997b, ; Yue and Mao,, 1999; Owen,, 2008). The logarithmic powers here cannot ‘set in’ until they are small enough to obey the $\Gamma^{1/2}n^{-1/2}$ upper bound. Some forms of scrambled net sampling satisfy a central limit theorem (Loh,, 2003; Basu and Mukherjee,, 2017).

Very recently, a strong law of large numbers

has been proved for scrambled net sampling assuming only that $f\in L^{1+\delta}[0,1]^{s}$ for some $\delta>0$ (Owen and Rudolf,, 2021). The motivation for this result was that Balandat et al., (2020) needed a strong law of large numbers to prove consistency for their use for scrambled nets in Bayesian optimization.

Figure 1 graphically compares MC, QMC and RQMC points for $s=2$. The underlying QMC method is a Sobol’ sequence using ‘direction numbers’ from Joe and Kuo, (2008). We can see that MC points leave voids and create clumps. The QMC points are spread out more equally and show a strong diagonal structure. The RQMC points satisfy the same discrepancy bounds as the QMC points do but have broken up some of the structure.

In favorable settings the empirical behaviour of QMC and RQMC for realistic $n$ can be as good as their asymptotic rates. In less favorable settings RQMC can be like MC with some reduced variance. The favorable integrands are those where $f$ is nearly additive or at least dominated by sums of only a few of their inputs at a time. See Caflisch et al., (1997) for a definition of functions of ‘low effective dimension’ and Dick et al., (2013) for a survey of work on reproducing kernel Hilbert spaces of favorable integrands.

We propose to combine the stochastic quasi-Newton method with RQMC samples to create a randomized quasi-stochastic quasi-Newton (RQSQN) algorithm. At the $k$’th iteration, we draw an independently scrambled refreshing sample ${\mathcal{Z}}_{k}=(\boldsymbol{z}_{k,1},\dots,\boldsymbol{z}_{k,n_{g}})$ of size $n=n_{g}$ via RQMC to compute the gradient estimator $\bar{g}({\mathcal{Z}}_{k};\theta_{k})$. Then we find the descent direction $H_{k}\bar{g}({\mathcal{Z}}_{k};\theta_{k})$ using an L-BFGS two-loop recursion. Then we update the solution by

Here $\alpha_{k}$ may be found by line search with the Wolfe condition when using L-BFGS. See Chapter 3.1 in Nocedal and Wright, (2006).

Algorithm 1 shows pseudo-code for an RQMC version of SQN based on L-BFGS. It resembles the SQN algorithm in Byrd et al., (2016), except that the random samples are drawn by using RQMC instead of being sampled without replacement from a finite data set. Note that we don’t compute the Hessian directly. We either compute the Hessian-vector product $\nabla^{2}f({\mathcal{Z}}_{t};\bar{\theta}_{t})s_{t}$ or use gradient differences $\nabla f(\bar{\theta}_{t})-\nabla f(\bar{\theta}_{t-1})$.

In this section, we study the convergence rate of a general quasi-Newton iteration based on $n$ sample points $\boldsymbol{z}_{k1},\dots,\boldsymbol{z}_{kn}\sim p$ at stage $k$. The algorithm iterates as follows

Here $H_{k}$ is an approximate inverse Hessian. We let $F(\theta)={\mathbb{E}}(f(\boldsymbol{z}_{ki};\theta))$. We assume that the gradient estimator $g({\mathcal{Z}}_{k};\theta_{k})$ is unbiased conditionally on $\theta_{k}$, i.e., ${\mathbb{E}}(g({\mathcal{Z}}_{k};\theta_{k})\mid\theta_{k})=\nabla F(\theta_{k})$. The $\theta_{k}$ are random because they depend on ${\mathcal{Z}}_{k^{\prime}}$ for $k^{\prime}<k$.

The Hessian estimates $H_{k}$ for $k>1$ are also random because they depend directly on some additional inputs $\tilde{\mathcal{Z}}_{t}$ for those Hessian update epochs $t$ which occur prior to step $k$. They also depend on ${\mathcal{Z}}_{k^{\prime}}$ for $k^{\prime}<k$ because those random variables affect $\theta_{k^{\prime}}$. In our algorithms $H_{k}$ is independent of $\theta_{k}$. For RQMC this involves fresh rerandomizations of the underlying QMC points at step $k$. The alternative of ‘going deeper’ into a given RQMC sequence using points $(k-1)n+1$ through $kn$ would not satisfy independence. While it might possibly perform better it is harder to analyze and we saw little difference empirically in doing that. Our theory requires regularity of the problem as follows.

These are standard assumptions in the study of SGD. For example, see Assumptions 4.1 and 4.5 of Bottou et al., (2018). Strong convexity implies that there exists a unique minimizer $\theta^{*}$ to estimate. We write $F^{*}=F(\theta^{*})$, for the best possible value of $F$. We must have some smoothness and Lipshitz continuity is a mild assumption. The quantity $M$ will prove to be important below. We can get a better $M$ from RQMC than from MC. The other way to reduce $M$ is to increase $n$. When RQMC has an advantage it is because it gets a smaller $M$ for the same $n$, or to put it another way, it can get comparable $M$ with smaller $n$.

For two symmetric matrices $A$ and $B$, $A\preccurlyeq B$ means that $B-A$ is positive semi-definite. We denote the spectral norm of $A$ by $\|A\|$.

The following theorem states the convergence rate of $\|\theta_{k}-\theta^{*}\|$.

In this section we investigate quasi-Newton quasi-Monte Carlo optimization for some VB problems. Variational Bayes begins with a posterior distribution $p(\boldsymbol{z}\!\mid\!\boldsymbol{x})$ that is too inconvenient to work with. This usually means that we cannot readily sample $\boldsymbol{z}$. We turn instead to a distribution $q_{\theta}$ for $\theta\in\Theta$ from which we can easily sample $\boldsymbol{z}\sim q_{\theta}$. We now want to make a good choice of $\theta$ and in VB the optimal value $\theta^{*}$ is taken to be the minimizer of the KL divergence between distributions:

In this section we are using the symbol $p(\cdot)$ as a generic probability distribution. It is the distribution of whatever random variable’s symbol is inside the parentheses and not necessarily the same $p$ from the introduction. We use ${\mathbb{E}}_{\theta}(\cdot)$ to denote expectation with respect to $\boldsymbol{z}\sim q_{\theta}$.

We suppose that $\boldsymbol{z}$ has a prior distribution $p(\boldsymbol{z})$. Then Bayes rule gives

The last term does not depend on $\theta$ and so we may minimize $\mathcal{D}_{\mathrm{KL}}(\cdot\,\|\,\cdot)$ by maximizing

This ${\mathcal{L}}(\cdot)$ is known as the evidence lower bound (ELBO). The first term ${\mathbb{E}}_{\theta}({\log p(\boldsymbol{x}\!\mid\!\boldsymbol{z})})$ expresses a preference for $\theta$ having a large value of the likelihood of the observed data $\boldsymbol{x}$ given the latent data $\boldsymbol{z}$. The second term $-\mathcal{D}_{\mathrm{KL}}(q_{\theta}(\boldsymbol{z}\,\|\,p(\boldsymbol{z}))$ can be regarded as a regularization, penalizing parameter values for which $q_{\theta}(\boldsymbol{z})$ is too far from the prior distribution $p(\boldsymbol{z})$.

To optimize ${\mathcal{L}}(\theta)$ we need $\nabla{\mathcal{L}}(\theta)$. It is usual to choose a family $q_{\theta}$ for which $\mathcal{D}_{\mathrm{KL}}(\cdot\,\|\,\cdot)$ and its gradient are analytically tractable. We still need to estimate the gradient of the first term, i.e.,

One method is to use the score function $\nabla\log p_{\theta}(z)$ and Fisher’s identity

with $f(\cdot)=\log(p(\boldsymbol{x}\!\mid\!\cdot))$. Unfortunately an MC strategy based on this approach can suffer from large variance.

The most commonly used method is to write the parameter $\theta$ as a function of some underlying common random variables. This is known as the reparameterization trick (Kingma and Welling,, 2014). Suppose that there is a base distribution $p_{0}$ and a transformation $T(\cdot;\theta)$, such that if $\boldsymbol{z}\sim p_{0}$, then $T(\boldsymbol{z};\theta)\sim q_{\theta}$. Then

It is often easy to sample from the base distribution $p_{0}$, and thus to approximate the expectation by MC or RQMC samples. This is the method we use in our examples.

RQMC methods are finding uses in simulation optimization problems in machine learning, especially in first order SGD algorithms. We have looked at their use in a second order, L-BFGS algorithm. RQMC is known theoretically and empirically to improve the accuracy in integration problems compared to both MC and QMC. We have shown that improved estimation of expected gradients translates directly into improved optimization for quasi-Newton methods. There is a small burden in reprogramming algorithms to use RQMC instead of MC, but that is greatly mitigated by the appearance of RQMC algorithms in tools such as BoTorch (Balandat et al.,, 2020) and the forthcoming scipy 1.7 (scipy.stats.qmc.Sobol) and QMCPy at https://pypi.org/project/qmcpy/.

Our empirical examples have used VB. The approach has potential value in Bayesian optimization (Frazier,, 2018) and optimal transport (El Moselhy and Marzouk,, 2012; Bigoni et al.,, 2016) as well.

The examples we chose were of modest scale where both first and second order methods could be used. In these settings, we saw that second order methods improve upon first order ones. For the autoencoder problem the second order methods converged faster than the first order ones did. This also happenend for the crossed random effects problem where the second order methods found better ELBOs than the first order ones did and RQMC-based quasi-Newton algorithm found a better ELBO than the MC-based quasi-Newton did without increasing the wall clock time.

It is possible that RQMC will bring an advantage to conjugate gradient approaches as they have some similarities to L-BFGS. We have not investigated them.

This work was supported by the National Science Foundation under grant IIS-1837931.