Title: Markovian Flow Matching: Accelerating MCMC with Continuous Normalizing Flows
Venue: NeurIPS 2024
Accepted: yes
Keywords: Sampling algorithms

Bayesian inference

MCMC

flow matching

continuous normalizing flows
Abstract: Continuous normalizing flows (CNFs) learn the probability path between a reference distribution and a target distribution by modeling the vector field generating said path using neural networks. Recently, Lipman et al. (2022) introduced a simple and inexpensive method for training CNFs in generative modeling, termed flow matching (FM). In this paper, we repurpose this method for probabilistic inference by incorporating Markovian sampling methods in evaluating the FM objective, and using the learned CNF to improve Monte Carlo sampling. Specifically, we propose an adaptive Markov chain Monte Carlo (MCMC) algorithm, which combines a local Markov transition kernel with a non-local, flow-informed transition kernel, defined using a CNF. This CNF is adapted on-the-fly using samples from the Markov chain, which are used to specify the probability path for the FM objective. Our method also includes an adaptive tempering mechanism that allows the discovery of multiple modes in the target distribution. Under mild assumptions, we establish convergence of our method to a local optimum of the FM objective. We then benchmark our approach on several synthetic and real-world examples, achieving similar performance to other state-of-the-art methods but often at a significantly lower computational cost.

Markovian Flow Matching: Accelerating MCMC with Continuous Normalizing Flows

Abstract

Continuous normalizing flows (CNFs) learn the probability path between a reference distribution and a target distribution by modeling the vector field generating said path using neural networks. Recently, Lipman et al. [45] introduced a simple and inexpensive method for training CNFs in generative modeling, termed flow matching (FM). In this paper, we repurpose this method for probabilistic inference by incorporating Markovian sampling methods in evaluating the FM objective, and using the learned CNF to improve Monte Carlo sampling. Specifically, we propose an adaptive Markov chain Monte Carlo (MCMC) algorithm, which combines a local Markov transition kernel with a non-local, flow-informed transition kernel, defined using a CNF. This CNF is adapted on-the-fly using samples from the Markov chain, which are used to specify the probability path for the FM objective. Our method also includes an adaptive tempering mechanism that allows the discovery of multiple modes in the target distribution. Under mild assumptions, we establish convergence of our method to a local optimum of the FM objective. We then benchmark our approach on several synthetic and real-world examples, achieving similar performance to other state-of-the-art methods, but often at a significantly lower computational cost.

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 [51], Bayesian inference [25], and molecular dynamics [43]. In particular, let π (d x ) be a target probability distribution on R d with density π ( x ) with respect to the Lebesgue measure of the form 1

where ˆ π : R d → R + is a continuously differentiable function which can be evaluated pointwise, and Z = ∫ R d ˆ π ( x )d x is an unknown normalizing constant. We are interested in generating samples from the target distribution π in order to approximate integrals of the form π [ f ] = E π [ f ( x )] , where d

A standard solution to this problem is Markov chain Monte Carlo (MCMC) [12, 64], which relies on the construction of a Markov process which admits the target π as its invariant distribution. One of the most broadly applicable and widely studied MCMC methods is the Metropolis-Hastings (MH) algorithm [32], which proceeds in two steps. First, given a current sample x , a new sample y is proposed according to some proposal distribution q ( ·| x ) . Then, this sample is accepted with probability α ( x, y ) = min { 1 , π ( y ) q ( x | y ) π ( x ) q ( y | x ) } . 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 [65]. 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., 48].

Naturally, the choice of proposal distribution q ( ·| 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 [38], protein folding [41], astrophysics [22], and sensor network localization [37]. On the one hand, local proposals, such as those employed in the Metropolis-Adjusted Langevin Algorithm (MALA) [66] or Hamiltonian Monte Carlo (HMC) [20, 56] struggle to transition between regions of high-probability, resulting in very long decorrelation times and few effective independent samples [e.g., 49]. On the other hand, global proposal distributions must be very carefully designed in order to avoid high rejection rates, particularly in high dimensions [17, 47].

Another popular approach to sampling is variational inference (VI) [10, 34, 61, 79], which obtains a parametric approximation π θ ∗ ( x ) ≈ π ( x ) to the target by minimising the Kullback-Leibler (KL) divergence to the target over a parameterized family of distributions D θ = { π θ : θ ∈ Θ } . State-of-theart VI methods use normalizing flows (NFs), which consist of a sequence of invertible transformations between a reference and a target distribution, to define a flexible variational family [62]. There has also been growing interest in the use of continuous normalizing flows (CNFs), which define a path between distributions using ordinary differential equations [15, 27, 45]. 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 utilize NFs to enhance the performance of MCMC algorithms; see, e.g., [28] 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., 33, 39, 59, 68]. 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., 24, 67].

Our contributions In this paper, we continue this promising line of work, introducing 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 Lipman et al. [45]. 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 [15]. Let v t be a time-dependent vector field that runs continuously in the unit interval. Under mild conditions, this vector field can be used to construct a time-dependent diffeomorphic map

called a flow ϕ : [0 , 1] × R d → R d , defined via the ordinary differential equation (ODE):

Given a reference density p 0 : R d → R + , and the flow ϕ , we can generate a probability density path p : [0 , 1] × R d → R + as the pushforward of p 0 under ϕ , viz p t := [ ϕ t ] ♯ p 0 , for t ∈ [0 , 1] . This yields, via the instantaneous change-of-variables formula [e.g., 14]

where x s := ϕ s ( x ) , and where ∇ 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 , in which case the ODE in (2) is referred to as a neural ODE [15]. In turn, this yields a deep parametric model ϕ θ t for the flow ϕ t , known as a CNF [27].

Flow Matching One would typically like to learn a CNF which maps between a given reference density p 0 and a target density π . Given samples from the target, one approach is to maximize the log-likelihood E x ∼ π [ log p θ 1 ( x ) ] . 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) provides an alternative, simulation-free method for training CNFs [45]. Let p t ( x ) be a target probability density path such that p 0 = p is a simple reference distribution, and p 1 ≈ π 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 v θ t ( x ) is defined as

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. [45, Theorem 2], it is equivalent to minimize the conditional flow-matching (CFM) loss

