A Complete Recipe for Stochastic Gradient MCMC

Introduction

Markov chain Monte Carlo (MCMC) has become a defacto tool for Bayesian posterior inference. However, these methods notoriously mix slowly in complex, high-dimensional models and scale poorly to large datasets. The past decades have seen a rise in MCMC methods that provide more efficient exploration of the posterior, such as Hamiltonian Monte Carlo (HMC) and its Reimann manifold variant . This class of samplers is based on defining a potential energy function in terms of the target posterior distribution and then devising various continuous dynamics to explore the energy landscape, enabling proposals of distant states. The gain in efficiency of exploration often comes at the cost of a significant computational burden in large datasets.

Recently, stochastic gradient variants of such continuous-dynamic samplers have proven quite useful in scaling the methods to large datasets . At each iteration, these samplers use data subsamples—or minibatches—rather than the full dataset. Stochastic gradient Langevin dynamics (SGLD) innovated in this area by connecting stochastic optimization with a first-order Langevin dynamic MCMC technique, showing that adding the “right amount” of noise to stochastic gradient ascent iterates leads to samples from the target posterior as the step size is annealed. Stochastic gradient Hamiltonian Monte Carlo (SGHMC) builds on this idea, but importantly incorporates the efficient exploration provided by the HMC momentum term. A key insight in that paper was that the naïve stochastic gradient variant of HMC actually leads to an incorrect stationary distribution (also see ); instead a modification to the dynamics underlying HMC is needed to account for the stochastic gradient noise. Variants of both SGLD and SGHMC with further modifications to improve efficiency have also recently been proposed .

In the plethora of past MCMC methods that explicitly leverage continuous dynamics—including HMC, Riemann manifold HMC, and the stochastic gradient methods—the focus has been on showing that the intricate dynamics leave the target posterior distribution invariant. Innovating in this arena requires constructing novel dynamics and simultaneously ensuring that the target distribution is the stationary distribution. This can be quite challenging, and often requires significant physical and geometrical intuition . A natural question, then, is whether there exists a general recipe for devising such continuous-dynamic MCMC methods that naturally lead to invariance of the target distribution. In this paper, we answer this question to the affirmative. Furthermore, and quite importantly, our proposed recipe is complete. That is, any continuous Markov process (with no jumps) with the desired invariant distribution can be cast within our framework, including HMC, Riemann manifold HMC, SGLD, SGHMC, their recent variants, and any future developments in this area. That is, our method provides a unifying framework of past algorithms, as well as a practical tool for devising new samplers and testing the correctness of proposed samplers.

The recipe involves defining a (stochastic) system parameterized by two matrices: a positive semidefinite diffusion matrix, $\mathbf{D}(\mathbf{z})$, and a skew-symmetric curl matrix, $\mathbf{Q}(\mathbf{z})$, where $\mathbf{z}=(\theta,r)$ with $\theta$ our model parameters of interest and $r$ a set of auxiliary variables. The dynamics are then written explicitly in terms of the target stationary distribution and these two matrices. By varying the choices of $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$, we explore the space of MCMC methods that maintain the correct invariant distribution. We constructively prove the completeness of this framework by converting a general continuous Markov process into the proposed dynamic structure.

For any given $\mathbf{D}(\mathbf{z})$, $\mathbf{Q}(\mathbf{z})$, and target distribution, we provide practical algorithms for implementing either full-data or minibatch-based variants of the sampler. In Sec. 3.1, we cast many previous continuous-dynamic samplers in our framework, finding their $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$. We then show how these existing $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$ building blocks can be used to devise new samplers; we leave the question of exploring the space of $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$ well-suited to the structure of the target distribution as an interesting direction for future research. In Sec. 3.2 we demonstrate our ability to construct new and relevant samplers by proposing stochastic gradient Riemann Hamiltonian Monte Carlo, the existence of which was previously only speculated. We demonstrate the utility of this sampler on synthetic data and in a streaming Wikipedia analysis using latent Dirichlet allocation .

A Complete Stochastic Gradient MCMC Framework

Marginalizing the auxiliary variables gives us the desired distribution on $\theta$. In this paper, we generically consider $\mathbf{z}$ as the samples we seek to draw; $\mathbf{z}$ could represent $\theta$ itself, or an augmented state space in which case we simply discard the auxiliary variables to perform the desired marginalization.

As in HMC, the idea is to translate the task of sampling from the posterior distribution to simulating from a continuous dynamical system which is used to define a Markov transition kernel. That is, over any interval $h$, the differential equation defines a mapping from the state at time $t$ to the state at time $t+h$. One can then discuss the evolution of the distribution $p(\mathbf{z},t)$ under the dynamics, as characterized by the Fokker-Planck equation for stochastic dynamics or the Liouville equation for deterministic dynamics . This evolution can be used to analyze the invariant distribution of the dynamics, $p^{s}(\mathbf{z})$. When considering deterministic dynamics, as in HMC, a jump process must be added to ensure ergodicity. If the resulting stationary distribution is equal to the target posterior, then simulating from the process can be equated with drawing samples from the posterior.

If the stationary distribution is not the target distribution, a Metropolis-Hastings (MH) correction can often be applied. Unfortunately, such correction steps require a costly computation on the entire dataset. Even if one can compute the MH correction, if the dynamics do not nearly lead to the correct stationary distribution, then the rejection rate can be high even for short simulation periods $h$. Furthermore, for many stochastic gradient MCMC samplers, computing the probability of the reverse path is infeasible, obviating the use of MH. As such, a focus in the literature is on defining dynamics with the right target distribution, especially in large-data scenarios where MH corrections are computationally burdensome or infeasible.

