# Markovian Flow Matching: Accelerating MCMC with Continuous Normalizing Flows

Alberto Cabezas
Department of Mathematics and Statistics
Lancaster University, UK
<EMAIL_ADDRESS>
&Louis Sharrock
Department of Mathematics and Statistics
Lancaster University, UK
<EMAIL_ADDRESS>
&Christopher Nemeth
Department of Mathematics and Statistics
Lancaster University, UK
<EMAIL_ADDRESS>
###### Abstract

Continuous normalizing flows (CNFs) learn the probability path between a
reference and a target density by modeling the vector field generating said
path using neural networks. Recently, Lipman et al. [42] introduced a simple
and inexpensive method for training CNFs in generative modeling, termed flow
matching (FM). In this paper, we re-purpose this method for probabilistic
inference by incorporating Markovian sampling methods in evaluating the FM
objective and using the learned probability path to improve Monte Carlo
sampling. We propose a sequential method, which uses samples from a Markov
chain to fix the probability path defining the FM objective. We augment this
scheme with an adaptive tempering mechanism that allows the discovery of
multiple modes in the target. Under mild assumptions, we establish convergence
to a local optimum of the FM objective, discuss improvements in the
convergence rate, and illustrate our methods on synthetic and real-world
examples.

##  1 Introduction

The task of sampling from a probability distribution known only up to a
normalization constant is a fundamental problem arising in a wide variety of
fields, including statistical physics [48], Bayesian inference [22], and
molecular dynamics [40]. In particular, let $\smash{\pi(\mathrm{d}x)}$ be a
target probability distribution on $\smash{\mathbb{R}^{d}}$ with density
$\smash{\pi(x)}$ with respect to the Lebesgue measure of the form111In a
slight abuse of notation, we use $\pi$ to denote both the target distribution
and its density.

$\pi(x)=\frac{\hat{\pi}(x)}{Z},$ (1)

where $\hat{\pi}:\mathbb{R}^{d}\rightarrow\mathbb{R}_{+}$ is a continuously
differentiable function which can be evaluated pointwise, and
$\smash{Z=\int_{\mathbb{R}^{d}}\hat{\pi}(x)\mathrm{d}x}$ is an unknown
normalizing constant. We are interested in generating samples from $\pi$ in
order to approximate integrals of the form $\pi[f]=\mathbb{E}_{\pi}[f(x)]$,
where $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$.

A standard solution to this problem is Markov chain Monte Carlo (MCMC) [61,
10], which relies on the construction of a Markov process which admits the
target $\pi$ as its invariant distribution. One of the most broadly applicable
and widely studied MCMC methods is the Metropolis-Hastings (MH) algorithm
[29], which proceeds in two steps. First, given a current sample $x$, a new
sample is proposed according to some proposal distribution $q(\cdot|x)$. Then,
this sample is accepted with probability
$\smash{\alpha(x,y)=\min\big{\\{}1,\tfrac{\pi(y)q(x|y)}{\pi(x)q(y|x)}\big{\\}}}$.
This strategy generates a Markov chain with the desired stationary
distribution and, under mild conditions on the proposal and the target, also
ensures that the Markov chain is ergodic [62]. However, for high-dimensional,
multi-modal settings, such methods can easily get stuck in local modes, and
suffer from very slow mixing times [e.g., 45].

Naturally, the choice of proposal distribution $q(\cdot|x)$ is critical to
ensuring that MH MCMC algorithms explore the target distribution within a
reasonable number of iterations. A key goal is to obtain proposal
distributions with fast mixing times, which can be applied generically to any
target distribution. This is particularly challenging in the face of complex,
multi-modal (or metastable) distributions, which commonly arise in
applications such as genetics [35], protein folding [38], astrophysics [19],
and sensor network localisation [34]. On the one hand, local proposals, such
as those employed in the Metropolis-Adjusted Langevin Algorithm (MALA) [63] or
Hamiltonian Monte Carlo (HMC) [17, 53] struggle to transition between regions
of high-probability, resulting in very long decorrelation times and few
effective independent samples [e.g., 46]. On the other hand, global proposal
distributions must be very carefully designed in order to avoid high rejection
rates, particularly in high dimensions [14, 44].

Another popular approach to sampling is variational inference (VI) [75, 31,
58, 8], which obtains a parametric approximation
${\pi}_{\theta^{*}}(x)\approx\pi(x)$ to the target by minimising the Kullback-
Leibler (KL) divergence to the target over a parameterized family of
distributions $\mathcal{D}_{\theta}=\\{\pi_{\theta}:\theta\in\Theta\\}$.
State-of-the-art VI methods use normalizing flows (NFs), consisting of a
sequence of invertible transformations between a reference and a target
distribution, to define a flexible variational family [59]. There has also
been growing interest in the use of continuous normalizing flows (CNFs), which
define a path between distributions using ordinary differential equations [13,
24, 42]. CNFs avoid the need for strong constraints on the flow but, until
recently, have been hampered by expensive maximum likelihood training.

In recent years, several works have sought hybrid methods which utilise NFs to
enhance the performance of MCMC algorithms; see [25] for a recent review. For
example, NFs have been successfully used to precondition complex Bayesian
posteriors, significantly improving the performance of existing MCMC methods
[e.g., 56, 30, 36, 65]. The synergy between local MCMC proposals and global,
flow-informed proposals has also been explored, leading to enhanced mixing
rates and effective estimation of multimodal targets [e.g., 21, 64]. In this
paper, we continue this promising line of work.

#### Our contributions

We introduce a new probabilistic inference scheme which integrates CNFs with
MCMC sampling techniques. Our approach utilizes flow matching (FM), a
scalable, simulation-free training objective for CNFs recently introduced by
[42]. This enables, for the first time, the incorporation of CNFs into an
adaptive MCMC algorithm. Concretely, our approach augments a local, gradient-
based Markov transition kernel with a non-local, flow-informed transition
kernel, defined using a CNF. This CNF, and the corresponding transition
kernel, are adapted on-the-fly using samples from the chain, which are used to
define the probability path for the FM objective. Our scheme also includes an
adaptive tempering mechanism, which is essential for discovering multiple
modes in complex target distributions. Under mild assumptions, we establish
that the flow-network parameters output by our method converge to a local
optimum of the FM objective. We then demonstrate empirically the performance
of our approach on several synthetic and real-world examples, illustrating
comparable or superior performance to other state-of-the-art sampling methods.

##  2 Preliminaries

#### Continuous Normalizing Flows

A continuous normalizing flow (CNF) is a continuous-time generative model
which is trained to map samples from a base distribution $p_{0}$ to a given
target distribution [13]. Let $v_{t}$ be a time-dependent vector field that
runs continuously in the unit interval. This vector field can be used to
construct a time-dependent diffeomorphic map called a flow
$\phi:[0,1]\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$, defined via the
ordinary differential equation (ODE):

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\phi_{t}(x)=v_{t}(\phi_{t}(x)),\quad\phi_{0}(x)=x.$
(2)

Given a reference density $p_{0}$, and the flow $\phi$, we can generate a
probability density path
$p:[0,1]\times\mathbb{R}^{d}\rightarrow\mathbb{R}_{+}$ as the pushforward of
$p_{0}$ under $\phi$, viz $p_{t}:=[\phi_{t}]_{\sharp}p_{0}$, for $t\in[0,1]$.
This yields, via the instantaneous change–of-variables formula [e.g., 12]

$\displaystyle\log p_{t}(x_{t})=\log p_{0}(x)-\int_{0}^{t}\nabla\cdot
v_{s}(x_{s})\mathrm{d}s,$ (3)

where $x_{s}:=\phi_{s}(x)$, and where $\nabla$ is the divergence operator,
i.e. the trace of the Jacobian matrix. In modern applications, the vector
field $v_{t}$ is often parameterized using a neural network $v_{t}^{\theta}$,
in which case the ODE in (2) is referred to as a neural ODE [13]. In turn,
this yields a deep parametric model $\phi_{t}^{\theta}$ for the flow
$\phi_{t}$, known as a CNF [24].

#### Flow Matching

We would like to learn a CNF which maps between a given reference density
$p_{0}$ and a target density $\pi$. Given samples from the target, one
approach is to maximize the log-likelihood
$\smash{\mathbb{E}_{x\sim\pi}\left[\log p_{1}^{\theta}(x)\right]}$. In
practice, however, maximum likelihood training is very slow as both sampling
and likelihood evaluation require multiple network passes to solve the ODE in
(2). Flow Matching (FM) Lipman et al. [42] provides an alternative,
simulation-free method for training CNFs. Let $p_{t}(x)$ be a target
probability density path such that $p_{0}=p$ is a simple reference
distribution, and $p_{1}\approx\pi$ is approximately equal to the target
distribution. Let $v_{t}(x)$ be a vector field which generates this
$p_{t}(x)$. Then the FM objective for the CNF vector field
$\smash{v_{t}^{\theta}(x)}$ is defined as

$\displaystyle\mathcal{L}(\theta;\pi)=\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{x\sim
p_{t}}\left[\|v_{t}^{\theta}(x)-v_{t}(x)\|^{2}_{2}\right].$ (4)

In practice, we do not have direct access to the target vector field,
$v_{t}(x)$, and so we cannot minimize (4) directly. However, as shown in
Lipman et al. [42, Theorem 2], it is equivalent to minimize the conditional
flow-matching (CFM) loss

$\displaystyle\mathcal{J}(\theta;\pi)=\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{x_{1}\sim\pi}\mathbb{E}_{x\sim
p_{t}(\cdot|x_{1})}\left[||v_{t}^{\theta}(x)-v_{t}(x|x_{1})||^{2}_{2}\right],$
(5)

where $p_{t}(\cdot|x_{1})$ is a conditional probability density path
satisfying $p_{0}(x|x_{1})=p_{0}$ and $p_{1}(x|x_{1})\approx\delta_{x_{1}}$,
and $v_{t}(\cdot|x_{1}):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ is a
conditional vector field that generates $p_{t}(\cdot|x_{1})$. There are
various choices for $p_{t}(\cdot|x_{1})$ and $v_{t}(\cdot|x_{1})$. For
simplicity, we here assume that the conditional probability path is Gaussian,
viz
$p_{t}(x|x_{1})=\mathcal{N}(x|m_{t}(x_{1}),s_{t}(x_{1})^{2}\mathbb{I}_{d})$,
where $m:[0,1]\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ denotes a time-
dependent mean, and $s:[0,1]\times\mathbb{R}\rightarrow\mathbb{R}_{+}$ a time-
dependent scalar standard deviation. For our experiments, we further adopt the
optimal transport conditional probability path introduced in [42], setting
$m_{t}(x_{1})=tx_{1}$ and $s_{t}(x_{1})=1-(1-\sigma_{\mathrm{min}})t$ for some
$\sigma_{\mathrm{min}}\ll 1$. In this case, the conditional vector field
assumes the particularly simple form
$\smash{v_{t}(x|x_{1})=\tfrac{x_{1}-(1-\sigma_{\mathrm{min}})x}{1-(1-\sigma_{\mathrm{min}})t}}$.

##  3 Markovian Flow Matching

In this section, we present our main contribution, an adaptive MCMC algorithm
which combines a non-local, flow-informed transition kernel trained via FM, a
local Markov transition kernel, and an adaptive annealing schedule. We begin
by describing how CNFs can be used within a MH MCMC algorithm.

###  3.1 MCMC with Flow Matching

Suppose, for now, that we have access to a CNF
$\smash{(\phi_{t}^{\theta})_{t\in[0,1]}}$, trained (e.g.) via flow-matching,
with corresponding vector field $\smash{(v_{t}^{\theta})_{t\in[0,1]}}$, which
generates a probability path $\smash{(p_{t}^{\theta}})_{t\in[0,1]}$ between a
reference density $p_{0}$ and an approximation of the target density $\pi$.
Given a point $\smash{x_{0}\in\mathbb{R}^{d}}$ on the reference space, we can
evaluate the log-density of the pullback of the target distribution $\pi$ as

$\displaystyle\log[\phi_{1}^{\theta}]^{\sharp}\pi(x_{0})$
$\displaystyle=\log\pi(\phi_{1}^{\theta}(x_{0}))-\int_{1}^{0}\nabla\cdot
v_{t}^{\theta}(\phi_{t}^{\theta}(x_{0}))\mathrm{d}t.$ (6)

Under the assumption that the CNF approximately transports samples from
$p_{0}$ to $\pi$, we expect that
$\smash{[\phi_{1}^{\theta}]_{\sharp}p_{0}\approx\pi}$ in the target space, and
that $\smash{[\phi_{1}^{\theta}]^{\sharp}\pi\approx p_{0}}$ in the reference
space. Given that the reference distribution $p_{0}$ is chosen such that it is
easy to sample from, this suggests the following strategy, sometimes referred
to as _neural trasport MCMC_ or _neutraMCMC_ [56, 41, 30, 25]. First,
transform initial positions $x_{1}$ from the target space to the reference
space by solving

