Second order ensemble Langevin method for sampling and inverse problems

Consider sampling the distribution

where $\Phi:\mathbb{R}^{N}\to\mathbb{R}$ is termed the potential function and $Z_{q}$ the normalization constant. A broad family of problems can be cast into this formulation; the Bayesian approach to inverse problems provides a particular focus for our work [34]. The point of departure for the algorithms considered in this paper is the following mean-field Langevin equation:

where $D$ stands for the gradient operator, $W$ is an $N$-dimensional standard Brownian motion, $\rho$ is the density associated to the law of $q$, and $\mathcal{C}(\rho)$ is the covariance under this density. This constitutes a mean-field generalization [23] of the standard Langevin equation [44]. Applying a particle approximation to the mean-field model results in an interacting particle system, and coupled Langevin dynamics [37, 23, 42, 24]. The benefit of preconditioning using the covariance is that it leads to mixing rates independent of the problem, provably for quadratic $\Phi$ and empirically beyond this setting, because of the affine invariance [26] of the resulting algorithms [24].

In order to further accelerate mixing and achieve sampling efficiency, we introduce an auxiliary variable $p\in\mathbb{R}^{N}$ and consider the Hamiltonian

where the new state variable $z:=(q^{\top},p^{\top})^{\top}\in\mathbb{R}^{2N}$. Define measure on $\mathbb{R}^{2N}$ by

where $Z_{q,p}$ is the normalization constant. The marginal distribution of $\Pi$ in the $q$ variable gives the desired distribution $\pi$, i.e. $\int\Pi(z)\mathrm{d}p=\pi(q)$. We now aim at sampling the joint distribution. To this end, consider the following underdamped Langevin dynamics in $\mathbb{R}^{2N}$:

with the choices

Here $W_{0}$ is a standard Brownian motion in $\mathbb{R}^{2N}$ with components $W^{\prime},W\in\mathbb{R}^{N}$. Then we have the following, proved in Subsection 8.1:

In practice, to simulate from such a mean-field model, it will be necessary to consider a particle approximation of the form

for the set of $I$ particles $Z=\{z^{(i)}\}_{i=1}^{I},$ and where $M(Z),K(Z),J(Z)$ are appropriate empirical approximations of $\mathcal{M}(\rho),\mathcal{K}(\rho),\mathcal{J}(\rho)$ based on the approximation

and the Hamiltonian is given by

We use the notation $H_{Z}$ to emphasize the dependence on the ensemble $Z$, but when the context is clear we also use the abbreviated notation $H$. Note that $H$ is the appropriate finite particle approximation of $\mathcal{H},$ given the particle approximation $M$ of $\mathcal{M}.$

In this paper we will concentrate on a specific choice of mean-field operators within the above general construction, which we now describe. Let $\mathcal{C}_{q}(\rho)$ denote the $q$-marginal in the covariance under the law of (1.3). We make the choices $\mathcal{K}_{1}=0$, $\mathcal{C}=\mathcal{M}=\mathcal{C}_{q}(\rho)$, and $\mathcal{K}_{2}=\gamma\mathcal{C}_{q}(\rho)$, for a scalar damping parameter $\gamma>0$. Then the underdamped Langevin dynamics yields

To implement a particle approximation of the mean-field dynamics (1.7) we introduce particles in the form $z^{(i)}(t)=\bigl{(}(q^{(i)}(t))^{\top},(p^{(i)}(t))^{\top}\bigr{)}$ and we use the ensemble covariance and mean approximations

In order to obtain affine invariance, we take the generalized square root of the ensemble covariance $C_{q}(Z)$, similarly to [24]. We introduce the $N\times I$ matrix

which allows us to define the empirical covariance and generalized (nonsymmetric) square root via

Now with $I$ independent standard Brownian motions $\{W^{(i)}\}_{i=1}^{I}\in\mathbb{R}^{I}$, a natural particle approximation of (1.7) is, for $i=1,\dots,I$,