Generically, all continuous Markov processes that one might consider for sampling can be written as a stochastic differential equation (SDE) of the form:

where $\mathbf{f}(\mathbf{z})$ denotes the deterministic drift and often relates to the gradient of $H(\mathbf{z})$, ${\mathbf{W}(t)}$ is a $d$-dimensional Wiener process, and $\mathbf{D}(\mathbf{z})$ is a positive semidefinite diffusion matrix. Clearly, however, not all choices of $\mathbf{f}(\mathbf{z})$ and $\mathbf{D}(\mathbf{z})$ yield the stationary distribution $p^{s}(\mathbf{z})\propto\exp(-H(\mathbf{z}))$.

When $\mathbf{D}(\mathbf{z})=0$, as in HMC, the dynamics of Eq. (2) become deterministic. Our exposition focuses on SDEs, but our analysis applies to deterministic dynamics as well. In this case, our framework—using the Liouville equation in place of Fokker-Planck—ensures that the deterministic dynamics leave the target distribution invariant. For ergodicity, a jump process must be added, which is not considered in our recipe, but tends to be straightforward (e.g., momentum resampling in HMC).

To devise a recipe for constructing SDEs with the correct stationary distribution, we propose writing $\mathbf{f}(\mathbf{z})$ directly in terms of the target distribution:

Here, $\mathbf{Q}(\mathbf{z})$ is a skew-symmetric curl matrix representing the deterministic traversing effects seen in HMC procedures. In contrast, the diffusion matrix $\mathbf{D}(\mathbf{z})$ determines the strength of the Wiener-process-driven diffusion. Matrices $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$ can be adjusted to attain faster convergence to the posterior distribution. A more detailed discussion on the interpretation of $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$ and the influence of specific choices of these matrices is provided in the Supplement.

Importantly, as we show in Theorem 1, sampling the stochastic dynamics of Eq. (2) (according to Itô integral) with $\mathbf{f}(\mathbf{z})$ as in Eq. (3) leads to the desired posterior distribution as the stationary distribution: $p^{s}(\mathbf{z})\propto\exp(-H(\mathbf{z}))$. That is, for any choice of positive semidefinite $\mathbf{D}(\mathbf{z})$ and skew-symmetric $\mathbf{Q}(\mathbf{z})$ parameterizing $\mathbf{f}(\mathbf{z})$, we know that simulating from Eq. (2) will provide samples from $p(\theta\mid\mathcal{S})$ (discarding any sampled auxiliary variables $r$) assuming the process is ergodic.

$p^{s}(\mathbf{z})\propto\exp(-H(\mathbf{z}))$ is a stationary distribution of the dynamics of Eq. (2) if $\mathbf{f}(\mathbf{z})$ is restricted to the form of Eq. (3), with $\mathbf{D}(\mathbf{z})$ positive semidefinite and $\mathbf{Q}(\mathbf{z})$ skew-symmetric. If $\mathbf{D}(\mathbf{z})$ is positive definite, or if ergodicity can be shown, then the stationary distribution is unique.

The equivalence of $p^{s}(\mathbf{z})$ and the target $p(\mathbf{z}\mid\mathcal{S})\propto\exp(-H(\mathbf{z}))$ can be shown using the Fokker-Planck description of the probability density evolution under the dynamics of Eq. (2) :

Eq. (4) can be further transformed into a more compact form :

We can verify that $p(\mathbf{z}\mid\mathcal{S})$ is invariant under Eq. (5) by calculating $\left[e^{-H(\mathbf{z})}\nabla H(\mathbf{z})+\nabla e^{-H(\mathbf{z})}\right]=0$. If the process is ergodic, this invariant distribution is unique. The equivalence of the compact form was originally proved in ; we include a detailed proof in the Supplement for completeness. ∎

2 Completeness of the Framework

An important question is what portion of samplers defined by continuous Markov processes with the target invariant distribution can we define by iterating over all possible $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$? In Theorem 2, we show that for any continuous Markov process with the desired stationary distribution, $p^{s}(\mathbf{z})$, there exists an SDE as in Eq. (2) with $\mathbf{f}(\mathbf{z})$ defined as in Eq. (3). We know from the Chapman-Kolmogorov equation that any continuous Markov process with stationary distribution $p^{s}(\mathbf{z})$ can be written as in Eq. (2), which gives us the diffusion matrix $\mathbf{D}(\mathbf{z})$. Theorem 2 then constructively defines the curl matrix $\mathbf{Q}(\mathbf{z})$. This result implies that our recipe is complete. That is, we cover all possible continuous Markov process samplers in our framework. See Fig. 1.

For the SDE of Eq. (2), suppose its stationary probability density function $p^{s}(\mathbf{z})$ uniquely exists, and that \left[\mathbf{f}_{i}(\mathbf{z})p^{s}(\mathbf{z})-\sum_{j=1}^{d}\dfrac{\partial}{\partial\theta_{j}}\Big{(}{\mathbf{D}}_{ij}(\mathbf{z})p^{s}(\mathbf{z})\Big{)}\right] is integrable with respect to the Lebesgue measure, then there exists a skew-symmetric $\mathbf{Q}(\mathbf{z})$ such that Eq. (3) holds.

