Irreversible Langevin MCMC on Lie Groups

TL;DR

This paper introduces an irreversible MCMC algorithm tailored for Lie groups, leveraging geometric structures to enhance convergence speed, and demonstrates its effectiveness with the example of SO(3).

Contribution

It develops a novel irreversible Hamiltonian Monte Carlo-like method on Lie groups using geometric and dissipation structures, generalizing HMC.

Findings

Algorithm converges faster than reversible methods.

Recovers HMC as a special case with long OU process simulation.

Numerical experiments on SO(3) validate the approach.

Abstract

It is well-known that irreversible MCMC algorithms converge faster to their stationary distributions than reversible ones. Using the special geometric structure of Lie groups $\mathcal G$ and dissipation fields compatible with the symplectic structure, we construct an irreversible HMC-like MCMC algorithm on $\mathcal G$, where we first update the momentum by solving an OU process on the corresponding Lie algebra $\mathfrak g$, and then approximate the Hamiltonian system on $\mathcal G \times \mathfrak g$ with a reversible symplectic integrator followed by a Metropolis-Hastings correction step. In particular, when the OU process is simulated over sufficiently long times, we recover HMC as a special case. We illustrate this algorithm numerically using the example $\mathcal G = SO(3)$.

Figures

Equations

$$\mathrm{d}Z_{t}=\Big{(}X_{H}(Z_{t})-\frac{\beta}{2}\sum_{i=1}^{m}\{H_{i},H\}X_{H_{i}}(Z_{t})\Big{)}\mathrm{d}t+\sum_{i=1}^{m}X_{H_{i}}(Z_{t})\circ\mathrm{d}W^{i}_{t}\,,$$

$$\mathcal{L}f=-\{f,H\}-\frac{\beta }{2}\sum _{i=1}^m\{H,H_i\}\{f,H_i\}+\frac{1}{2}\sum _{i=1}^m\{\{f,H_i\},H_i\}.$$

$${\int }_{\mathcal{M}}\{f,g\}{\omega }^n=0.$$

$${\mathbb{P}}_{\infty }(z)=p_{\infty }{\omega }^n:=\frac{1}{Z}{e}^{-\beta H(z)}{\omega }^n,Z={\int }_{\mathcal{M}}{e}^{-\beta H(z)}{\omega }^n$$

$$g\{f,H_i\}\{H,H_i\}$$

$$=\cdots =\{gf\{H,H_i\},H_i\}-f\{g\{H,H_i\},H_i\}$$

$$g\{\{f,H_i\},H_i\}=\{g\{f,H_i\},H_i\}-\{f\{g,H_i\},H_i\}+f\{\{g,H_i\},H_i\}.$$

$${\int }_{\mathcal{M}}g(\mathcal{L}f){\omega }^n={\int }_{M}g\left(-\{f,H\}-\frac{\beta }{2}\sum _{i=1}^m\{f,H_i\}\{H,H_i\}+\frac{1}{2}\sum _{i=1}^m\{\{f,H_i\},H_i\}\right){\omega }^n$$

$$={\int }_{\mathcal{M}}\left(-\{fg,H\}-\frac{\beta }{2}\{fg\{H,H_i\},H_i\}+\frac{1}{2}\left(\{g\{f,H_i\},H_i\}-\{f\{g,H_i\},H_i\}\right)\right){\omega }^n$$

$$+{\int }_{M}f\left(\{g,H\}+\frac{\beta }{2}\{g\{H,H_i\},H_i\}+\frac{1}{2}\{\{g,H_i\},H_i\}\right){\omega }^n$$

$$={\int }_{\mathcal{M}}f\left(\{g,H\}+\frac{\beta }{2}\{g\{H,H_i\},H_i\}+\frac{1}{2}\{\{g,H_i\},H_i\}\right){\omega }^n={\int }_{\mathcal{M}}f({\mathcal{L}}^{\ast }g){\omega }^n,$$

$${\mathcal{L}}^{\ast }g=\{g,H\}+\frac{\beta }{2}\sum _{i=1}^m\{g\{H,H_i\},H_i\}+\frac{1}{2}\sum _{i=1}^m\{\{g,H_i\},H_i\}.$$