where p t ( ·| x 1 ) is a conditional probability density path satisfying p 0 ( x | x 1 ) = p 0 and p 1 ( x | x 1 ) ≈ δ x 1 , and v t ( ·| x 1 ) : R d → R d is a conditional vector field that generates p t ( ·| x 1 ) . There are various choices for p t ( ·| x 1 ) and v t ( ·| x 1 ) . For simplicity, we here assume that the conditional probability path is Gaussian, viz p t ( x | x 1 ) = N ( x | m t ( x 1 ) , s t ( x 1 ) 2 I d ) , where m : [0 , 1] × R d → R d denotes a timedependent mean, and s : [0 , 1] × R → R + a time-dependent scalar standard deviation. For our experiments, we further adopt the optimal transport conditional probability path introduced in [45], setting m t ( x 1 ) = tx 1 and s t ( x 1 ) = 1 -(1 -σ min ) t for some σ min ≪ 1 . In this case, the conditional vector field assumes the particularly simple form v t ( x | x 1 ) = x 1 -(1 -σ min ) x 1 -(1 -σ 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, gradient-based Markov transition kernel; and an adaptive annealing schedule. We begin by describing how CNFs can be used within a MHMCMCalgorithm.

3.1 MCMCwith Flow Matching

Suppose, for now, that we have access to a CNF ( ϕ θ t ) t ∈ [0 , 1] , trained (e.g.) via flow-matching, with corresponding vector field ( v θ t ) t ∈ [0 , 1] , which generates a probability path ( p θ t ) t ∈ [0 , 1] between a reference density p 0 and an approximation of the target density π . Given a point x 0 ∈ R d on the reference space, we can evaluate the log-density of the pullback of the target distribution π as

Under the assumption that the CNF approximately transports samples from p 0 to π , we expect that [ ϕ θ 1 ] ♯ p 0 ≈ π in the target space, and that [ ϕ θ 1 ] ♯ π ≈ 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 [28, 33, 44, 59]. First, transform initial positions x 1 from the target space to the reference space by solving

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

This corresponds to using a transformation-informed proposal in a Markov transition step, an approach which has been successfully applied using (discrete) normalizing flows [28, 33, 44, 59].

There are various possible choices for the proposal distribution on the reference space (see Appendix A). For example, [23] 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 [39]. Concretely, our flow-informed random-walk transition kernel, summarized in Algorithm 2 (see Appendix A), can be written as

where ρ θ (d y | x ) is the distribution defined by the transition

and α ( x, y ) = min { 1 , π ( y ) ρ θ ( x | y ) π ( x ) ρ θ ( y | x ) } and b ( x ) = ∫ R d α ( x, y ) ρ θ (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 π . Clearly, however, the CFM objective is not immediately applicable in our setting, since we do not have access to samples from the target π .

Tong et al. [72] 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.

In this paper, we adopt a different approach, updating the parameters of a CNF based on a dynamic estimate of the CFM objective obtained via an adaptive MCMC algorithm. This is similar in spirit to other recent flow-informed MCMC algorithms [24, 35, 67], and the Markovian score climbing algorithm in [54].

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 π , 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

where q (d y | x ) is given by

and where, as elsewhere, α ( x, y ) = min { 1 , π ( y ) q ( x | y ) π ( x ) q ( y | x ) } and b ( x ) = ∫ R d α ( x, y ) q (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 nonlocal 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 R := P ◦ Q , given explicitly by R ( x, d y ; π, θ ) = ∫ Z Q ( x, d z ; π ) P ( z, d y ; π, θ ) . 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 P are updated based on a new estimate of J ( θ ; π ) . To be precise, suppose we write µ t := µ 0 R k ( · , · ; π, θ ) for the distribution of the Markov chain with kernel R ( · , · ; π, θ ) after k ∈ N steps, starting from initialization µ 0 , where R k = R ◦ R · · · ◦ R . Our objective function is then given by

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

where v t ( x | x 1 ) = x 1 -(1 -σ min ) x 1 -(1 -σ min ) t and ϕ t ( x | x 1 ) = (1 -(1 -σ min t ) x + tx 1 . In practice, we will optimize a Monte Carlo estimate of this objective, namely,

where { x i ( k ) } N i =1 are the samples from N chains of our MCMC algorithm after k ∈ N iterations, x i 0 i . i . d . ∼ p 0 , and t i ∼ 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 vectororiented 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 π . 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., 24, Section IV.C]. To alleviate this problem, one approach is to iteratively target a sequence of annealed densities { π k ( x ) } k =0: K , which smoothly interpolate between a simple base distribution π 0 ( x ) (e.g., a standard Gaussian), and the target distribution π K ( x ) := π ( x ) . This idea is central to other Monte Carlo sampling methods such as Sequential Monte Carlo (SMC) [18] and Annealed Importance Sampling (AIS) [55], as well as sampling methods used in score-based generative modelling [e.g., 70]. In our case, the annealed targets act as intermediary steps within the flow-informed MCMC scheme.

A standard way in which to construct the sequence { π k ( x ) } k =0: K is to use a geometric interpolation, defining

Algorithm 1 Markovian Flow Matching

1: Input: target π , base π 0 , vector field v θ t , reference p 0 , concentration σ min , initial parameters θ 0 , target ESS fraction α , iterations K , number of particles N , local kernel Q , flow-informed kernel P , MCMCsteps per resampling step k Q , step sizes ε 1: K . ).

where β 0: K is a sequence of temperatures which satisfies 0 = β 0 < β 1 < · · · < β K = 1 [e.g., 55]. In practice, it can be difficult to choose a good sequence of temperatures that provides a smooth transition between densities. One heuristic for adaptively setting this sequence is based on the effective sample size (ESS). In particular, by setting the ESS to a user-specified percentage α of the number of particles N , the next temperature β k in the schedule can be determined by solving the recursive equation [9]

where w β k -1 i ( β ) = [ π K ( x i ) β π 0 ( x i ) 1 -β ] / [ π K ( x i ) β k -1 π 0 ( x i ) 1 -β k -1 ] = [ π K ( x i ) /π 0 ( x i )] β -β k -1 are new importance weights given the current temperature β k -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 .

Convergence The output of Algorithm 1 is a vector of parameters θ K which defines a CNF ( ϕ θ K t ) t ∈ [0 , 1] . Under the assumption that the parameter estimate converges, that is, θ K → θ ∗ global as K →∞ , where θ ∗ global = arg min θ ∈ Θ J ( θ ; π ) is the global minimizer of the CFM objective J ( θ ; π ) in (5), this CNF is guaranteed to generate a probability path ( p θ t ) t ∈ [0 , 1] which transports samples from the reference p 0 to the true target π [e.g., 45].

In practice, the objective J ( θ ; π ) 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 ( v θ t ) t ∈ [0 , 1] . This being said, it is reasonable to ask whether θ 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 θ K → θ ∗ almost surely as K →∞ , where θ ∗ denotes a local minimum of the CFM objective. This proposition closely mirrors [54, Proposition 1]. Its proof, which relies on a classical result in [6, 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 ( θ K ) K ∈ N is a bounded sequence, which almost surely visits a compact subset of the domain of attraction of θ ∗ infinitely often. Then θ K → θ ∗ almost surely.

4 Related work

In recent years, a number of works have proposed algorithms which combine MCMC techniques with NFs; see, e.g., [2, 28] for recent surveys. Broadly speaking, these algorithms fall into two distinct 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 [59], and since been investigated in a number of different works [e.g., 13, 33, 44, 54, 58, 68, 82, 84]. Flow MCMC methods, meanwhile, utilize the pushforward of the base distribution through the NF as an (independent) proposal within an MCMC scheme. This approach was first studied by [3], and further extended in [4, 31, 57].

More recently, [24, 67] 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 [35]. 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, [5, 50] consider augmenting SMC with NFs, while [19, 52] use NFs (or diffusion models) within AIS.

Although less directly comparable to our own approach, several other recent works have proposed to use (controlled) diffusion processes to sample from unnormalized probability distributions. Such works include Zhang and Chen [85], who introduce the path integral sampler , Vargas et al. [75], who propose the denoising diffusion sampler , and Zhang et al. [83], who introduce generative flow samplers . Some other relevant contributions in this direction include [1, 8, 16, 60, 63, 69, 73, 74, 75, 76, 77, 78].

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 four relevant methods. The Denoising Diffusion Sampler [DDS; 75] 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; 24] is an augmented MCMC scheme which uses a mixture of MALA and adaptive transition kernels learned using discrete NFs. Flow Annealed Importance Sampling Bootstrap [FAB; 52] is an augmented AIS scheme minimizing the mass-covering α -divergence with α = 2 . Finally, Adaptive Tempered SMC (AT-SMC), i.e. the SMC algorithm described in [18] using a MALA transition kernel and a sequence of annealed distributions chosen adaptively by solving (17).

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 [85], we parameterize the vector field as

where the neural networks are standard MLPs with 2 hidden layers, using a Fourier feature augmentation for t [71], and where NN ∗ outputs a real value that reweights the vector field output using the time component. This architecture is also used by DDS [75, Section 4]. Meanwhile, FAB and NF-MCMC use rational quadratic splines [21]. Flows are trained using Adam [40] with a linear decay schedule terminating at ε K = 0 . We report results for all methods averaged over 10 independent runs with varying random seeds. Code to reproduce the experiments is provided at https://github.com/albcab/mfm.

5.1 4-mode Gaussian mixture

Our first example is a mixture of four Gaussians, evenly spaced and equally weighted, in twodimensional 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., 29]; 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 Figure 1. Further algorithmic details and results are provided in Appendix C.2.

In this experiment, only our method (Figure 1a) and DDS (Figure 1c) learn the fully separated modes, reflecting the greater expressivity of CNFs in comparison to the discrete NFs used in, e.g., NF-MCMC

(Figure 1d). 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).