$\displaystyle\begin{bmatrix}x_{0}\\\ \log p_{1}^{\theta}(x_{1})-\log
p_{0}(x_{0})\end{bmatrix}=\begin{bmatrix}x_{1}\\\
0\end{bmatrix}+\int^{0}_{1}\begin{bmatrix}v_{t}^{\theta}(x_{t})\\\
-\nabla\cdot v_{t}^{\theta}(x_{t})\end{bmatrix}\mathrm{d}t,$ (7)

which integrates the combined dynamics of $x_{t}$ and the log-density of the
sample backwards in time. Then, generate MCMC proposals $y_{0}$ in the
reference space using any standard MCMC scheme which targets the pullback of
the target distribution, as defined in (6). Finally, transform accepted
proposals back to target space using the forward dynamics, viz

$\begin{bmatrix}y_{1}\\\ \log p_{1}^{\theta}(y_{1})-\log
p_{0}(y_{0})\end{bmatrix}=\begin{bmatrix}y_{0}\\\
0\end{bmatrix}+\int_{0}^{1}\begin{bmatrix}v_{t}^{\theta}(y_{t})\\\
-\nabla\cdot v_{t}^{\theta}(y_{t})\end{bmatrix}\mathrm{d}t.$ (8)

This corresponds to using a transformation-informed proposal in a Markov
transition step, an approach which has been successfully applied using
(discrete) normalizing flows [56, 41, 30, 25]. There are various possible
choices for the proposal distribution on the reference space (see Appendix A).
For example, [20] consider an independent MH (IMH) proposal, where i.i.d.
samples are drawn from the reference distribution. Here we focus on a flow-
informed random-walk, motivated largely by its superior empirical performance
in numerical experiments. This proposal performs particularly well on high-
dimensional problems, where overfitting of the CNF can be corrected with
stochastic steps, while exacerbated by independent proposals [36]. Concretely,
our flow-informed random-walk transition kernel, summarized in Algorithm 2
(see Appendix A), can be written as

$\displaystyle
P(x,\mathrm{d}y;\pi,\theta)=\alpha(x,y)\rho_{\theta}(\mathrm{d}y|x)+(1-b(x))\delta_{x}(\mathrm{d}y),$
(9)

where $\rho_{\theta}(\mathrm{d}y|x)$ is the distribution defined by the
transition

$\displaystyle
x_{0}=x+\int_{1}^{0}v_{t}^{\theta}(\phi_{t}^{\theta}(x))\mathrm{d}t,\quad
y_{0}\sim\mathcal{N}(x_{0},\sigma_{\mathrm{opt}}^{2}),\quad
y=y_{0}+\int_{0}^{1}v_{t}^{\theta}(\phi_{t}^{\theta}(y_{0}))\mathrm{d}t,$ (10)

and
$\smash{\alpha(x,y)=\min\big{\\{}1,\tfrac{\pi(y)\rho_{\theta}(x|y)}{\pi(x)\rho_{\theta}(y|x)}\big{\\}}}$
and $b(x)=\int_{\mathbb{R}^{d}}\alpha(x,y)\rho_{\theta}(\mathrm{d}y|x)$.

#### Training the CNF

Thus far, we have assumed that it is possible to train a CNF which maps
samples from the reference distribution $p_{0}$ to (an approximation of) the
target distribution $\pi$. Clearly, however, the CFM objective is not
immediately applicable in our setting, since we do not have access to samples
from the target $\pi$. [68] propose two alternatives in this case: (i) use an
importance sampling reweighted objective function, or (ii) use samples from a
long-run MCMC algorithm (e.g., MALA) as approximate target samples. Both of
these approaches, however, have limitations. The former is unlikely to succeed
when the proposal distribution differs significantly from the target
distribution, while the latter will only perform well when the chosen MCMC
method mixes well. We will adopt a different approach, similar in spirit to
other flow-informed MCMC algorithms [21, 64, 32], and the _Markovian score
climbing_ algorithm in [51], which updates the parameters of a CNF based on a
dynamic estimate of the CFM objective obtained via an adaptive MCMC algorithm.

###  3.2 Adaptive MCMC with Flow Matching

#### Overview

Our adaptive MCMC scheme combines a non-local, flow-informed transition kernel
(e.g., a flow informed random-walk) and a local transition kernel (e.g.,
MALA), which generate new samples from a sequence of annealed target
distributions. These new samples are used to define a new estimate of the CFM
objective in (5), which is optimized to define a new CNF. These steps are
repeated until the samples converge in distribution to the target $\pi$, and
the flow-network parameters converge to a local minima of the flow matching
objective (see Proposition 3.1). This scheme, which we refer to as _Markovian
Flow Matching_ (MFM), is summarized in Algorithm 1.

#### Sampling

There is significant freedom regarding the choice of both the local and the
non-local MCMC algorithms. In our experiments, we adopt the Metropolis-
Adjusted Langevin Algorithm (MALA) as the local algorithm. Thus, the local
Markov kernel $Q$ is given by

$\displaystyle Q(x,\mathrm{d}y;\pi)$
$\displaystyle=\alpha(x,y)q(\mathrm{d}y|x)+(1-b(x))\delta_{x}(\mathrm{d}y),$
(11)

where $q(\mathrm{d}y|x)$ is given by

$q(\mathrm{d}y|x)\propto\exp\left(-\frac{1}{4\tau}\|y-x-\tau\nabla\log\pi(x)\|^{2}\right)\mathrm{d}y,$
(12)

and where, as elsewhere,
$\smash{\alpha(x,y)=\min\big{\\{}1,\tfrac{\pi(y)q(x|y)}{\pi(x)q(y|x)}\big{\\}}}$
and $b(x)=\int_{\mathbb{R}^{d}}\alpha(x,y)q(\mathrm{d}y|x)$. In principle,
however, other choices such as HMC could also be used. Meanwhile, for the non-
local MCMC algorithm, we adopt the flow-informed random walk with non-local
Markov kernel $P$ defined in (9). Together, assuming alternate local and non-
local steps, these two kernels define a Markov chain with Markov transition
kernel $\smash{R:=P\circ Q}$, given explicitly by
$\smash{R(x,\mathrm{d}y;\pi,\theta)=\int_{\mathcal{Z}}Q(x,\mathrm{d}z;\pi)P(z,\mathrm{d}y;\pi,\theta)}$.
In practice, the balance between local and non-local moves is controlled by
the hyperparameter $k_{Q}$, which sets the number of local steps before a
global step.

#### Training

Following each MCMC step, the parameters of the flow-informed Markov
transition kernel $\smash{P}$ are updated based on a new estimate of
$\smash{\mathcal{J}(\theta;\pi)}$. To be precise, suppose we write
$\smash{\mu_{t}:=\mu_{0}R^{k}(\cdot,\cdot;\pi,\theta})$ for the distribution
of the Markov chain with kernel $\smash{R(\cdot,\cdot;{\pi,\theta})}$ after
$k\in\mathbb{N}$ steps, starting from initialisation $\mu_{0}$, where
$\smash{R^{k}=R\circ R\cdots\circ R}$. Our objective function is then given by

$\displaystyle\mathcal{J}(\theta;\mu_{k})=\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{x_{1}\sim\mu_{k}}\mathbb{E}_{x\sim
p_{t}(\cdot|x_{1})}\left[||v_{t}^{\theta}(x)-v_{t}(x|x_{1})||^{2}_{2}\right].$
(13)

For our choice of conditional probability path (i.e., the optimal transport
path), we can in fact rewrite this objective as [42, Section 4.1]

$\displaystyle\mathcal{J}(\theta;\mu_{k},\sigma_{\mathrm{min}})=\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{x_{1}\sim\mu_{k}}\mathbb{E}_{x_{0}\sim
p_{0}}\left[||v_{t}^{\theta}(\phi_{t}(x_{0}|x_{1}))-v_{t}(\phi_{t}(x_{0}|x_{1})|x_{1})||^{2}_{2}\right],$
(14)