$${\mathcal{L}}^{\ast }{p}_{\infty }=p_{\infty }^{\prime }(H)\{H,H\}+$$

$$-\frac{\beta }{2}\sum _{i=1}^m\{p_{\infty }(H)\{H,H_i\},H_i\}=0.$$

$${\int }_A{\mathbb{P}}_{\infty }^1={\int }_{T^{\ast }A}{\iota }^{\ast }{\mathbb{P}}_{\infty },$$

$${\mathbb{P}}_{\infty }^1(\mathrm{d}q)=\frac{1}{Z_1}{e}^{-V(q)}\sqrt{|g|}\mathrm{d}q,Z_1={\int }_{\mathcal{Q}}{e}^{-V(q)}\sqrt{|g|}\mathrm{d}q,$$

$$\mathrm{d}{Z}_t$$

$$=-\frac{\beta }{2}\sum _{i=1}^m{v}^k{e}_k(U_i)e_j(U_i){\partial }_{v^j}(Z_t)\mathrm{d}t-\sum _{i=1}^m{e}_j(U_i){\partial }_{v^j}(Z_t)\circ \mathrm{d}{W}_t^i.$$

$$\frac{\mathrm{d}{Z}_t}{\mathrm{d}t}=X_H(Z_t).$$

$$\dot{q}=0,\mathrm{d}{v}_t=-\frac{\beta }{2}{\nabla }_q{U}_i(q){\nabla }_q{U}_i(q{)}^T{v}_t\mathrm{d}t-{\nabla }_q{U}_i(q)\circ \mathrm{d}{W}_t^i.$$

$$X_{H_i}(g,v)=-e_j{|}_g(U_i){\xi }_j=:{\sigma }_{ji}(g){\xi }_j.$$

$$\displaystyle X_{H_{i}}(g,v)=-\text{Tr}\big{(}\nabla U_{i}^{T}g\xi_{j}\big{)}\xi_{j}=\sigma_{ji}(g)\xi_{j},$$

$$\dot{g}=0,\mathrm{d}{v}_t=-\frac{\beta }{2}(\sigma (g)\sigma (g{)}^T{)}_{jk}{v}_t^k{\xi }_j\mathrm{d}t+{\sigma }_{ji}(g){\xi }_j\mathrm{d}{W}_t^i,$$

$$\mathrm{d}{\mathbf{v}}_t=-\frac{\beta }{2}\sigma {\sigma }^T{\mathbf{v}}_t\mathrm{d}t+\sigma \mathrm{d}{\mathbf{W}}_t,$$

$$v_{t+h}=e^{-\frac{\beta }{2}Dh}{v}_t+\sigma {\int }_t^{t+h}{e}^{-\frac{\beta }{2}D(t+h-s)}\mathrm{d}{\mathbf{W}}_s,$$

$$\begin{array}{rl}p(v_0,v)& =\frac{1}{(2\pi {)}^{n/2}|\det {\Sigma }_h|}\exp \left(-\frac{1}{2}{\Vert v-e^{-\frac{\beta }{2}Dh}{v}_0\Vert }_{{\Sigma }_h^{-1}}^2\right),\\ \mathrm{where}{\Sigma }_h& =\frac{1}{\beta }\left(\mathrm{Id}-e^{-\beta Dh}\right).\end{array}$$

$$v^{\ast }\simeq \mathcal{N}\left(e^{-\frac{\beta }{2}Dh}{v}_0,{\Sigma }_h\right)$$

$${\bar{v}}_{k+\frac{1}{2}}$$

$${\bar{g}}_{k+1}$$

$${\bar{v}}_{k+1}$$

$$\alpha =\min \left\{1,\exp \left(-H({\bar{g}}_N,{\bar{v}}_N)+H({\bar{g}}_0,{\bar{v}}_0)\right)\right\},$$

$$U_i(g)=ϵ\mathrm{Tr}(e^{-{\xi }_i}g)\mathrm{for}\mathrm{span}({\xi }_i)=\mathfrak{so}(3).$$

Full text

1 Introduction