The integrability condition is usually satisfied when the probability density function uniquely exists. A constructive proof for the existence of ${\mathbf{Q}}(\mathbf{z})$ is provided in the Supplement.

3 A Practical Algorithm

In practice, simulation relies on an $\epsilon$-discretization of the SDE, leading to a full-data update rule

Calculating the gradient of $H(\mathbf{z})$ involves evaluating the gradient of $U(\theta)$. For a stochastic gradient method, the assumption is that $U(\theta)$ is too computationally intensive to compute as it relies on a sum over all data points (see Sec. 2). Instead, such stochastic gradient algorithms examine independently sampled data subsets $\widetilde{\mathcal{S}}\subset\mathcal{S}$ and the corresponding potential for these data:

The specific form of Eq. (7) implies that $\widetilde{U}(\theta)$ is an unbiased estimator of $U(\theta)$. As such, a gradient computed based on $\widetilde{U}(\theta)$—called a stochastic gradient —is a noisy, but unbiased estimator of the full-data gradient. The key question in many of the existing stochastic gradient MCMC algorithms is whether the noise injected by the stochastic gradient adversely affects the stationary distribution of the modified dynamics (using $\nabla\widetilde{U}(\theta)$ in place of $\nabla U(\theta)$). One way to analyze the impact of the stochastic gradient is to make use of the central limit theorem and assume

resulting in a noisy Hamiltonian gradient $\nabla\widetilde{H}(\mathbf{z})=\nabla H(\mathbf{z})+[\mathcal{N}(0,\mathbf{V}(\theta)),\mathbf{0}]^{T}$. Simply plugging in $\nabla\widetilde{H}(\mathbf{z})$ in place of $\nabla H(\mathbf{z})$ in Eq. (6) results in dynamics with an additional noise term ({\mathbf{D}}(\mathbf{z}_{t})+{\mathbf{Q}}(\mathbf{z}_{t})\big{)}[\mathcal{N}(0,\mathbf{V}(\theta)),\mathbf{0}]^{T}. To counteract this, assume we have an estimate $\hat{\mathbf{B}}_{t}$ of the variance of this additional noise satisfying $2\mathbf{D}(\mathbf{z}_{t})-\epsilon_{t}\hat{\mathbf{B}}_{t}\succeq 0$ (i.e., positive semidefinite). With small $\epsilon$, this is always true since the stochastic gradient noise scales down faster than the added noise. Then, we can attempt to account for the stochastic gradient noise by simulating

This provides our stochastic gradient—or minibatch— variant of the sampler. In Eq. (9), the noise introduced by the stochastic gradient is multiplied by $\epsilon_{t}$ (and the compensation by $\epsilon_{t}^{2}$), implying that the discrepancy between these dynamics and those of Eq. (6) approaches zero as $\epsilon_{t}$ goes to zero. As such, in this infinitesimal step size limit, since Eq. (6) yields the correct invariant distribution, so does Eq. (9). This avoids the need for a costly or potentially intractable MH correction. However, having to decrease $\epsilon_{t}$ to zero comes at the cost of increasingly small updates. We can also use a finite, small step size in practice, resulting in a biased (but faster) sampler. A similar bias-speed tradeoff was used in to construct MH samplers, in addition to being used in SGLD and SGHMC.

Applying the Theory to Construct Samplers

We explicitly state how some recently developed MCMC methods fall within the proposed framework based on specific choices of $\mathbf{D}(\mathbf{z})$, $\mathbf{Q}(\mathbf{z})$ and $H(\mathbf{z})$ in Eq. (2) and (3). For the stochastic gradient methods, we show how our framework can be used to “reinvent” the samplers by guiding their construction and avoiding potential mistakes or inefficiencies caused by naïve implementations.

The key ingredient in HMC is Hamiltonian dynamics, which simulate the physical motion of an object with position $\theta$, momentum $r$, and mass $\mathbf{M}$ on an frictionless surface as follows (typically, a leapfrog simulation is used instead):

Eq. (12) is a special case of the proposed framework with $\mathbf{z}=(\theta,r)$, $H(\theta,r)=U(\theta)+\frac{1}{2}r^{T}M^{-1}r$, {\mathbf{Q}}(\theta,r)=\left(\begin{array}[]{ll}0&-\mathbf{I}\\ \mathbf{I}&0\end{array}\right) and $\mathbf{D}(\theta,r)=\mathbf{0}$.

As discussed in , simply replacing $\nabla U(\theta)$ by the stochastic gradient $\nabla\widetilde{U}(\theta)$ in Eq. (12) results in the following updates:

where the $\approx$ arises from the approximation of Eq. (8). Careful study shows that Eq. (15) cannot be rewritten into our proposed framework, which hints that such a naïve stochastic gradient version of HMC is not correct. Interestingly, the authors of proved that this naïve version indeed does not have the correct stationary distribution. In our framework, we see that the noise term $\mathcal{N}(0,2\epsilon_{t}\mathbf{D}(\mathbf{z}))$ is paired with a $\mathbf{D}(\mathbf{z})\nabla H(\mathbf{z})$ term, hinting that such a term should be added to Eq. (15). Here, {\mathbf{D}}(\theta,r)=\left(\begin{array}[]{ll}0&0\\ 0&\epsilon\mathbf{V}(\theta)\end{array}\right), which means we need to add $\mathbf{D}(\mathbf{z})\nabla H(\mathbf{z})=\epsilon\mathbf{V}(\theta)\nabla_{r}H(\theta,r)=\epsilon\mathbf{V}(\theta)\mathbf{M}^{-1}r$. Interestingly, this is the correction strategy proposed in , but through a physical interpretation of the dynamics. In particular, the term $\epsilon\mathbf{V}(\theta)\mathbf{M}^{-1}r$ (or, generically, $\mathbf{C}\mathbf{M}^{-1}r$ where $\mathbf{C}\succeq\epsilon\mathbf{V}(\theta)$) has an interpretation as friction and leads to second order Langevin dynamics:

Here, $\hat{\mathbf{B}}_{t}$ is an estimate of $\mathbf{V}(\theta_{t})$. This method now fits into our framework with $H(\theta,r)$ and $\mathbf{Q}(\theta,r)$ as in HMC, but with {\mathbf{D}}(\theta,r)=\left(\begin{array}[]{ll}0&0\\ 0&\mathbf{C}\end{array}\right). This example shows how our theory can be used to identify invalid samplers and provide guidance on how to effortlessly correct the mistakes; this is crucial when physical intuition is not available. Once the proposed sampler is cast in our framework with a specific $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$, there is no need for sampler-specific proofs, such as those of .

SGLD proposes to use the following first order (no momentum) Langevin dynamics to generate samples

This algorithm corresponds to taking $\mathbf{z}=\theta$ with $H(\theta)=U(\theta)$, ${\mathbf{D}}(\theta)=\mathbf{D}$, ${\mathbf{Q}}(\theta)=\mathbf{0}$, and $\hat{\mathbf{B}}_{t}=\mathbf{0}$. As motivated by Eq. (9) of our framework, the variance of the stochastic gradient can be subtracted from the sampler injected noise to make the finite stepsize simulation more accurate. This variant of SGLD leads to the stochastic gradient Fisher scoring algorithm .

SGLD can be generalized to use an adaptive diffusion matrix ${\mathbf{D}}(\theta)$. Specifically, it is interesting to take ${\mathbf{D}}(\theta)=\mathbf{G}^{-1}(\theta)$, where $\mathbf{G}(\theta)$ is the Fisher information metric. The sampler dynamics are given by

Taking $\mathbf{D}(\theta)={\mathbf{G}}(\theta)^{-1}$, ${\mathbf{Q}}(\theta)=\mathbf{0}$, and $\hat{\mathbf{B}}_{t}=\mathbf{0}$, this SGRLD method falls into our framework with correction term $\Gamma_{i}(\theta)=\sum\limits_{j}\dfrac{\partial{\mathbf{D}}_{ij}(\theta)}{\partial\theta_{j}}$. It is interesting to note that in earlier literature , $\Gamma_{i}(\theta)$ was taken to be $2\ |\mathbf{G}(\theta)|^{-1/2}\sum\limits_{j}\dfrac{\partial}{\partial\theta_{j}}\left(\mathbf{G}_{ij}^{-1}(\theta)|\mathbf{G}(\theta)|^{1/2}\right)$. More recently, it was found that this correction term corresponds to the distribution function with respect to a non-Lebesgue measure ; for the Lebesgue measure, the revised $\Gamma_{i}(\theta)$ was as determined by our framework . Again, we have an example of our theory providing guidance in devising correct samplers.

Finally, the SGNHT method incorporates ideas from thermodynamics to further increase adaptivity by augmenting the SGHMC system with an additional scalar auxiliary variable, $\xi$. The algorithm uses the following dynamics:

We can take $\mathbf{z}=(\theta,r,\xi)$, $H(\theta,r,\xi)=U(\theta)+\dfrac{1}{2}r^{T}r+\dfrac{1}{2d}(\xi-A)^{2}$, {\mathbf{D}}(\theta,r,\xi)=\left(\begin{array}[]{ccc}0&0&0\\ 0&A\cdot\mathbf{I}&0\\ 0&0&0\end{array}\right), and {\mathbf{Q}}(\theta,r,\xi)=\left(\begin{array}[]{ccc}0&-\mathbf{I}&0\\ \mathbf{I}&0&r/d\\ 0&-r^{T}/d&0\end{array}\right) to place these dynamics within our framework.

In our framework, SGLD and SGRLD take $\mathbf{Q}(\mathbf{z})=0$ and instead stress the design of the diffusion matrix $\mathbf{D}(\mathbf{z})$, with SGLD using a constant $\mathbf{D}(\mathbf{z})$ and SGRLD an adaptive, $\theta$-dependent diffusion matrix to better account for the geometry of the space being explored. On the other hand, HMC takes $\mathbf{D}(\mathbf{z})=0$ and focuses on the curl matrix $\mathbf{Q}(\mathbf{z})$. SGHMC combines SGLD with HMC through non-zero $\mathbf{D}(\theta)$ and $\mathbf{Q}(\theta)$ matrices. SGNHT then extends SGHMC by taking $\mathbf{Q}(\mathbf{z})$ to be state dependent. The relationships between these methods are depicted in the Supplement, which likewise contains a discussion of the tradeoffs between these two matrices. In short, $\mathbf{D}(\mathbf{z})$ can guide escaping from local modes while $\mathbf{Q}(\mathbf{z})$ can enable rapid traversing of low-probability regions, especially when state adaptation is incorporated. We readily see that most of the product space $\mathbf{D}(\mathbf{z})\times\mathbf{Q}(\mathbf{z})$, defining the space of all possible samplers, has yet to be filled.

2 Stochastic Gradient Riemann Hamiltonian Monte Carlo