where
$\smash{v_{t}(x|x_{1})=\tfrac{x_{1}-(1-\sigma_{\mathrm{min}})x}{1-(1-\sigma_{\mathrm{min}})t}}$
and $\phi_{t}(x|x_{1})=(1-(1-\sigma_{\mathrm{min}}t)x+tx_{1}$. In practice, we
will optimize a Monte Carlo estimate of this objective, namely,

$\displaystyle\mathcal{J}(\theta;\\{x^{i}(k)\\}_{i=1}^{N},\sigma_{\mathrm{min}})$
$\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\|v_{t_{i}}^{\theta}(\phi_{t_{i}}(x^{i}_{0}|x^{i}(k)))-v_{t_{i}}(\phi_{t_{i}}(x^{i}_{0}|x^{i}(k))|x^{i}(k))\|_{2}^{2},$
(15)

where $\smash{\\{x^{i}(k)\\}_{i=1}^{N}}$ are the samples from $N$ chains of
our MCMC algorithm after $k\in\mathbb{N}$ iterations,
$\smash{x^{i}_{0}\overset{\mathrm{i.i.d.}}{\sim}p_{0}}$, and
$\smash{t_{i}\sim\mathcal{U}(0,1)}$. The use of $N$ particles allow the
state’s mutation to $N$ computing cores running in parallel at each iteration.
Sampling steps can be run in parallel using modern vector-oriented libraries,
before each particle is used to approximate the loss and update the
parameters. Thus, the speedup gained by using more than one core scales
linearly with the number of cores as long as there are as many cores as there
are particles.

#### Annealing

For complex (e.g., multimodal) target distributions, it can be challenging to
learn a CNF that successfully maps between the reference $p_{0}$ and the
target $\pi$. For example, if the locations of the modes of the target are not
known a priori, and the MCMC chains are initialized far from one or more of
the modes, it is unlikely that the local MCMC kernel, and therefore the
trained flow, will ever discover these modes [e.g., 21, Section IV.C]. To
alleviate this problem, one approach is to iteratively target a sequence of
annealed densities $\smash{\\{\eta_{t}(x)\\}_{t=0:T}}$, which smoothly
interpolate between a simple base distribution $\smash{\eta_{0}(x)}$ (e.g., a
standard Gaussian), and the target distribution $\smash{\eta_{T}(x):=\pi(x)}$.
This idea is central to other Monte Carlo sampling methods such as Sequential
Monte Carlo (SMC) [15] and Annealed Importance Sampling (AIS) [52], as well as
sampling methods used in score-based generative modelling [e.g., 66]. In our
case, the annealed targets act as intermediary steps within the flow-informed
MCMC scheme.

In the setting of Bayesian inference, where the target distribution takes the
form $\pi(x)\propto\mathcal{L}(\mathcal{D}|x)\pi_{0}(x)$, for a given
likelihood $\mathcal{L}(\mathcal{D}|x)$ and prior $\pi_{0}(x)$, a standard way
in which to construct $\smash{\\{\eta_{t}(x)\\}_{t=0:T}}$ is to use a
geometric interpolation from the prior to the posterior [52], viz
$\eta_{t}(x)\propto\mathcal{L}(\mathcal{D}|x)^{\beta_{t}}\pi_{0}(x)$, where
$\beta_{0:T}$ is a sequence of temperatures satisfying
$0=\beta_{0}<\beta_{1}<\dots<\beta_{T}=1$. In practice, it can be difficult to
choose a good sequence of temperatures. One heuristic for adaptively setting
this sequence is the effective sample size (ESS). In particular, by setting
the ESS to a user-specified percentage $\alpha$ of the number of particles
$N$, the next temperature $\beta_{t}$ in the schedule can be determined
through the recursive equation [7]

$\displaystyle\beta_{t}=\inf\bigg{\\{}\beta_{t-1}<\beta\leq
1:\frac{\left[\frac{1}{N}\sum_{i=1}^{N}w_{i}^{\beta_{t-1}}(\beta)\right]^{2}}{\frac{1}{N}\sum_{i=1}^{N}w_{i}^{\beta_{t-1}}(\beta)^{2}}=\alpha\bigg{\\}},$
(16)

where
$w_{i}^{\beta_{t-1}}(\beta)=\mathcal{L}(\mathcal{D}|x)^{\beta}\pi_{0}(x)/\mathcal{L}(\mathcal{D}|x)^{\beta_{t-1}}\pi_{0}(x)=\mathcal{L}(\mathcal{D}|x)^{\beta-\beta_{t-1}}$
are new importance weights given the current temperature $\beta_{t-1}$. In
practice, we find that the inclusion of this adaptive tempering scheme is
essential in the presence of highly multimodal target distributions, enabling
the discovery of modes which are not known _a priori_.

Algorithm 1 Markovian Flow Matching

1:Input: $\pi_{0}$, likelihood $\mathcal{L}$, vector field $v_{t}^{\theta}$,
reference $p_{0}$, concentration $\sigma_{min}$, initial parameters
$\theta_{0}$, target ESS fraction $\alpha$, iterations $K$, number of
particles $N$, local kernel $Q$, flow-informed kernel $P$, MCMC steps per
resampling step $k_{Q}$, step sizes $\varepsilon_{1:K}$.

2:Output: flow-network parameters $\theta_{K}$

3:Sample $x^{i}_{0}\sim\pi_{0}$ for $i=1,\dots,N$ (_initialise samples_)

4:for $k=1:K$ do

5: if $\beta_{k-1}<1$ then

6: $\beta_{k}=$ solve (16) for $\beta_{k-1}<\beta\leq 1$ (_update annealing
temperature_)

7: end if

8: if $k\mod k_{Q}+1=0$ then

9: $x^{i}_{k}\sim
P(x^{i}_{k-1},\cdot;\pi_{0}\mathcal{L}^{\beta_{k}},\theta_{k-1})$ for
$i=1,\dots,N$ (_flow-informed Markov transition_).

10: else

11: $x^{i}_{k}\sim Q(x^{i}_{k-1},\cdot;\pi_{0}\mathcal{L}^{\beta_{k}})$ for
$i=1,\dots,N$ (_local Markov transition_).

12: end if

13:
$\theta_{k}=\theta_{k-1}+\varepsilon_{k}\nabla_{\theta}\mathcal{J}(\theta_{k-1}|\\{x^{i}_{k}\\}_{i=1}^{N},\sigma_{\mathrm{min}})$
(_update flow-network parameters_)

14:end for

#### Convergence

The output of Algorithm 1 is a vector of parameters $\theta_{K}$ which defines
a CNF $\smash{(\phi_{t}^{\theta_{K}})_{t\in[0,1]}}$. Under the assumption that
the parameter estimate converges, that is,
$\smash{\theta_{K}\rightarrow\theta^{*}_{\mathrm{global}}}$ as
$K\rightarrow\infty$, where
$\smash{\theta^{*}_{\mathrm{global}}=\operatorname*{arg\,min}_{\theta\in\Theta}\mathcal{J}(\theta;\pi)}$
denotes the global minimizer of the CFM objective $\mathcal{J}(\theta;\pi)$ in
(5), this CNF is guaranteed to generate a probability path
$(p_{t}^{\theta})_{t\in[0,1]}$ which transports samples from the reference
$p_{0}$ to the true target $\pi$ [e.g., 42].

In practice, the objective $\smash{\mathcal{J}(\theta;\pi)}$ is highly non-
convex, and thus it is not possible to establish a convergence result of this
type without imposing unreasonably strong assumptions on the vector field
$\smash{(v_{t}^{\theta})_{t\in[0,1]}}$. This being said, it is reasonable to
ask whether $\theta_{K}$ converges to a local optimum of the CFM objective. We
now answer this question in the affirmative. In particular, under mild
regularity conditions, Proposition 3.1 guarantees that
$\theta_{K}\rightarrow\theta^{*}$ almost surely as $K\rightarrow\infty$, where
$\theta^{*}$ denotes a local minimum of the CFM objective. The proof of this
result, which relies on a classical result in [5, Theorem 3.17], is provided
in Appendix B.

######  Proposition 3.1.

Assume that Assumptions B.1 \- B.6 hold (see Appendix B). Assume also that
$\smash{(\theta_{K})_{K\in\mathbb{N}}}$ is a bounded sequence, which almost
surely visits a compact subset of the domain of attraction of
$\smash{\theta^{*}}$ infinitely often. Then
$\smash{\theta_{K}\rightarrow\theta^{*}}$ almost surely.

##  4 Related work

In recent years, a number of works have proposed algorithms which combine MCMC
techniques with NFs; see, e.g., [25, 1] for recent surveys. Broadly speaking,
these algorithms fall into two categories. _NeutraMCMC_ methods leverage NFs
as reparameterization maps which simplify the geometry of the target
distribution, before running (local) MCMC samplers in the latent space. This
technique was pioneered in the seminal paper [56], and since been investigated
in a number of different works [e.g., 30, 51, 55, 41, 79, 76, 65, 11]. _Flow
MCMC_ methods, meanwhile, utilise the pushforward of the base distribution
through the NF as an (independent) proposal within an MCMC scheme. This
approach was first studied by [2], and further extended in [54, 3, 28]. More
recently, [21, 64] have introduced adaptive MCMC schemes which combine local
MCMC samplers (e.g., MALA or HMC), with a non-local, flow-informed proposal
(IMH or i-SIR); see also [32]. Our algorithm combines aspects of both
_neutraMCMC_ and _flow MCMC_ methods and, unlike any existing approach, make
use of a CNF (as opposed to a discrete NF), by leveraging the conditional flow
matching objective.

The use of NFs within other Monte Carlo algorithms has also been the subject
of recent interest. For example, [4, 47] consider augmenting SMC with NFs,
while [16, 49] use NFs (or diffusion models) within AIS. Finally, although not
directly comparable to our own approach, several recent works have proposed to
use (controlled) diffusion processes to sample from unnormalized probability
distributions. Such methods include [80], who introduce the _path integral
sampler_ , [71], who propose the _denoising diffusion sampler_ , and [77] who
introduce _generative flow samplers_. Other relevant contributions in this
direction include [69, 6, 60, 72, 71, 73, 70, 74, 57].

##  5 Experiments

In this section, we evaluate the performance of MFM (Algorithm 1) on two
synthetic and two real data examples. Our method is benchmarked against three
relevant methods. The Denoising Diffusion Sampler [DDS; 71] is a VI method
which approximates the reversed diffusion process from a reference
distribution to an extended target distribution by minimizing the KL
divergence. Adaptive Monte Carlo with Normalizing Flows [NF-MCMC; 21] is an
augmented MCMC scheme which uses a mixture of MALA and adaptive transition
kernels learned using discrete NFs. Finally, Flow Annealed Importance Sampling
Bootstrap [FAB; 49] is an augmented AIS scheme minimizing the mass-covering
$\alpha$-divergence with $\alpha=2$.

For each experiment, all MALA kernels use the same step size, targeting an
acceptance rate of close to 1 since we estimate expectations, e.g. in (14),
using the current ensemble of particles, rather than a single long chain.
Following [78], we parameterize the vector field as

$\displaystyle
v_{t}^{\theta}(x)=\text{NN}(x,t;\theta_{1})+\text{NN}(t;\theta_{2})\times\nabla\log\pi(x),$
(17)

where the neural networks are standard MLPs with 2 hidden layers, using a
Fourier feature augmentation for $t$ [67]. This architecture is also used by
DDS [71, Section 4]. Meanwhile, FAB and NF-MCMC use rational quadratic splines
[18]. Flows are trained using Adam [37] with a linear decay schedule
terminating at $\varepsilon_{K}=0$. We report results for all methods averaged
over 10 independent runs with varying random seeds. The code to reproduce the
experiments is provided at albcab/mfm.

###  5.1 4-mode Gaussian mixture

Our first example is a mixture of four Gaussians, evenly spaced and equally
weighted, in two-dimensional space. The four mixture components have means
$(8,8)$, $(-8,8)$, $(8,-8)$, $(-8,-8)$, and all have identity covariance. This
ensures that the modes are sufficiently separated to mean that jumping between
modes requires trajectories over sets with close to null probability. Given
the synthetic nature of the problem, we can measure approximation quality
using the Maximum Mean Discrepancy (MMD) [e.g., 26]; see Appendix C.1 for
details. We can also include, as a benchmark, the results for an approximation
learned using FM with true target samples. Diagnostics for all models are
presented in Table 1, and learned flow samples in Fig. 1. Further algorithmic
details and results are provided in Appendix C.2.

| 4-mode | 16-mode
---|---|---
|  | MMD | seconds |  | MMD | seconds
FM w/ $\pi$ samples |  |  $5.45$e-4$\pm 3.51$e-4 | $20.5\pm 0.41$ |  |  $2.16$e-3$\pm 3.95$e-4 | $22.8\pm 0.22$
MFM $k_{Q}=10^{2}$ |  |  $1.39$e-3$\pm 8.22$e-4 | $38.1\pm 1.29$ |  |  $3.57$e-3$\pm 7.68$e-4 | $38.5\pm 1.75$
DDS |  |  $1.76$e-4$\pm 2.32$e-4 | $114.\pm 0.68$ |  |  $1.02$e-1$\pm 4.10$e-2 | $115.\pm 0.64$
NF-MCMC |  |  $5.85$e-3$\pm 3.91$e-3 | $72.0\pm 11.7$ |  |  $8.05$e-3$\pm 1.42$e-2 | $67.0\pm 12.3$
FAB |  |  $2.69$e-4$\pm 2.06$e-4 | $101.\pm 3.24$ |  |  $1.51$e-3$\pm 1.06$e-3 | $102.\pm 4.32$
Table 1: Diagnostics for the two synthetic examples. MMD is the Maximum Mean
Discrepancy between real samples from the target and samples generated from
the learned flow. Results are averaged and empirical 95% confidence intervals
over 10 independent runs.

In this experiment, only our method (Fig. 1(a)) and DDS (Fig. 1(c)) learn the
fully separated modes, reflecting the greater expressivity of CNFs in
comparison to the discrete NFs used in, e.g., NF-MCMC (Fig. 1(d)). It is worth
noting that DDS provides a closer approximation to the real target than MFM
and, notably, even FM trained using true target samples (top row). Given that
both methods use the same network architecture but a different learning
objective, this suggests a potential limitation with the FM objective, at
least when using this network architecture. This being said, MFM is notably
more efficient than DDS (as well as the other methods) in terms of total
computation time. While this is not a critical consideration in this
synthetic, low-dimensional setting, it is a significant advantage of MFM in
higher-dimensional settings involving real data (e.g., Section 5.3 and Section
5.4).

(a) MFM $k_{Q}=10^{2}$

(b) FAB

(c) DDS

(d) NF-MCMC

Figure 1: Comparison between MFM, FAB, DDS, and NF-MCMC. Samples from the
target density for the 4-mode Gaussian mixture example.

###  5.2 16-mode Gaussian mixture

The second experiment is a mixture of bivariate Gaussians with 16 mixture
components. This is a modification of the 4-mode example, with contrasting
qualities that illustrate other characteristics of each of the presented
methods. In this case, the modes are evenly distributed on $[-16,16]^{2}$,
with random log-normal variances. The number of modes reduces the size of sets
of (near) null probability between the modes, making jumping between them
easier. To increase the difficulty of this model, all methods are initialized
on a concentrated region of the sampling space. Diagnostics are presented in
Table 1 and learned flow samples in Fig. 2. Further details are provided in
Appendix C.3.

In this example, DDS collapses to the modes closest to the initial positions
while our method captures the whole target. Since the modes are no longer
separated by areas of near-zero probability, the discrete NF methods are now
able to accurately capture the target density. In this case, DDS marginally
outperforms MFM as measured by the MMD, but this slight improvement in
performance comes at the cost of a much higher run-time.

(a) MFM $k_{Q}=10^{2}$

(b) FAB

(c) DDS

(d) NF-MCMC

Figure 2: Comparison between MFM, FAB, DDS, and NF-MCMC. Samples from the
target density for the 16-mode Gaussian mixture example.

###  5.3 Field system

Our first real-world example considers the Allen–Cahn equation, used as a
benchmark in [21], and described in Appendix C.4. This fundamental reaction-
diffusion equation is central to the study of phase transitions in condensed
matter systems. Incorporating random forcing terms or thermal fluctuations
allows for a stochastic treatment of the dynamics, capturing the inherent
randomness and uncertainties in physical systems. This system leads to a
discretized target density of the form

$\displaystyle\log\pi(x)=-\beta\bigg{(}\frac{a}{2\Delta
s}\sum_{i=1}^{d+1}(x_{i}-x_{i-1})^{2}+\frac{\Delta
s}{4a}\sum_{i=1}^{d}(1-x_{i}^{2})^{2}\bigg{)},$ (18)

where $\Delta s=\frac{1}{d}$, and with boundary conditions $x_{0}=x_{d+1}=0$.
In our experiments, we set the dimensionality $d=64$. Meanwhile, parameters
values are chosen to ensure bimodality at $x=\pm 1$: $a=0.1$ and inverse
temperature $\beta=20$. The bimodality induced by the two global minima
complicates mixing when using traditional MCMC updates. Learning the global
geometry of the target and using that information to propose transitions
facilitates movement between modes. Unlike previous work [e.g., 21], we
deliberately choose not to employ an informed base measure. Instead, we opt
for a standard Gaussian with no additional information, making the problem
significantly more challenging. This choice illustrates the robustness of our
approach.

(a) MFM $k_{Q}=10^{2}$

(b) FAB

(c) DDS

(d) NF-MCMC

Figure 3: Comparison between MFM, FAB, DDS, and NF-MCMC. Representative
samples from the target density for the Field system example.

Numerical diagnostics for each method are presented in Table 2. In this case,
we use the Kernelized Stein Discrepancy (KSD) as a measure of sample quality
[e.g., 43, 23]; see Appendix C.1 for details. While this is not a perfect
metric, it does allow us to qualitatively compare the different methods
considered.

In this case, the tempering mechanism of our method is crucial for ensuring
that the learned flow does not collapse on one of the modes and instead
explores both global minima. This is confirmed when plotting the samples
generated in the grid in Fig. 3. This experiment demonstrates the capability
of our method to capture complex multi-modal densities, even without an
informed base measure, at a _significantly_ lower computational cost (e.g.,
10-25x faster) than competing methods. It is worth noting that MFM (and the
other two benchmarks) significantly outperformed NF-MCMC in this example,
despite the similarities between MFM and NF-MCMC. While we tested various
hyperparameter configurations for NF-MCMC, we were not able to find a setting
that achieved comparbale results.

###  5.4 Log-Gaussian Cox point process

Bayesian inference for high-dimensional spatial models is known to be
challenging. One such model is the log-Gaussian Cox point process, which here
we use to model the locations of 126 Scots pine saplings in a natural forest
in Finland. See Appendix C.5 for full details. The target space is discretized
to an $M=40\times 40$ regular grid, rendering the target dimension $d=1600$.
In Table 2, we report diagnostics for each algorithm.

In this case, the lack of multimodality in the target makes it a good fit for
non-tempered schemes. Similar to the previous example, NF-MCMC is unable to
obtain an accurate approximation to the target distribution. We suspect that
this may be a result of non-convergence: due to memory issues, it was not
possible to run NF-MCMC (or FAB) for more than $K=10^{3}$ iterations. This
also explains the (relatively) smaller run times of these algorithms in this
example. By a small margin, DDS provides the best approximation of the target,
slightly outperforming MFM and FAB. Meanwhile, MFM provides a good
approximation to the target at a lower computational cost with respect to its
counterparts.

| Field system | Log-Gaussian Cox |  |
---|---|---|---|---
$K=10^{4}$ |  | KSD U-stat. | KSD V-stat. |  | seconds |  | KSD U-stat. | KSD V-stat. |  | seconds
MFM |  | $2.17\pm 2.19$ | $20.4\pm 2.30$ |  | $146\pm 1.93$ |  |  $1.05$e-1$\pm 0.05$ | $31.2\pm 0.16$ |  | $910\pm 47.1$
DDS |  | $15.2\pm 35.9$ | $18.0\pm 36.9$ |  | $2400\pm 8.65$ |  |  $7.59$e-2$\pm 0.02$ | $24.7\pm 0.08$ |  | $3260\pm 8.41$
NF-MCMC |  | $548\pm 325$ | $549\pm 325$ |  | $2000\pm 15.6$ |  | $11.8\pm 7.55$ | $89.0\pm 238$ |  | $215\pm 46.4$
FAB |  | $0.14\pm 0.42$ | $1.78\pm 0.42$ |  | $3880\pm 7.19$ |  |  $1.55$e-1$\pm 0.06$ | $52.3\pm 2.02$ |  | $1040\pm 2.78$
Table 2: Real examples diagnostics where Diagonstic for the two real data
examples. KSD U-stat and V-stat are the Kernel Stein Discrepancy U- and
V-statistics between the target and samples generated from the learned flow.
Results are averaged and empirical 95% confidence intervals over 10
independent runs.

##  6 Conclusion

Summary. In this paper, we introduced Markovian Flow Matching, a new approach
to sampling from unnormalized probability distributions that augments MCMC
with CNFs. Our method combines a local Markov kernel with a non-local, flow-
informed Markov kernel, which is adaptively learned during sampling using FM.
It also incorporates an adaptive tempering mechanism, which allows for the
discovery of multiple target modes. Under mild assumptions, we established
convergence of the flow network parameters to a local optimum of the FM
objective. Finally, we benchmarked the performance of our algorithm on several
examples, illustrating comparable performance to other state-of-the-art
methods, often at a fraction of the computational cost.

Limitations. We highlight three limitations of our work. First, our
theoretical result established convergence of the flow network parameters
obtained via MFM to a _local_ minimum of the FM objective. Further work is
required to understand how well these local minima generalize, in order to
accurately quantify how accurately the corresponding CNF captures the target
posterior. Second, we did not establish non-asymptotic convergence rates for
our method. Finally, since it was not the main focus of this work, we did not
explore in great detail other choices of architecture for the flow network. We
expect that, for certain targets, this could have a significant impact on the
performance of MFM. This remains an interesting avenue for future research.

## Acknowledgments and Disclosure of Funding

LS and CN were supported by the Engineering and Physical Sciences Research
Council (EPSRC), grant number EP/V022636/1. CN acknowledges further support
from the EPSRC, grant numbers EP/S00159X/1 and EP/Y028783/1.

## References

  * Albergo and Vanden-Eijnden [2023] M. S. Albergo and E. Vanden-Eijnden.  Learning to sample better.  _arXiv preprint arXiv:2310.11232_ , 2023.
  * Albergo et al. [2019] M. S. Albergo, G. Kanwar, and P. E. Shanahan.  Flow-based generative models for Markov chain Monte Carlo in lattice field theory.  _Physical Review D_ , 100(3):034515, 2019.
  * Albergo et al. [2021] M. S. Albergo, G. Kanwar, S. Racanière, D. J. Rezende, J. M. Urban, D. Boyda, K. Cranmer, D. C. Hackett, and P. E. Shanahan.  Flow-based sampling for fermionic lattice field theories.  _Physical Review D_ , 104:114507, 2021.  doi: 10.1103/PhysRevD.104.114507.
  * Arbel et al. [2021] M. Arbel, A. Matthews, and A. Doucet.  Annealed flow transport Monte Carlo.  In _Proceedings of the 38th International Conference on Machine Learning (ICML)_ , 2021.
  * Benveniste et al. [1990] A. Benveniste, M. Metivier, and P. Priouret.  _Adaptive Algorithms and Stochastic Approximations_.  Springer-Verlag, Berlin, Heidelberg, 1990.  doi: 10.1007/978-3-642-75894-2.
  * Berner et al. [2024] J. Berner, L. Richter, and K. Ullrich.  An optimal control perspective on diffusion-based generative modeling.  _Transaction on Machine Learning Research (TMLR)_ , 2024.
  * Beskos et al. [2016] A. Beskos, A. Jasra, N. Kantas, and A. Thiery.  On the convergence of adaptive sequential Monte Carlo methods.  _Annals of Applied Probability_ , 26(2):1111–1146, 2016.  doi: 10.1214/15-AAP1113.
  * Blei et al. [2017] D. M. Blei, A. Kucukelbir, and J. D. McAuliffe.  Variational inference: a review for statisticians.  _Journal of the American Statistical Association_ , 112(518):859–877, 2017.  doi: 10.1080/01621459.2017.1285773.
  * Bradbury et al. [2018] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang.  JAX: composable transformations of Python+NumPy programs, 2018\.  URL http://github.com/google/jax.
  * Brooks et al. [2011] S. Brooks, A. Gelman, G. Jones, and X.-L. Meng.  _Handbook of Markov Chain Monte Carlo_.  Chapman and Hall/CRC, 1st edition, 2011.  doi: 10.1201/b10905.
  * Cabezas and Nemeth [2023] A. Cabezas and C. Nemeth.  Transport elliptical slice sampling.  In _Proceedings of the 26th International Conference on Artificial Intelligence and Statistics (AISTATS)_ , 2023.
  * Chen and Lipman [2024] R. Chen and Y. Lipman.  Flow matching on general geometries.  In _Proceedings of the 12th International Conference on Learning Representations (ICLR)_ , 2024.
  * Chen et al. [2018] R. T. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud.  Neural ordinary differential equations.  In _Proceedings of the 32nd Annual Conference on Neural Information Processing Systems (NeurIPS)_ , 2018.
  * Del Debbio et al. [2021] L. Del Debbio, J. Marsh Rossney, and M. Wilson.  Efficient modeling of trivializing maps for lattice $\phi$ 4 theory using normalizing flows: a first look at scalability.  _Physical Review D_ , 104(9):094507, 2021.  doi: 10.1103/PhysRevD.104.094507.
  * Del Moral et al. [2006] P. Del Moral, A. Doucet, and A. Jasra.  Sequential Monte Carlo samplers.  _Journal of the Royal Statistical Society Series B: Statistical Methodology_ , 68(3):411–436, 2006.
  * Doucet et al. [2022] A. Doucet, W. Grathwohl, A. G. Matthews, and H. Strathmann.  Score-based diffusion meets annealed importance sampling.  In _Proceedings of the 36th Annual Conference on Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Duane et al. [1987] S. Duane, A. D. Kennedy, B. J. Pendleton, and D. Roweth.  Hybrid Monte Carlo.  _Physics letters B_ , 195(2):216–222, 1987.  doi: 10.1016/0370-2693(87)91197-X.
  * Durkan et al. [2019] C. Durkan, A. Bekasov, I. Murray, and G. Papamakarios.  Neural spline flows.  In _Proceedings of the 33rd Annual Conference on Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Feroz et al. [2009] F. Feroz, M. Hobson, and M. Bridges.  Multinest: an efficient and robust bayesian inference tool for cosmology and particle physics.  _Monthly Notices of the Royal Astronomical Society_ , 398(4):1601–1614, 2009.
  * Gabrié et al. [2021] M. Gabrié, G. M. Rotskoff, and E. Vanden-Eijnden.  Efficient bayesian sampling using normalizing flows to assist Markov chain Monte Carlo methods.  In _Proceedings of the 38th International Conference on Machine Learning (ICML): 3rd Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models_ , 2021.
  * Gabrié et al. [2022] M. Gabrié, G. M. Rotskoff, and E. Vanden-Eijnden.  Adaptive Monte Carlo augmented with normalizing flows.  _Proceedings of the National Academy of Sciences_ , 119(10):e2109420119, 2022.  doi: 10.1073/pnas.2109420119.
  * Gelman et al. [1995] A. Gelman, J. B. Carlin, H. S. Stern, and D. B. Rubin.  _Bayesian data analysis_.  Chapman and Hall/CRC, 1995.
  * Gorham and Mackey [2017] J. Gorham and L. Mackey.  Measuring sample quality with kernels.  In _Proceedings of the Proceedings of the 34th International Conference on Machine Learning (ICML)_ , 2017.
  * Grathwohl et al. [2018] W. Grathwohl, R. T. Chen, J. Bettencourt, I. Sutskever, and D. Duvenaud.  FFJORD: Free-form continuous dynamics for scalable reversible generative models.  In _Proceedings of the 7th International Conference on Learning Representations (ICLR)_ , 2018.
  * Grenioux et al. [2023] L. Grenioux, A. Durmus, É. Moulines, and M. Gabrié.  On sampling with approximate transport maps.  In _Proceedings of the 40th International Conference on Machine Learning (ICML)_ , 2023.
  * Gretton et al. [2012] A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schöolkopf, and A. Smola.  A kernel two-sample test.  _Journal of Machine Learning Research (JMLR)_ , 13:723–773, 2012.
  * Gu and Kong [1998] M. G. Gu and F. H. Kong.  A stochastic approximation algorithm with Markov chain Monte-Carlo method for incomplete data estimation problems.  _Proceedings of the National Academy of Sciences_ , 95(13):7270–7274, 1998.  doi: 10.1073/pnas.95.13.7270.
  * Hackett et al. [2021] D. C. Hackett, C.-C. Hsieh, M. S. Albergo, D. Boyda, J.-W. Chen, K.-F. Chen, K. Cranmer, G. Kanwar, and P. E. Shanahan.  Flow-based sampling for multimodal distributions in lattice field theory.  _arXiv preprint arXiv:2107.00734_ , 2021.
  * Hastings [1970] W. K. Hastings.  Monte Carlo sampling methods using Markov chains and their applications.  _Biometrika_ , 57(1):97–109, 1970.  doi: 10.1093/biomet/57.1.97.
  * Hoffman et al. [2019] M. Hoffman, P. Sountsov, J. V. Dillon, I. Langmore, D. Tran, and S. Vasudevan.  Neutra-lizing bad geometry in Hamiltonian Monte Carlo using neural transport.  In _Proceedings of the 1st Symposium on Advances in Approximate Bayesian Inference (AABI)_ , 2019.
  * Hoffman et al. [2013] M. D. Hoffman, D. M. Blei, C. Wang, and J. Paisley.  Stochastic variational inference.  _Journal of Machine Learning Research_ , 14(1):1303–1347, 2013.
  * Hunt-Smith et al. [2024] N. T. Hunt-Smith, W. Melnitchouk, F. Ringer, N. Sato, A. W. Thomas, and M. J. White.  Accelerating Markov chain Monte Carlo sampling with diffusion models.  _Computer Physics Communications_ , 296:109059, 2024.  doi: 10.1016/j.cpc.2023.109059.
  * Hutchinson [1989] M. F. Hutchinson.  A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines.  _Communications in Statistics-Simulation and Computation_ , 18(3):1059–1076, 1989.
  * Ihler et al. [2005] A. Ihler, J. Fisher III, R. Moses, and A. Willsky.  Nonparametric belief propagation for self-localization of sensor networks.  _IEEE Journal on Selected Areas in Communications_ , 23(4), 2005.  doi: 10.1109/JSAC.2005.843548.
  * Jensen et al. [2004] S. T. Jensen, X. S. Liu, Q. Zhou, and J. S. Liu.  Computational discovery of gene regulatory binding motifs: a bayesian perspective.  _Statistical Science_ , 19(1):188–204, 2004.  doi: 10.1214/088342304000000107.
  * Karamanis et al. [2022] M. Karamanis, F. Beutler, J. A. Peacock, D. Nabergoj, and U. Seljak.  Accelerating astronomical and cosmological inference with preconditioned Monte Carlo.  _Monthly Notices of the Royal Astronomical Society_ , 516(2):1644–1653, 2022.  doi: 10.1093/mnras/stac2272.
  * Kingma and Ba [2015] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization.  In _Proceedings of the 3rd International Conference on Learning Representations (ICLR)_ , 2015.
  * Kou et al. [2006] S. Kou, J. Oh, and W. H. Wong.  A study of density of states and ground states in hydrophobic-hydrophilic protein folding models by equi-energy sampling.  _The Journal of Chemical Physics_ , 124(24), 2006.  doi: 10.1063/1.2208607.
  * Lee [1990] A. Lee.  _U-Statistics: Theory and Practice_.  Taylor & Francis, 1st edition, 1990.  doi: 10.1201/9780203734520.
  * Leimkuhler and Matthews [2015] B. Leimkuhler and C. Matthews.  Molecular dynamics.  _Interdisciplinary applied mathematics_ , 39:443, 2015.  doi: 10.1007/978-3-319-16375-8.
  * Li and Wang [2018] S.-H. Li and L. Wang.  Neural network renormalization group.  _Physical Review Letters_ , 121:260601, 2018.  doi: 10.1103/PhysRevLett.121.260601.
  * Lipman et al. [2023] Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le.  Flow matching for generative modeling.  In _Proceedings of the 11th International Conference on Learning Representations (ICLR)_ , 2023.
  * Liu et al. [2016] Q. Liu, J. Lee, and M. Jordan.  A kernelized Stein discrepancy for goodness-of-fit tests.  In _Proceedings of the 33rd International Conference on Machine Learning (ICML)_ , 2016.
  * Mahmoud et al. [2022] A. H. Mahmoud, M. Masters, S. J. Lee, and M. A. Lill.  Accurate sampling of macromolecular conformations using adaptive deep learning and coarse-grained representation.  _Journal of Chemical Information and Modeling_ , 62(7):1602–1617, 2022.  doi: 10.1021/acs.jcim.1c01438.
  * Mangoubi et al. [2018] O. Mangoubi, N. S. Pillai, and A. Smith.  Does Hamiltonian Monte Carlo mix faster than a random walk on multimodal densities?  _arXiv preprint arXiv:1808.03230_ , 2018.
  * Mangoubi et al. [2021] O. Mangoubi, N. Pillai, and A. Smith.  Simple conditions for metastability of continuous Markov chains.  _Journal of Applied Probability_ , 58(1):83–105, 2021.  doi: 10.1017/jpr.2020.83.
  * Matthews et al. [2022] A. Matthews, M. Arbel, D. J. Rezende, and A. Doucet.  Continual repeated annealed flow transport Monte Carlo.  In _Proceedings of the 39th International Conference on Machine Learning (ICML)_ , 2022.
  * Metropolis et al. [1953] N. Metropolis, A. W. Rosenbluth, M. N. Rosenbluth, A. H. Teller, and E. Teller.  Equation of state calculations by fast computing machines.  _The journal of chemical physics_ , 21(6):1087–1092, 1953.
  * Midgley et al. [2023] L. I. Midgley, V. Stimper, G. N. Simm, B. Schölkopf, and J. M. Hernández-Lobato.  Flow annealed importance sampling bootstrap.  In _Proceedings of the 11th International Conference on Learning Representations (ICLR)_ , 2023.
  * Møller et al. [1998] J. Møller, A. R. Syversveen, and R. P. Waagepetersen.  Log Gaussian Cox processes.  _Scandinavian Journal of Statistics_ , 25(3):451–482, 1998.  doi: 10.1111/1467-9469.00115.
  * Naesseth et al. [2020] C. Naesseth, F. Lindsten, and D. Blei.  Markovian score climbing: variational inference with KL(p||q).  _Proceedings of the 34th Annual Conference on Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Neal [2001] R. M. Neal.  Annealed importance sampling.  _Statistics and computing_ , 11:125–139, 2001.  doi: 10.1023/A:1008923215028.
  * Neal et al. [2011] R. M. Neal et al.  MCMC using Hamiltonian dynamics.  _Handbook of Markov chain Monte Carlo_ , 2(11):2, 2011.
  * Nicoli et al. [2020] K. A. Nicoli, S. Nakajima, N. Strodthoff, W. Samek, K.-R. Müller, and P. Kessel.  Asymptotically unbiased estimation of physical observables with neural samplers.  _Physical Review E_ , 101(2):023304, 2020.  doi: 10.1103/PhysRevE.101.023304.
  * Noé et al. [2019] F. Noé, S. Olsson, J. Köhler, and H. Wu.  Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning.  _Science_ , 365(6457), 2019.  doi: 10.1126/science.aaw1147.
  * Parno and Marzouk [2018] M. D. Parno and Y. M. Marzouk.  Transport map accelerated Markov chain Monte Carlo.  _SIAM/ASA Journal on Uncertainty Quantification_ , 6(2):645–682, 2018.  doi: 10.1137/17M1134640.
  * Phillips et al. [2024] A. Phillips, H.-D. Dau, M. J. Hutchinson, V. D. Bortoli, G. Deligiannidis, and A. Doucet.  Particle denoising diffusion sampler.  In _Proceedings of the 41st International Conference on Machine Learning (ICML)_ , 2024.
  * Ranganath et al. [2014] R. Ranganath, S. Gerrish, and D. Blei.  Black box variational inference.  In _Proceedings of the 17th International Conference on Artificial Intelligence and Statistics (AISTATS)_ , 2014.
  * Rezende and Mohamed [2015] D. Rezende and S. Mohamed.  Variational inference with normalizing flows.  In _Proceedings of the 32nd International Conference on Machine Learning (ICML)_ , 2015.
  * Richter et al. [2024] L. Richter, J. Berner, and G.-H. Liu.  Improved sampling via learned diffusions.  In _Proceedings of the 12th International Conference on Learning Representations (ICLR)_ , 2024.
  * Robert and Casella [2004] C. P. Robert and G. Casella.  _Monte Carlo Statistical Methods_.  Springer-Verlag, New York, 2nd edition, 2004.  doi: 10.1007/978-1-4757-4145-22.
  * Roberts and Rosenthal [2004] G. O. Roberts and J. S. Rosenthal.  General state space Markov chains and MCMC algorithms.  _Probability Surveys_ , 1:20 – 71, 2004.  doi: 10.1214/154957804100000024.
  * Roberts and Tweedie [1996] G. O. Roberts and R. L. Tweedie.  Exponential convergence of Langevin distributions and their discrete approximations.  _Bernoulli_ , 2(4):341–363, 1996.
  * Samsonov et al. [2022] S. Samsonov, E. Lagutin, M. Gabrié, A. Durmus, A. Naumov, and E. Moulines.  Local-global MCMC kernels: the best of both worlds.  In _Proceedings of the 36th Annual Conference on Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Schönle and Gabrié [2023] C. Schönle and M. Gabrié.  Optimizing Markov chain Monte Carlo convergence with normalizing flows and gibbs sampling.  In _Proceedings of the 37th Annual Conference on Neural Information Processing Systems (NeurIPS): AI for Science Workshop_ , 2023.
  * Song and Ermon [2019] Y. Song and S. Ermon.  Generative modeling by estimating gradients of the data distribution.  In _Proceedings of the 33rd Annual Conference on Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Tancik et al. [2020] M. Tancik, P. Srinivasan, B. Mildenhall, S. Fridovich-Keil, N. Raghavan, U. Singhal, R. Ramamoorthi, J. Barron, and R. Ng.  Fourier features let networks learn high frequency functions in low dimensional domains.  In _Proceedings of the 34th Annual Conference on Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Tong et al. [2024] A. Tong, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, K. Fatras, G. Wolf, and Y. Bengio.  Improving and generalizing flow-based generative models with minibatch optimal transport.  _Transactions on Machine Learning Research (TMLR)_ , 2024.
  * Tzen and Raginsky [2019] B. Tzen and M. Raginsky.  Theoretical guarantees for sampling and inference in generative models with latent diffusions.  In _Proceedings of the 32nd Annual Conference on Learning Theory (COLT)_ , 2019.
  * Vargas and Nüsken [2023] F. Vargas and N. Nüsken.  Transport, variational inference and diffusions: with applications to annealed flows and Schrödinger bridges.  _Proceedings of the 40th International Conference on Machine Learning (ICML): Workshop on New Frontiers in Learning, Control, and Dynamical Systems_ , 2023.
  * Vargas et al. [2023a] F. Vargas, W. Grathwohl, and A. Doucet.  Denoising diffusion samplers.  In _Proceedings of the 11th International Conference on Learning Representations (ICLR)_ , 2023a.
  * Vargas et al. [2023b] F. Vargas, A. Ovsianas, D. Fernandes, M. Girolami, N. D. Lawrence, and N. Nüsken.  Bayesian learning via neural Schrödinger–Föllmer flows.  _Statistics and Computing_ , 33(1):3, 2023b.  doi: 10.1007/s11222-022-10172-5.
  * Vargas et al. [2023c] F. Vargas, T. Reu, and A. Kerekes.  Expressiveness remarks for denoising diffusion models and samplers.  In _Proceedings of the 5th Symposium on Advances in Approximate Bayesian Inference (AABI)_ , 2023c.
  * Vargas et al. [2024] F. Vargas, S. Padhy, D. Blessing, and N. Nüsken.  Transport meets variational inference: controlled Monte Carlo diffusions.  In _Proceedings of the 12th International Conference on Learning Representations (ICLR)_ , 2024.
  * Wainwright et al. [2008] M. J. Wainwright, M. I. Jordan, et al.  Graphical models, exponential families, and variational inference.  _Foundations and Trends® in Machine Learning_ , 1(1–2):1–305, 2008.  doi: 10.1561/2200000001.
  * Zhang et al. [2023a] B. J. Zhang, Y. M. Marzouk, and K. Spiliopoulos.  Transport map unadjusted Langevin algorithms.  _arXiv preprint arXiv:2302.07227_ , 2023a.
  * Zhang et al. [2024] D. Zhang, R. T. Q. Chen, C.-H. Liu, A. Courville, and Y. Bengio.  Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization.  In _Proceedings of the 12th International Conference on Learning Representations (ICLR)_ , 2024.
  * Zhang et al. [2022] L. Zhang, B. Carpenter, A. Gelman, and A. Vehtari.  Pathfinder: Parallel quasi-Newton variational inference.  _Journal of Machine Learning Research_ , 23(306):1–49, 2022.
  * Zhang et al. [2023b] L. Zhang, D. M. Blei, and C. A. Naesseth.  Transport score climbing: variational inference using forward kl and adaptive neural transport.  _Transactions on Machine Learning Research (TMLR)_ , 2023b.
  * Zhang and Chen [2022] Q. Zhang and Y. Chen.  Path integral sampler: a stochastic control approach for sampling.  In _Proceedings of the 10th International Conference on Learning Representations (ICLR)_ , 2022.

##  Appendix A Flow-Informed Markov Chain Monte Carlo Methods

Algorithm 2 Flow-informed Random Walk Metropolis Hastings

1:Input: initial $x$, target $\pi$, vector field $v_{t}^{\theta}$, flow
parameters $\theta$

2:Output: $x^{\prime}$

3:$\sigma_{opt}\leftarrow 2.38/\sqrt{d}$

4:$\phi_{1}(x)=x_{t=1}\leftarrow x$

5:$\begin{bmatrix}x_{0}\\\ \Delta\log
p(x_{0})\end{bmatrix}\leftarrow\begin{bmatrix}x\\\
0\end{bmatrix}+\displaystyle\int_{1}^{0}\begin{bmatrix}v_{t}^{\theta}(\phi_{t}(x))\\\
-\nabla\cdot v_{t}^{\theta}(\phi_{t}(x))\end{bmatrix}\mathrm{d}t$

6:$y_{0}\sim\mathcal{N}(\cdot|x_{0},\sigma_{opt}^{2})$

7:$\begin{bmatrix}y_{1}\\\ \Delta\log
p(y_{1})\end{bmatrix}\leftarrow\begin{bmatrix}y_{0}\\\
0\end{bmatrix}+\displaystyle\int_{0}^{1}\begin{bmatrix}v_{t}^{\theta}(\phi_{t}(y))\\\
-\nabla\cdot v_{t}^{\theta}(\phi_{t}(y))\end{bmatrix}\mathrm{d}t$

8:$\alpha\leftarrow\min\left\\{1,\frac{\pi(y_{1})\exp(-\Delta\log
p(y_{1}))}{\pi(x_{1})\exp(\Delta\log p(x_{0}))}\right\\}$

9:With probability $\alpha$ make $x^{\prime}\leftarrow y_{1}$ else
$x^{\prime}\leftarrow x$

Algorithm 3 Flow-informed Independent Metropolis Hastings

1:Input: initial $x$, target density $\pi$, vector field $v_{t}^{\theta}$,
reference density $p_{0}$, flow parameters $\theta$.

2:Output: $x^{\prime}$

3:$\phi_{1}(u)=u_{t=1}\leftarrow x$

4:$\begin{bmatrix}u_{0}\\\ \Delta\log
p(u_{0})\end{bmatrix}\leftarrow\begin{bmatrix}u_{1}\\\
0\end{bmatrix}+\displaystyle\int_{1}^{0}\begin{bmatrix}v_{t}^{\theta}(\phi_{t}(u))\\\
-\nabla\cdot v_{t}^{\theta}(\phi_{t}(u))\end{bmatrix}dt$

5:$\phi_{0}(x)=x_{t=0}\sim p_{0}$

6:$\begin{bmatrix}x_{1}\\\ \log
p(x_{1})\end{bmatrix}\leftarrow\begin{bmatrix}x_{0}\\\ \log
p_{0}(x_{0})\end{bmatrix}+\displaystyle\int_{0}^{1}\begin{bmatrix}v_{t}^{\theta}(\phi_{t}(x))\\\
-\nabla\cdot v_{t}^{\theta}(\phi_{t}(x))\end{bmatrix}dt$

7:$\alpha\leftarrow\min\left\\{1,\frac{\pi(x_{1})p_{0}(u_{0})\exp(-\Delta\log
p(u_{0}))}{\exp(\log p(x_{1}))\pi(u_{1})}\right\\}$

8:With probability $\alpha$ make $x^{\prime}\leftarrow x_{1}$ else
$x^{\prime}\leftarrow x$

Algorithm 4 Flow-informed Conditional Importance Sampling

1:Input: initial $x$, target density $\pi$, vector field $v_{t}^{\theta}$,
reference density $q_{0}$, flow parameters $\theta$, number of importance
samples $K$.

2:Output: $x^{\prime}$

3:$\phi_{1}(u)=u_{t=1}\leftarrow x$

4:$\begin{bmatrix}u_{0}\\\ \Delta\log
p(u_{0})\end{bmatrix}\leftarrow\begin{bmatrix}u_{1}\\\
0\end{bmatrix}+\displaystyle\int_{1}^{0}\begin{bmatrix}v_{t}^{\theta}(\phi_{t}(u))\\\
-\nabla\cdot v_{t}^{\theta}(\phi_{t}(u))\end{bmatrix}dt$

5:$w_{0}\leftarrow\frac{\pi(u_{1})}{p_{0}(u_{0})\exp(-\Delta\log p(u_{0}))}$

6:$x^{(0)}\leftarrow x$

7:for $k=1:K$ do

8: $\phi_{0}(x)=x_{t=0}\sim q_{0}$

9: $\begin{bmatrix}x_{1}\\\ \log
p(x_{1})\end{bmatrix}\leftarrow\begin{bmatrix}x_{0}\\\ \log
p_{0}(x_{0})\end{bmatrix}+\displaystyle\int_{0}^{1}\begin{bmatrix}v_{t}^{\theta}(\phi_{t}(x))\\\
-\nabla\cdot v_{t}^{\theta}(\phi_{t}(x))\end{bmatrix}dt$

10: $w_{k}\leftarrow\frac{\pi(x_{1})}{\exp(\log p(x_{1}))}$

11: $x^{(k)}\leftarrow x_{1}$

12:end for

13:Choose $k^{\prime}$ with probability $P(k^{\prime}=k)\propto w_{k}$, then
make $x^{\prime}\leftarrow x^{(k^{\prime})}$

##  Appendix B Proof of Proposition 3.1

Our proof follows closely the proof of [51, Proposition 1]. Let $\theta^{*}$
be a minimiser of the CFM objective in (4), which we recall is given by

$\mathcal{J}(\theta;\pi)=\mathbb{E}_{x_{1}\sim\pi}\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{x\sim
p_{t}(\cdot|x_{1})}\left[||v_{t}^{\theta}(x)-v_{t}(x|x_{1})||^{2}\right]:=\mathbb{E}_{x_{1}\sim\pi}\left[j(\theta,x_{1})\right]$
(19)

where we have defined
$j(\theta,x_{1})=\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{x\sim
p_{t}(\cdot|x_{1})}\left[||v_{t}^{\theta}(x)-v_{t}(x|x_{1})||^{2}\right]$.
Now, consider the ordinary differential equation (ODE) given by

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\theta(t)=\nabla\mathcal{J}_{\theta}(\theta(t);\pi),\quad\theta(0)=\theta_{0},\quad
t\geq 0.$ (20)

We say that $\smash{\hat{\theta}}$ is stability point of this ODE if, given
the initial condition $\smash{\theta(0)=\hat{\theta}}$, the ODE admits the
unique solution $\smash{\theta(t)=\hat{\theta}}$ for all $t\geq 0$. Naturally,
the minimiser $\theta^{*}$ is a stability point of this ODE, since
$\smash{\left.\nabla_{\theta}\mathcal{J}(\theta;\pi)\right|_{\theta=\theta^{*}}=0}$.
Meanwhile, we call $\Theta$ the domain of attraction of $\theta^{*}$ if, given
the initial condition $\theta(0)\in\Theta$, the solution $\theta(t)\in\Theta$
for all $t\geq 0$, and $\theta(t)$ converges to $\theta^{*}$ as
$t\rightarrow\infty$.

Let $x_{k}\in\mathbb{R}^{d_{x}}$, and $\mathcal{X}\subseteq\mathbb{R}^{d_{x}}$
be an open subset of $\mathbb{R}^{d_{x}}$. Let
$\Theta\subseteq\mathbb{R}^{d_{\theta}}$ be an open set in
$\mathbb{R}^{d_{\theta}}$, and $\Theta_{c}\subseteq\Theta$ be a compact subset
of $\Theta$. Consider the Markov transition kernel $M:=P\circ Q^{k}$ given by
a cycle of $k_{Q}$ repeated transitions of a MALA transition kernel and a
flow-informed RWMH transition kernel, viz

$\displaystyle M_{\pi,\theta}(x,\mathrm{d}y)=\int\cdots\int
Q(x,\mathrm{d}x_{1};\theta)Q(x_{1},\mathrm{d}x_{2})\dots
Q(x_{k_{Q}-1},\mathrm{d}x_{k_{Q}};\theta)P(x_{k_{Q}},\mathrm{d}y;\pi,\theta).$
(21)

This transition kernel is $\pi$-invariant since both $P$ and $Q$ are
$\pi$-invariant. In addition, let $M^{k}_{\pi,\theta}(x,\mathrm{d}y)$ be the
repeated application of this Markov transition kernel, namely,

$\displaystyle M^{k}_{\pi,\theta}(x,\mathrm{d}y)=\int\cdots\int
M_{\pi,\theta}(x,\mathrm{d}x_{1})M_{\pi,\theta}(x_{1},\mathrm{d}x_{2})\cdots
M_{\pi,\theta}(x_{k-2},\mathrm{d}x_{k-1})M_{\pi,\theta}(x_{k-1},\mathrm{d}y).$
(22)

Following [27, 51], we impose the following assumptions, for some sufficiently
large positive real number $q>1$.

######  Assumption B.1 (Robbins-Monro Condition).

The step size sequence $(\varepsilon_{k})_{k=1}^{\infty}$ satisfies the
following requirements:

$\sum_{k=1}^{\infty}\varepsilon_{k}=\infty,\quad\sum_{k=1}^{\infty}\varepsilon_{k}^{2}<\infty.$
(23)

######  Assumption B.2 (Integrability).

There exists a constant $C_{1}>0$ such that for, any $\theta\in\Theta$,
$x\in\mathcal{X}$ and $k\geq 1$,

$\displaystyle\int(1+|y|^{q})M_{\pi,\theta}^{k}(x,\mathrm{d}y)\leq
C_{1}(1+|x|^{q}).$ (24)

######  Assumption B.3 (Convergence of the Markov Chain).

For each $\theta\in\Theta$, it holds that

$\displaystyle\lim_{k\rightarrow\infty}\sup_{x\in\mathcal{X}}\frac{1}{1+|x|^{q}}\int(1+|y|^{q})|M_{\pi,\theta}^{k}(x,\mathrm{d}y)-\pi(\mathrm{d}y)|=0.$
(25)

######  Assumption B.4 (Continuity in $\theta$).

There exists a constant $C_{2}$ such that for all
$\theta,\theta^{\prime}\in\Theta_{c}$,

$\displaystyle\left|\int(1+|y|^{q})(M_{\pi,\theta}^{k}(x,\mathrm{d}y)-M_{\pi,\theta^{\prime}}^{k}(x,\mathrm{d}y))\right|\leq
C_{2}|\theta-\theta^{\prime}|(1+|x|^{q}).$ (26)

######  Assumption B.5 (Continuity in $x$).

There exists a constant $C_{3}$ such that for all $x_{1},x_{2}\in\mathcal{X}$,

$\displaystyle\sup_{\theta\in\Theta}\left|\int(1+|y|^{q+1})(M_{\pi,\theta}^{k}(x_{1},\mathrm{d}y)-M_{\pi,\theta}^{k}(x_{2},\mathrm{d}y))\right|\leq
C_{3}|x_{1}-x_{2}|(1+|x_{1}|^{q}+|x_{2}|^{q}).$ (27)

######  Assumption B.6 (Conditions on the Objective Function).

For any compact subset $\Theta_{c}\subset\Theta$, there exist positive
constants $p,K_{1},K_{2},K_{3}$ and $v>1/2$ such that for all
$\theta,\theta^{\prime}\in\Theta_{c}$ and $x,x_{1},x_{2}\in\mathcal{X}$,

$\displaystyle\left|\nabla_{\theta}j(\theta,x_{1})\right|$ $\displaystyle\leq
K_{1}(1+|x_{1}|^{p+1}),$ (28)
$\displaystyle|\nabla_{\theta}j(\theta,x_{1})-\nabla_{\theta}j(\theta,x_{1}^{\prime})$
$\displaystyle\leq K_{2}|x_{1}-x_{1}^{\prime}|(1+|x_{1}|^{p}+|x_{2}|^{p}),$
(29)
$\displaystyle|\nabla_{\theta}j(\theta,x_{1})-\nabla_{\theta}j(\theta^{\prime},x_{1})$
$\displaystyle\leq K_{3}|\theta-\theta^{\prime}|^{v}(1+|x_{1}|^{p+1}).$ (30)

With the above assumptions, the result follows from Theorem 1 of [27] by
setting $x\rightarrow x_{1}$, $\Pi_{\theta}\rightarrow M_{\pi,\theta}$,
$H(\theta,x)\rightarrow\nabla_{\theta}j(\theta,x_{1})$.

##  Appendix C Additional Experimental Details

Code for the numerical experiments is written in Python with array
computations handled by JAX [9]. The implementation of relevant methods for
comparison is sourced from open source repositories: DDS using
franciscovargas/denoising_diffusion_samplers, NF-MCMC using kazewong/flowMC,
and FAB using lollcat/fab-jax. All experiments are run on an NVidia V100 GPU
with 32GB of memory. In the following subsections, we will give more details
on the modelling and hyperparameter choices for each experiment, along with
additional results.

###  C.1 Diagnostics

Let $\pi$ and $\nu$ be two probability measures. Let $\mathcal{F}$ denote the
unit ball in a reproducing kernel Hilbert space (RKHS) $\mathcal{H}$,
associated with the positive definite kernel
$k:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow\mathbb{R}$. Then the maximum
mean discrepancy (MMD) between $\pi$ and $\nu$ is defined as [26, Section 2.2]

$\displaystyle\text{MMD}_{k}^{2}(\pi,\nu)$
$\displaystyle=\|m_{\pi}-m_{\nu}\|_{\mathcal{F}}^{2},$ (31)

where $m_{\pi}$ is the mean embedding of $\pi$, defined via
$\mathbb{E}_{\pi}[f]=\langle f,m_{\pi}\rangle_{\mathcal{H}}$ for all
$f\in\mathcal{H}$. Using standard properties of the RKHS, the squared MMD can
be written as [26, Lemma 6]

$\mathrm{MMD}_{k}^{2}(\pi,\nu)=\mathbb{E}_{x,x^{\prime}\sim\pi}\left[k(x,x^{\prime})\right]-2\mathbb{E}_{x\sim\pi,y\sim\nu}\left[k(x,y)\right]+\mathbb{E}_{y\sim\nu,y^{\prime}\sim\nu}\left[k(y,y^{\prime})\right].$
(32)

Thus, given samples $(x_{i})_{i=1}^{m}\sim\pi$ and $(y_{i})_{i=1}^{m}\sim\nu$,
an unbiased estimate of the squared MMD can be computed as

$\displaystyle\widehat{\text{MMD}}_{k}^{2}(\pi,\nu)$
$\displaystyle=\frac{1}{m(m-1)}\sum_{i=1}^{m}\sum_{i\neq
j}^{m}k(x_{i},x_{j})-\frac{2}{m^{2}}\sum_{i=1}^{m}\sum_{j=1}^{m}k(x_{i},y_{j})$
$\displaystyle+\frac{1}{m(m-1)}\sum_{i=1}^{m}\sum_{j\neq
i}^{m}k(y_{i},y_{j}).$ (33)

For a kernel $k$, the kernel Stein discrepancy (KSD) between $\pi$ and $\nu$
is defined as the MMD between $\pi$ and $\nu$, using the Stein kernel
$k_{\pi}$ associated with $k$, which is defined as

$\displaystyle k_{\pi}(x,x^{\prime})$
$\displaystyle=\nabla_{x}\cdot\nabla_{x^{\prime}}k(x,x^{\prime})+\nabla_{x}k(x,x^{\prime})\cdot\nabla_{x^{\prime}}\log\pi(x^{\prime})$
$\displaystyle+\nabla_{x^{\prime}}k(x,x^{\prime})\cdot\nabla_{x}\log\pi(x)+k(x,x^{\prime})\nabla_{x}\log\pi(x)\cdot\nabla_{x^{\prime}}\log\pi(x),$
(34)

and satisfies the Stein identity
$\mathbb{E}_{\pi}\left[k_{\pi}(x,\cdot)\right]=0$. We thus have that

$\displaystyle\mathrm{KSD}_{k}^{2}(\pi,\nu)$
$\displaystyle=\mathrm{MMD}_{k_{\pi}}^{2}(\pi,\nu)$ (35)
$\displaystyle=\mathbb{E}_{x,x^{\prime}\sim\pi}\left[k_{\pi}(x,x^{\prime})\right]-2\mathbb{E}_{x\sim\pi,y\sim\nu}\left[k_{\pi}(x,y)\right]+\mathbb{E}_{y\sim\nu,y^{\prime}\sim\nu}\left[k_{\pi}(y,y^{\prime})\right]$
(36)
$\displaystyle=\mathbb{E}_{y\sim\nu,y^{\prime}\sim\nu}\left[k_{\pi}(y,y^{\prime})\right].$
(37)

We can obtain estimates of the KSD by using U-statistics or V-statistics. In
particular, an unbiased estimate of $\smash{\mathrm{KSD}_{k}^{2}(\pi,\nu)}$ is
given by the U-statistic [39]

$\widehat{\mathrm{KSD}}^{2}_{k,U}(\pi,\nu)=\frac{1}{n(n-1)}\sum_{i=1}^{n}\sum_{i\neq
j}^{n}k_{\pi}(y_{i},y_{i}^{\prime}).$ (38)

Alternatively, we can estimate $\smash{\mathrm{KSD}_{k}^{2}(\pi,\nu)}$ using a
biased (but non-negative) V-statistic of the form [43, Section 4]

$\widehat{\mathrm{KSD}}^{2}_{k,V}(\pi,\nu)=\frac{1}{n^{2}}\sum_{i=1}^{n}\sum_{j=1}^{n}k_{\pi}(y_{i},y_{i}^{\prime}).$
(39)

In all of our numerical experiments, we calculate the U- and V- statistics
using the inverse multi-quadratic kernel
$k(x,x^{\prime})=(1+(x-x^{\prime})^{T}(x-x^{\prime}))^{\beta}$ due to its
favourable convergence properties [23, Theorem 8], setting
$\beta=-\frac{1}{2}$.

###  C.2 4-mode Gaussian mixture

For this experiment, all methods use $N=128$ parallel chains for training and
$128$ hidden dimensions for all neural networks. Methods with a MALA kernel
use a step size of $0.2$, and methods with splines use 4 coupling layers with
8 bins and range limited to $[-16,16]$.

In Table 3, we present results for $K=10^{3}$ iterations.Since MFM is much
more efficient than other methods, we also report results for a great number
of total iterations. Table 4 contains results for $K=5\cdot 10^{3}$ iterations
for MFM. In the main text, we present results for $K=5\cdot 10^{3}$ learning
iterations for MFM, and $K=10^{3}$ iterations for the other algorithms, since
this renders the total computational cost of all algorithms somewhat
comparable.

For both choices of $K$, we also present results using Hutchinson’s trace
estimator (HTE) [33, 24] to calculate the MH acceptance probability in the
flow-informed Markov transition kernel. As expected, its effect on sample
quality becomes more apparent as $k_{Q}$ increases. However, its effect on
computation time is less significant than in larger dimensional examples.

$K=10^{3}$ | $\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ | KSD U-stat. | KSD V-stat. | MMD | seconds
---|---|---|---|---|---
FM w/ $\pi$ samples | $-4.39\pm 0.10$ |  $6.51$e-3$\pm 2.62$e-3 |  $6.85$e-3$\pm 2.62$e-3 |  $1.39$e-3$\pm 6.00$e-4 | $5.88\pm 0.30$
MFM $k_{Q}=K$ | $-4.65\pm 0.06$ |  $8.32$e-3$\pm 8.95$e-3 |  $8.69$e-3$\pm 8.95$e-3 |  $3.05$e-3$\pm 2.51$e-3 | $9.31\pm 0.58$
MFM $k_{Q}=10^{2}$ | $-4.62\pm 0.08$ |  $6.69$e-3$\pm 4.45$e-3 |  $7.07$e-3$\pm 4.45$e-3 |  $2.05$e-3$\pm 1.39$e-3 | $11.1\pm 0.58$
– w/ HTE | $-4.63\pm 0.10$ |  $6.47$e-3$\pm 4.74$e-3 |  $6.84$e-3$\pm 4.74$e-3 |  $2.06$e-3$\pm 1.92$e-3 | $12.5\pm 0.64$
MFM $k_{Q}=10$ | $-4.59\pm 0.07$ |  $7.24$e-3$\pm 2.58$e-3 |  $7.61$e-3$\pm 2.58$e-3 |  $1.55$e-3$\pm 8.30$e-4 | $26.1\pm 1.21$
– w/ HTE | $-4.67\pm 0.10$ |  $1.17$e-2$\pm 8.44$e-3 |  $1.21$e-2$\pm 8.44$e-3 |  $2.57$e-3$\pm 1.27$e-3 | $31.7\pm 1.51$
DDS | $-4.22\pm 0.03$ |  $9.89$e-4$\pm 1.05$e-3 |  $1.30$e-3$\pm 1.05$e-3 |  $1.76$e-4$\pm 2.32$e-4 | $114.\pm 0.68$
NF-MCMC | $-4.37\pm 0.21$ |  $1.80$e-2$\pm 1.44$e-2 |  $1.83$e-2$\pm 1.44$e-2 |  $5.85$e-3$\pm 3.91$e-3 | $72.0\pm 11.7$
FAB | $-4.67\pm 0.16$ |  $2.31$e-3$\pm 1.19$e-3 |  $2.69$e-3$\pm 1.21$e-3 |  $2.69$e-4$\pm 2.06$e-4 | $101.\pm 3.24$
Table 3: Diagnostics for the 4-mode Gaussian mixture with $K=10^{3}$. $\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ is the Monte Carlo approximation of the log-target density using the learned flow to generate samples; KSD U-stat and V-stat are the Kernel Stein Discrepancy U- and V-statistics between the target and samples generated from the learned flow; MMD is the Maximum Mean Discrepancy between real samples from the target and samples generated from the learned flow. Results are averaged and empirical 95% confidence intervals over 10 independent runs. $K=5\cdot 10^{3}$ | $\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ | KSD U-stat. | KSD V-stat. | MMD | seconds
---|---|---|---|---|---
FM w/ $\pi$ samples | $-4.22\pm 0.05$ |  $2.20$e-3$\pm 1.19$e-3 |  $2.52$e-3$\pm 1.19$e-3 |  $5.45$e-4$\pm 3.51$e-4 | $20.5\pm 0.41$
MFM $k_{Q}=K$ | $-4.48\pm 0.06$ |  $3.62$e-3$\pm 1.51$e-3 |  $3.97$e-3$\pm 1.51$e-3 |  $2.49$e-3$\pm 2.41$e-3 | $25.4\pm 0.94$
MFM $k_{Q}=10^{2}$ | $-4.47\pm 0.07$ |  $3.97$e-3$\pm 1.59$e-3 |  $4.32$e-3$\pm 1.58$e-3 |  $1.39$e-3$\pm 8.22$e-4 | $38.1\pm 1.29$
– w/ HTE | $-4.49\pm 0.06$ |  $4.00$e-3$\pm 1.85$e-3 |  $4.35$e-3$\pm 1.85$e-3 |  $2.06$e-3$\pm 1.69$e-3 | $43.6\pm 1.77$
MFM $k_{Q}=10$ | $-4.44\pm 0.07$ |  $3.13$e-3$\pm 1.25$e-3 |  $3.47$e-3$\pm 1.26$e-3 |  $7.92$e-4$\pm 5.82$e-4 | $145.\pm 2.59$
– w/ HTE | $-4.48\pm 0.05$ |  $5.33$e-3$\pm 3.38$e-3 |  $5.68$e-3$\pm 3.38$e-3 |  $1.19$e-3$\pm 5.11$e-4 | $176.\pm 4.20$
Table 4: Diagnostics for the 4-mode Gaussian mixture with $K=5\cdot 10^{3}$.
$\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ is the Monte Carlo approximation
of the log-target density using the learned flow to generate samples; KSD
U-stat and V-stat are the Kernel Stein Discrepancy U- and V-statistics between
the target and samples generated from the learned flow; MMD is the Maximum
Mean Discrepancy between real samples from the target and samples generated
from the learned flow. Results are averaged and empirical 95% confidence
intervals over 10 independent runs.

###  C.3 16-mode Gaussian mixture

Like the 4-mode example, all methods use $N=128$ parallel chains for training
and $128$ hidden dimensions for all neural networks. Methods with a MALA
kernel use a step size of $0.2$, and methods with splines use 4 coupling
layers with 8 bins and range limited to $[-16,16]$. In Table 5 we present
results for $K=10^{3}$ iterations. In Table 6, we provide results for MFM for
$K=5\times 10^{3}$ learning iterations. In the main text, we present results
for $K=5\cdot 10^{3}$ learning iterations for MFM and $K=10^{3}$ iterations
for all other algorithms, which yields a more comparable total computation
time.

$K=10^{3}$ | $\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ | KSD U-stat. | KSD V-stat. | MMD | seconds
---|---|---|---|---|---
FM w/ $\pi$ samples | $-7.19\pm 0.18$ |  $3.16$e-2$\pm 6.47$e-3 |  $3.22$e-2$\pm 6.49$e-3 |  $9.88$e-3$\pm 1.14$e-3 | $6.22\pm 0.24$
MFM $k_{Q}=K$ | $-7.89\pm 0.16$ |  $2.60$e-2$\pm 4.18$e-3 |  $2.66$e-2$\pm 4.20$e-3 |  $2.32$e-2$\pm 3.63$e-3 | $10.8\pm 0.67$
MFM $k_{Q}=10^{2}$ | $-7.20\pm 0.30$ |  $1.72$e-2$\pm 3.17$e-3 |  $1.77$e-2$\pm 3.19$e-3 |  $1.59$e-2$\pm 4.62$e-3 | $12.9\pm 0.65$
– w/ HTE | $-7.25\pm 0.29$ |  $1.89$e-2$\pm 4.10$e-3 |  $1.94$e-2$\pm 4.13$e-3 |  $1.65$e-2$\pm 7.56$e-3 | $14.3\pm 0.53$
MFM $k_{Q}=10$ | $-7.32\pm 0.08$ |  $3.06$e-2$\pm 4.70$e-3 |  $3.12$e-2$\pm 4.72$e-3 |  $1.15$e-2$\pm 8.73$e-4 | $26.0\pm 0.62$
– w/ HTE | $-7.23\pm 0.24$ |  $3.40$e-2$\pm 1.07$e-2 |  $3.46$e-2$\pm 1.07$e-2 |  $1.31$e-2$\pm 1.80$e-3 | $29.6\pm 0.65$
DDS | $-5.86\pm 0.20$ |  $6.65$e-3$\pm 6.69$e-3 |  $6.94$e-3$\pm 6.69$e-3 |  $1.02$e-1$\pm 4.10$e-2 | $115.\pm 0.64$
NF-MCMC | $-5.74\pm 0.35$ |  $1.23$e-2$\pm 1.71$e-2 |  $1.26$e-2$\pm 1.72$e-2 |  $8.05$e-3$\pm 1.42$e-2 | $67.0\pm 12.3$
FAB | $-5.89\pm 0.28$ |  $4.22$e-3$\pm 3.31$e-3 |  $4.58$e-3$\pm 3.34$e-3 |  $1.51$e-3$\pm 1.06$e-3 | $102.\pm 4.32$
Table 5: Diagnostics for the 16-mode Gaussian mixture with $K=10^{3}$. $\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ is the Monte Carlo approximation of the log-target density using the learned flow to generate samples; KSD U-stat and V-stat are the Kernel Stein Discrepancy U- and V-statistics between the target and samples generated from the learned flow; MMD is the Maximum Mean Discrepancy between real samples from the target and samples generated from the learned flow. Results are averaged and empirical 95% confidence intervals over 10 independent runs. $K=5\cdot 10^{3}$ | $\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ | KSD U-stat. | KSD V-stat. | MMD | seconds
---|---|---|---|---|---
FM w/ $\pi$ samples | $-5.90\pm 0.08$ |  $4.74$e-3$\pm 1.11$e-3 |  $5.11$e-3$\pm 1.11$e-3 |  $2.16$e-3$\pm 3.95$e-4 | $22.8\pm 0.22$
MFM $k_{Q}=K$ | $-6.18\pm 0.06$ |  $3.78$e-3$\pm 9.90$e-4 |  $4.13$e-3$\pm 9.97$e-4 |  $1.96$e-2$\pm 3.24$e-3 | $29.4\pm 1.38$
MFM $k_{Q}=10^{2}$ | $-6.00\pm 0.10$ |  $7.01$e-3$\pm 3.07$e-3 |  $7.39$e-3$\pm 3.09$e-3 |  $3.57$e-3$\pm 7.68$e-4 | $38.5\pm 1.75$
– w/ HTE | $-6.00\pm 0.11$ |  $7.38$e-3$\pm 2.62$e-3 |  $7.76$e-3$\pm 2.64$e-3 |  $4.87$e-3$\pm 1.67$e-3 | $40.9\pm 1.84$
MFM $k_{Q}=10$ | $-6.13\pm 0.11$ |  $7.95$e-3$\pm 1.80$e-3 |  $8.35$e-3$\pm 1.81$e-3 |  $3.68$e-3$\pm 7.55$e-4 | $119\pm 5.76$
– w/ HTE | $-6.00\pm 0.11$ |  $9.85$e-3$\pm 3.63$e-3 |  $1.03$e-2$\pm 3.64$e-3 |  $8.81$e-3$\pm 1.28$e-3 | $134\pm 3.58$
Table 6: Diagnostics for the 16-mode Gaussian mixture with $K=5\cdot 10^{3}$.
$\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ is the Monte Carlo approximation
of the log-target density using the learned flow to generate samples; KSD
U-stat and V-stat are the Kernel Stein Discrepancy U- and V-statistics between
the target and samples generated from the learned flow; MMD is the Maximum
Mean Discrepancy between real samples from the target and samples generated
from the learned flow. Results are averaged and empirical 95% confidence
intervals over 10 independent runs.

###  C.4 Field system

The Allen–Cahn equation is defined in terms of a random field
$\phi:[0,1]\rightarrow\mathbb{R}$ satisfying the following stochastic partial
differential equation:

$\frac{\partial\phi}{\partial t}=a\frac{\partial^{2}\phi}{\partial
s^{2}}+a^{-1}(\phi-\phi^{3})+\sqrt{2\beta^{-1}}\eta(t,s),$ (40)

where $a>0$ is a parameter, $\beta$ is the inverse temperature, $s\in[0,1]$
denotes the spatial variable, and $\eta$ is spatiotemporal white noise. The
imposition of Dirichlet boundary conditions ensures $\phi(s=0)=\phi(s=1)=0$.

The associated Hamiltonian, reflecting a spatial coupling term penalizing
changes in $\phi$, takes the form:

$U_{*}[\phi]=\beta\int_{0}^{1}\left[\frac{a}{2}\left(\frac{\partial\phi}{\partial
s}\right)^{2}+\frac{1}{4a}\left(1-\phi^{2}(s)\right)^{2}\right]ds.$ (41)

At low temperatures, this coupling induces alignment of the field in either
the positive or negative direction, leading to two global minima, $\phi_{+}$
and $\phi_{-}$, with typical values of $\pm 1$.

For this example, all methods use $N=1024$ parallel chains for training and
$256$ hidden dimensions for all neural networks. Methods with a MALA kernel
use a step size of $0.0001$, and methods with splines use 8 coupling layers
with 8 bins and range limited to $[-5,5]$. Results for $K=10^{4}$ learning
iterations are presented in Table 7. In the main text, we present results for
$k_{Q}=10^{2}$ for MFM.

Interestingly, in high-dimensional problems such as this and the following,
the version of the algorithm using Hutchinson’s trace estimator [33, 24] to
calculate the MH acceptance probability has little apparent effect on the
approximation quality. It does, however, have a significant impact on the
computation time.

$K=10^{4}$ | $\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ | KSD U-stat. | KSD V-stat. |  | seconds
---|---|---|---|---|---
MFM $k_{Q}=K$ | $-73.2\pm 1.04$ | $2.22\pm 1.73$ | $20.2\pm 1.73$ |  | $45.2\pm 0.76$
MFM $k_{Q}=K/10$ | $-68.7\pm 1.29$ | $2.32\pm 1.87$ | $20.4\pm 1.88$ |  | $53.0\pm 0.70$
– w/ HTE | $-72.6\pm 1.52$ | $2.03\pm 1.61$ | $20.1\pm 1.63$ |  | $46.0\pm 0.86$
MFM $k_{Q}=K/100$ | $-56.8\pm 0.65$ | $2.17\pm 2.19$ | $20.4\pm 2.30$ |  | $146\pm 1.93$
– w/ HTE | $-68.7\pm 2.30$ | $2.42\pm 1.74$ | $21.0\pm 1.80$ |  | $61.6\pm 0.95$
DDS | $-76.3\pm 17.4$ | $15.2\pm 35.9$ | $18.0\pm 36.9$ |  | $2400\pm 8.65$
NF-MCMC | $-26.9\pm 9.62$ | $548\pm 325$ | $549\pm 325$ |  | $2000\pm 15.6$
FAB | $-50.4\pm 0.14$ | $0.14\pm 0.42$ | $1.78\pm 0.42$ |  | $3880\pm 7.19$
Table 7: Diagnostics for the field system with $K=10^{4}$.
$\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ is the Monte Carlo approximation
of the log-target density using the learned flow to generate samples; KSD
U-stat and V-stat are the Kernel Stein Discrepancy U- and V-statistics between
the target and samples generated from the learned flow. Results are averaged
and empirical 95% confidence intervals over 10 independent runs.

###  C.5 Log-Gaussian Cox point process

The original $10\times 10$ square meter plot is standardized to the unit
square, as depicted in Fig. 4. We discretize the unit square $[0,1]^{2}$ into
a $M=40\times 40$ regular grid. The latent intensity process
$\Lambda=\\{\Lambda_{m}\\}_{m\in M}$ is specified as
$\Lambda_{m}=\exp(X_{m})$, where $X=\\{X_{m}\\}_{m\in M}$ is a Gaussian
process with a constant mean $\mu_{0}\in\mathbb{R}$ and exponential covariance
function $\Sigma_{0}(m,n)=\sigma^{2}\exp\left(-|m-n|/(30\beta)\right)$ for
$m,n\in M$, i.e. $X\sim\mathcal{N}(\mu_{0}1_{d},\Sigma_{0})$ for
$1_{d}=[1,\dots,1]^{T}\in\mathbb{R}^{d}$ with dimension $d=900$. The chosen
parameter values are $\sigma^{2}=1.91$, $\beta=1/33$, and
$\mu_{0}=\log(126)-\sigma^{2}/2$, estimated by [50]. The number of points in
each grid cell $Y=\\{Y_{m}\\}_{m\in M}\in\mathbb{N}^{40\times 40}$ are
modelled as conditionally independent and Poisson distributed with means
$a\Lambda_{m}$,

$\displaystyle\mathcal{L}(Y|X)=\prod_{m\in[1:30]^{2}}\exp(x_{m}y_{m}-a\exp(x_{m})),$
(42)

where $a=1/40^{2}$ represents the area of each grid cell.

Figure 4: Visualization of the standardized 10 $\times$ 10 square meter plot
and the locations of 126 Scots pine saplings in Finland.

For this example, all methods use $N=128$ parallel chains for training and
$1024$ hidden dimensions for all neural networks. Methods with a MALA kernel
use a step size of $0.01$, and methods with splines use 8 coupling layers with
8 bins and range limited to $[-10,10]$. Results for $K=10^{4}$ learning
iterations are presented in Table 8. In the main text, we present results for
$k_{Q}=10^{3}$ for MFM. We were unable to run NF-MCMC and FAB for $K=10^{4}$
iterations because of memory issues; instead, we present results for
$K=10^{3}$ iterations only for the models using discrete normalizing flows.

$K=10^{4}$ | $\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ | KSD U-stat. | KSD V-stat. |  | seconds
---|---|---|---|---|---
MFM $k_{Q}=K$ | $-2110\pm 11.3$ |  $1.07$e-1$\pm 4.74$e-2 | $31.2\pm 0.20$ |  | $109\pm 4.15$
MFM $k_{Q}=10^{3}$ | $-2110\pm 7.39$ |  $1.05$e-1$\pm 5.11$e-2 | $31.2\pm 0.16$ |  | $910\pm 47.1$
– w/ HTE | $-2110\pm 8.35$ |  $1.07$e-1$\pm 4.73$e-2 | $31.2\pm 0.17$ |  | $122\pm 0.64$
MFM $k_{Q}=10^{2}$ | $-2110\pm 7.61$ |  $1.12$e-1$\pm 5.22$e-2 | $31.2\pm 0.22$ |  | $8470\pm 193$
– w/ HTE | $-2120\pm 8.50$ |  $1.21$e-1$\pm 4.96$e-2 | $31.3\pm 0.18$ |  | $250\pm 3.65$
DDS | $-1850\pm 8.59$ |  $7.59$e-2$\pm 2.24$e-2 | $24.7\pm 0.08$ |  | $3260\pm 8.41$
NF-MCMC | $-1410\pm 53.8$ | $11.8\pm 7.55$ | $89.0\pm 238$ |  | $215\pm 46.4$
FAB | $-3070\pm 80.7$ |  $1.55$e-1$\pm 6.12$e-2 | $52.3\pm 2.02$ |  | $1040\pm 2.78$
Table 8: Log Gaussian Cox point process diagnostics for $K=10^{4}$ where
$\mathbb{E}_{[\phi_{1}]_{\\#}p_{0}}\log\pi$ is the Monte Carlo approximation
of the log-target density using the learned flow to generate samples; KSD
U-stat and V-stat are the Kernelized Stein discrepancy’s U- and V-statistics
between the target and samples generated from the learned flow. Results are
averaged and empirical 95% confidence intervals over 10 independent runs.