In this work, we construct an irreversible MCMC algorithm on Lie groups, which generalises the standard Hamiltonian Monte Carlo (HMC) algorithm on $\mathbb{R}^{n}$. The HMC method [duane1987hybrid] generates samples from a probability density (with respect to an appropriate reference measure) known up to a constant factor by generating proposals using Hamiltonian mechanics, which is approximated by a reversible symplectic numerical integrator and followed by a Metropolis-Hastings step to correct for the bias introduced during the numerical approximations. The resulting time-homogeneous Markov chain is thus reversible, and allow distant proposals to be accepted with high probability, which decreases the correlations between samples (for a basic reference on HMC see [neal2011mcmc], and for a geometric description see [barp2018geometry, betancourt2017geometric]). However, it is well-known that ergodic irreversible diffusions converge faster to their target distributions [duncan2016variance, rey2015irreversible], and several irreversible MCMC algorithms based on Langevin dynamics have been proposed [ottobre2016markov, ottobre2016function].

From a mechanical point of view, diffusions on Lie groups are important since they form the configuration space of many interesting systems, such as the free rigid body. For example in [chen2015constrained] Euler-Poincaré reduction of group invariant symplectic diffusions on Lie groups are considered in view of deriving dissipative equations from a variational principle, and in [arnaudon2018noise] Langevin systems on coadjoint orbits are constructed by adding noise and dissipation to Hamiltonian systems on Lie groups. The phase transitions of this system were analysed using a sampling method [arnaudon2018networks]. In lattice gauge theory one typically uses the HMC algorithm for semi-simple compact Lie groups which was originally presented in [kennedy88b] and extended to arbitrary Lie groups in [barp2019hamiltonianB], see also [Kennedy:2012, clark2008tuning, clark2007accelerating]. In [barp2019homogeneous], it was shown how to construct HMC on homogeneous manifolds using symplectic reduction, which includes sampling on Lie groups as a special case.

To construct an irreversible algorithm on Lie groups, we first extend Langevin dynamics to general symplectic manifolds $\mathcal{M}$ based on Bismut’s symplectic diffusion process [bismut1981mecanique]. Our generalised Langevin dynamics with multiplicative noise and nonlinear dissipation has the Gibbs measure as the invariant measure, which allows us to design MCMC algorithms that sample from a Lie group $\mathcal{G}$ when we take $\mathcal{M}=T^{*}\mathcal{G}$. In our Langevin system the irreversible component is determined by Hamiltonian vector fields which are compatible with the symplectic structure, thus avoiding the appearance of divergence terms associated to the volume distortion. We are then free to choose the noise-generating Hamiltonians to best suit the target distribution. Choosing Hamiltonians that only depend on position allows us to proceed with a Strang splitting of the dynamics into a position-dependent OU process in the fibres which can be solved exactly, and a Hamiltonian part which is approximated using a leapfrog scheme, followed by a Metropolis-Hastings acceptance/rejection step in a similar fashion to [ottobre2016markov, ottobre2016function, bou2010patch]. Ideally one wants to choose these Hamiltonians to achieve the fastest convergence to stationarity.

On a general manifold, it would be necessary to introduce local coordinates in order to solve the OU process on the fibres, making it difficult to implement. However, since our base manifold is a Lie group, the Maurer-Cartan form defines an isomorphism between the cotangent bundle $T^{*}\mathcal{G}$ and the trivial bundle $\mathcal{G}\times\mathfrak{g}^{*}$, which, given an inner-product on $\mathfrak{g}$, may further be identified with $\mathcal{G}\times\mathfrak{g}$. As a result, one may pull back the OU process on $T_{g}^{*}\mathcal{G}$ to an OU process on $\mathfrak{g}$ for any $g\in\mathcal{G}$, thus avoiding the problem of having to choose appropriate charts. Hence on Lie groups, we obtain a practical irreversible MCMC algorithm which generalises the $\mathbb{R}^{n}$-version of the irreversible algorithm considered in [ottobre2016markov, ottobre2016function].

Finally, we simulate this algorithm in the special case $\mathcal{G}=SO(3)$ and perform a Maximum Mean Discrepancy (MMD) test to show that on average, the irreversible algorithm converges faster to the stationary measure than the corresponding reversible HMC on $SO(3)$.

2 Diffusions on Symplectic Manifolds