k

Q

= 10

2

(b) FAB

(c) DDS

(d) NF-MCMC

15

15

15

15

10

O

0

10

O

10

O

10

O

0

5

5

5

Q

O

Q

0

O

5

O

O

-10

-10

-10

-10

15

15

15

15

15

-10

-5

0

10

15

15

-10

0

5

10

15

15

10

-5

V+

0

5

10

15

-15

-10

5

O

5

10

15

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 Figure 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, FAB marginally outperforms MFM as measured by the MMD, but this slight improvement in performance comes at the cost of a much higher run-time.

5.3 Field system

Our first real-world example considers the stochastic Allen-Cahn model [7], used as a benchmark in [24], and described in Appendix C.5. 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 model leads to a discretized target density

k

Q

= 10

2

(b) FAB

(c) DDS

(d) NF-MCMC

15

15

15

15

10

10

10

10

5

5

D

0

0

-5

-5

-10

-10

10

10

15

15

15

15

15

-10

5

Q

5

10

15

15

-10

-5

0

5

10

15

15

-10

-5

D

5

10

15

15

-10

5

10

15

which takes the form

with ∆ s = 1 d , and boundary conditions x 0 = x d +1 = 0 . In our experiments, we take d = 64 . Meanwhile, following [24], other parameter values are chosen to ensure bimodality at x = ± 1 : a = 0 . 1 , b = 1 /a = 10 , and β = 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., 24], 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.

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., 26, 46]; 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 Figure 3. This experiment demonstrates the ability 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. Indeed, while FAB was the best performing method in this experiment as measured by the KSD, it failed to capture both of the modes in the target distribution, and required a much greater total computation time (see Table 2).

It is worth noting that MFM (and the other two benchmarks, DDS and FAB) 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 comparable results in the absence of an informed base measure.

k

Q

= 10

2

(b) FAB

(c) DDS

(d) NF-MCMC

1.5

0.5

0.5

0.0

1.0

0.0

0.0

0.2

0.5

0.4

0.0

0.5

0.5

0.6

0.5

-1.0

1.0

0.8

1.0

1.5

1.5

1.0

1.5

0

10

20

30

40

50

60

0

10

20

30

40

50

60

0

10

20

30

40

50

60

0

10

20

30

40

50

60

5.4 Log-Gaussian Cox process