In later sections, we will add corrections ensuring preservation of the desired $q$ marginals of the invariant measure. These corrections will account for both finite particle approximation of the mean field limit, and the fact that (1.9) does not account for the dependence of $H_{Z}(z)$ on $Z$; note that this latter dependence is accounted for in (1.5). To establish the form of the correction we follow the exposition of the well-trodden subject of determining a diffusion with given invariant measure from [40][Theorem 1].

In subsequent sections we will employ an ensemble approximation of $C_{q}(Z)D\Phi(\cdot)$, as in [23], thereby avoiding the need to compute adjoints of the forward model; we note also that in the linear case this approximation is exact. We will show that the resulting interacting particle system has the potential to provide accurate derivative-free inference for certain classes of inverse problems. In the remainder of this section we provide a literature review, we highlight our contributions, and we outline the structure of the paper.

The overdamped Langevin equation is the canonical SDE that is invariant with respect to given target density distribution. Many sampling algorithms are built upon this idea, and in particular, it is shown to govern a large class of Monte Carlo Markov Chain (MCMC) methods; see [47, 43, 40]. To enhance mixing and accelerate convergence, a second-order method, Hybrid Monte Carlo (HMC, also referred to as Hamiltonian Monte Carlo) [19, 41] has been proposed, leading to underdamped Langevin dynamics. There have been many attempts to justify the empirically observed fast convergence speed of second-order methods in comparison with first-order methods [4]. Recently a quantitative convergence rate is established in [10], showing that the underdamped Langevin dynamics converges faster than the overdamped Langevin dynamics when the log of the target $\pi$ has a small Poincaré constant; see also [21].

The idea of introducing preconditioners within the context of interacting particle systems used for sampling is developed in in [37]. Preconditioning via ensemble covariance is shown to boost convergence and numerical performance [23]. Other choices of preconditioners in sampling will lead to different forms of SDEs and associated Fokker-Planck equations with different structures, which can result in different sampling methods effective in different scenarios; see for example [38]. Affine invariance is introduced in [26] where it is argued that this property leads to desirable convergence properties for interacting particle systems used for sampling: methods that satisfy affine invariance are invariant under an affine change of coordinates, and are thus uniformly effective for problems that can be rendered well-conditioned under an affine transformation. An affine invariant version of the mean-field underdamped Langevin dynamics [23] is proposed in [42, 24].

Kalman methods have shown wide success in state estimation problems since their introduction by Evensen; see [22] for an overview of the field, and the papers [45, 31] for discussion of their use in inverse problems. Using the ensemble covariance as a preconditioner leads to affine invariant [24] and gradient-free [23] approximations of Langevin dynamics; this is desirable in practical computations in view of the intractability of derivatives in many large-scale models arising in science and engineering [15, 28, 35, 30]. See [2, 49] for analysis of these methods and, in the context of continuous data-assimilation, see [17, 3, 51]. There are other derivative-free methods that can be derived from the mean-field perspective, and in particular consensus-based methods show promise for optimization [12] and have recently been extended to sampling in [11].

Recent work has established the convergence of ensemble preconditioning methods to mean field limits; see for example [18]. For other works on rigorous derivation of mean field limits of interacting particle systems, see [50, 13, 32]. For underpinning theory of Hamiltonian-based sampling, see [7, 8, 39, 5].

The following contributions are made in this paper:

1. 1.

   We introduce an underdamped second order mean field Langevin dynamics, with a covariance-based preconditioner.

2. 2.

   In the case of Bayesian inverse problems defined by a linear forward map, we show that that this mean field model preserves Gaussian distributions under time evolution and, if initialized at a Gaussian, converges to the desired target at rate independent of the linear map.

3. 3.

   We introduce finite particle approximations of the mean field model, and introduce finite size corrections to preserve invariance of the desired target measure, resulting in an affine invariant method.

4. 4.

   For Bayesian inverse problems, we introduce a gradient-free approximation of the algorithm, based on ensemble Kalman methodology.

5. 5.

   In the context of Bayesian inverse problems we provide numerical examples to demonstrate that the algorithm resulting from the previous considerations has desirable sampling properties.