We consider diffusion processes on symplectic manifolds $(\mathcal{M},\omega)$, where we have a natural volume form $\omega^{n}$, and define the canonical Poisson bracket $\{g,f\}:=\omega(X_{f},X_{g})=X_{g}f$, where $X_{g}$ is the Hamiltonian vector field associated to the Hamiltonian $g:\mathcal{M}\to\mathbb{R}$, i.e. $\mathrm{d}g=\iota_{X_{g}}\omega$. Given arbitrary functions $H$ and $H_{i}$ for $i=1,\ldots,m$ on $\mathcal{M}$, we consider the SDE

[TABLE]

which has the generator, or forward Kolmogorov operator

[TABLE]

To show that the Gibbs measure is an invariant measure for (1), we will need the following lemma:

Lemma 1

For a symplectic manifold $(\mathcal{M},\omega)$ and two functions $f,g\in C^{\infty}(\mathcal{M})$ such that either $\partial\mathcal{M}=\varnothing$ or $g|_{\partial M}=0$, we have the following identity

[TABLE]

For the proof, see [cruzeiro2018momentum], Section 4.3. Hence (1) enables us to build MCMC algorithms on any symplectic manifold, and in particular the cotangent bundle of Lie groups, that converge to the Gibbs measure:

Theorem 2.1

Given a symplectic manifold $(\mathcal{M},\omega)$ without boundary, equation (1) on $\mathcal{M}$ has the Gibbs measure

[TABLE]

as its stationary measure for any choice of $H_{i}:\mathcal{M}\to\mathbb{R}$ where $i=1,\ldots,m$.

Proof

Using the Leibniz rule $\{fg,h\}=f\{g,h\}+g\{f,h\}$, we have

[TABLE]

and similarly

[TABLE]

Hence one can compute the $L^{2}(\mathcal{M},\omega^{n})$-adjoint of the operator $\mathcal{L}$ as follows

[TABLE]

where we have used Lemma 1 to integrate the Poisson brackets to [math]. Hence, we obtain the Fokker-Planck operator

[TABLE]

Now, by the derivation property of the Poisson bracket, $\{f\circ g,h\}=f^{\prime}\circ g\{g,h\}$ and noting that $p_{\infty}^{\prime}(H)=-Z^{-1}\beta e^{-\beta H}=-\beta p_{\infty}(H)$, one can check that

[TABLE]

Therefore $\mathbb{P}_{\infty}(z)=p_{\infty}(z)\omega^{n}$ is indeed an invariant measure for (1).

If $\mathcal{M}=T^{*}\mathcal{Q}$ is the cotangent bundle of a manifold without boundary $\mathcal{Q}$, we define the marginal measure $\mathbb{P}^{1}_{\infty}$ on $\mathcal{Q}$ by

[TABLE]

for any measurable set $A\subset\mathcal{Q}$, where $\iota:T^{*}A\rightarrow T^{*}\mathcal{Q}$ is the inclusion map. In addition, if $(\mathcal{Q},\gamma)$ is a Riemannian manifold, we can consider the Hamiltonian function $H(q,p)=\frac{1}{2}\gamma_{q}(p,p)+V(q)$, for $(q,p)\in T^{*}\mathcal{Q}$, and the marginal invariant measure $\mathbb{P}_{\infty}^{1}(\mathrm{d}q)$ of the process (1) is simply

[TABLE]

where $\sqrt{|g|}\mathrm{d}q$ is the Riemannian volume form.

The MCMC algorithm which we will derive in section §3 is based on a Strang splitting of the dynamics (1) into a Hamiltonian part and a Langevin part. Hereafter, we identify $T^{*}\mathcal{Q}$ with $T\mathcal{Q}$ through the metric and just consider the dynamics on $T\mathcal{Q}$ instead of $T^{*}\mathcal{Q}$.

3 Irreversible Langevin MCMC on Lie Groups