Bayesian inference for high-dimensional spatial models is known to be challenging. One such model is the log-Gaussian Cox process (LGCP) introduced in [53], which is used to model the locations of 126 Scots pine saplings in a natural forest in Finland. See Appendix C.6 for full details. The target space is discretized to a M = 40 × 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, MFMprovides a good approximation to the target at a lower computational cost with respect to its competitors.

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 output by our algorithm to a local optimum of the FM objective. We also 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 and Future Work . 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. Indeed, a promising avenue for further research lies in developing tailored CNFs designed for particular posterior distributions. This approach would go beyond the current practice of including the gradient of the log-posterior and instead exploit unique characteristics intrinsic to each model when constructing the flow.

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

T. Akhound-Sadegh, J. Rector-Brooks, A. J. Bose, S. Mittal, P. Lemos, C.-H. Liu, M. Sendera, S. Ravanbakhsh, G. Gidel, Y. Bengio, et al. Iterated denoising energy matching for sampling from Boltzmann densities. In Proceedings of the 41st International Conference on Machine Learning (ICML) , 2024. 7

M. S. Albergo and E. Vanden-Eijnden. Learning to sample better. arXiv preprint arXiv:2310.11232 , 2023. 6

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. 7

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. 7

M. Arbel, A. Matthews, and A. Doucet. Annealed flow transport Monte Carlo. In Proceedings of the 38th International Conference on Machine Learning (ICML) , 2021. 7

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. 6

N. Berglund, G. D. Gesù, and H. Weber. An Eyring-Kramers law for the stochastic Allen-Cahn equation in dimension two. Electronic Journal of Probability , 22:1 - 27, 2017. doi: 10.1214/ 17-EJP60. 8

J. Berner, L. Richter, and K. Ullrich. An optimal control perspective on diffusion-based generative modeling. Transaction on Machine Learning Research (TMLR) , 2024. 7

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. 6

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. 2

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 . 18

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. 2

A. Cabezas and C. Nemeth. Transport elliptical slice sampling. In Proceedings of the 26th International Conference on Artificial Intelligence and Statistics (AISTATS) , 2023. 7

R. Chen and Y. Lipman. Flow matching on general geometries. In Proceedings of the 12th International Conference on Learning Representations (ICLR) , 2024. 3

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. 2, 3

V. De Bortoli, M. Hutchinson, P. Wirnsberger, and A. Doucet. Target score matching. arXiv preprint arXiv:2402.08667 , 2024. 7

L. Del Debbio, J. Marsh Rossney, and M. Wilson. Efficient modeling of trivializing maps for lattice ϕ 4 theory using normalizing flows: a first look at scalability. Physical Review D , 104(9): 094507, 2021. doi: 10.1103/PhysRevD.104.094507. 2

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. 5, 7

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. 7

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. 2

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. 7

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. 2

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. 4

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. 2, 4, 5, 7, 8, 9, 20

A. Gelman, J. B. Carlin, H. S. Stern, and D. B. Rubin. Bayesian data analysis . Chapman and Hall/CRC, 1995. 1

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. 9, 18

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. 2, 3, 19, 21

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. 2, 4, 6

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. 7, 18

M. G. Gu and F. H. Kong. A stochastic approximation algorithm with Markov chain MonteCarlo 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. 17

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. 7

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. 2

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. 2, 4, 7

M. D. Hoffman, D. M. Blei, C. Wang, and J. Paisley. Stochastic variational inference. Journal of Machine Learning Research , 14(1):1303-1347, 2013. 2

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. 4, 7

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):10591076, 1989. 19, 21

A. Ihler, J. Fisher III, R. Moses, and A. Willsky. Nonparametric belief propagation for selflocalization of sensor networks. IEEE Journal on Selected Areas in Communications , 23(4), 2005. doi: 10.1109/JSAC.2005.843548. 2

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. 2

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. 2, 4

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In Proceedings of the 3rd International Conference on Learning Representations (ICLR) , 2015. 7

S. Kou, J. Oh, and W. H. Wong. A study of density of states and ground states in hydrophobichydrophilic protein folding models by equi-energy sampling. The Journal of Chemical Physics , 124(24), 2006. doi: 10.1063/1.2208607. 2

A. Lee. U-Statistics: Theory and Practice . Taylor & Francis, 1st edition, 1990. doi: 10.1201/ 9780203734520. 18

B. Leimkuhler and C. Matthews. Molecular dynamics. Interdisciplinary Applied Mathematics , 39:443, 2015. doi: 10.1007/978-3-319-16375-8. 1

S.-H. Li and L. Wang. Neural network renormalization group. Physical Review Letters , 121: 260601, 2018. doi: 10.1103/PhysRevLett.121.260601. 4, 7

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. 1, 2, 3, 5, 6

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. 9, 18

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. 2

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. 2

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. 2

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. 7

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. 1

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. 7, 18, 20

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. 10, 22

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. 4, 6, 7, 17

R. M. Neal. Annealed importance sampling. Statistics and Computing , 11:125-139, 2001. doi: 10.1023/A:1008923215028. 5, 6

R. M. Neal et al. MCMC using Hamiltonian dynamics. Handbook of Markov chain Monte Carlo , 2(11):2, 2011. 2

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. 7

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. 7, 20

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. 2, 4, 7

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. 7

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. 2

D. Rezende and S. Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on Machine Learning (ICML) , 2015. 2

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. 7

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. 2

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. 2

G. O. Roberts and R. L. Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli , 2(4):341-363, 1996. 2

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. 2, 4, 7

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. 2, 7

M. Sendera, M. Kim, S. Mittal, P. Lemos, L. Scimeca, J. Rector-Brooks, A. Adam, Y. Bengio, and N. Malkin. Improved off-policy training of diffusion samplers. arXiv preprint arXiv:2402.05098 , 2024. 7

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. 5

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. 7

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. 4

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. 7

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. 7

F. Vargas, W. Grathwohl, and A. Doucet. Denoising diffusion samplers. In Proceedings of the 11th International Conference on Learning Representations (ICLR) , 2023. 7

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, 2023. doi: 10.1007/s11222-022-10172-5. 7

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) , 2023. 7

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. 7

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. 2

K. W. K. Wong, M. Gabrié, and D. Foreman-Mackey. flowMC: Normalizing flow enhanced sampling package for probabilistic inference in JAX. Journal of Open Source Software , 8(83): 5021, 2023. 18

H. Wu, J. Köhler, and F. Noé. Stochastic normalizing flows. Advances in Neural Information Processing Systems , 33:5933-5944, 2020. 20

