Ensemble sampler for inﬁnite-dimensional inverse problems
AABI
Abstract
We introduce a new Markov chain Monte Carlo sampler for inﬁnite-dimensional
Bayesian inverse problems.
The new sampler is based on the aﬃne invariant ensemble sampler, which we extend for the ﬁrst time to function spaces.
The new sampler is more eﬃcient than preconditioned Crank-Nicolson, yet it requires no gradient information or posterior covariance information, making the sampler broadly applicable.
Introduction
In Bayesian inverse problems, Markov chain Monte Carlo (MCMC) methods are needed to approximate distributions on inﬁnite-dimensional function spaces.
Yet designing an eﬃcient and broadly applicable MCMC method for this context has proved challenging.
The ﬁrst MCMC sampler for function spaces was preconditioned Crank-Nicolson (pCN, [2]).
This early approach can be used for a wide variety of Bayesian inverse problems, yet pCN is ineﬃcient, requiring a very large number of iterations to estimate statistics of the posterior distribution [3].
In contrast, newly developed MCMC methods sample the posterior distribution more eﬃciently [3, 4, 1, 8, 7], yet the new samplers require gradient or posterior covariance information which is challenging to obtain for many inverse problems.
This raises the question: what function space sampler is more eﬃcient than pCN without requiring gradient or posterior covariance information?
A gradient-free sampler is needed when the forward model has a legacy code base and obtaining derivatives is diﬃcult or impossible.
A covariance-free sampler is needed when estimating the posterior covariance matrix is a challenging procedure that would use up excessive computational resources.
In the literature on ﬁnite-dimensional MCMC, there is a gradient-free sampler that quickly “learns” the posterior covariance structure through an ensemble of interacting walkers.
This sampler, called the aﬃne invariant ensemble sampler (AIES, [6]), is easy to tune, easy to parallelize, and eﬃcient at sampling spaces of moderate dimensionality ( d ≤ 20).
AIES is used extensively due to its implementation in the well-known emcee package for python [5].
The main contribution of this work is to extend the AIES to the setting of inﬁnitedimensional inverse problems with a Gaussian prior distribution.
We ﬁrst identify the
Karhunen–Lo`eve (KL) expansion for the prior distribution.
Then, we alternate between
AIES and pCN sampling steps.
We use AIES to sample the posterior distribution on the ﬁrst few KL coeﬃcients, while we use pCN to sample the posterior distribution on the higher KL coeﬃcients.
By combining AIES with pCN, we obtain an ensemble sampler that is stable, eﬃcient, and broadly applicable across a large class of Bayesian inverse problems.
Background
In this section we brieﬂy review the pCN and AIES samplers.
Firstly, pCN is a dimension-independent random walk sampler and has the following proposal: starting from a position U , the pCN update takes the form
[Equation]
where ξ ∼ N (0 , C ) is a random draw from the Gaussian prior distribution and ω ∈ (0 , 1] is the step size parameter.
This proposal is independent of dimension and has acceptance probability
[Equation]
In contrast, AIES is a ﬁnite-dimensional sampler is deﬁned by an ensemble of walkers −→ X =
( X 1 , ...
X L ) to sample from a product density π ( x 1 ) · · · π ( x L ) with each walker taking values in R M .
The proposal for walker X i to the new position is:
[Equation]
where Z ∈ [1 /a, a ] is a random number with density g ( Z ) ∝ 1 / √ z ,.
This proposal has acceptance probability
[Equation]
New ensemble sampler
Under the posterior distribution π ( du ) ∝ exp ( φ ( u ))
π 0 ( du ), the KL coordinates have an unknown distribution, which must be approximated through sampling.
However, the prior distribution restricts the posterior distribution in substantial ways.
The higher KL coeﬃcients (ie: the coeﬃcients of the high frequency basis elements) in particular, are already highly constrained under the the prior distribution, so the prior and posterior distributions must be similar.
As a consequence, the pCN approach, which is designed for eﬃcient sampling from the prior distribution, works well for sampling the posterior distribution on the higher KL coordinates.
The most diﬃcult part of sampling the Bayesian posterior distribution is sampling the ﬁrst few KL coeﬃcients.
These coordinates are barely constrained under the prior distribution, so the posterior distribution can take an arbitrary shape, which may be poorly scaled or multimodal.
As a consequence, the pCN sampler may work very slowly on these coordinates, in which case an improved sampling strategy is required.
To address the sampling bottlenecks in Bayesian inverse problems, we propose a Metropoliswithin-Gibbs sampler that alternates between using pCN on the higher KL coeﬃcients and using the AIES sampler on the ﬁrst few KL coeﬃcients.
The proposed function space ensemble sampler robustly improves the sampling performance of pCN, as we show through two examples.
Remarkably, this is the ﬁrst known sampler that improves on the performance of pCN without requiring gradients and without requiring posterior covariance information.
Algorithm
To sample a Bayesian posterior distribution π ( du ) ∝ exp ( φ ( u ))
π 0 ( du ) with a
Gaussian prior π (0 , C ), we perform the following steps:
Identify the matrix J whose columns are the ﬁrst M eigenvectors of the covariance operator.
Set P =
JJ T and Q
= I − JJ T .
Initialise an ensemble of walkers −→ X =
( X 1 , ...
X L ).
Repeat the following sampling steps: (a) For each walker X
i , propose the update
[Equation]
where ξ ∼ N (0 , C ).
Accept the proposal with probability
[Equation]
(b) For each walker X
i , randomly choose a diﬀerent walker X j (cid:54) =
X i .
Propose the update
[Equation]
where Z ∈ [1 /a, a ] has density g ( z ) ∝ 1 / √ z .
Accept with probability
[Equation]
Numerical experiments
In this section we will consider two examples to illustrate the performance of the FES sampler on problems that have both scalar and functional parameters.
The ﬁrst example - the advection equation - shows how the FES sampler can increase mixing speed by two orders of magnitude over pCN.
We also illustrate the eﬀect of varying the main tuning parameter M.
The second example - path reconstruction for langevin dynamics - shows the robustness of FES to mild multimodality and model misﬁt.
We compare FES to pCN as well as to another gradient-free sampler (which is main competitor of our method) and ﬁnd that FES performs better.
For both examples we ﬁx the stretch move step size to a = 2, as recommended in [5] and tune the pCN proposal ω to give an acceptance rate of 20%.
4. Advection equation
Our ﬁrst example involves parameter estimation for the advection equation, which is linear hyperbolic partial diﬀerential equation (PDE).
This simple model is the special case of more complicated nonlinear hyperbolic PDEs arising in ﬂuid dynamics, acoustics, motorway traﬃc ﬂow, and many other applications.
[Equation]
We simultaneously estimate the advection speed c and the initial condition ρ 0 from a set of noisy observations.
In this example we set a uniform prior on the wave speed π 0 ( c ) ∼ U (0 , 1 . 4), and set a Gaussian process prior on the initial condition on [0 , 10].
We use a squared exponential kernel, with l = 1 and σ 0 = 130:
[Equation]
We discretise the initial condition using 200 points.
We generate observations we set c true = 0 .
5 and use a sample from the GP prior as initial condition.
We consider 9 equally spaced observation in the x-t plane and add Gaussian observational noise with standard deviation σ obs = 0 . 1.
We run FES with L = 100 walkers and M = 0 , 1 , 5 , 10 , 20 for 20 , 20 , 10 , 2 , 2 million iterations respectively.
The sampler uses Metropolis-with-Gibbs updates as in algorithm (1); namely it alternates proposing the ﬁnite space and the complementary subspace.
We also run a ”vanilla” sampler that proposes a joint update for the wave speed and initial condition.
The proposals are a Gaussian random walk for the wave speed and pCN for the initial condition.
We show the ACF plots in ﬁgure (1): these illustrate how FES can show a two orders of magnitude speedup compared to pCN.
We also see how M = 10 is optimal, and that choosing M too large causes mixing to slow down again.
4. Path reconstruction for Langevin dynamics
In this example we compare the performance of pCN, FES, and the hybrid sampler developed in [8].
This hybrid sampler also considers a truncation of the KL expansion, but has adaptive
Gaussian proposals for the ﬁnite-dimensional space.
We also compare the performance of FES for diﬀerent number of walkers as well as joint and MwG updates.
We ﬁnd that MwG updates are preferable, and that using a large number of walkers is better (as recommended in [5]).
We consider Langevin dynamics on a single well potential giving the evolution of position
X t and momentum P t :
Our second example involves parameter estimation for the noisy harmonic oscillator
[Equation]
We use as true path X ( t ) = sin(4 t ) and generate noisy observations from it.
Note that this model misﬁt makes the sampling problem more challenging.
We simultaneously estimate the drift parameter α , the diﬀusion parameter σ , and the driving Brownian motion W from noisy observations.
We compare the performance of pCN, the new ensemble sampler, and another gradient-free sampler in the literature.
This posterior exhibits multimodality in the alpha parameter which change the frequency of the X t paths (as seen in ﬁgure 2).
[Figure]
We show the ACF plots of the diﬀerent samplers in ﬁgure (3).
We ﬁnd that pCN works slowly because of the multimodality of the α parameter.
However, the new ensemble sampler freely traverses the modes, leading to faster sampling than pCN and the other alternative approach that explicitly estimates the posterior covariance matrix.
FES is therefore more robust than the other samplers to mild multimodality as well as model misﬁt (recall that the data was generated from X ( t )
= sin(4 t )).
We also ﬁnd that as the hybrid sampler takes a large number of iterations for the adaptation of the Gaussian proposal to stabilise (around 1 million iterations).
This long adaptation can be a problem when the forward model is computationally expensive, as is often the case with ODE and PDE models.
Conclusion
We introduce a new gradient-free sampler for function spaces which is shown to outperform a random walk proposal (pCN) as well as a specialised gradient-free sampler in the literature.
In terms of limitations, there are two aspects to our work.
Both work well but both can potentially be improved further.
First, we approximate the subspace where pCN is not good enough for sampling using the KL expansion and keeping the ﬁrst few basis elements.
It would beneﬁcial to ﬁnd a better way to estimate this low-dimensional subspace without gradients.
However, estimating the full covariance is diﬃcult, especially for high discretisation.
After isolating the subspace where pCN is not good enough for sampling, the
AIES is a good gradient-free sampler, but there might be other better methods that can be proposed, speciﬁcally, in the case of extreme multimodality, which is acknowledged [6] to be the main limitation of AIES.