Consider a $n$ dimensional Lie group $\mathcal{G}$ and let $e_{i},\theta^{i},i=1,\ldots,n$ be an orthonormal basis of left-invariant vector fields and dual one-forms respectively. We consider $H=V\circ\pi+T:T\mathcal{G}\rightarrow\mathbb{R}$, where $T$ is the kinetic energy associated to a bi-invariant metric on $\mathcal{G}$ and $V\propto\log\chi:\mathcal{G}\rightarrow\mathbb{R}$ is the potential energy, where $\chi$ is the distribution we want to sample from on $\mathcal{G}$. We let $v^{i}:T\mathcal{G}\to\mathbb{R}$ be the fibre coordinate functions with respect to the left-invariant vector fields, $v^{i}(g,u_{g}):=\theta^{i}_{g}(u_{g})$.

Vector fields tangent to $T\mathcal{G}$ (i.e., elements of $\Gamma(TT\mathcal{G})$) can be expanded in terms of left-invariant vector fields $e_{i}$ and the fibre-coordinate vector fields $\partial_{v^{i}}$, (i.e., $\Gamma(TT\mathcal{G})\cong\Gamma(T\mathcal{G}\oplus T\mathfrak{g})$). We consider noise Hamiltonians that depend only on position, $H_{i}=U_{i}\circ\pi$ where $U_{i}:\mathcal{G}\to\mathbb{R}$, so the corresponding Hamiltonian vector fields can be written as $X_{H_{i}}=-e_{j}(U_{i})\partial_{v^{j}}$, (see [barp2019hamiltonianB]). Hence the stochastic process (1) on $T\mathcal{G}$ can be split up into a Langevin part

[TABLE]

and a Hamiltonian part

[TABLE]

Note the geodesics are given by the one-parameter subgroups, with Hamiltonian vector field $X_{T}=v^{k}e_{k}$. Since $X_{H_{i}}$ only has components in the fibre direction $\partial_{v^{i}}$ (i.e., it has no $e_{i}$ components along $\mathcal{G}$) the diffusion starting at any point $(g,v)\in T\mathcal{G}$ remains in $T_{g}\mathcal{G}$, i.e., with the same base point $g$. When $\mathcal{G}=\mathbb{R}^{n}$ and we use the standard kinetic energy $T=\frac{1}{2}\|v\|_{\mathbb{R}^{n}}^{2}$, then vector fields become gradients, i.e. $e_{j}=\partial_{q^{j}}$, and equation (4) becomes the space dependent Langevin equation for $(q,v)\in T\mathbb{R}^{n}$,

[TABLE]

Now let $\xi_{i}:=e_{i}(1)$ be a basis of the Lie algebra $\mathfrak{g}$, where $1$ is the identity. Then $e_{i}(g)=\partial_{1}L_{g}\xi$ and we can identify $T\mathcal{G}$ with $\mathcal{G}\times\mathfrak{g}$ through the relation $(g,v^{i}e_{i}(g))\sim(g,v^{i}\xi_{i})$. In other words, we may now think of $v^{i}$ as the Lie algebra coordinate functions $v^{i}:\mathcal{G}\times\mathfrak{g}\to\mathbb{R}$ with $v^{i}(g,u)=\theta^{i}_{1}(u)$, and since $\mathfrak{g}$ is a vector space, we can identify $\partial_{v^{i}}\sim\xi_{i}$ and write

[TABLE]

for $i=1,\ldots,m$ and $j=1,\ldots,n$. For matrix Lie groups, this becomes

[TABLE]

where $(\nabla U_{i})_{ab}:=\partial_{x_{ab}}U_{i}$, where $x_{ab}$ are the matrix coordinates of $g\in\mathcal{G}$. The Langevin equation (4) can then be written as

[TABLE]

and if we identify $\mathfrak{g}\sim\mathbb{R}^{n}$, $v^{i}\xi_{i}\sim\mathbf{v}\in\mathbb{R}^{n}$, we get a standard OU process on $\mathbb{R}^{n}$

[TABLE]

for the vector-valued Wiener process $\mathbf{W}_{t}=(W_{t}^{1},\ldots,W_{t}^{m})\in\mathbb{R}^{m}$. Note that the noise term can be interpreted as an Itô integral since the diffusion coefficient $\sigma$ does not depend on $\mathbf{v}$. We can solve this Langevin equation explicitly if the matrix $D=\sigma\sigma^{T}$ is invertible. This is the case if the vectors $\nabla U_{i}$ form an orthonormal basis of $\mathbb{R}^{n}$, or more generally if they satisfy the Hörmander condition. The explicit solution is then given by