In Sec. 3.1, we have shown how our framework unifies existing samplers. In this section, we now use our framework to guide the development of a new sampler. While SGHMC inherits the momentum term of HMC, making it easier to traverse the space of parameters, the underlying geometry of the target distribution is still not utilized. Such information can usually be represented by the Fisher information metric , denoted as $\mathbf{G}(\theta)$, which can be used to precondition the dynamics. For our proposed system, we consider $H(\theta,r)=U(\theta)+\frac{1}{2}r^{T}r$, as in HMC/SGHMC methods, and modify the $\mathbf{D}(\theta,r)$ and $\mathbf{Q}(\theta,r)$ of SGHMC to account for the geometry as follows:

We refer to this algorithm as stochastic gradient Riemann Hamiltonian Monte Carlo (SGRHMC). Our theory holds for any positive definite $\mathbf{G}(\theta)$, yielding a generalized SGRHMC (gSGRHMC) algorithm, which can be helpful when the Fisher information metric is hard to compute.

A naïve implementation of a state-dependent SGHMC algorithm might simply (i) precondition the HMC update, (ii) replace $\nabla U(\theta)$ by $\nabla\widetilde{U}(\theta)$, and (iii) add a state-dependent friction term on the order of the diffusion matrix to counterbalance the noise as in SGHMC, resulting in:

However, as we show in Sec. 4.1, samples from these dynamics do not converge to the desired distribution. Indeed, this system cannot be written within our framework. Instead, we can simply follow our framework and, as indicated by Eq. (9), consider the following update rule:

which includes a correction term $\nabla_{\theta}\left({\mathbf{G}}(\theta)^{-1/2}\right)$, with $i$-th component $\sum\limits_{j}\dfrac{\partial}{\partial\theta_{j}}\left({\mathbf{G}}(\theta)^{-1/2}\right)_{ij}$. The practical implementation of gSGRHMC is outlined in Algorithm 1.

Experiments

In Sec. 4.1, we show that gSGRHMC can excel at rapidly exploring distributions with complex landscapes. We then apply SGRHMC to sampling in a latent Dirichlet allocation (LDA) model on a large Wikipedia dataset in Sec. 4.2. The Supplement contains details on the specific samplers considered and the parameter settings used in these experiments.

In this section we aim to empirically (i) validate the correctness of our recipe and (ii) assess the effectiveness of gSGRHMC. In Fig. 2(left), we consider two univariate distributions (shown in the Supplement) and compare SGLD, SGHMC, the naïve state-adaptive SGHMC of Eq. (27), and our proposed gSGRHMC of Eq. (30). See the Supplement for the form of $\mathbf{G}(\theta)$. As expected, the naïve implementation does not converge to the target distribution. In contrast, the gSGRHMC algorithm obtained via our recipe indeed has the correct invariant distribution and efficiently explores the distributions. In the second experiment, we sample a bivariate distribution with strong correlation. The results are shown in Fig. 2(right). The comparison between SGLD, SGHMC, and our gSGRHMC method shows that both a state-dependent preconditioner and Hamiltonian dynamics help to make the sampler more efficient than either element on its own.

2 Online Latent Dirichlet Allocation

We also applied SGRHMC (with ${\mathbf{G}}(\theta)={\rm diag}(\theta)^{-1}$, the Fisher information metric) to an online latent Dirichlet allocation (LDA) analysis of topics present in Wikipedia entries. In LDA, each topic is associated with a distribution over words, with $\beta_{kw}$ the probability of word $w$ under topic $k$. Each document is comprised of a mixture of topics, with $\pi^{(d)}_{k}$ the probability of topic $k$ in document $d$. Documents are generated by first selecting a topic $z_{j}^{(d)}\sim\pi^{(d)}$ for the $j$th word and then drawing the specific word from the topic as $x_{j}^{(d)}\sim\beta_{z_{j}^{(d)}}$. Typically, $\pi^{(d)}$ and $\beta_{k}$ are given Dirichlet priors.

The goal of our analysis here is inference of the corpus-wide topic distributions $\beta_{k}$. Since the Wikipedia dataset is large and continually growing with new articles, it is not practical to carry out this task over the whole dataset. Instead, we scrape the corpus from Wikipedia in a streaming manner and sample parameters based on minibatches of data. Following the approach in , we first analytically marginalize the document distributions $\pi^{(d)}$ and, to resolve the boundary issue posed by the Dirichlet posterior of $\beta_{k}$ defined on the probability simplex, use an expanded mean parameterization shown in Figure 3(upper left). Under this parameterization, we then compute $\nabla\log p(\theta|\mathbf{x})$ and, in our implementation, use boundary reflection to ensure the positivity of parameters $\theta_{kw}$. The necessary expectation over word-specific topic indicators $z_{j}^{(d)}$ is approximated using Gibbs sampling separately on each document, as in . The Supplement contains further details.

For all the methods, we report results of three random runs. When sampling distributions with mass concentrated over small regions, as in this application, it is important to incorporate geometric information via a Riemannian sampler . The results in Fig. 3(right) indeed demonstrate the importance of Riemannian variants of the stochastic gradient samplers. However, there also appears to be some benefits gained from the incorporation of the HMC term for both the Riemmannian and non-Reimannian samplers. The average runtime for the different methods are similar (see Fig. 3(lower left)) since the main computational bottleneck is the gradient evaluation. Overall, this application serves as an important example of where our newly proposed sampler can have impact.

Conclusion