B. J. Zhang, Y. M. Marzouk, and K. Spiliopoulos. Transport map unadjusted Langevin algorithms. arXiv preprint arXiv:2302.07227 , 2023. 7

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. 7

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) , 2023. 7

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. 7

A Flow-Informed Markov Chain Monte Carlo Methods

Algorithm 2 Flow-informed Random Walk Metropolis Hastings

1: Input: initial x , target π , vector field v θ t , flow parameters θ

2: Output:

x

′

3: σ opt ← 2 . 38 / √ d

9: With probability α make x ′ ← y 1 else x ′ ← x

Algorithm 3 Flow-informed Independent Metropolis Hastings

1: Input: initial x , target density π , vector field v θ t , reference density p 0 , flow parameters θ .

2: Output:

x

′

3: ϕ 1 ( u ) = u t =1 ← x

5: ϕ 0 ( x ) = x t =0 ∼ p 0

8: With probability α make x ′ ← x 1 else x ′ ← x

Algorithm 4 Flow-informed Conditional Importance Sampling

1: Input: initial x , target density π , vector field v θ t , reference density q 0 , flow parameters θ , number of importance samples K .

2: Output: x

′

3: ϕ 1 ( u ) = u t =1 ← x

7: for k = 1 : K do

12: end for

13: Choose k ′ with probability P ( k ′ = k ) ∝ w k , then make x ′ ← x ( k ′ )

B Proof of Proposition 3.1

Our proof follows closely the proof of [54, Proposition 1]. Let θ ∗ be a minimizer of the CFM objective in (4), which we recall is given by

where we have defined j ( θ, x 1 ) = E t ∼U (0 , 1) E x ∼ p t ( ·| x 1 ) [ || v θ t ( x ) -v t ( x | x 1 ) || 2 ] . Now, consider the ordinary differential equation (ODE) given by

We say that ˆ θ is stability point of this ODE if, given the initial condition θ (0) = ˆ θ , the ODE admits the unique solution θ ( t ) = ˆ θ for all t ≥ 0 . Naturally, the minimizer θ ∗ is a stability point of this ODE, since ∇ θ J ( θ ; π ) | θ = θ ∗ = 0 . Meanwhile, we call Θ the domain of attraction of θ ∗ if, given the initial condition θ (0) ∈ Θ , the solution θ ( t ) ∈ Θ for all t ≥ 0 , and θ ( t ) converges to θ ∗ as t →∞ .

Let x k ∈ R d x , and X ⊆ R d x be an open subset of R d x . Let Θ ⊆ R d θ be an open set in R d θ , and Θ c ⊆ Θ be a compact subset of Θ . Consider the Markov transition kernel M := P ◦ Q k given by a cycle of k Q repeated transitions of a MALA transition kernel and a flow-informed RWMH transition kernel, viz

This transition kernel is π -invariant since both P and Q are π -invariant. In addition, let M k π,θ ( x, d y ) be the repeated application of this Markov transition kernel, namely,

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

Assumption B.1 (Robbins-Monro Condition) . The step size sequence ( ε k ) ∞ k =1 satisfies the following requirements:

Assumption B.2 (Integrability) . There exists a constant C 1 > 0 such that for, any θ ∈ Θ , x ∈ X and k ≥ 1 ,

Assumption B.3 (Convergence of the Markov Chain) . For each θ ∈ Θ , it holds that

1

1 +

k

→∞

sup

x

∈X

|

x

|

q

k

π,θ

Assumption B.4 (Continuity in θ ) . There exists a constant C 2 such that for all θ, θ ′ ∈ Θ c ,

Assumption B.5 (Continuity in x ) . There exists a constant C 3 such that for all x 1 , x 2 ∈ X ,

Assumption B.6 (Conditions on the Objective Function) . For any compact subset Θ c ⊂ Θ , there exist positive constants p, K 1 , K 2 , K 3 and v > 1 / 2 such that for all θ, θ ′ ∈ Θ c and x, x 1 , x 2 ∈ X ,

With the above assumptions, the result follows from Theorem 1 of [30] by setting x → x 1 , Π θ → M π,θ , H ( θ, x ) →∇ θ j ( θ, x 1 ) .

∫

q

lim

(1 +

|

y

|

)

|

M

(

x,

d

y

)

-

π

(d

y

)

|

= 0

.

(26)

C Additional Experimental Details

Code for the numerical experiments is written in Python with array computations handled by JAX [11]. The implementation of relevant methods for comparison is sourced from open source repositories: DDS using franciscovargas/denoising_diffusion_samplers, NF-MCMC using kazewong/flowMC [80], and FAB using lollcat/fab-jax [52]. 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 π and ν be two probability measures. Let F denote the unit ball in a reproducing kernel Hilbert space (RKHS) H , associated with the positive definite kernel k : R d × R d → R . Then the maximum mean discrepancy (MMD) between π and ν is defined as [29, Section 2.2]

where m π is the mean embedding of π , defined via E π [ f ] = ⟨ f, m π ⟩ H for all f ∈ H . Using standard properties of the RKHS, the squared MMD can be written as [29, Lemma 6]

Thus, given samples ( x i ) i m =1 ∼ π and ( y i ) i m =1 ∼ ν , an unbiased estimate of the squared MMD can be computed as

̸

̸

For a kernel k , the kernel Stein discrepancy (KSD) between π and ν is defined as the MMD between π and ν , using the Stein kernel k π associated with k , which is defined as

and satisfies the Stein identity E π [ k π ( x, · )] = 0 . We thus have that

We can obtain estimates of the KSD by using U-statistics or V-statistics. In particular, an unbiased estimate of KSD 2 k ( π, ν ) is given by the U-statistic [42]

̸

Alternatively, we can estimate KSD 2 k ( π, ν ) using a biased (but non-negative) V-statistic of the form [46, Section 4]

In all of our numerical experiments, we calculate the U- and V- statistics using the inverse multiquadratic kernel k ( x, x ′ ) = (1 + ( x -x ′ ) T ( x -x ′ )) β due to its favourable convergence properties [26, Theorem 8], setting β = -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 · 10 3 iterations for MFM and AT-SMC. In the main text, we present results for K = 5 · 10 3 learning iterations for MFM and AT-SMC, 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) [27, 36] 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.

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 and AT-SMC for K = 5 × 10 3 learning iterations. In the main text, we present results for K = 5 · 10 3 learning iterations for MFM and AT-SMC and K = 10 3 iterations for all other algorithms, which yields a more comparable total computation time.