In Section 2, we introduce the inverse problems context that motivates us. In Section 3 we discuss the equilibrium distribution of the mean field model; and then we introduce a finite particle approximation of the mean field model, identifying corrections which ensure that the equilibrium distribution is an independent product (of dimension equal to the particle cardinality) of copies of the desired equilibrium distribution. Section 4 introduces the ensemble Kalman approximation of the finite particle system; and in that section we also demonstrate affine invariance of the resulting method. Section 5 presents analysis of the finite particle system in the case of linear inverse problems, where the ensemble Kalman approximation is exact; we demonstrate that the relaxation time to equilibrium is independent of the specific linear inverse problem considered, a consequence of affine invariance. In Section 6 we provide numerical results which demonstrate the efficiency and potential value of our method, and in Section 7 we draw conclusions. Proofs of the propositions are given in the appendix, Section 8.

The mean field underdamped Langevin equation (1.3) has an associated nonlinear and nonlocal Fokker-Planck equation giving the evolution of the law of particles $z(t)$, denoted $\rho(z,t)$. The equation for this law is (see Proof 8.1.)

By Proposition 1.1 this Fokker-Planck equation has $\Pi(z)$ as its equilibrium; this follows as for standard linear Fokker-Planck equations [44] since the dependence of the sample paths on $\rho$ involves only the mean and covariance; see Proof 8.1 and see also [20, 27, 33, 46, 44] for discussions on how to derive Fokker-Planck equations with a prescribed stationary distribution. In the specific case (1.7), recall that $\mathcal{J}$ and $\mathcal{K}$ are given by

These choices satisfy the assumption of Proposition 1.1.

The approximation of (1.7) by the interacting particle system (1.9) does not result in an equilibrium distribution with marginals on the position coordinates given by independent products $\pi.$ To address this issue we return to the formulation in (1.5), (1.6), with the following empirical approximations to $\mathcal{K}$ and $\mathcal{J}$:

where $C_{q}(Z)$ is given in (1.8); note that in defining $H_{Z}(z)$ to be consistent with choices leading to (1.7) we make the approximation $\mathcal{M}(\rho)\approx M(Z)=C_{q}(Z),$ leading to

We would like the resulting finite particle system (1.5), (3.14), (3.15) to have $\pi^{\otimes I}$ as the marginal on the position coordinates of its equilibrium solution; in general, however, it does not. We now determine a drift correction term which, when added to the system results in invariance of $\Pi^{I}$; here, for normalization constant $Z_{I}$,

Note that $\Pi^{I}$ does not have product form, but its marginal on the $q$ coordinates does, and is $\pi^{\otimes I}.$

To achieve the desired invariance of ${\Pi^{I}}$ we add the correction term $\nabla_{z^{(i)}}\cdot\bigl{(}K(Z)-J(Z)\bigr{)}$ to the interacting particle system (1.5) to obtain the following coupled system of coupled diffusions in $\mathbb{R}^{N},$ for $i=1,\dots,I:$

By Theorem 1 in both [40] and [20] this will result in a linear Fokker-Planck equation (3.18) with $\Pi^{I}$ as its equilibrium. This Fokker-Planck equation has the form

where $\rho^{I}$ stands for the joint probability of the $I$ particles, $H^{\otimes I}$ is the vector containing $I$ copies of $H$, and $K^{\otimes I},J^{\otimes I}$ are block-diagonal matrices with $I$ copies of $K,J$ on the diagonal entry respectively. Note that in the case of finite particle approximation, the preconditioner is no longer dependent of the density $\rho$, and thus the Fokker-Planck equation is indeed linear. For the specific choices of $K(Z),J(Z),M(Z)$ and $H_{Z}(z)$ made in (3.14) and (3.15), equation (3.17) simplifies to the following system of diffusions in $\mathbb{R}^{N}$, holding for $i=1,\dots,I:$

where we have defined $\hat{q}^{(i)}:=q^{(i)}-\bar{q}$. Concerning these equations, we establish the following proposition. We postpone the proof to Subsection 8.2; it is based on direct calculations of the terms $\nabla_{z^{(i)}}\cdot(K-J)$ and $D_{z^{(i)}}{H_{Z}}(z^{(i)})$.

We now make the ensemble Kalman approximation to approximate the gradient term by differences, as in [23]:

where $\bar{\mathcal{G}}:=\frac{1}{I}\sum_{k=1}^{I}\mathcal{G}\left(q^{(k)}\right)$. Invoking this approximation within (3.19), using the specific form of $\Phi$, yields the following system of interacting particles in $\mathbb{R}^{N}$, for $i=1,\dots,I:$

We will use this system as the basis of all our numerical experiments.

In this subsection, we show the affine invariance property [26, 24, 37] for the Fokker-Planck equations in the mean-field regime (3.12) and for the ensemble approximation (3.18); the particle equation in the mean-field regime (1.7), for the ensemble approximation (1.9), and its correction (3.19); and the gradient-free approximation (4.20). For simplicity of presentation, we only state the results in the case of ensemble approximation, and the mean field case is a straightforward analogy upon dropping all of the particle superscripts.

Then we have the observation that all of our systems are affine invariant.

We defer the proof in Subsection 8.3. The significance of affine invariance is that it implies that the rate of convergence is preserved under affine transformations. The proposed methodology is thus uniformly effective for problems that become well-conditioned under an affine transformation. Proposition 5.1, which follows in the next section, illustrates this property in the setting of linear forward map $\mathcal{G}(\cdot).$

We employ a splitting method to integrate the stochastic dynamical system given by equation (4.20): the first capturing the Hamiltonian evolution, including the finite-size correction, and the second capturing an OU process in momentum space. The Hamiltonian evolution follows the equation:

The OU process follows the equation:

We implement a symplectic Euler integrator [48] for the Hamiltonian part, i.e., take a half step $\epsilon/2$ of momentum updates, then a full step $\epsilon$ of position updates, and finally a half step $\epsilon/2$ of momentum updates.²²2With the ensemble approximation the system is only approximately Hamiltonian; the splitting used in symplectic Euler is still well-defined, however. Let $Z_{j}$ be the collection of all position and momentum particles at time $j$: $\{q_{j}^{(i)},p_{j}^{(i)}\}_{i=1}^{I}.$ Starting from time $j$ this symplectic Euler integration gives map $Z_{j}\mapsto\hat{Z}_{j}.$ We set $q_{j+1}^{(i)}$ to be the $i^{\rm th}$ position coordinates of $\hat{Z}_{j}$ and then update the momentum coordinates using the OU process; the damping coefficient $\gamma>0$ is treated as a hyperparameter of EKHMC. Vector $Z$ is set at the value given by output of the preceding symplectic Euler integrator, denoted by $\hat{Z}_{j}$. The update of the $i^{\rm th}$ momentum coordinate given by the OU process is then:

where $\hat{p}^{(i)}_{j}$ are the momentum coordinates from $\hat{Z}_{j}$. Within the OU process the damping coefficient $\gamma>0$ is treated as a hyperparameter of EKHMC.

Similarly to [23], and as there with the goal of improving stability and convergence speed towards posterior distribution, we implement an adaptive step size, i.e., the true step size is rescaled by the magnitude of the “force field” $F_{H}$:

We follow the example presented in Section 4.3 of the paper [23]. We start by defining the forward map which is given by the one-dimensional elliptic boundary value problem

The solution is given explicitly by

The forward model operator $\mathcal{G}$ is then defined by

Here $u=\bigl{(}u_{1},u_{2}\bigr{)}^{T}$ is a constant vector that we want to find and we assume that we are given noise measurements $y$ of $p(\cdot)$ at locations $x_{1}=0.25$ and $x_{2}=0.75$. Parameters are chosen according to [23], but we summarize them here for completeness:

• •

  noise $\eta\sim\mathsf{N}(0,\Gamma)$, $\Gamma=0.1^{2}I_{2}$;

• •

  prior $\pi_{0}=\mathsf{N}(0,\Gamma_{0})$, $\Gamma_{0}=\sigma^{2}I_{2}$, $\sigma=10;$

• •

  measurement $y=(27.5,79.7)^{T};$

• •

  number of particles $J=10^{3};$

• •

  initialization: $(u_{1},u_{2})\sim\mathsf{N}(-3.5,0.1^{2})\times\mathsf{U}(70,110).$