We presented a general recipe for devising MCMC samplers based on continuous Markov processes. Our framework constructs an SDE specified by two matrices, a positive semidefinite $\mathbf{D}(\mathbf{z})$ and a skew-symmetric $\mathbf{Q}(\mathbf{z})$. We prove that for any $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$, we can devise a continuous Markov process with a specified stationary distribution. We also prove that for any continuous Markov process with the target stationary distribution, there exists a $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$ that cast the process in our framework. Our recipe is particularly useful in the more challenging case of devising stochastic gradient MCMC samplers. We demonstrate the utility of our recipe in “reinventing” previous stochastic gradient MCMC samplers, and in proposing our SGRHMC method. The efficiency and scalability of the SGRHMC method was shown on simulated data and a streaming Wikipedia analysis.

This work was supported in part by ONR Grant N00014-10-1-0746, NSF CAREER Award IIS-1350133, and the TerraSwarm Research Center sponsored by MARCO and DARPA. We also thank Mr. Lei Wu for helping with the proof of Theorem 2 and Professors Ping Ao and Hong Qian for many discussions.

References

Supplementary Material: A Complete Recipe for Stochastic Gradient MCMC

Appendix A Proof of Stationary Distribution

In this section, we provide a proof for Theorem 1. To prove the theorem, we first show when ${\mathbf{f}}(\mathbf{z})$ satisfies Eq. (3) in the main paper, then the following Fokker-Planck equation of the dynamics:

is equivalent to the following compact form :

We can further decompose the second term as follows

The second equality follows because $\sum_{ij}\dfrac{\partial}{\partial\mathbf{z}_{i}}\dfrac{\partial}{\partial\mathbf{z}_{j}}[\mathbf{Q}_{ij}(\mathbf{z})p(\mathbf{z},t)]=0$ due to anti-symmetry of $\mathbf{Q}$. Putting these back into the formula, we get

We can then verify that $p(\mathbf{z}\mid\mathcal{S})\propto e^{-H(\mathbf{z})}$ is invariant under the compact form by calculating

The above proof follows directly from and is provided here for readers’ convenience.

Appendix B Proof of Completeness

In this section, we provide a constructive proof for Theorem 2, the existence of ${\mathbf{Q}}(\mathbf{z})$.

We first rewrite Eq. (3) in the main paper and notice that finding matrix $\mathbf{Q}(\mathbf{z})$ is equivalent to finding the matrix $\mathbf{Q}(\mathbf{z})p^{s}(\mathbf{z})$ such that

where the right hand side is a divergence-free vector.

We transform the above equation and its constraint into the frequency domain and obtain a set of linear equations.

Then we construct a solution to the linear equations and use inverse Fourier transform to obtain $\mathbf{Q}(\mathbf{z})$.

Multiplying $p^{s}(\mathbf{z})$ on both sides of Eq. (3) in the main paper, and noting that:

The equation for ${\mathbf{Q}}_{ij}(\mathbf{z})$ can now be written as:

Recall that the Fokker-Planck equation for the stochastic process, Eq. (2), is:

We can immediately observe that the right hand side of Eq. (S.4) has a divergenceless property by substituting the stationary probability density function $p^{s}(\mathbf{z})$ into Eq. (S.5):

The nice forms of Eqs. (S.4) and (S.6) imply that the questions can be transformed into a linear algebra problem once we apply a Fourier transform to them. Denote the Fourier transform of ${\mathbf{Q}}(\mathbf{z})p^{s}(\mathbf{z})$ as $\hat{{\mathbf{Q}}}({\mathbf{k}})$; and Fourier transform of {\mathbf{f}}_{i}(\mathbf{z})p^{s}(\mathbf{z})-\sum\limits_{j}\dfrac{\partial}{\partial\mathbf{z}_{j}}\Big{(}{\mathbf{D}}_{ij}(\mathbf{z})p^{s}(\mathbf{z})\Big{)} as $\hat{{\mathbf{F}}}_{i}({\mathbf{k}})$, where ${\mathbf{k}}=({\mathbf{k}}_{1},\cdots,{\mathbf{k}}_{n})^{T}$ is the set of the spectral variables. That is:

Then, \dfrac{\partial}{\partial\mathbf{z}_{j}}\Big{(}{\mathbf{Q}}_{ij}(\mathbf{z})p^{s}(\mathbf{z})\Big{)} is transformed to $2\pi{\rm i}\ \hat{{\mathbf{Q}}}_{ij}{\mathbf{k}}_{j}$, and Eq. (S.4) becomes the following equivalent form in Fourier space:

Hence, it is clear that matrix $\hat{{\mathbf{Q}}}$ must be a skew-symmetric projection matrix from the span of ${\mathbf{k}}$ to the span of $\hat{{\mathbf{F}}}$, where ${\mathbf{k}}$ and $\hat{{\mathbf{F}}}$ are always orthogonal to each other. We thereby construct $\hat{{\mathbf{Q}}}$ as combination of two rank $1$ projection matrices:

We arrive at the final result that matrix ${\mathbf{Q}}(\mathbf{z})$ is equal to $p^{s}(\mathbf{z})^{-1}$ times the inverse Fourier transform of $\hat{{\mathbf{Q}}}({\mathbf{k}})$:

Thus, if \left({\mathbf{f}}_{i}(\mathbf{z})p^{s}(\mathbf{z})-\sum\limits_{j}\dfrac{\partial}{\partial\mathbf{z}_{j}}\Big{(}{\mathbf{D}}_{ij}(\mathbf{z})p^{s}(\mathbf{z})\Big{)}\right) belongs to the space of $L^{1}$, then any continuous time Markov process, Eq. (2), can be turned into this new formulation. ∎