C.4 Many Well

We also present a synthetic problem approximating the 32-dimensional Many Well distribution given by the product of 16 copies of the 2-dimensional Double Well distribution [e.g., 52, 58, 81],

where each copy of the Double Well is evaluated on a different pair of the 32 inputs. The 32dimensional Many Well has 2 16 = 65536 modes, one for each possible choice of mode in each of the 16 copies of the double well. We can obtain exact samples from the Many Well by sampling each independent copy of the Double Well.

Like previous synthetic examples, 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 . 1 , and methods with splines use 4 coupling layers with 8 bins and a range limited to [ -16 , 16] . Table 7 presents results for K = 10 3 iterations. In Table 8, we provide results for MFM and AT-SMC for K = 5 × 10 3 learning iterations. In the main text, we present results for K = 5 · 10 3 learning iterations for MFM and AT-SMC and K = 10 3 iterations for all other algorithms, which yields a more comparable total computation time.

C.5 Field system

The stochastic Allen-Cahn equation is defined in terms of a random field ϕ : [0 , 1] → R satisfying the following stochastic partial differential equation [e.g., 24, Section V]:

where a > 0 is a parameter, β is the inverse temperature, s ∈ [0 , 1] denotes the spatial variable, and η is spatiotemporal white noise. We impose Dirichlet boundary conditions throughout, so that ϕ ( s = 0) = ϕ ( s = 1) = 0 .

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

At low temperatures, this coupling induces alignment of the field in either the positive or negative direction, leading to two global minima, ϕ + and ϕ -, with typical values of ± 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 9. 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 [27, 36] 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.

C.6 Log-Gaussian Cox process

The original 10 × 10 square meter plot is standardized to the unit square. We discretize the unit square [0 , 1] 2 into a M = 40 × 40 regular grid. The latent intensity process Λ = { Λ m } m ∈ M is specified as Λ m = exp( X m ) , where X = { X m } m ∈ M is a Gaussian process with a constant mean µ 0 ∈ R and exponential covariance function Σ 0 ( m,n ) = σ 2 exp( -| m -n | / (40 β )) for m,n ∈ M ,

i.e. X ∼ N ( µ 0 1 d , Σ 0 ) for 1 d = [1 , . . . , 1] T ∈ R d with dimension d = 1600 . The chosen parameter values are σ 2 = 1 . 91 , β = 1 / 33 , and µ 0 = log(126) -σ 2 / 2 , corresponding to the values estimated in [53]. The number of points in each grid cell Y = { Y m } m ∈ M ∈ N 40 × 40 are modelled as conditionally independent and Poisson distributed with means a Λ m ,