[TABLE]

which has the transition probability

[TABLE]

3.1 MCMC algorithm on Lie groups

From the Langevin system considered in the previous section, we can construct the following MCMC algorithm to sample from the distribution $\chi:=e^{-\frac{\beta}{2}V}$ on $\mathcal{G}$. Starting from $(g_{0},v_{0})\in T\mathcal{G}$, we iterate the following

Solve equation (9) exactly until time $h$ by sampling

[TABLE]

to obtain $(\bar{g}_{0},\bar{v}_{0})=(g_{0},v^{*})$; 2. 2.

Approximate the Hamiltonian system (5) using $N$ Leapfrog trajectories with step size $\delta>0$. Starting at $(\bar{g}_{0},\bar{v}_{0})=(g_{0},v^{*})$:

 For $k=0,\ldots,N-1$: 111For a non-matrix group, simply replace $\mathrm{Tr}\left(\partial_{x}V^{T}\bar{g}_{k}\xi_{i}\right)$ with $e_{i}|_{g}(V)$.

[TABLE]

to obtain $(\bar{g}_{N},\bar{v}_{N})$. The time step $\delta$ and number of steps $N$ are to be tuned appropriately by the users. 3. 3.

Accept or reject the proposal by a Metropolis-Hastings step. We accept the proposal $(\bar{g}_{N},\bar{v}_{N})$ with probability

[TABLE]

and set $(g_{1},v_{1})=(\bar{g}_{N},\bar{v}_{N})$. On the other hand, if the proposal is rejected, we set $(g_{1},v_{1})=(\bar{g}_{0},-\bar{v}_{0})$.

Notice also that in the limit $h\to\infty$, this algorithm become the standard HMC algorithm, which is reversible, but for finite $h$, the algorithm is irreversible (see [ottobre2016markov] for a detailed discussion).

4 Example on $SO(3)$

As an example we will pick the rotation group $\mathcal{G}=SO(3)$, where the Lie algebra $\mathfrak{so}(3)$ consists of $3\times 3$ anti-symmetric matrices so that the kinetic energy on the Lie algebra is $T(v)=\frac{1}{2}\mathrm{Tr}(v^{T}v)$, for $v\in\mathfrak{so}(3)$. The potentials on $SO(3)$ will be defined by functions $V,U_{i}:\mathbb{R}^{3\times 3}\to\mathbb{R}$ on matrices. We then pick the potential to be of the form $V(g)=e^{\alpha Tr(g)}$ and isotropic noise with

[TABLE]

We then obtain samples $g_{t}$ on $SO(3)$, which we can project onto the sphere by simply letting the group act on a vector $z=(0,0,1)$, to get $x_{t}=g_{t}z$, see panel (a) of figure 1. From these samples $g_{t}$, we can also estimate the convergence rate of the MCMC algorithm by computing the maximum mean discrepancy (MMD) between the set of first $N$ samples and the whole sequence, using the values on the diagonal of the matrices $g_{t}$.

We see that in figure 1, small values for $h$ give MCMC algorithms with a faster convergence rate than the HMC limit, i.e., $h\to\infty$. This is a direct consequence from the irreversibility of the algorithm, as explained in [duncan2016variance, rey2015irreversible]. Even if a faster convergence is desirable, one has to ensure that the correct distribution is sampled, and if $h$ is taken too small, the algorithm will be close to a pure Hamiltonian dynamics, with additional reversal steps $v\to-v$ when the proposed state is rejected. We observe this effect already for $h=0.01$ where the convergence is as fast as $h=0.1$ for the first steps of the chain, but then later slows down, as the distribution is not sampled correctly.

Acknowledgements The authors thank A. Duncan for the useful insights that helped improve this work. AA acknowledges EPSRC funding through award EP/N014529/1 via the EPSRC Centre for Mathematics of Precision Healthcare. ST acknowledges the Schrödinger scholarship scheme for funding during this work. AB was supported by a Roth Scholarship funded by the Department of Mathematics, Imperial College London, by EPSRC Fellowship (EP/J016934/3) and by The Alan Turing Institute under the EPSRC grant [EP/N510129/1].