Entries in the skew-symmetric projector ${\mathbf{Q}}_{ij}(\mathbf{z})$ constructed here are real.

Denote ${\mathbf{a}}_{i}^{2}=\sum\limits_{l\neq i}{\mathbf{k}}_{l}^{2}$, then the inverse Fourier transform of $\dfrac{{\mathbf{k}}_{i}}{(2\pi{\rm i})\cdot\sum\limits_{l}{\mathbf{k}}_{l}^{2}}$ along the partial variable ${\mathbf{k}}_{i}$ is equal to:

where $H[x]$ is the Heaviside function. Because ${\mathbf{g}}_{i}(\mathbf{z})$ is an even function in $k_{l}$, $l\neq i$, its total inverse Fourier transform is real.

Therefore, the inverse Fourier transform of $\dfrac{{\mathbf{k}}_{i}\hat{{\mathbf{F}}}_{j}({\mathbf{k}})}{(2\pi{\rm i})\cdot\sum\limits_{l}{\mathbf{k}}_{l}^{2}}$ is the convolution of two real functions.

Appendix C 2-D Case as a Simple Intuitive Example of the Construction

Appendix D Previous MCMC Algorithms in the Form of Continuous Markov Processes as Elements in the Current Recipe

This section parallels that of Sec. 3.1 of the main paper, but in terms of the continuous dynamics underlying the samplers. This allows us to rapidly draw connections with our SDE framework of Sec. 2.1. Fig. S.4 provides a cartoon visualization of the portion of the product space $\mathbf{D}(\mathbf{z})\times\mathbf{Q}(\mathbf{z})$ already covered by past methods, after casting these methods in our framework below. Our proposed gSGRHMC method covers a portion of this space previously not explored.

The continuous dynamics underlying Eq. (12) in the main paper are

Again, we see Eq. (S.16) is a special case of our proposed framework with $\mathbf{z}=(\theta,r)$, $H(\theta,r)=U(\theta)+\frac{1}{2}r^{T}M^{-1}r$, {\mathbf{Q}}(\theta,r)=\left(\begin{array}[]{ll}0&-I\\ I&0\end{array}\right) and $\mathbf{D}(\theta,r)=\mathbf{0}$.

As described in , replacing $\nabla U(\theta)$ by the stochastic gradient $\nabla\widetilde{U}(\theta)$ in the $\epsilon$-discretized HMC system of Eq. (12) (resulting in Eq. (15)) has a continuous-time representation as:

Analogously to Sec. 3.1, these dynamics do not fit into our framework. Instead, in our framework we see that the noise term $\sqrt{2\mathbf{D}(\mathbf{z})}\textrm{d}\mathbf{W}(t)$ is paired with a $\mathbf{D}(\mathbf{z})\nabla H(\mathbf{z})$ term, hinting that such a term must be added to the dynamics of Eq. (S.19). Here, {\mathbf{D}}(\theta,r)=\left(\begin{array}[]{ll}0&0\\ 0&\epsilon\mathbf{V}\end{array}\right), which means we need to add a term of the form $\mathbf{D}(\mathbf{z})\nabla H(\mathbf{z})=\epsilon\mathbf{V}\nabla_{r}H(\theta,r)=\epsilon\mathbf{V}\mathbf{M}^{-1}r$. Interestingly, this is the correction strategy proposed in , but through a physical interpretation of the dynamics. In particular, the term $\epsilon\mathbf{V}\mathbf{M}^{-1}r$ (or, generically, $\mathbf{C}\mathbf{M}^{-1}r$) has an interpretation as a friction term and leads to second order Langevin dynamics:

This method now fits into our framework with $H(\theta,r)$ and $\mathbf{Q}(\theta,r)$ as in HMC, but here with {\mathbf{D}}(\theta,r)=\left(\begin{array}[]{ll}0&0\\ 0&C\end{array}\right).

SGLD proposes to use the following first order (no momentum) Langevin dynamics to generate samples

This algorithm corresponds to taking $\mathbf{z}=\theta$ with $H(\theta)=U(\theta)$, ${\mathbf{D}}(\theta)=\mathbf{D}$, ${\mathbf{Q}}(\theta)=0$. As in the case of SGHMC, the variance of the stochastic gradient can be subtracted from the sampler injected noise $\sqrt{2\mathbf{D}}\mathbf{W}(t)$ to make the finite stepsize simulation more accurate. This variant of SGLD leads to the stochastic gradient Fisher scoring algorithm .

SGLD can be generalized to use an adaptive diffusion matrix ${\mathbf{D}}(\theta)$. Specifically, it is interesting to take ${\mathbf{D}}(\theta)=\mathbf{G}^{-1}(\theta)$, where $\mathbf{G}(\theta)$ is the Fisher information metric. The sampler dynamics is given by

Taking $\mathbf{D}(\theta)={\mathbf{G}}^{-1}(\theta)$ and ${\mathbf{Q}}(\theta)=\mathbf{0}$, the SGRLD method falls into our current framework with the correction term $\Gamma_{i}(\theta)=\sum\limits_{j}\dfrac{\partial{\mathbf{D}}_{ij}(\theta)}{\partial\theta_{j}}$.

Finally, the continuous dynamics underlying the SGNHT algorithm in Sec. 3.1 are