Here $\mathsf{U}$ is the uniform distribution.

We choose $a=0.01$, $\epsilon=0.2$ in both EKS and EKHMC. We find choosing $\gamma=1$ causes overshooting, a phenomenon which can be ameliorated by increasing to $\gamma=100$; indeed this latter value appears, empirically, to be close to optimal in terms of convergence speed. We conjecture this difference from the linear case, where the optimal value is $\gamma=1.83$ (see Remark 5.2), is due to the non-Gaussianity of the desired target: the particles have accumulated considerable momentum when entering the linear convergence regime, and so extra damping is required to counteract this and avoid overshooting. Despite the desirable problem independence of the optimal $\gamma$ in the linear case, we expect case-by-case optimization to be needed for nonlinear problems.

We evolve the ensemble of particles for 200 iterations, and record their positions in the last iteration as the approximation of the posterior distribution, shown in Figure 1. The EKHMC samples are quite similar to EKS samples, both of which are reasonable approximations of the samples obtained from the gold standard HMC [19] simulations – but both approximations of the gold standard miss the full spread of the true distribution, because of the ensemble Kalman approximation they invoke. We also compute four ensemble quantities to characterize the evolution of EKHMC and EKS: $u_{1}$ mean, $u_{2}$ mean, two eigenvalues $\sigma_{1}$ and $\sigma_{2}$ of the covariance matrix $C_{q}(Z)$. EKHMC has faster convergence towards equilibrium than EKS, benefiting from the second-order dynamics.

This example follows the problem setting detailed in [23]. We summarize the essential problem specification here for completeness.The forward problem of porous medium flow, defined by permeability $a(\cdot)$ and source term $f(\cdot)$, is to find the pressure field $p(\cdot)$ where $p(\cdot)$ is solution of the following elliptic PDE:

We assume that the permeability $a(x)=a(x;u)$ depends on some unknown parameters $u\in\mathbb{R}^{d}$. The resulting inverse problem is, given (noisy) pointwise measurements of $p(x)$ on a grid, to infer $a(x;u)$ or $u$. We model $a(x;u)$ as a log-Gaussian field. The Gaussian underlying this log-Gaussian model has mean zero and has precision operator defined as $\mathcal{C}^{-1}=(-\triangle+\tau^{2}\mathcal{I})^{\alpha}$; here $\triangle$ is equipped with Neumann boundary conditions on the spatial-mean zero functions. We set $\tau=3$ and $\alpha=2$ in the experiments. Such parametrization yields as Karhunen-Loeve (KL) expansion

where $\varphi_{l}(x)={\rm cos}(\pi\langle l,x\rangle)$, $\lambda_{l}=(\pi^{2}|\ell|^{2}+\tau^{2})^{-\alpha}$, $K\equiv\mathbb{Z}^{2}$. Draws from this random field are Hölder with any exponent less than one. In practice, we truncate $K$ to have dimension $d=2^{8}$. We generate a truth random field by sampling $u\sim\mathsf{N}(0,I_{d})$. We create data $y$ with $\eta\sim\mathsf{N}(0,0.1^{2}I_{K})$. We choose the prior covariance $\Gamma_{0}=10^{2}I_{d}$.

To characterize the evolution of EKHMC and compare it with EKS, we compute two metrics:

For these metrics, we use norms

where the first is defined in the negative Sobolev space $H^{-2}$, whilst the second in the $L^{2}$ space.

We set $a=0.01$ and $\epsilon=1.0$ for both EKHMC and EKS. We set $\gamma=1$ in EKHMC; unlike the previous example, this choice does not lead to over-shooting. We simulate the particles for $100$ iterations, and use their positions in the last iteration as the approximation to the posterior distribution. We compare the samples obtained from MCMC (we use the pCN variant on RWMH [16]), EKS ($J=128,512,2048$) and EKHMC ($J=128,512,2048$) in Figure 2. The evolution of $||u||_{H^{-2}}$ and $||u||_{L^{2}}$ are plotted to show that convergence is achieved, and that EKHMC converges faster than EKS.