Deterministic Langevin Monte Carlo with Normalizing Flows for Bayesian Inference
Abstract
We propose a general purpose Bayesian inference algorithm for expensive likelihoods, replacing the stochastic term in the Langevin equation with a deterministic density gradient term.
The particle density is evaluated from the current particle positions using a Normalizing Flow (NF), which is differentiable and has good generalization properties in high dimensions.
We take advantage of NF preconditioning and NF based Metropolis-Hastings updates for a faster convergence.
We show on various examples that the method is competitive against state of the art sampling methods.
Introduction
The task of Bayesian inference is to determine the posterior p ( x | y ) of d parameters x given the likelihood p ( y | x ) of some data y and given the prior p ( x ) , such that the posterior is given by
Bayes theorem as p ( x | y )
=
p ( y | x ) p ( x ) /p
( y ) .
While we have access to the joint p ( x , y )
=
p ( y | x ) p ( x ) , the normalization p ( y ) is generally unknown.
Bayesian posterior inference can be related to the general task of ﬁnding a stationary distribution in an external potential by deﬁning
U ( x ) = − ln p ( y | x ) − ln p ( x ) =
L ( x )+ P ( x ) , where L ( x ) is the negative log likelihood and P ( x ) the negative log prior.
A common approach to this task is to use a particle based Langevin equation, which solves a stochastic differential equation (SDE) for the particle evolution (Langevin Monte Carlo, LMC). A related method is Hamiltonian Monte Carlo (HMC) [Duane et al., 1987] which uses particle positions and velocities to evolve the particles.
In these cases the long time equilibrium solution is that of the target distribution, the posterior p ( x | y ) .
LMC and HMC are two of the most popular gradient based implementations of Markov Chain Monte Carlo (MCMC), but for ﬁnite step size they require Metropolis-Hastings (MH) adjustment to become unbiased.
In this case they have theoretical convergence guarantees to the stationary distribution (under mild ergodicity assumptions), but suffer from chain element correlations, which may render the convergence to be very slow.
In this paper we propose a Deterministic Langevin Monte Carlo (DLMC) with Normalizing Flows (NF).
It uses a random initialization from the prior, but the subsequent evolution is deterministic, replacing the stochastic velocity term in the Langevin equation with a deterministic density gradient term.
To avoid the need to solve the Fokker-Planck equation directly we evaluate the density deﬁned by the current particle positions via an NF, and then use its gradient to update the ﬁrst order particle Langevin dynamics.
NFs are differentiable and have good generalization properties in high dimensions.
Deterministic sampling methods have until now been limited to particle based methods. Two examples are Stein Variational Gradient Descent (SVGD)
[Liu and Wang, 2016], and interacting particle solutions of Maoutsa et al. [2020], which proposes particle based scores for the gradient of the log-density.
However, these particle based methods scale poorly to high dimensions.
The main novelty of our contribution is the introduction of NFs to the class of deterministic methods for evaluating stationary distributions: we use NFs for density estimation and gradient evaluation, as well as for MH adjustment to correct for imperfect NF density estimation and accelerate convergence.
Our primary target application is to expensive likelihoods (wall clock time of seconds or more).
An example are inverse problems in scientiﬁc applications, where often one must solve an expensive
Ordinary or Partial Differential Equation (ODE or PDE) forward model, where the computational cost can be minutes or hours.
In such applications the computational cost of the NF (which is of order seconds in our experiments) is not dominant.
For fast model evaluations standard MCMC methods sufﬁce and our target are problems where the cost of standard methods is prohibitive.
The computational cost of a forward model is unrelated to the complexity of the posterior distribution: often the posteriors are simple Gaussian distributions even when the computational cost is high.
In such settings the main goal of posterior analysis is to minimize the number of likelihood calls to reach some prescribed precision on the posterior.
DLMC has the following advantages in comparison to several other MCMC samplers:
• DLMC takes advantage of gradient information of the target distribution, but without stochastic noise, which slows down the particle evolution towards the target in LMC and HMC.
As a result, DLMC gradient based optimization can rapidly move the particles into the region of high posterior mass for a fast convergence.
• DLMC is not a sequential Markov Chain, and DLMC particles can be updated in parallel at each time step, which can take advantage of trivial machine parallelization of likelihood evaluations, reducing the wall-clock time compared to standard MCMC.
• DLMC particles are initialized as a random draw from the prior, which covers the posterior.
Each particle produces an independent sample in the limit of a large number of particles ( N ).
• DLMC uses NF for density estimation, and takes advantage of NF preconditioning and NF based MH updates, for a faster convergence and to help correct for imperfect density estimation. • DLMC can handle multimodal posteriors, in contrast to standard MCMC samplers.
Langevin and Fokker-Planck equations
The overdamped Langevin equation is a stochastic differential equation describing particle motion in an external potential and subject to a random force with zero mean, [Equation]
[Equation]
) where (cid:104) η ( t ) (cid:105) = 0 and (cid:104)
η
i ( t ) η j ( t (cid:48) )
(cid:105) =
δ ij δ ( t − t (cid:48) ) .
For simplicity we set the diffusion coefﬁcient, temperature and mobility to unity.
The Langevin equation is a ﬁrst order velocity equation which has a deterministic velocity − ∇ U ( x ) and a stochastic velocity η .
Here the gradient operator is with respect to the parameters of U ( x ) , ∇ U ( x ) =
dU ( x ) /d x .
In practice we need to discretize the Langevin equation, and for any ﬁnite step size ∆ t the result is a biased distribution.
An example is the Ornstein-Uhlenbeck process, where U ( x ) is a harmonic potential, exp( − U ( x ))
=
N ( x ; µ , Σ ) .
Standard Langevin algorithm updates lead to the solution q ( x ) =
N ( µ , Σ
( I − ∆ t Σ − 1 / 2) − 1 )
[Wibisono, 2018].
We can see that for a ﬁnite ﬁxed stepsize ∆ t the solution converges to a biased answer regardless of the number of time integration steps taken, or the number of samples used.
A solution to this is to supplement the Langevin updates with the Metropolis-Hastings (MH)
Adjustment Langevin Algorithm (MALA), where MH acceptance guarantees detailed balance [Roberts and Tweedie, 1996].
A similar MH adjustment is also required for HMC [Duane et al., 1987].
The Langevin equation can be viewed as a particle implementation of the evolution of the particle probability density q ( x ( t )) , which is governed by the deterministic Fokker-Planck equation, a continuity equation for the density,
[Equation]
We deﬁned V ( x ( t ))
= − ln q ( x ( t )) and expressed current as density times velocity, where the two terms in the probability current J correspond to the two velocity terms in the Langevin equation.
One can see from equation 2 that the stationary distribution ˙
q ( x , t ) = 0 is given by p ( x | y ) =
exp( − U ( x )) /p
( y ) .
The Fokker-Planck equation is a PDE, while the Langevin equation is stochastic (SDE).
Langevin diffusion provides a particle description of the density q ( x , t ) and in the large time limit both equations lead to a stationary distribution equal to the posterior p ( x | y ) .
If we replace the stochastic velocity in the Langevin equation 1 with the deterministic velocity in equation 2, we obtain the deterministic Langevin equation [Maoutsa et al., 2020, Song et al., 2020],
[Equation]
Equation 3 gives the dynamics of a particle representation of the density q ( x ( t )) , which in the large time limit t → ∞ converges to the same distribution as the solution of equation 1.
This follows from the fact that the stationary solution of equation 3, ˙ x ( t ) = 0 is given by V ( x ( t ))
=
U ( x ( t )) + c , with c being a constant independent of x , which must thus be equal to log p ( y ) , since q ( x ( t )) is normalized.
There is no stochastic noise in Equation 3, so it can be interpreted as a deterministic analog of the Langevin equation, which is a ﬁrst order equation for velocity.
The main difﬁculty in solving equation 3 is the instantaneous density term q ( x ( t )) needed to obtain
V ( x ( t )) .
Here we solve this by evolving N particles together, using NFs to evaluate their particle density.
While this makes the particles interacting, in the limit of large N the particle evolution is independent.
Furthermore, we can use the NF to do MH adjustment to correct for imperfect NF density estimation (Section 3.2).
Upon time discretization of Equation 3 we obtain Algorithm 1.
A detailed discussion of the algorithm implementation is given in Appendix B. We initialize at t =by providing initial samples drawn at random from the known prior distribution p ( x ) , similar to Sequential Monte Carlo [Del Moral et al., 2006].
This has several positive features: since the posterior is proportional to the prior times the likelihood, the prior always covers the posterior.
Moreover, when the likelihood is weakly informative the convergence of DLMC is fast.
We then apply DLMC gradient updates to obtain the particle position updates.
When initializing from the prior the initial density distribution cancels out the prior in the target, and the initial gradient update is the log likelihood term ∇ L ( x ) .
Equation 3 can be interpreted as a gradient based minimization of the time dependent objective U ( x ( t ))
− V ( x ( t )) .
In this view we can optimize the objective using any optimization method.
Initially, if we start the particles from the prior p ( x ) , DLMC is simply performing optimization of the target likelihood p ( y | x ) .
Later, as q ( x ( t )) ∼ p ( x | y ) , the gradient of the optimization objective becomes zero, and the particles stop moving.
One can use any optimization method to move the particles, so if the gradient of U ( x ) is not available we can use gradient free optimization methods to do so.
For gradient based optimization we can use many of the standard optimization methods [Nocedal and Wright, 2006], such as ﬁrst order (stochastic or deterministic) gradient descent, momentum based methods, or (quasi) second order methods (e.g., L-BFGS, Newton, conjugate gradient etc.).
The algorithm can be stopped when the particle distribution becomes stationary.
For example, we can stop when the estimated ﬁrst and second moments of the posterior stop changing.
Given the formulation above, the performance of DLMC depends on the ability of the NF to approximate the particle density and the corresponding gradient at each time step. The addition of MH updates (Section 3.2) help to correct for imperfect density estimation. However, in situations where the NF ﬁt is a poor match for the target the MH acceptance can be very low.
Furthermore, in high dimensions ( O (10 3 ) ) we are limited by the training cost to use simple NF approximations.
However, in this scenario DLMC can still be used to move particles quickly towards the typical set, acting as an initialization for MCMC in the NF latent space.
This is discussed further in Section 5.5.
Normalizing Flows provide a powerful framework for density estimation and sampling [Dinh et al., 2017, Papamakarios et al., 2017, Kingma and Dhariwal, 2018, Dai and Seljak, 2021].
These models map the d -dimensional data x to d -dimensional latent variables z through a sequence of invertible transformations f = f 1 ◦ f 2 ◦ ... ◦
f L , such that z = f ( x ) and z is mapped to a base distribution π ( z ) , which is chosen to be a standard Normal distribution N (0 , I ) .
The probability density of data x can be evaluated using the change of variables formula:
[Equation]
The Jacobian determinant of each transform J l = | det( ∂f l ( x ) ∂ x ) | must be easy to compute for evaluating the density, and the transformation f l should be easy to invert for efﬁcient sampling.
In this paper we use Sliced Iterative Normalizing Flow (SINF) [Dai and Seljak, 2021], which has been shown to achieve better performance for small training data (below a few thousand particles), and is considerably faster, than alternatives.
3.1 Normalizing Flow as a preconditioner
While the core of DLMC is the NF based density gradient, the same NF can be used for additional convergence acceleration towards the target.
The proposed algorithm can be slow when the condition number of the problem is high.
A typical solution is to precondition with a covariance matrix, which mimics the global curvature of the problem.
This becomes more difﬁcult if the curvature is position dependent.
NFs can address this problem by warping the space in which we perform the optimization. This warping maps spatially varying curvature to the latent space where the curvature is constant and isotropic [Parno and Marzouk, 2018, Hoffman et al., 2019].
We can implement DLMC in the NF latent space z , where the target distribution p ( x | y ) is transformed into
[Equation]
where J ( z ) = (cid:12)(cid:12)(cid:12)(cid:12) det (cid:18)
∂ f − 1 ( z )
∂ z (cid:19)(cid:12)(cid:12)(cid:12)(cid:12) is the absolute value of the Jacobian determinant.
We can deﬁne the potential in latent space as
[Equation]
If we can evaluate ∇ U ( x ) using methods such as auto-differentiation and if the NF Jacobian is also differentiable we can also evaluate the gradient in latent space z as
[Equation]
In latent space the latent variables are distributed as π ( z ) =
exp( [Equation] V
[Equation]
z ))
=
N (0 , I ) , so V ( z ) =
z T z / 2 + ( d/ 2) ln(2 π ) and ∇ V ( z ) =
z .
This gives the deterministic Langevin gradient descent updates as
[Equation]
It is interesting to compare this deterministic update to the standard stochastic Langevin equation update, which in latent space is
[Equation]
where η (cid:48) ∼ N (0 , 2 I ∆ t ) .
Both updates have the potential gradient term, which is the deterministic drift velocity.
The difference between the two is that DLMC also uses the current position of the sample z for a deterministic update of the density gradient z ( t )∆ t , which moves the particles radially outwards from the center of the latent space, while the stochastic Langevin algorithm adds a Gaussian random variable η (cid:48) to the update.
We ﬁnd that performing DLMC updates in the NF latent space accelerates convergence.
3.2 Normalizing Flow for Metropolis-Hastings adjustment
The main requirement of DLMC is that q ( x ( t )) properly describes the density of samples x ( t ) at time t .
This can only be ensured in the large N limit.
Moreover, for non-convex optimization problems it is possible that the samples get stuck at local minima far from the global minimum.
One can add a Metropolis-Hastings (MH) adjustment to correct for imperfect density estimation, 1: Input: initial samples ( t at random prior p ( x ) size N access x )
P
2: Initial update: 3: for i = 1 to N do 4: x i (∆ t ) =
x i ( t = 0) − ∇ L ( x i ( t )) ∆ t 5: end for 6: while not converged do 7: t ← t + ∆ t 8: Run Normalizing Flow on current particle positions to obtain q ( x ( t ))
=
exp( − V ( x ( t )))
and map z ( t ) =
f ( x ( t )) 9: for i = 1 to N do 10: if latent space update then 11:
Map to latent space:
z i ( t ) =
f ( x i ( t )) 12:
DL update: z i ( t + ∆ t ) =
z i ( t ) − [ ∇ U ( z i ( t )) − z i ( t )]
∆ t 13:
Inverse map: x i ( t + ∆ t ) =
f − 1 ( z i ( t + ∆ t )) 14: else 15: DL update: x i ( t + ∆ t ) = x i ( t ) − ∇ [ U ( x i ( t )) − V ( x i ( t ))]
∆ t 16: end if 17:
MH update: draw a sample ˜ x i ∼ q ( x ( t )) , replace x i with ˜ x
i with probability r ( x i , ˜ x i ) 18: end for 19: end while 20: Output : a set of particles x i ( t ) distributed as a target distribution p ( x | y ) .
eliminating samples that get stuck far from the global minimum. This step is particularly necessary for disconnected multi-modal peaks of the target to be properly equilibrated. We note that whilst LMC (i.e., MALA) and HMC utilize local MH adjustment, they typically cannot equilibriate between separate posterior peaks without the use of methods such as annealing.
Speciﬁcally, at a given t we can draw independent particles ˜ x i ∼ q ( x ( t )) , one for each existing particle x i .
We compare the new particle against an existing particle to decide whether we accept or reject it.
Since these new samples are independent the MH acceptance rate is (e.g., Albergo et al.
[Equation]
Note that the normalization constant cancels out in this ratio, so we can evaluate this with knowledge of U ( x ) only.
Initially, when q ( x ( t )) is very broad compared to p ( x | y ) , the acceptance rate of MH will be low, since most of the samples drawn will have low p ( x | y ) , so this process may only eliminate the worst performers, and replace them with samples in the higher density region of p ( x | y ) .
As DLMC progresses q ( x ( t )) becomes closer to p ( x | y ) , and the acceptance increases and reaches r ( x , ˜ x )
=if the NF learns the target perfectly.
In high dimensions the MH acceptance can remain low. However, the MH update is still necessary for proper equilibration of multi-modal peaks.
Related work
NFs and transport maps for Bayesian posterior analysis have been explored in Parno and Marzouk [2018], Hoffman et al. [2019], Arbel et al.
These approaches use NFs as latent space preconditioners, using standard MCMC methods such as HMC, LMC or MH deployed in latent space, where the geometry may be more favorable for fast mixing. In Albergo et al.
[2019] NFs are used as an MH transition proposal, giving independent samples.
DLMC goes beyond these NF applications in that it also uses NFs for the density gradient in the deterministic Langevin equation.
DLMC is an ensemble method, in that all of the particles inform the NF density evaluation, which is in turn used to update each particle position.
There have been many other ensemble methods proposed in the literature, such as the Afﬁne Invariant Ensemble Sampler [Goodman and Weare, 2010] or Differential Evolution MCMC [Ter Braak, 2006]. DLMC differs from these methods in
[Figure]
[Figure]
that it evaluates the particle density explicitly and uses the density and target gradient to update the particle evolution.
Two deterministic algorithms related to DLMC that have recently been proposed are Stein Variational
Gradient Descent (SVGD) [Liu and Wang, 2016], which uses direct pair summation over all particles, and score based interacting particles [Maoutsa et al., 2020], inspired by generative score based models [Song et al., 2020], which use the particles to directly estimate the density gradient term.
DLMC differs from these in using NFs instead of direct pair interactions, which we ﬁnd requires fewer particles to accurately represent the particle density.
DLMC is related to Variational Inference (e.g., Blei et al. [2017]) in that both approximate the posterior with a parametrized normalized function q ( x ) .
However, DLMC does not use standard variational optimization such as KL minimization to optimize q ( x ) , and does not use q ( x ) for the posterior, which is reported by the particles instead.
In our experiments we observe that the posterior distribution of DLMC particles agrees better with the exact posteriors than the samples from the NF density.
This is expected, since we use the MH step to correct the NF.
DLMC thus contains components of both sampling and optimization: sampling as optimization in the space of distributions has been pursued in many directions since Jordan et al. [Figure]
Examples
The main DLMC application is for expensive likelihoods.
We can simulate this by synthetically creating expensive model evaluations, but will also compare the performance in terms of the number of likelihood evaluations.
The implementation of DLMC requires NF and particle evolution.
A discussion of implementation details can be found in Appendix B. For NF we use SINF, which has very few hyper-parameters [Dai and Seljak, 2021], is fast and iterative. The number of layers L can be chosen based on cross-validation, where we set aside 20% of the samples and iterate until validation data start to diverge from the training data. However, for the d = 1000 Gaussian (Section 5.5) we ﬁx L = 5 .
Typical SINF training time is of order seconds on a CPU.
The implementation of particle evolution requires us to specify the learning rate for the particle updates.
At each iteration, we take Adagrad updates in the ∇ ( U ( x ( t )) − V ( x ( t ))) direction [Duchi et al., 2011].
We use learning rates between 0 . 001 and 0 .
1 , with smaller learning rates being more robust for targets with complicated geometries such as funnel distributions.
In our plots we focus on the quality of marginal posteriors, because that is typically the primary goal of Bayesian inference
[2019] and show the squared bias of the second moment versus wall clock time or number of gradient evaluations.
The precision of this number is limited by the effective number of samples (ESS).
Here we choose bias-squared of 0.01 as sufﬁcient in terms of posteriors, and corresponds for a Gaussian distribution to an ESS of 200.
The number of particles N is a hyperparameter that must exceed this
ESS.
Overall we ﬁnd there is a tradeoff between number of particles and number of iterations, such that starting with a larger number of particles does not always imply a higher overall computational cost, due to the improved NF density estimation resulting in fewer DLMC iterations.
The main baseline we compare against is the No-U-Turn Sampler (NUTS) [Hoffman et al., 2014], an adaptive HMC variant implemented in the NumPyro library [Phan et al., 2019].
It requires tuning and burn-in, which need to be included in the computational budget.
NUTS is a standard baseline since it typically outperforms other samplers such as LMC, SVGD and SMC.
Where we include NUTS as a baseline, we use 500 tuning steps.
NUTS has the number of leapfrog steps as a tunable parameter, and when it is one HMC becomes LMC/MALA
[Girolami and Calderhead, 2011]. In some of our experiments we compare default DLMC (denoted serial, S-DLMC) to parallel DLMC (P-DLMC), where for the latter we show the wall clock time based on the assumption that we can evaluate the likelihood gradient in parallel on N CPU cores.
We assume a likelihood gradient cost of 1 minute, and the cost of the NF itself (seconds) is negligible.
5.1 Gaussian mixture
Multimodal distributions are a failure mode for many standard MCMC samplers as they cannot move from one peak to another when the potential barrier between them is too high.
They need to be augmented with an annealing procedure, where one slowly morphs the target distribution from the prior to the posterior.
Annealed Importance Sampling [Neal, 2001] and Sequential MC (SMC) [Del Moral et al., 2006] are examples of such procedures.
They require equilibrating the distribution at every temperature level using standard MCMC, and as a result are signiﬁcantly slower than standard MCMC.
DLMC uses a different strategy from annealing, but can also handle multi-modal distributions, as long as an NF can approximate it.
As an example we take a Gaussian mixture with 1/3 of the posterior mass in the ﬁrst peak and 2/3 in the second peak, in d = 100 .
We choose N = 500 , drawn from a broad initial prior ( − 2 , 2) d .
The log of the prior to posterior volume ratio is 415, which is a challenge for non-gradient based MCMCs, as they become very slow in such settings.
We apply logit transforms to all the variables such that we sample in an unconstrained space.
DLMC is gradient based and makes rapid progress towards the region of high posterior mass.
The resulting DLMC posterior after 90 iterations for the ﬁrst two variables is shown in ﬁgure 1, and is in a near perfect agreement with the true distribution.
5.2 Rosenbrock function
The 32-d example comes from a popular variant of the multidimensional Rosenbrock function [Rosenbrock, 1960], which is composed of 16 uncorrelated 2-d bananas.
The model log-likelihood is ln L = − (cid:80) n/ 2 i =1 (cid:2) ( g ( x [-
-]2 i − 1 ) 2 − x 2 i ) 2 /Q + ( x 2 i − 1 − 1) 2 (cid:3) , Q = 0 .
1 , n = 32 , with Gaussian prior N (0 , 6 2 ) on all the parameters.
We assume g ( x ) = x and use an initial burn-in phase with N =for 50 iterations, before upsampling to N = 1000 .
The posterior for the ﬁrst two variables is shown in ﬁgure 1.
We see that DLMC perfectly matches NUTS.
The quality of the posterior versus wall clock time is shown in ﬁgure 3, where we are idealizing a data assimilation process by assigning a high computational cost of 1 minute to the evaluation of g ( x ) .
This allows us to compare serial and parallel DLMC.
S-DLMC achieves results that are competitive with NUTS.
P-DLMC can improve the wall clock time by N = 1000 .
5.3 Hierarchical Bayesian analysis
Hierarchical Bayesian analyses are very common in many applications.
They often lead to funnel type posteriors, which are challenging for standard samplers.
Moreover, forward models relating the variance of the latent space to some parameters of interest can be expensive.
As an example, in cosmology data analysis applications we measure Fourier modes and want to determine their power spectrum, which in turn constrains cosmological parameters.
The evaluation of the power spectrum as a function of cosmological parameters requires an expensive PDE evolution of structure in the universe as a function of time [Lewis et al., 2000].
Here we choose a simpliﬁed version where the power variance P = g ( θ ) is determined by cosmological parameter θ and controls the variance of parameters
z i .
Parameters z i are not directly observable, and instead we observe their noisy version with a likelihood of the data y , so a combined prior and likelihood are
[Equation]
where σ 2 is the measurement noise and g ( θ ) is an expensive function that evaluates the power spectrum P as a function of θ .
To connect to Neal’s funnel we assume the result of the PDE is that the power spectrum depends exponentially on θ , with an approximate relation P =
g ( θ ) =
exp( θ ) , with a 1 minute computational cost of performing g ( θ ) .
The goal is the posterior of parameter θ marginalized over z i .
We use N = 500 .
In the limit σ 2 → ∞ the setup corresponds to Neal’s funnel, which is deemed problematic for many samplers such as HMC [Neal et al., 2011].
Note however that the DLMC solution in the σ 2 → ∞ is trivial: the data likelihood is non-informative, and after drawing samples from the prior of equationthe likelihood updates of equation 11 are zero, so the DLMC stopping criterion (stationary posterior) is satisﬁed immediately.
To make the problem closer to actual applications in cosmology we analyze it with a ﬁnite noise σ 2 , such that the data are informative of the parameter θ .
We create a simulated data set with the true value θ true = 0 and use d = 101 .
The resulting posterior for σ = 5 and σ = 0 .
1 is compared to
NUTS in ﬁgure 2.
We also show the prior (i.e., Neal’s funnel) to highlight that the posterior is narrow compared to the prior.
For σ = 5 we used 15,000 likelihood calls, which is orders of magnitude fewer likelihood evaluations than the NUTS baseline, which struggles to converge at all (ﬁgure 3) without the help of a reparametrization or preconditioner [Hoffman et al., 2019].
To highlight the advantage of DLMC with parallelization we further show the performance of (embarrassingly) parellel DLMC, where we evaluate the likelihood update for all 500 particles in parallel, so each iteration is ∼ 1 minute.
This converges to the required precision b 2 = {+
+} − 2 inminutes.
Sequential MCMC methods such as NUTS cannot be so efﬁciently parallelized because they require a hyperparameter tuning and burn-in phase (each of order 10 4 likelihood calls).
This leads to 3-4 orders of magnitude larger wall clock time.
SMC can take advantage of parallelization, but needs a lot of temperature levels and a lot of steps to equilibriate at each temperature level [Del Moral et al., 2006, Wu et al., 2017], such that the overall wall clock time is still orders of magnitude larger than parallel DLMC.
5.4 Sparse logistic regression
Hierarchical logistic regression with a sparse prior applied to the German credit dataset is a popular benchmark for sampling methods [Dua and Graff, 2017].
We use the example from Hoffman et al. [2019] with d = 51 , applying a log-transform to bounded variables such that we sample in an unconstrained space.
This is a challenging example for NUTS, with samples requiring on average ∼ 125 leapfrog steps.
We run an extended burn-in phase with 10 particles to shrink from the prior to posterior volume, using 500 iterations.
We then upsample to 5000 samples, with DLMC converging to the correct solution within 10 iterations.
Upsampling to lower sample sizes, such as 1000 particles, also converges to the correct solution, but requires more DLMC iterations such that the computational cost is not reduced.
The resulting marginals are shown in ﬁgure 4 for three representative variables showing good agreement.
5.5 Ill conditioned Gaussian
To asses the performance of DLMC in high dimensions we consider an ill conditioned Gaussian target in d = 1000 .
We construct the target covariance Σ by sampling its eigenvalues from Gamma(1 , 1) , giving a condition number of ∼ 4000 .
The prior and likelihood are then deﬁned as
[Equation] where the prior and likelihood covariance are chosen such that the posterior covariance is Σ . In this high dimensional setting we are limited by the computational cost of the NF, and the ability of the NF to approximate the current particle density. However, we can use a simple ﬂow with 5 layers, to move particles quickly towards the typical set with DLMC.
In this experiment we evolvedparticles until we satisﬁed the condition V = (cid:80) di
=1 (cid:104) x · ∇ U ( x ) (cid:105) i ≤ 1 .
1 d , where the sum is over the target dimensions.
This virial condition is chosen because initially we have V (cid:29) d , with the particle ensemble satisfying V =
d at equilibrium [Leimkuhler and Matthews, 2015].
For this problem we are unable to achieve low bias with DLMC alone.
However, we can treat
DLMC as an initialization for HMC i.e., we run DLMC until we reach the virial threshold and then perform parallel HMC with each particle in the NF latent space. DLMC reaches the virial threshold in 28 steps. Running parallel HMC chains in latent space with 10 leapfrog steps and the step size targeting an acceptance rate of 0.65, we reach a bias-squared of 0.01 after 5 HMC iterations. This corresponds to ∼ 10 5 serial likelihood evaluations, and ∼ 100 parallel evaluations.
The resulting particle distributions from DLMC and the subsequent latent space HMC are shown in ﬁgure 1. By comparison, NUTS requires ∼ 10 5 serial likelihood evaluations to reach a bias-squared of 0.01. For this example the serial cost of DLMC+HMC and NUTS is comparable, with DLMC+HMC providing a performance improvement of ∼ 10 3 when implemented in parallel. In this experiment we have used a very simple implementation of HMC, and the convergence rate could likely be improved by using some ensemble adaptation scheme e.g., Hoffman and Sountsov [2022].
In Appendix C we compare to kernel density estimation based DLMC and SVGD, both of which give inferior results, conﬁrming that the NF is the key ingredient for the success of deterministic methods.
In Appendix D we present results of ablation studies. We show DLMC without MH, DLMC without NF preconditioning, and pure NF MH without DLMC.
In all cases the results are inferior to DLMC with preconditioning and with MH adjustment.
We also show NF with MH without DL is inferior to DLMC.
Conclusions
We present a new general purpose Bayesian inference algorithm we call Deterministic Langevin
Monte Carlo with Normalizing Flows.
The method uses the deterministic Langevin equation, which requires the density gradient, which we propose to be evaluated using a Normalizing Flow determined by the positions of all the particles at a given time.
Particle positions are then updated in the direction of log target density minus log current density gradient.
To correct for imperfect NF density estimation, we add an MH step.
The process is repeated until the particle distribution becomes stationary, which happens when the particle density equals the target density.
In our experiments we ﬁnd that DLMC signiﬁcantly outperforms baselines such as LMC, HMC and
SVGD, and even more in terms of wall-clock time when likelihoods can be evaluated in parallel on a multi-core computer.
A possible limitation is that NF training may not learn the true target distribution for small N .
While DLMC with MH does offer some robustness against imperfect density estimation, it can fail in situations where the MH acceptance rate is very low, such that it cannot reach the stationary distribution in a ﬁnite number of steps.
This is particularly apparent in O (10 3 ) dimensions, where we are limited by the NF training cost and the ﬁdelity of its gradient approximation. Here, convergence on the correct target can be achieved by running DLMC until the particle ensemble reaches the typical set, at which point we perform a short MCMC run in latent space.
Indeed, treating DLMC as an initialization for a latent space MCMC in this way can provide some additional robustness through the associated asymptotic guarantees, whilst avoiding the expensive tuning and burn-in process typically required for MCMC.
It would nonetheless be interesting to explore what NF architectures, regularization, and training procedures safeguard against low MH acceptance rates.
Acknowledgements
This material is based upon work supported by the U.S. Department of Energy, Ofﬁce of Science,
Ofﬁce of Advanced Scientiﬁc Computing Research under Contract No. DE-AC02-05CH11231 at Lawrence Berkeley National Laboratory to enable research for Data-intensive Machine Learning and
Analysis.
We thank Minas A Checklist For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[No] German credit data taken from public UCI repository.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] Code included as supplemental material.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]