Again, we see we can take $\mathbf{z}=(\theta,r,\xi)$, $H(\theta,r,\xi)=U(\theta)+\dfrac{1}{2}r^{T}r+\dfrac{1}{2d}(\xi-A)^{2}$, {\mathbf{D}}(\theta,r,\xi)=\left(\begin{array}[]{ccc}0&0&0\\ 0&A\cdot\mathbf{I}&0\\ 0&0&0\end{array}\right), and {\mathbf{Q}}(\theta,r,\xi)=\left(\begin{array}[]{ccc}0&-\mathbf{I}&0\\ \mathbf{I}&0&r/d\\ 0&-r^{T}/d&0\end{array}\right) to place these dynamics within our framework.

Appendix E Discussion of Choice of D and Q

A lot of choices of $\mathbf{D}(\mathbf{z})$ and $\mathbf{Q}(\mathbf{z})$ could potentially result in faster convergence of the samplers than those previously explored. For example, $\mathbf{D}(\mathbf{z})$ determines how much noise is introduced. Hence, an adaptive diffusion matrix $\mathbf{D}(\mathbf{z})$ can facilitate a faster escape from a local mode if $||\mathbf{D}(\mathbf{z})||$ is larger in regions of low probability, and can increase accuracy near the global mode if $||\mathbf{D}(\mathbf{z})||$ is smaller in regions of high probability. Motivated by the fact that a majority of the parameter space is covered by low probability mass regions where less accuracy is often needed, one might want to traverse these regions quickly. As such, an adaptive curl matrix $\mathbf{Q}(\mathbf{z})$ with $2$-norm growing with the level set of the distribution can facilitate a more efficient sampler. We explore an example of this in the gSGRHMC algorithm of the synthetic experiments (see Supp. F.1).

Appendix F Parameter Settings in Synthetic and Online Latent Dirichlet Allocation Experiments

In the synthetic experiment using gSGRHMC, we specifically consider $\mathbf{G}(\theta)^{-1}=D\sqrt{|\widetilde{U}(\theta)+C|}$. The constant $C$ ensures that ${\widetilde{U}(\theta)+C}$ is positive in most cases so that the fluctuation is indeed smaller when the probability density function is higher. Note that we define $\mathbf{G}(\theta)$ in terms of $\widetilde{U}(\theta)$ to avoid a costly full-data computation. We choose $D=1.5$ and $C=0.5$ in the experiments. The design of $\mathbf{G}$ is motivated by the discussion in Supp. E, taking $\mathbf{Q}(\theta)$ to have 2-norm growing with the level sets of the potential function can lead to faster exploration of the posterior.

Comparison of SGLD, SGHMC, the naïve implementation of SGRHMC (Eq. (27)), and the gSGRHMC methods is shown in Fig. S.5, indicating the incorrectness of the naïve SGRHMC.

F.2 Online Latent Dirichlet Allocation Experiment

In the online latent Dirichlet allocation (LDA) experiment, we used minibatches of $50$ documents and $K=50$ topics. Similar to , the stochastic gradient of the log posterior of the parameter $\theta$ on a minibatch $\widetilde{\mathcal{S}}$ is calculated as

where $\alpha$ is the hyper-parameter for the Gamma prior of per-topic word distributions, and $\gamma$ for the per-document topic distributions. Here, $n_{dkw}$ is the count of how many times word $w$ is assigned to topic $k$ in document $d$ (via $z_{j}^{(d)}=k$ for $x_{j}=w$). The $\cdot$ notation indicates $n_{dk\cdot}=\sum_{w}n_{dkw}$. To calculate the expectation of the latent topic assignment counts $n_{dkw}$, Gibbs sampling is used on the topic assignments in each document separately, using the conditional distributions

where $\backslash j$ represents a count excluding the topic assignment variable $z_{j}^{(d)}$ being updated. See for further details.

We follow the experimental settings in for Riemmanian samplers (SGRLD and SGRHMC), taking the hyper-parameters of Dirichlet priors to be $\gamma=0.01$ and $\alpha=0.0001$. Since the non-Riemmanian samplers (SGLD and SGHMC) do not handle distributions with mass concentrated over small regions as well as the Riemmanian samplers, we found $\gamma=0.1$ and $\alpha=0.01$ to be optimal hyper-parameters for them and use these instead for SGLD and SGHMC. In doing so, we are modifying the posterior being sampled, but wished to provide as good of performance as possible for these baseline methods for a fair comparison. For the SGRLD method, we keep the stepsize schedule of \epsilon_{t}=\left(a\cdot\bigg{(}1+\dfrac{t}{b}\bigg{)}\right)^{-c} and corresponding optimal parameters $a,b,c$ used in the experiment of . For the other methods, we use a constant stepsize because it was easier to tune. (A constant stepsize for SGRLD performed worse than the schedule described above, so again we are trying to be as fair to baseline methods as possible when using non-constant stepsize for SGRLD.) A grid search is performed to find $\epsilon_{t}=0.02$ for the SGRHMC method; $\epsilon_{t}=0.01$, $\mathbf{D}=I$ (corresponding to Eq. (19) in the main paper) for the SGLD method; and $\epsilon_{t}=0.1$, $\mathbf{C}=\mathbf{M}=I$ (corresponding to Eq. (18) in the main paper) for the SGHMC method.

For a randomly selected subset of topics, in Table S.1 we show the top seven most heavily weighted words in the topic learned with the SGRHMC sampler.