where a = 1 / 40 2 represents the area of each grid cell. 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 10. 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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims made in the abstract and the introduction provide an accurate reflection of the theoretical, methodological, and experiment contributions in the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The final section includes a discussion of the limitations of this work, including the computational efficiency of the proposed method, situations in which the proposed method may be outperformed by other baselines, and hyperparameters which may strongly influence the performance of the proposed method.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Our theoretical result (Proposition 3.1 in Section 3) is accompanied by a full set of assumptions and a complete proof (in Appendix B). Any results upon which the proof relies are properly cited. In addition, all of the theorems, formulas, and proofs in the paper are numbered and cross-referenced.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The paper includes full details of the algorithm used to obtain the experimental results in the paper, including all of the hyper-parameter settings for each of the different numerical experiments. Our code is also made publicly available on GitHub (see Section 5).

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Our code is made publicly available on GitHub (see Section 5).

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide all of the training and test details, including hyperparameter settings, train-test splits, choice of optimizer, etc., necessary to understand and reproduce our experimental results.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We report error bars for all of our numerical experiments (see Section 5). In the text, we provide full details regarding how these error bars are calculated, as well as the factors which vary between each experimental run.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We include full details of the type of compute workers used in the main text (see Appendix C). We also include details of the running time for each of the main numerical experiments, for both our proposed method and the relevant baselines.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conforms with NeurIPS Code of Ethics in every respect, including with regards to Potential Harms Caused by the Research Process , Societal Impact and Potential Harmful Consequences , and Impact Mitigation Measures . All of the authors have read and reviewed the NeurIPS Code of Ethics to ensure that the research conducted in this paper conforms with the guidelines.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper contains foundational research on a new method for approximate statistical inference. It is not tied to any particular applications and, as such, these are not directly discussed in the main text.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper does not include any datasets or models associated with a high risk of misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The creators and original owners of assets used in the paper (e.g., code for relevant baselines) are properly credited (see Section 5). We both cite the original papers that produced these assets, and include URLs to the code.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release any new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing experiments or any research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing experiments or any research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe proposed method falls in the wider category of models that integrate MCMC methods with normalizing flows. Specifically, the method expands on the neutra-MCMC model, where the key idea is to run MCMC in a simpler reference space rather than in the (more complicated) target one. The main difference is that previous work uses discrete normalizing flows while the proposed approaches employs the Flow Matching framework and hence relies on a continuous velocity field, which is provided with a continuous Normalizing Flow.\n\nSTRENGTHS:\n- the authors propose an interesting way to use the Flow Matching framework in the context of Flow-enhanced MCMC samplers\n- I also found interesting the use of a sample-based annealing scheme (which is however not novel). The authors claim it works particularly well for multi-modal distributions, which is definitely a very relevant setting\n\nWEAKNESSES:\n- the novelty aspect of the proposed method seem limited to extend current formalism of MCMC - Normalizing Flows (NFs) methods (specifically neutra-MCMC) to the specific architecture of Continuous Normalizing Flows (CNF). However, I don't see a priori why the (continuous) CNF should be preferable discrete NFs (which are even more expressive than CNF [1]) if not because the flow matching objective requires access to the velocity field provided by CNF\n\n[1] Perugachi-Diaz et al., Invertible DenseNets with Concatenated LipSwish, NeurIPS 2021\n\nQUESTIONS:\n1. Some work analyzes the performance of methods combining MCMC with NFs methods in terms of (i) how easy is it to sample from multi modal distributions and (ii) scalability with dimensions [1]. They found that neutra-MCMC (similarly to the proposed method) is efficient in multi-modal settings but struggles more in high dimensions, compared to methods that use Flows as proposal for MH. Did you observe a consistent behaviour with their analysis? For instance I would have expected NF-MCMC to perform better in the Log-Gaussian Cox experiment. Or did the proposed method outperform \"flow-MCMC\" methods also when inflating the dimensions even more? (e.g. in the experiments in Section 5.3 and 5.4)\n2. The use of CNF are motivated by the availability of the velocity field, which is required by the Flow Matching objective. However, other neutra-MCMC methods work very well also with discrete NFs. CMF are definitely more expressive than standard discrete flows that are commonly used in the literature (mainly real NVP), but one could also use more powerful discrete NFs like [2]. I was wondering whether the authors could discuss the benefits of the proposed approach compared to neutra-MCMC with discrete NFs.\n3. The proposed method uses a hybrid between local and global updates, similarly to [3] (which is cited in the paper). The main difference it the use of a CNF (for the Flow Matching objective) instead of NFs. So I think it would be very relevant to compare the proposed method against [3]. Also the authors of [3] use very naive layers (real NVP, which consist of basic shift and scale transforms), while more flexible layers like [2] might exclude that the difference in performance is to be connected to the expressiveness of the flows. Similarly, the NF-MCMC uses very simple NFs layers (always real NVP).\n4. Even though the Flow Matching objective could allow to learn complicated distributions, in the end the authors use a very simple velocity field to transform the conditional distribution. Doesn't this limit the expressiveness of the learnt distribution?\n5. In line 218 do the authors mean FAB instead of DDS?\n\n\n[1] Grenioux et al, On Sampling with Approximate Transport Maps, ICML 2023\n\n[2] Perugachi-Diaz et al., Invertible DenseNets with Concatenated LipSwish, NeurIPS 2021\n\n[3] Samsonov, Local-Global MCMC kernels: the best of both worlds, NeurIPS 2022\n\nLIMITATIONS:\nLimitations are clearly stated.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nA method called Markovian flow matching (MFM) for training neural ODEs (continuous normalizing flows) by flow matching to sample from distributions given as unnormalized density functions is proposed. The method obtains samples on which to perform FM training by an MCMC procedure that alternates two kinds of kernels: Metropolis-adjusted Langevin in the target space (as has been done in past work) and Gaussian proposals in the latent space, i.e., the pullback of the target by the learned ODE. An annealing scheme is additionally used. MFM is evaluated on toy low-dimensional tasks, a synthetic density from statistical physics, and a high-dimensional data-derived density (log-Gaussian Cox process).\n\nSTRENGTHS:\n- The idea to do MCMC in the latent space of a CNF by pullback along the flow map is original, as far as I know.\n  - There are two general families methods for training continuous-time models to sample Boltzmann densities: using maximum likelihood on (possibly reweighted) approximate samples (e.g., FAB and the proposed MFM) and using distribution-matching objectives (e.g., DDS, PIS), as well as hybrid approaches.\n  - It's interesting to see progress on making methods of the first kind efficient, and to see them combined with CNFs. In the past, e.g., in [Tong et al. \"Improving and generalizing flow-based generative models\", D.2], FM has been used with MCMC, but only in the target space.\n- No serious problems with the writing, including the global structure of the paper and the exposition of preliminaries.\n\nWEAKNESSES:\n- Some math bugs (imprecise exposition):\n  - L72: $v_t$ being **any** time-dependent vector field does not imply existence of a diffeomorphic flow map (even if it is continuous in $t$, if that is how we interpret \"runs continuously in the unit interval\"). One also needs integrability conditions.\n  - L75: $p_0$ has to be strictly positive for the path to take values in $\\mathbb{R}^+$ and for the discussion the discussion that follows to make sense.\n- Comparison with SDE models: Approximate (MCMC) samples can also be used to train diffusion models (neural SDEs):\n  - First, this could be done either by minimising a variational bound on the log-likelihood, similar to what is done here for CNFs -- is it possible to do such a comparison? At present the SDE baseline (DDS) is not using MCMC, only \"on-policy\" forward exploration.\n  - Second, methods such as that in [Sendera et al. \"Improved off-policy training of diffusion samplers\"] use MCMC on the target density to obtain samples for training with a distribution-matching, not maximum-likelihood, objective.\n- The main weakness is that the experiments are not very comprehensive or convincing:\n  - The 2D GMMs are toy problems; all methods, including DDS, should find the modes when appropriately tuned on such problems, and so it is hard to draw any conclusions from the results.\n  - On the field system, MFM, although much faster than FAB, is performing **much** worse in the metrics used.\n  - On LGCP: \n    - The use of this task for benchmarking samplers of Boltzmann densities is questionable in the first place, since there is no known ground truth (even the exact normalising constant is diagreed upon) -- see, for example, the aforementioned [Sendera et al., B.1] for a discussion of inconsistent or irreproducible evaluation in past work on this dataset.\n    - Why use the nonstandard metric (KSD)? Most past work has used estimates of the log-partition function, which here is reported in the appendix.\n    - On that more standard metric, MFM is significantly underperforming (Table 8).\n- On both the field system and LGCP, how are we to know if the proposed method is useful on this problem without a point of comparison? I would strongly suggest a long-run HMC, MALA, and SMC (for example, run for the same wall time as the shortest and longest of the ML methods evaluated) as baselines.\n- I would also strongly suggest evaluations on a more comprehensive set of benchmarks that are common in past work on samplers of Boltzmann densities (as in the FAB and DDS papers, among others).\n\nQUESTIONS:\nMinor:\n- L25: You probably want to specify \"a new sample $y$\" and in the subsequent lines clarify that the \"mild conditions\" are also needed for the target to be the *only* stationary distribution.\n- L77: Dash between \"change\" and \"of\" instead of hyphen.\n- L86: Wrong placement of parentheses.\n- L246: I do not believe [78] introduced this parametrisation -- it was in [80] (and possibly in earlier work as well).\n- I don't understand equation (18), which involves $x_{65}$ in the first sum although the dimensionality is 64. Is there is mistake, or should the indices be interpreted modulo $d$?\n- L308: Typo in \"comparable\".\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a novel method for incorporating flow matching with MCMC; entailing constructing a Markov kernel as a mixture of regular MCMC step and flow step (from data space (\\pi_1) to prior (\\pi_0, typically Gaussian) and back with some added noise. The flow guarantees a likelihood and can hence be used for accept/ reject corrections. \n\nIn this sampling setting, one does not have access to samples from the data and hence the flow is trained from MCMC samples jointly within sampling.\n\nSTRENGTHS:\nI think the authors did a great job on this paper:\n- The method is simple and well explained\n- Experiments are thorough with appropriate baselines including recent DDS, good performance and code is available\n- The authors clearly articulate limitations, which I appreciate\n- Prior and related work is discussed and not hidden\n\nWEAKNESSES:\nI worry how long it takes the flow network to converge. Simulating the flow back and forward to generate samples at regular frequency during training could be quite slow if involving a large number of network evaluations for the flow. And if the MCMC locally does not mix then the flow will not see many samples. I understand tempering may help resolve this but at what cost?\n\nSimilarly, I imagine the samples from the flow could be quite far from the target distribution if the flow was not pre-trained or initialized to be close to some nearby distribution.\n\nThe idea itself is quite straightforward. This is not necessarily a weakness in itself, I am surprised something similar has not been considered which am unaware of. I understand there are many similar schemes that are cited using non-local Markov updates with normalizing flows (discrete), just not using flow matching (e.g. https://arxiv.org/abs/2105.12603) and other works from the same authors. I appreciate the authors have discussed and cited these. This however limits the novelty of the methodological contribution slightly in my opinion.\n\n\nMinor presentation:\n- I struggle to understand from Figure  what the ground truth samples should look like. I assume like the samples from the author's method. It would help readers to show the density to know what the ground truth should be, similar to the other experiments.\n- Similarly Table 2, show in bold the top performing items to help readers\n\nSome other relevant works using diffusion / flow for sampling:\n- Target score matching, Bortoli et al 2024 https://arxiv.org/abs/2402.08667\n- Iterated Denoising Energy Matching for Sampling from Boltzmann Densities, Akhound-Sadegh et al 2024\nand although cited the work of Phillips et al 2024 (already cited) could be used as a baseline given code is available (also in jax)\n- CRAFT and AFT  (Annealed flow transport Monte Carlo, 2021) may also be good baselines to include, code is available in jax for those too.\n\nQUESTIONS:\nMy understanding is that the flow update pushes samples into a simpler prior / latent type space then back to the data space. Is it possible to include Markov transitions in the prior space i.e. in the simpler (Gaussian) distribution, which are easier and encourage non local updates (in data space)? The addition of Gaussian noise is a bit like this I suppose but has anything else been considered?\n\n\nDoes this actually require flow matching? I feel any diffusion model would also have be usable here, it can be trained in the same way and one can use the probability flow ODE to have tractable likelihoods and if needed similar to flow matching. There are many similarities between the two models, especially for Gaussian marginal, but diffusions seem to be better performing generative models still, see e.g. Lipman 2022 (cited) vs Karras 2022 (https://arxiv.org/abs/2206.00364) in terms of generative performance. If just using Langevin dynamics without accept/reject correction then the added stochasticity of diffusion samplers may also help, plus the benefit of Langevin dynamics corrector steps along the diffusion paths could possibly be used (Song 2021).\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to use continuous normalizing flows (CNFs) to define the proposal distribution in a MCMC framework. While the use of flow models for MCMC proposals is not entirely new, this paper introduces an interesting training procedure that iteratively updates the learned CNF model while performing MCMC. Existing works either (i) make use of an importance weighted training objective, or (ii) first perform long-run MCMC to obtain high-quality samples. Instead, this work makes use of the Flow Matching training objective to fit to its current MCMC samples, where samples are obtained from a MCMC involving the CNF itself.\n\nSTRENGTHS:\n- Overall, I feel this paper is a combination of good solid ideas that individually are not entirely new, but packaged in a way that makes sense.\n- The writing is also clear and to the point.\n- Experiments seem to be carried out on some standard benchmark potential functions.\n- The paper repeatedly mentions that the proposed method is favorable to competing neural MCMC approaches based on wallclock time, being significantly faster in some cases.\n\nWEAKNESSES:\n- I feel the set of experiments is \"too standard\". Given that Bayesian inference is such a long field, it would be make a paper much stronger if it can show meaningful improvement on real unsolved problems where exploration is key.\n- There is a need to add ablation experiments that provide better intuition regarding what benefit the flow-informed Markov transition is doing.\n\nQUESTIONS:\nGenerally, while the proposed MCMC algorithm as a whole makes sense, it is composed of multiple components. In order to better understand each proposed component, it would be good to have ablation experiments analyzing the behavior of each component.\n\n- Given that there is a series of local MCMC updates, it would be good to compare against MCMC with just the Q transition kernel. This can help answer what exactly is the change in improvement from including a flow-informed transition kernel.\n\n- I think what would be really interesting is a plot of the acceptance rate between the P (flow-informed) and Q (local) kernels. I would imagine that at the beginning P has poor acceptance rates because the CNF is poorly trained, while after the MCMC chain is run for long enough times, its acceptance rates are much higher. Meanwhile, the acceptance rate of Q should be constant but perhaps lower than the acceptance rate of P? Regardless, it would be interesting to see such a plot to better understand the behavior of the proposed algorithm.\n\n- If you turn off the annealing, how much worse does the approach perform?\n\n### Minor comments\n- The likelihood notation L(D|x)  when discussing annealing (Eq 16) is not used anywhere else (and also not used in the experiments since they directly define pi(x). I felt this to be confusing because as it is currently written, it may seem that the annealing trick is restricted to only Bayesian inference settings. I would imagine the annealing simply takes L(x) = pi(x) / pi_0(x) for some pi_0(x). It might be good to clarify this?\n- There are 3 instances of \"NeutraMCMC\" which should probably be \"Neural MCMC\" ?\n\nLIMITATIONS:\nThe paper is written clearly with objective analyses, with no excessive beautification of the proposed method. It's nice.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-17"}