Importance Corrected Neural JKO Sampling
Authors: Johannes Hertrich
Robert Gruhlke
Venue: ICML 2025 poster (2025)
Keywords: Sampling
Wasserstein Gradient Flows
Normalizing Flows
Rejection Sampling
Primary area: probabilistic_methods->monte_carlo_and_sampling_methods
TL;DR: We accurately sample from unnormalized densities by alternating local flow steps with non-local rejection steps.

Abstract: In order to sample from an unnormalized probability density function, we propose to combine continuous normalizing flows (CNFs) with rejection-resampling steps based on importance weights. We relate the iterative training of CNFs with regularized velocity fields to a JKO scheme and prove convergence of the involved velocity fields to the velocity field of the Wasserstein gradient flow (WGF). The alternation of local flow steps and non-local rejection-resampling steps allows to overcome local minima or slow convergence of the WGF for multimodal distributions. Since the proposal of the rejection step is generated by the model itself, they do not suffer from common drawbacks of classical rejection schemes. The arising model can be trained iteratively, reduces the reverse Kullback-Leibler (KL) loss function in each step, allows to generate iid samples and moreover allows for evaluations of the generated underlying density. Numerical examples show that our method yields accurate results on various test distributions including high-dimensional multimodal targets and outperforms the state of the art in almost all cases significantly.

Importance Corrected Neural JKO Sampling
Johannes Hertrich 12
Robert Gruhlke
Abstract
the underlying stochastic differential equation follows the Wasserstein-2 gradient flow of the reverse KL. divergence F(u) = KL v/). Over the last years, generative models like normalizing flows (Rezende & Mohamed, 2015) or dif- fusion models (Ho et al., 2020; Song et al., 2021) became more popular for sampling. see, e-g., (Phillips al., 2024; Vargas al., 2023a). Also these methods are based on the reverse KL divergence as loss function. While genera- tive models have successfully been applied in data-driven setups, their application to the problem of sampling from arbitrary unnormalized densities is not straightforward. This difficulty arises from the significantly harder nature of the problem, even in moderate dimensions, particularly when dealing with target distributions that exhibit phenomena such as concentration effects, multimodalities, heavy tails, or other issues related to the curse of dimensionality.
In order sample from unnormalized probabil- ity density function, propose to combine con- tinuous normalizing flows (CNFs) with rejection- resampling steps based on importance weights. We relate the iterative training of CNFs with reg- ularized velocity fields to a JKO scheme and prove convergence of the involved velocity fields to the velocity field of the Wasserstein gradient flow (WGF). The alternation of local flow steps and non-local rejection-resampling steps allows to overcome local minima or slow convergence of the WGF for multimodal distributions. Since the proposal of the rejection step is generated by the model itself, they do not suffer from common drawbacks of classical rejection schemes. The arising model can be trained iteratively, reduces the reverse Kullback-Leibler (KL) loss function in each step, allows to generate iid samples and moreover allows for evaluations of the generated underlying density. Numerical examples show that our method yields accurate results on various test distributions including nigh-dimensional mul- timodal targets and outperforms the state of the art in almost all cases significantly.
In particular, the reverse KL is non-convex in the Wasser- stein space as soon as the target density v not log-coneave which is for example the case when consists of multi- ple modes. In this case generative models often collapse to one small number of modes. We observe that for continuous normalizing flows (CNFs, Chen et al., 2018 8; Grathwohl et al., 2019) this can be prevented by regular- izing the norm of the velocity field as proposed under the name OT-flow by (Onken et al., 2021). In particular, this regularization converts the objective functional into a convex one. However, the minimizer of the regularized loss function is no longer given by the target measure but by the Wasserstein proximal mapping of the objective func- tion applied onto the latent distribution. Considering that the
1. Introduction
We consider the problem of sampling from an unnormal- ized probability density function. That is, we are given an integrable function g: Rd + R> and we aim to generate samples from the probability distribution given by the density q(x) = 22) where the normalizing constant 7,y Jg+ (1:)dæ: is unknown. Many classical sampling methods are based on Markov chain Monte Carlo (MCMC) methods like the overdamped Langevin sampling, see, e.g., (Welling& & Teh, 2011). The generated probability path of
(JKO) scheme (Jordan et al., 1998) iteratively applies the Wasserstein proximal mapping and converges to the Wasserstein gradient flow, several pa- pers proposed approximate the steps of the scheme by generative models, see (Altekrüger et al., 2023; Alvarez- Melis et al., 2022; Fan et al., 2022; Lambert et al., 2022; Mokrov et al., 2021; Vidal et al., 2023; Xu et al., 2024). We will refer to this class of method by the name neural JKO. Even though this approximates the same gradient flow of the Langevin dynamics these approaches have usually the advantage of faster inference (once they are trained) and additional possibly allow for density evaluations.
'Equal contribution Unnivessité Paris Dauphine PSL Uniiersiity College London 'FU Berlin. Correspondence to: Johannes Hertrich jjohannes.hertriche dawichhh@uuupphhpps.eeeeu Robert Grublke <r.ruulke @ fu- du-berllin.de.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
However, already the time-continuous Wasserstein gradi- ent flow suffers from the non-convexity of the reverse KL

---

Importance Corrected Neural JKO Sampling
loss function by getting stuck in local minima or by very slow convergence times. In particular, it is well-known that Langevin-basel sampling methods often do not distribute the mass correctly onto multimodal target distributions. As a remedy, (Neal, 2001) proposed importance sampling, i.e., to reweight the sample based on quotients of the target dis- tribution and its current approximation which leads to an unbiased estimator of the Monte Carlo integral. However, this estimator may lead to highly imbalanced weights be- tween these samples and the resulting estimator might have a large variance. Moreover, the density of the current ap- proximation has to be known up to a possible multiplicative constant, which is not the case for many MCMC methods like Langevin sampling or Hamiltonian Monte Carlo. In (Del Moral al., 2006) the authors propose a scheme for al- ternating importance sampling steps with local Monte Carlo steps. However, the corresponding importance sampling step generates non-iid samples, that de-correlate over time by construction.
malized densities are often based on Markov Chain Monte
Carlo (MCMC) methods, see e.g. (Gilks et al., 1995). In particular first order based variants, such as the Hamilto- nian Monte Carlo (HMC) (Betancourt, 2017; Hoffman & Gelman, 2014) and the Metropolis Adjusted Langevin Algo- rithm (MALA Girolami & Calderhead, 201 I; Rossky et al., 1978; Roberts & Tweedie, 1996) are heavily used in practice, see (Andrieu et al., 2003) for an overview. The viewpoint of these samplers sample space description of gradient flows defined in a metricized probability space then allows for ex- tensions such as interacting particle systems (Chen et al., 2023; Eigel et al., 2024; Garbuno- Inigo et al., 2020; Wang & Li, 2022). However, since these algorithms are based on local transformations of the samples, they are unable to dis- tribute the mass correctly among different modes, which can partially be corrected by importance sampling (Neal, 2001) and sequential Monte Carlo samplers (SMC) (Del Moral et al., 2006) as described above. In contrast to our model, SMC approximates the density of the approximation by as- signing "inverse Markov kernels" to certain MCMC kemels, which might lead to propagating errors. Furthermore, the generation of additional samples requires to rerun the whole procedure which can be very costly. Other approaches ap- proximate the target density by combination of transport maps and low-rank models such as tensor trains (Figel et al., 2022) yielding efficient access to posterior statistics.
Contributions In this paper, we propose a sampling method which combines neural JKO steps based on CNFs with importance based rejection steps. While the CNFs adjust the position of the generated samples locally, the rejection steps readjust the inferred distribution non-locally based on the quotient of generated and target distribution. Then, in each rejection step we resample the rejected points based on the current constructed generative model. We illustrate this procedure in Figure 1.
In the last years, generative models became very popular, including VAES (Kingma & Welling, 2014), normalizing flows (Rezende & Mohamed, 2015), diffusion models (Ho et al., 2020) flow-matching (Lipman al., 2022) which is also known rectified flow (Liu et al., 2022). In con- trast to our setting, they initially consider the modeling task, i.e., they assume that they are given samples from the target measure instead of an unnormalized density. How- ever, there are several papers, which adapt these algorithms for the sampling task. For normalizing flows, this mostly amounts to changing the loss function (Hoffman et al., 2019; Marzouk et al., 2016; Qiu & Wang, 2024). Very recently, there appeared also flow-matching variant for the sampling task (Woo & Ahn, 2024) based on (Akhound- -Saheh et al., 2024). For diffusion (and stochastic control) models this was done based on variational approaches (Blessing et al., 2024; Phillips et al., 2024; Vargas et al., 2023a;b; Zhang & Chen, 2021) or by computing the score by solving a PDE (Richter & Berner, 2024; Sommer et al., 2024). These methods usually provide much faster sampling times than MCMC methods and are often used in combination with some conditioning parameter for inverse problems, where a (generative) prior is combined with a known likelihood term (Ardizzone et al., 2019; Altekrüger & Hertrich, 2023; Andrle et al., 2021: Denker et al., 2024). Combinations of generative models with stochastic sampling steps were con- sidered in the literature for generative modeling under the
Our method generates independent samples and allows to evaluate the density of the generated distribution. In our numerical examples, we apply our method to common test distributions up to the dimension d = 1600 which are par- tially highly multimodal. We show that our importance corrected neural JKO sampling (neural JKO IC) achieves significantly better results than the comparisons'.
From theoretical side, we prove that the velocity fields from sequence of neural JKO steps strongly converge to the velocity field of the corresponding Wasserstein gradi- ent flow and that the reverse KL loss function decreases throughout the importance-bant based rejection steps.
Outline In Section 2, we recall the fundamental concepts which will be required. Afterwards, we consider neural JKO schemes more detail in Section 3. We introduce our importance-basee rejection steps in Section 4. Finally, we evaluate our model numerically and compare it to existing methods in Section 5. Conclusions are drawn in Section 6. Additionally, proofs, further numerical and technical details are presented in Appendix A- E.
Related Work Common methods for sampling of unnor-
'The code is available at https:///ithub.... pphertrich/neure __o__ii
2

---

Importance Corrected Neural JKO Sampling
Lasent distribution
After Step
After Stop 2
After Step
PNI
CNB
-2
After Step
After Step
Siep
Approsimation
onjectoon
maarping
soarming
-1
-2
-2
-2
Figure I: Iterative application of neural JKO steps and rejection steps for a shifted mixture target distribution. The orange samples are rejected in the next following rejection step and the blue samples are the resampled points. The latter approach enables for the correction of wrong mode weights introduced by the underlying WGF. See Figure 3 for more steps.
name stochastic normalizing flows (Hagemann et al., 2023; 2022; Noé al., 2019; Wu et al., 2020) and for sampling under the name annealed flow transport Monte Carlo (Arbel al., 2021; Matthews al., 2022). (Gabrié al., 2022) use normalizing flows to Iearn proposal distributions in Metropolis- Hastings algorithm. These generative models can be adapted to follow a Wasserstein gradient flow, by mimicking a JKO scheme with generative models or directly following the velocity field a kernel-based functional. Such approaches were proposed for generative modeling (Fan al., 2022; Hagemann et al., 2024; Hertrich et al., 2024; Liutkus et al., 2019; Vidal et al., 2023; Xu et al., 2024), sampling (Fan et al., 2022: Lambert et al., 2022; Liu & Wang, 2016a; Mokrov et al., 2021) or other tasks (Altekrüger et al., 2023; Arbel et al., 2019; Alvarez-Melis et al., 2022). Gradient flows the reverse KL divergence with respect different metric were considered in (Liu & Wang, 2016b) under the name Stein variational gradient de- scent. Moreover, Lu et al. (2019) study gradient flows in the Wasserstein-Fisher Rao metric which can be implemented via birth-death processes.
2.1. Curves in Wasserstein Spaces
Wasserstein Distance Let P(RR) be the space prob- ability measures on Rd and denote by P2(Rd) == {o P(R Sa || [lpx) < 2o} the subspace of proba- bility measures with finite second moment. Let P be the subspace of absolutely continuous measures from (2(R)). Moreover, we denote for μ. (IR by I(p, v) == { P2 Rd): P1 μt P2 v) the set of all transport plans with marginals μ and where P: Rd x Rd + Rd defined by (2( = I, is the pro- jection onto the i-th component for i = 1, 2. Then, we equip (R with the Wasserstein-2 metric defined by
(2( = inf
xEr(ar)) JRIXRd || y
If p (RR), the above problem has always unique minimizer.
Absolutely Continuous Curves A curve y: 1 + P2 (IRd) on the interval / S R is called absolutely continuous there exists a Borel velocity field v: Rd / + Rd with J [Iv (alee))lllaaa))))))) tion
After the first version of our paper appeared, several re- lated preprints were released, including (Albergo& & Vanden- Eijnden, 2024; Chen al., 2025; He al., 2025; Wu & Xie, 2024).
< oo such that the continuity equa-
(t)(1) + V (u(-t))(t)) =
is fulfilledd on I x Ed in a weak sense. Then, any velocity field solving the continuity equation (1) for fixed y char- acterizes 7 as yt) 2(,, ((to)) where 2 is the solution of the ODE á(x,, () v(=(2, (), () with =(x, 40) x and to E I. It can be shown that for an absolutely continuous curve there exists a unique solution of minimal norm which is equivalently characterized by the co-called regular tangent space T2 see Appendix for details. An abso-
2. Preliminaries
In this section, we provide rough overview of the required concepts for this paper. To this end, we first revisit the basic definitions of Wasserstein gradient flows, e.g., based on (Ambrosio al., 2005). Afterwards we recall continuous normalizing flows with regrguulaiizztinns..
3

---

Importance Corrected Neural JKO Sampling
lutely continuous curve is geodesic if there exists some >> 0 such that W2(7(+),>(t) = c|s - t|-
Theorem 2.2 (Ambrosio al., 2005, Thm 11.2.1). Let F: P2(RR) - (-00,+00 be proper, Isc, coercive, and X-convex along generalized geodesics, and let μ² E dom F. Then the curves P/r defined via the minimizing movement scheme (4) converge + O locally uniformly to locally Lipschitz curve y (0, +00) + P2((dd) which is the unique Wasserstein gradient flow of F with (0-h)/(00++-))))))) = np.
The following theorem formulates dynamic version of the Wasserstein distance based minimal energy curves in the Wasserstein space.
Theorem 2.1 (Benamou & Brenier, 2000). Assume that H. E (R)) Then, (1,, v) is equal to
2.2. Continuous Normalizing Flows and OT-Flows
inf Ra x= (0. I- ((2.))==(((
The concept of normalizing flows first appeared in (Rezende & Mohamed, 2015). follows the basic idea to approxi- mate a probability distribution v/ by considering a simple latent distribution μo (usually a standard Gaussian) and to construct diffeomorphism Tg: I - R depending on some parameters 0 such that v ~ To In practice, the diffeomorphism can be approximated by coupling- based neural networks (Dinh et al., 2016; Kingma & Dhariwal, 2018), residual architectures (Behrmann et al., 2019; Chen et al.. 2019; Hertrich, 2023) or autoregressive flows (De Cao et al., 2020; Durkan et al., 2019; Huang et al., 2018; Pa- pamakarios et al., 2017). In this paper, mainly fo- cus on continuous normalizing flows proposed by (Chen et al., 2018; Grathwohl et al., 2019), see also (Ruthotto & Haber, 2021) for an overview. Here the diffeomor- phism To is parameterized as neural ODE. To this end let 0g: TRd R - Rd be neural network with parameters 0 and let zæ: Rdd [0, + RR for fixed v 0 be the solution of èg = ug, with initial condition zg 2x(2.)) 0) Then, we define To via the solution zg as To(x)) = =0(a, r). The density pe of #o can be described by the change-of- variables formula po(x) = (x() (dtp///le(((// (VT((z)), where po is the density of the latent distribution po. In the case of con- tinuous normalizing flows, it can be shown that the denom- inator can be computed as log(Idet(VTT) (x)( = Eo r) with 0,6o 1) trace( trace(Vvo (vor 1)) and Eo(: 0) 0. In order to train a normalizing flow, one usually uses the Kullback-Leible divergence. If v is given by a density ((x) gxx)/2g. where Z, is an unknown normalizing constant, then this amounts to the reverse KL loss function
Moreover, there exists unique minimizing velocity field to and the curve defined by (t) = =(-,t)))) #1 with E [o, 1] and i(x, () = v(z(,, 1), z(x, 0) = x is a geodesic which fulfills the continuity equation (t) V V-(((,,t))(t) = O.
Let T > O. Then, by substitution of t by t/T and rescaling v in the time variable, this is equal to
T %/o l
inf v: x01(0, $(x. ((x.))= 0(=,0)=,, -
(2)
Wasserstein Gradient Flows An absolutely continuous curve y: (0, x)) + (2 with velocity field or E (2 (2( is a Wasserstein gradient with respect to
F: ) + (-20, ooj if o E (t)),,)), for for.e.t > 0, where (F(μ) denotes the reduced Fréchet subdiffential at JL, see Appendix A for a definition.
To compute Wasserstein gradient flows numerically, we can use the generalized minimizing movements or Jordan- Kinderlehrer-Oue (JKO) scheme (Jordan et al., 1998). To this end, we consider the Wasserstein proximal mapping defined as
( = arg min (Rd)
+ ++)
Then, define as μ for A E N the steps of the minimizing movements scheme, i.e.,
(()) = KL Mo, v)
= Ezwuul - log(g(To(x))) + log(po (x)) lo(a, + C,
Ez l log(g
1 12,
1 A+1 = (rox,=(1)))
(3)
for some constant C' independent of 8.
We denote the piecewise constant interpolations 7r: [0, oo) scheme by
(RR)) of the minimizing movement
In order to stabilize and accelerate the training, (Onken et al., 2021) propose to regularize the velocity field vs by its expected squared norm. More precisely, they propose to add the regularizer R()0) = 7 fo (0 l). R) to the loss function. This leads to straight trajectories in the ODE such that adaptive solvers only require very few steps to solve them. Following Theorem 2.1. the authors of (Onken et al., 2021) note that for B > the functional
+t) 1 t E (0.
(4)
Then, the following convergence result holds true. We recall the necessary definitions of coercivity and X-convexity in Wasserstein spaces in Appendix A.
4

---

Importance Corrected Neural JKO Sampling
-(0)+ + (R(0) has the same minimizer as the functional [())+ 2W2 To), which relates the JKO scheme as pointed out by (Vidal et al., 2023).
(3). This leads n (, u1 where
(Ur, 2r,k) E
arg min
E2 (5)
3. Neural JKO Scheme
where (, v) is given by the dynamic formulation (2) as
In the following, we learn the steps (3) of the JKO scheme by neural ODEs. While similar schemes were already sug- gested in several papers (Altekrüger et al., 2023; Alvarez- Melis et al., 2022: Fan et al., 2022; Lambert et al., 2022; Mokrov et al., 2021; Vidal et al., 2023; Xu et al., 2024), we are particularly interested in the convergence properties of the corresponding velocity fields. In Subsection 3.1, we introduce the general scheme and derive its properties. Af- terwards, in Subsection 3.2, we describe the corresponding neural network approximation. Throughout this section, we consider the following assumptions the objective func- tional FF and our initialization /40-
2 L. / (). + (((())))#1))
Finally, we concatenate the velocity fields of all steps by 0r (kr, = Ur and obtain the ODE
(x = u+ (x.
2- 2x 0) I.
As a straightforward observation, we obtain that the curve defined by the velocity field v, is the geodesic interpolation between the points from JKO scheme (3). We state a proof in Appendix B.2.
Assumption 3.1. Let F: P2 + RU {oo} be proper, lower semi-continuous with respect to narrow conver- gence, coercive, X-convex along generalized geodesics and bounded from below. Moreover, assume that dom( C PH and that F has finite metric derivative (H))( < 0o the initialization Ni E P2 (R.)).
Corollary 3.2. Under Ass. 3.1 the following holds true.
(i) It holds that
W2 2 (+++))==+ S- (1.
ie., Ur is the optimal velocity field from Theorem 2.
(ii) The curve 7+()) 2( #) fulfills 7 "KT + tr)
This assumption is fulfilled for many important divergences and loss functions F. We list some examples in Ap- pendix B.1. We will later pay particular attention the reverse Kullback-Leibler divergence F(u) (L,
((1 t)l + + for L E [0, 1J, where T is the optimal transport map between μ and ph+t
(iii Vy and "ir solve the continuity equation
S p(r) log p(xx) dz, where v is a fixed target measure and
a(w)
p and y are the densities of μs and respectively. This
Oct (t) + V (u+(,,))(t)) =
functional fulfills Assumption 3.1 if log(4) is A-convex.
We stress the fact that the X-convexity property is assumed for some x E IR, which explicitly includes negative values. Consequently, the theoretical results are also applicable for target densities which are not log-concave.
Analogously to Theorem 2.2 one can show that also the curves "/r are converging locally uniformly to the unique Wasserstein gradient flow (see, e.g., the proof of Ambrosio et al., 2005, Thm 11.1.6). In the next subsection, we will approximate the velocity fields Ur,k by neural networks. In order retain the stability of the resulting scheme, the next theorem states that also the velocity fields converge strongly towards the velocity field from the Wasserstein gradient flow. The proof is given in Appendix B.3.
Additionally, (Ambrosio et al., 2005, Lem 9.2.7) states that the functional G(p) 1 W22 (v. μ²² + F(u) (X+ + !)- convex along geodesics. In partieular, for functional G is strongly convex such that expect that optimizing it with generative model is much easier than optimizing F, see also Appendix F.1 for discussion how this can prevent mode collapse.
2r
x the
Theorem 3.3. Suppose that Assumption 3. is fulfilled and ler (π)) S (0, oc) with + 0. Then, (Va)), ): converges strongly to the velocity field 0 E L2(7, Rd x [0.T]) of
the Wasserstein gradient flow ry: (0, oo) + P2(R)) of F starting in no.
3.1. Piecewise Geodesic Interpolation
In order to represent the JKO scheme by neural ODEs, we first reformulate based on Theorem 2.1. To this end, we insert the dynamic formulation the Wasserstein distance in the Wasserstein proximal mapping defining the steps in
In some cases the limit velocity field v can be stated explic- itly, even though its direct computation is intractable. For details, we refer to Appendix B.
5

---

Importance Corrected Neural JKO Sampling
3.2. Neural JKO Sampling
4. Sypportance-Basede Rejection Steps
In the following, we learn the velocity fields vr, as neural ODEs in order to sample from target measure given by the density g(x) Z, (x) with unknown normalization constant 2n = Sg (())dx. To this end, we consider the Wasserstein gradient flow with respect to the reverse KL loss function F(u) = [I, 2) which has the unique minimizer " = v. Note that similar derivations were done in (Vidal al., 2023; Xu et al., 2024) for the application of generative modeling instead of sampling.
While a large number of existing sampling methods rely on Wasserstein gradient flows with respect to some divergences, it is well known that these loss functions are non-convex. This leads to very slow convergence or only convergence to suboptimal local minima. In particular, if the target distri- bution is multimodal the modes often do not have the right mass assigned. In this section, we derive novel rejection step that corrects such imbalanced mass assignments, Re- markably, we will see in Theorem 4.2, that it is possible to access the density after these rejection steps. This observa- tion is crucial, since it allows for the iterative combination of several rejection and neural JKO steps, leading to the proposed importance corrected neural JKO sampling.
Then, due to (5) the loss function [()) from the JKO steps for the training of the velocity field ny reads as
(,'',, / | (1o0( (100) + T(
Importance Sampling As a remedy, many sampling algo- rithms from the literature are based on importance weights, see, e.g., sequential Monte Carlo samplers (Del Moral et al., 2006) or annealed importance sampling (Neal, 2001). That is, we assign to each generated sample z; weight u, = g(x,)
where 2ø is the solution of (o( t) = v0(20(1, (x. t) with (2 0) = x. Now, following the derivations of continuous normalizing flows, cf. Section 2.2, the second term of L can be rewritten (up to an additive constant) as
p(z,)' where p is some proposal density and q is the density of the target distribution v. Then, for any integrable function f: IRd R it holds that Σ (x;) is an unbiased estimator of Jr (x))d(()). Note that impor- tance sampling is very sensitive with respect to the proposal which needs to be designed carefully and problem adapted.
F(zol; +)++) x Eenut logggloo r))) Lo (2. T),,
where to solves (o(or,)) trace( (iace(Vvo(z0(a) 1)) with Co(-,01 0) = O. Moreover, we can rewrite the first term of L based on
Rejection Steps Inspired by importance sampling, we propose to use importance based rejection steps. More pre- cisely, let μ be a proposal distribution where we can sample from with density p(æ) f and denote by the target distribution with density g(x;) = g(x))/z, Zy. In the fol- lowing, we assume that we have access to the unnormalized densities f and g. but not to the normalization constants Z and Z, Then, for a random variable X ~ ft, we now generate new random variable X by the following proce- dure: First, we compute the importance based acceptance probability (XX) min {1, x } min { ,(X) ) x(x)
= Enwut (wo (2, T),,
where wa is the solution of (n( ()- = (=0 1), 1212 Hence, we can represent C up to an additive constant as
[()) = Exnntt (og(((z0((( - fa(2,, 7) + wo(x, (6)
where (zo, wo) solves the ODE system
xf( S where > 0 is positive hyperparameter and ē eZyy Then, we set X X with probability a(X) and choose X = X' otherwise, where X" 1 and X are independent.
ēo(x,t) tx(z,t) wo(xx
) -
vo(2o(a, trace c0 (Z)(, [100(20 t), t)ll
/ ZZ
(7)
Remark 4.1. This is a one-step approximation of the classi- cal rejection sampling scheme (Von Neumann, 1951), see also (Andrieu et al., 2003) for an overview. More precisely, we arrive at the classical rejection sampling scheme by choosing ē > sup, <<x)/p(x) and redo the procedure when X is rejected instead of choosing X X'.
with initial conditions (x 0) = x, fz(2. 0) = 0 and wo(x, 0) 0. In particular, the loss function L can be eval- uated and differentiated based on samples from n.a Once the parameters are optimized, we can evaluate the JKO steps in the same way as standard continuous normalizing flows. We summarize training and evaluation of the JKO steps in Algorithm 3 and 4 in Appendix D. In practice, the density values are computed and stored in log-space for numerical stability. Additionally, note that the continuous normalizing flows can be replaced by other normalizing flow architectures, see Appendix E2 and Remark F.3 for details.
Similarly to importance sampling, the rejection sampling al- gorithm is highly sensitive towards the proposal distribution p. In particular, it suffers from the curse of dimensionality, in case of a non-tail non-tailored proposal p. We will tackle this problem later in the section by choosing already close to the target density q.
6

---

Importance Corrected Neural JKO Sampling
The following theorem describes the density of the distribu- tion of X. Moreover, it ensures that the KL divergence to the target distribution decreases. We include the proof in Appendix C.I
Remark 4.4 (Choice of c). We choose the hyperparameter such that a constant ratio 0 of the samples will be resampled, i.e. that [[a(X)] 1 - To this end, we assume that we are given samples 21. Ei..... NN from X and approximate
Theorem 4.2. Let ji be the distribution of X. Then, the following holds true.
N = N min {1, g(x.) cp(2;)
N
B[((())
1 R
N Σ o(s))
1
(i) i admits the density P given by p(x) = p()){((()) + 1 1-)[a(X])) In particular, we have p(x) [
Note that the right side of this formula depends mono- tonically on c such that we can find c > 0 such that E[a(X)] = 1 - r by bisection search. In our numeri- cal experiments, we set r 0.2. We summarize the choice of in Algorithm 5 in Appendix D.1.
with f(z) = f(x)(a(x) +1 +1-B|a(X)]]
(fi) It holds that KL KL(H,V).
The application of a single importance- rejection step is summarized in Algorithm I and illustrated in Figure 2 in the appendix. Note that the value [[a(X)] can easily be estimated based on samples during the training phase. Indeed, given N iid copies X1 XN of X, we obtain that [X) ~ Ei a(Xx) is an unbiased estimator fulfilling the eтror estimate from the following corollary. The proof is a direct consequence of Hoeffding's inequality and given in Appendix C.2.
Algorithm Importance corrected neural JKO sampling
unnormalized target density 9,
initial measure y with density po-
Input:
Number N N of samples,
Number K of total steps.
{a') Output: Sample generator {a') Mi=1 ~ U = z==
Let 2 ff..... 2 ~ 11p for k = 1...., K do
Corollary 4.3. Let X1,..,, XN be iid copies of X. Then, it holds
Define 1h with density p and draw samples af...., xby nº either by neural JKO step by solving the ODE (7) or importance based rejection by Alg. 1.
1 N Bla(X)]
N Σ [(N))
2m VN
EO
1 VN
(details in Alg. 3, 4 resp. Alg. 5) end for
Algorithm 1 Sampling and density propagation for importance-based rejection steps
Set i = pk with density pK
samples . of 1² with density vh, hyperparameter e (see Remark 4.4).
Input:
Neural JKO Sampling with Importance Correction Fi- nally, we combine the neural JKO scheme from the previ- ous section with our rejection steps obtain sampling algorithm. More precisely, start with simple latent dis- tribution μl0 with known density and which we can sample from. In our numerical experiments this will be a standard Gaussian distribution. Now, we iteratively generate distribu- tions μg. & = 1, .... K by applying either neural JKO steps as described in Algorithm 3 and or importance- jection steps as described in Algorithm 5 and 1. We call the resulting model importance corrected neural JKO model, which we summarize in Algorithm 2. During the sampling process we can maintain the density values p²(xx) of the density d of μı for the generated samples by the Algorithms and 1. Moreover, we can also use Theorem 4.2 to evaluate the density px at some arbitrary point x E Rdd We outline this density evaluation process in Appendix D.2. Remark 4.5 (Runtime Limitations). The sampling time of our importance corrected neural JKO sampling depends exponentially on the number of rejection steps since in each rejection step we resample constant fraction of the samples.
can draw samples from ,²²,,
Assume:
can evaluate the unnormalized density g.
for i 1,..., N' do
1. Compute au (=/²)) min {
(
cp (a
2. Draw uniformly from [0, 1] and of from m
3. Set A
k-1
x u S a(a)),
re-
x
if a(at).
4. Compute a (x (a)+1)= min {I,
giottyy I ept (attil))
5. Define the density value
, + A+1) (x. (ac(x)))) +1 +1 Elan(Xe)]
end for
Samples of ++1 ,..., x of 1
Output:
Density values (p+ (p++(ax/
(rtt1)) )
7

---

Importance Corrected Neural JKO Sampling
Table 1: Energy distance. We run each method 5 times and state the average value and corresponding standard deviations. The rightmost column shows the reference sampling error, i..., the lower bound magnitude of the average energy distance between two sets of different samples drawn from the ground truth. A smaller energy distance indicates a better result.
Sampler
Distrbution
Neuni (omrs)
Sampling limor
Mustache
46x 10x1
10-1
1.7 1
10
10-3
1
10"
stinedd Modes
11xx10-14:
12
82 10
13x
12x
=51
26 10-6
8hifle Podly
13x
10
122x 10
10
11
10
82x
10
10-2
133x
10
+22
34x 10x10-3-00
2.
Funnel
12x =511x
5.1x
*482%
10
16
10
74 10 22
10
4.6 10
1.
14x 10
+8.2%
3.1% 10
GMM-10
12x1002253x
122
12x101132200
10
3.7
18x
2.0
10
L1x
53% 15%
10
366 10-5
GMM-22
9.1 91x10-12288
1o-1
1ox10-142777
10-1
5.0
100x10-22
15
14x
10-1
1ax
10
10
10
11x 11x1100
6.4% 10-6
GMM-50
2.4 24x
10-1
2. 10
7.5
10-1
2.3
10-2 10
18x
10
1.7%
10-1
27 1ox10-+778 10
10x 10
4.0
10-5
L1 10-+
GMM-110
37 10 1.7 10
2
2.8
Lox
4.7
222
10
60x
10
10
6MM-210
10 18 10
98x10793110
39%
10
89
20%
Table 2: Estimated (Z). We run each method 5 times and state the average value and corresponding standard deviations. The DDS values for LGCP are from (Vargas et al., 2023a). Higher values of log ZZ) estimates correspond to better results.
Sumples
Dutribation
DDS
CRAFT
Neural IKO
Neural IKO IC (ours)
Ground Truth
Mustache
-1.5x 10 2.7
10
5.5
10
-30x 10 2.5
10
73 10
10
shilled Mods
-6.7 10- 2.0.
10-2
-12
114 10-1
-34 10-1 231
10-3
+5.1 x
10
224
10-3
shifted 8 Peaky
-122x10
1.2x10-1220
2.22
10
-18x10
262 79
-35x 1
10
-211 10
4.3
10
Funne
-1.8% 115x10-14-00
+6.8
10
-12x
-14 10
+.6x
10
-1.1 10
+11
GMM
-23 10
1.
-
1177
-43
10
13.5
10
1120
GMM-20
-6.1% 10x.10-1-6000
10-2
-15x
10
10-1
-6.3
10
10-2
40.4
10-1
10-3
GMM-50
-10x
10t1
-23
993
+11x
GMM-100
-30%
10-1
-23x
-18x 10
+98%
-39 10 7
GMM 200
-944 10
1
7.2
10
6.3
10 13
-22
100
2 10
-5.6
1113
LOOP
303.0 8
7.7
10
17.8322 10
1.9 10-1
5082110 10-1
DEE aonilable
However, due to the moderate exponential base of 1 + we will see in the numerical part that we are able to perform significant number of rejection steps in a tractable time.
the rejection steps by reporting values for the same neural JKO scheme without rejection steps (neural JKO).
For evaluating the quality of our results, we use two different metrics. First, we evaluate the energy distance (Székely, 2002). is a kernel metric which can be evaluated purely based on two sets of samples from the model and the ground truth. Morcover it encodes the geometry of the space such that a slight perturbation of the samples only leads to a slight change in the energy distance. Second, we estimate the log- normalizing constant which is equivalent to approximating the reverse KL loss of the model. A higher estimate of the log- normaliziiin constant corresponds to a smaller KL divergence between generated and target distribution and therefore higher similarity of the two measures. Since this requires the density of the model, this approach is not applicable for MALA and HMC.
In our numerics, we build our importance corrected neural JKO model by first applying n1 E N neural JKO steps followed by n2 E N blocks consisting out of one neural JKO step and three importance-basee rejection steps, where n and n2 are hyperparameters given for each model separately in Table 6. For the neural JKO steps, we choose an initial step size To > 0 as paperr-rramtterraeee and then increase the step size exponentially by Tk+1 4Tg. Note that one could alternatively use adaptive step sizes similar to (Xu et al., 2024). However, for our setting, we found that the simple step size rule is sufficient.
5. Numerical Results
The results are given in Table 1 and 2. Our importance corrected neural JKO sampling significantly outperforms the comparison for all test distributions. In particular, we observe that for shifted modes, shifted 8 peaky and GMM- d the energy distance between neural JKO IC and ground truth samples is in the same order of magnitude the energy distance between to different sets of ground truth samples. This implies that the distribution generated by neural JKO IC is indistinguishable from the target distribution in the energy distance. For these examples, the log( Z) esitmate is sometimes slightly larger than the ground truth, which can be explained by numerical effects, see Remark E. for detailed discussion. We point out to a precise description
We compare our method with classical Monte Carlo sam- plers like a Metropolis adjusted Langevin sampling (MALA) and Hamiltonian Monte Carlo (HMC), see c.g., (Betancourt, 2017; Roberts& & Tweedie, 1996). Additionally, we com- pare with two recent deep-learning based sampling algo- rithms, namely denoising diffusion samplers (DDS, Vargas et al., 2023a) and continual repeated annealed flow transport Monte Carlo (CRAFT, Matthews et al., 2022). We evaluate all methods on a set of common test distributions which is described in detail in Appendix E. Moreover, we report the error measures for our importance corrected neural JKO sampler (neural JKO IC), see Appendix E.5 for implementa- tion details. Additionally, we emphasize the importance of
8

---

Importance Corrected Neural JKO Sampling
of the metrics, the implementation details and additional experiments and figures in Appendix E.
project number 530824055 and by the EPSRC programme grant The Mathematics of Deep Learning with reference EP/V026259/1. RG acknowledges support by the DFG MATH+ project AA5-5 (was EF1-25)- Wasserstein Gradi- end Flows for Generalised Transport in Bayesian Inversion.
6. Conclusions
Methodology We proposed a novel and expressive gen- erative method that enables the efficient and accurate sam- pling from a prescribed unnormalized target density which is empirically confirmed in numerical examples. To this end, we combine local sampling steps, relying on piece- wise geodesic interpolations of the JKO scheme realized by CNFs, and non-local rejection and resampling steps based on importance weights. Since the proposal of the rejection step is generated by the model itself, they do not suffer from the curse of dimensionality as opposed to classical variants of rejection sampling. The proposed approach provides the advantage that can draw independent samples while correcting imbalanced mode weights, iteratively refine the current approximation and evaluate the density of the gener- ated distribution. This is a consequence of the density value propagation through CNFs and Theorem 4.2 and enables possible further post- pest-prcesssing steps that require density evaluations of the approximated sample process.
References
Akhound-Sadegh, T., Rector-Brooks, J., Joey Bose, A., Mit- tal, S., Lemos, P., Liu, C.-H., Sendera, M., Ravanbakhsh, S., Gidel, G., Bengio, Y., et al. Iterated denoising en- ergy matching for sampling from boltzmann densities. In Proceedings of the s International Conference on Machine Learning, pp. 760-786, 2024.
Albergo, M. S. and Vanden-Eijnden, E. Nets: A non-equilibrium transport sampler. arXiv preprint arXiv:2410.0271,, 2024.
Altekrüger, F. and Hertrich, J. WPPNets and WPPFlows:
The power of Wasserstein patch priors for superresolution. SIAM Journal on Imaging Sciences, 16(3):1033-100 1067, 2023.
Altekrüger. F., Hertrich, J., and Steidl, G. Neural Wasser- stein gradient flows for maximum mean discrepancies with Riesz kernels. In International Conference on Mu- chine Learning, pp. 664-690. PMLR, 2023.
Outlook Our method allows for the pointwise access to the approximated target density and the log normalization constant. These quantities can be used for the error moni- toring during the training and hence provide guidelines for the adaptive design of the emulator in terms of CNF -or rejection/resamplingg steps and provide a straightforward stopping criterion. The importance-basee rejection steps can also be used in other domains like fine-tuning of score- based diffusion models for downstream tasks, see Denker et al. (2025).
Alvarez-Melis, D., Schiff, Y., and Mrouch, Y. Optimizing functionals on the space of probabilities with input convex neural networks. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
Ambrosio, L.., Gigli, N., and Savare, G. Gradient Flows in Metric Spaces and in the Space of Probability Measures. Lectures in Mathematics ETH Zürich. Birkhäuser, Basel, 2005.
Limitations In the situation, when the emulator is realized through a stack of underlying rejection/resampling steps, the sample generation process time negatively affected, see Remark 4.5. In order to resolve the drawback we plan to uti- lize diffusion models for the sample generation. This is part of ongoing and future work by the authors. Finally, the use of continuous normalizing flows comes with computational challenges, which we discuss in detail in Appendix F.2.
Andrieu, C., De Freitas, N., Doucet, A., and Jordan, M. I. An introduction to MCMC for machine learning. Machine Learning, 50:5-43, 2003.
Andrle, A., Farchmin, N., Hagemann, P., Heidenreich, S., Soltwisch, V., and Steidl, G. Invertible neural networks versus MCMC for posterior reconstruction in grazing in- cidence X-ray fluorescence. In International Conference on Scale Space and Variational Methods in Computer Vision, pp. 528-539. Springer, 2021.
Impact Statement
This paper studies a general problem from the field of statis- ties and machine learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Arbel, M., Korba, A., Salim, A., and Gretton, A. Maximum mean discrepancy gradient flow. Advances in Neural Information Processing Systems, 32, 2019,
Acknowledgements
Arbel, M., Matthews, A., and Doucet, A. Annealed flow transport Monte Carlo. In International Conference on Machine Learning, pp. 318-330. PMLR, 2021.
JH acknowledges funding by the German Research Foun- dation (DFG) within the Walter-Benjamin Programme with

---

Importance Corrected Neural JKO Sampling
Ardizzone, L., Kruse, J., Wirkert, S., Rahner, D., Pellegrini, E. W., Klessen, R. S., Maier-Hein, L., Rother, C.,, and Köthe, U. Analyzing inverse problems with invertible neural networks. In International Conference on Learn- ing Representations, 2019.
Del Moral, P., Doucet, A., and Jasra, A. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3) 2006.
Denker, A., Vargas, F., Padhy, S., Didi, K., Mathis, S., Du- tordoir, V., Barbano, R., Mathieu, E., Komorowska, U. J., and Lio, P. DEFT: Efficient finetuning of conditional diffusion models by learning the generalised h--tranfform. arXiv preprint arXiv:2406.00 017721 2024.
Behrmann, J., Grathwohl, W., Chen, R. T. Q., Duvenaud, D., and Jacobsen, J.-H. Invertible residual networks. In International Conference on Machine Learning, pp. 573- 582. PMLR, 2019.
Denker, A., Padhy, S., Vargas, F., and Hertrich, J. Itera- tive importance fine-tuning of diffusion models. arXiv preprint arXiv:2502.0444: 2025.
Benamou, J.-D. and Brenier, Y. A computational fluid me- chanics solution to the Monge-Kantorovien mass transfer problem. Numerische Mathematik, 84(3)375-333, 2000.
Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density esti- mation using real NVP. In International Conference on Learning Representations, 2016.
Betancourt, M. A conceptual introduction to Hamiltonian Monte Carlo. arXiv preprint arXiv:/70) 724444444, 2017.
Durkan, C., Bekasov, A., Murray, I., and Papamakarios, G. Neural spline flows. Advances in Newral Information Processing Systems, 2019.
Blessing, D., Jia, X., Esslinger, J., Vargas, F., and Neumann, G. Beyond ELBOs: A large-scale evaluation of varia- tional methods for sampling. In International Conference on Machine Learning. 2024.
Eigel, M., Gruhlke, R., and Marschall, M. Low-rank tensor reconstruction of concentrated densities with application to bayesian inversion. Statistics and Computing, 32(2):
Brenier, Y. Décomposition polaire et réarrangement mono- tone des champs de vecteurs. Comptes Rendus de l'Académie des Sciences Paris Series I Mathematics, 305 (19):805-808, 1987.
27, 2022.
Eigel, M., Gruhlke, R., and Sommer, D. Less interaction with forward models Langevin dynamics: Enrichment and homotopy. SIAM Journal on Applied Dynamical Systems, 1370)1187- 2024.
Chen, J., Richter, L., Berner, J., Blessing, D., Neumann, G., and Anandkumar, A. Sequential controlled Langevin diffusions. In The Thirteenth International Conference on Learning Representations, 2025.
Fan, J., Zhang, Q., Taghvaei, A., and Chen, Y. Variational Wasserstein gradient flow. In International Conference Machine Learning, pp. 6185-6216 PMLR, 2022.
Chen, R. T. Q. torchdiffeq. 2018. URL https:// g thub. com/rtqicctt heaggchhe//tchhiie ffeq
Flamary, R., Courty, N., Gramfort, A., Alaya, M., Bois- bunon, A., Chambon, S., Chapel, L., Corenflos, A., Fa- tras, K., Fournier, N., et al. POT: Python optimal trans- port. Journal of Machine Learning Research, 22(78):1 2021.
Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 31, 2018.
Chen, R. T. Q., Behrmann, J., Duvenaud, D. K., and Ja- cobsen, J.-H. Residual flows for invertible generative modeling. Advances in Neural Information Processing Systems, 32, 2019.
Gabrié, M.. Rotskoff, G., and Hinden-Eijnden. E. Adaptive Monte Carlo augmented with normalizing flows. Pro- ceedings of the National Academy of Sciences, 119(10): e2 e2109420119 19, 2022.
Garbuno- Inigo, A., Nüsken, N., and Reich, S. Affine in- variant interacting Langevin dynamics for Bayesian infer- ence. SIAM Journal on Applied Dynamical Systems, 19 13):633316 2020.
Chen, Y., Huang, D. Z., Huang, J., Reich, S., and Stuart, A. M. Sampling via gradient flows in the space of proba- bility measures. arXiv preprint arXi::2310.03597, 2023.
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26, 2013.
Gilks, W. R., Richardson, S., and Spiegelhalter, D. Markov chain Monte Carlo in practice. CRC press, 1995.
Girolami, M. and Calderhead, B. Riemann manifold Langevin and Hamiltonian Monte Carlo methods. Jour- nal of the Royal Statistical Society: Series B (Statistical Methodology), 73(2): 2011.
De Cao, N., Titov, I., and Aziz, W. Block neural autore- gressive flow. In Uncertainty in Artificial Intelligence, pp. 1263-1273, PMLRR, 2020.
10

---

Importance Corrected Neural JKO Sampling
Grathwohl, W., Chen, R. T. Q., Bettencourt, J., Sutskever, I., and Duvenaud, D. Ffjord: Free-form continuous dynam- ics for scalable reversible generative models. In Interna- tional Conference on Learning Representations, 2019.
Hutchinson, M. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Com- munications in Statistics-Simulationn and Computation, 18(3):155--0776 1989.
Hagemann, P., Hertrich, J., and Steidl, G. Stochastic nor- malizing flows for inverse problems: Markov chains viewpoint. SIAM/ASA Journal on Uncertainty Quantifi- cation, (0(3):1162-1)) 11900 2022.
Jordan, R., Kinderlehrer, D., and Otto, F. The variational for- mulation of the Fokker-Planck equation. SIAM Journal on Mathematical Analysis, 29(1): 1998.
Kingma, D. P. and Dhariwal, P. Glow: Generative flow with invertible Ixl convolutions. Advances in Neural Information Processing Systems, 31, 2018.
Hagemann, P., Hertrich, J., and Steidl, G. Generalized nor- malizing flows via Markov chains. Cambridge University Press, 2023.
Kingma, P. and Welling, M. Auto-encoding variational Bayes. In International Conference on Learning Repre- sentations, 2014.
Hagemann, P., Hertrich, J., Altekrüger, F., Beinert, R., Chemseddine, J., and Steidl, G. Posterior sampling based on gradient flows of the MMD with negative distance kernel. In International Conference Learning Repre- sentations, 2024.
Lambert, M., Chewi, S., Bach, F., Bonnabel, S., and Rigol- let, P. Variational inference via Wasserstein gradient flows. Advances in Neural Information Processing Systems, 35: 14434-14447, 2022.
He, J., Du, Y., Vargas, F., Zhang, D., Padhy, S., Ou Yung R., Gomes, C. P., and Hernández-Lobato, M. No trick, no treat: Pursuits and challenges towards simulation-free training of neural samplers. In Frontiers in Probabilistic Inference: Learning meets Sampling, 2025.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Liu, Q. and Wang, D. Stein variational gradient descent: A general purpose Bayesian inference algorithm. Advances Neural Information Processing Systems, 29, 201
Hertrich, J. Proximal residual flows for Bayesian inverse problems. In International Conference Scale Space and Variational Methods in Computer Vision, pp. 210- 222. Springer, 2023.
Liu, Q. and Wang, D. Stein variational gradient descent: general purpose Bayesian inference algorithm. Advances in Neural Information Processing Systems. 29, 2016b.
Hertrich, J., Wald, C., Altekrüger, F., and Hagemann, Generative sliced MMD flows with Riesz kernels. In International Conference Learning Representations, 2024.
Liu, X., Gong, C.,, and Liu, Q Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint 2rXiv:2209.03300 2022.
Ho. J., Jain, A., and Abbeel, P. Denoising diffusion prob- abilistic models. In Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851. Curran Associates, Inc., 2020.
Liutkus, A., Simsekli, U., Majewski, S., Durmus, A., and
Stöter, F.-R. Sliced-wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. International Conference on Machine Learning, pp. 4104 4113. PMLR, 2019.
Hoeffding, W. Probability inequalities for sums of bounded random variables. Journal of the American Statistical Association, 58(301):13-30, 1963.
Lu, Y., Lu, J., and Nolen, J. Accelerating Langevin sampling with birth-death, arXiv preprint arXiv:1905.09863 2019.
Hoffman, M., Sountsov, P., Dillon, J., Langmore, I., Tran, D., and Vasudevan, S. Neutra-lizing lizing bad geometry in Hamiltonian Monte Carlo using neural transport. arXiv preprint arXiv: 1.Xiv:I903.00 119900033744.. 2019.
Marzouk, Y., Moselhy, T., Parno, M., and Spantini, A. Sam- pling via measure transport: An introduction. Handbook of uncertainty quantification, 1:2. 2016.
Matthews, A., Arbel, M., Rezende, D. J., and Doucet, A.
Hoffman, M. D. and Gelman, A. The no-u-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1): 1593-1623, 2014.
Continual repeated annealed flow transport Monte Carlo. In International Conference on Machine Learning, pp. 15 15196-15219. PMLR, 2022.
Midgley, L. I., Stimper, V., Simm, N. C., Schölkopf, B., and Homanddz-LLbbtee J. M. Flow annealed impor- tance sampling bootstrap. In The Eleventh International Conference on Learning Representations, 2023.
Huang, C.- W., Krueger, D., Lacoste, A., and Courville, A. Neural autoregressive flows. In International Conference on Machine Learning, pp. 2078-2087. PMLR, 2018.

---

Importance Corrected Neural JKO Sampling
Mokrov, P., Korotin, A., Li, L.., Genevay, A., Solomon, J. M., and Burnaev, E. Large-scale Wasserstein gradient flows. Advances in Neural Information Processing Systems, 34: 15243-15256, 2021.
Ruthotto, L. and Haber, E. An introduction deep genera- tive modeling. GAMM-Mittelllggnn, 44(2):e021000088, 2021.
Santambrogio, F. Optimal transport for applied mathemati- cians. Birküuser, NY, 55(58-63):94, 2015.
Møller, J., Syversveen, A. R., and Waagepetersen, R. Log gaussian cox processes. Scandinavian Journal of Statistics, 25(3):451-482 1998.
Sejdinovic, D., Sriperumbudur, B., Gretton, A., and Fuku- mizu, K. Equivalence of distance-based and RKHS-based statistics in hypothesis testing. The Annals of Statistics, 41(5):2263 - 1291,, 2013.
Neal, R. M. Annealed importance sampling. Statistics and computing, 11:125-139, 2001.
Sommer, D.. Gruhlke, R., Kirstein, M., Eigel, M., and Schillings, C. Generative modelling with tensor train approximations of Hamilton-Jacobi Jaco----laaaoiiillllla equations. arXiv preprint arXiv:2402.1528. 2024.
Neal, R. M. Slice sampling. The Annals of Statistics, 31(3): 705-767, 2003.
Noé, F., Olsson, S., Köhler, J.. and Wu, H. Boltzmann gener- ators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 147, 2019.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
Onken, D., Fung, S. W., Li, X., and Ruthotto, L. OT- flow: Fast and accurate continuous normalizing flows via optimal transport. In AAAI Conference on Artificial Intelligence, volume 35, pp. 9223 9232, 2021.
Székely, G. E-statistics: The energy of statistical samples. Techical Report, Bowling Green University, 2002.
Papamakarios, G., Pavlakou, T., and Murray, I. Masked autoregressive flow for density estimation. Advances in Neural Information Processing Systems, pp. 2338-2347, 2017.
Vargas, F., Grathwohl, W., and Doucet, A. Denoising dif- fusion samplers. International Conference on Learning Representations, 2023a.
Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. Trans- port meets variational inference: Controlled Monte Carlo diffusions. arXiv preprint arXiv:2307.01050, 2023b.
Peyré, G. and Cuturi, M. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 11(5-6):355-66 607, 2019.
Vidal, A., Wu Fung, S., Tenorio, L., Osher, S., and Nurbekyan, I.. Taming hyperparameter tuning in continu- ous normalizing flows using the JKO scheme. Scientific Reports, 13(1):4501, 2023.
Phillips, A., Dau, H.-D., Hutchinson, M. J., De Bortoli, V., Deligiannidis, G., and Doucet, A. Particle denoising dif- fusion sampler. In International Conference on Machine Learning, 2024.
Von Neumann, J. Various techniques used in connection with random digits. Applied Math Series, 12(36-38):1, 1951.
Qiu, Y. and Wang. X. Efficient multimodal sampling via tempered distribution flow. Journal of the American Sta- tistical Association, 1196446):4446444888 2024.
Wang, Y. and Li, W. Accelerated information gradient flow.
Rezende, D. and Mohamed, S. Variational inference with normalizing flows. In International Conference on Ma- chine Learning, pp. 1530-1538. PMLR, 2015.
Journal of Scientific Computing, 90:1-47, 2022.
Welling, M. and Teh, Y. W. Bayesian learning via stochastic gradient Langevin dynamics. In International Conference on Machine Learning, pp. 681-688, 2011.
Richter, L. and Berner, J. Improved sampling via learned diffusions. In International Conference on Learning Rep- resentations, 2024.
Woo, D. and Ahn, S. Iterated based flow matching for sampling from Boltzmann densities. arXiv preprint arXiv:2408$ 16249, 2024.
Roberts, G. O. and Tweedie, R. L. Exponential convergence of Langevin distributions and their discrete approxima- tions. Bernoulli, 2(4):341-363, 1996.
Wu, D. and Xie, Y. Annealing flow generative model to- wards sampling high high-dimensional and multi-mode distri- butions. arXiv preprint arXiv:2409.20544 2024.
Rossky, P. J., Doll, J. D., and Friedman, H. L. Brownian dynamics as smart Monte Carlo simulation. The Journal of Chemical Physics, 69(10)44628 0630446666663,,, 1978. doi: 10. 3063/1.433411
Wu, H., Köhler, J., and Noé, F. Stochastic normalizing flows. Advances in Neural Information Processing Systems, 33: 5933-5944, 2020.
12

---

Importance Corrected Neural JKO Sampling
Xu, C., Cheng, X., and Xie, Y. Normalizing flow neural net- works by JKO scheme. Advances in Neural Information Processing Systems, 36, 2024.
Zhang, Q. and Chen, Path integral sampler: a stochas- tic control approach for sampling. arXiv preprint arXiv:2 2021.
13

---

Importance Corrected Neural JKO Sampling
A. Background on Wasserstein Spaces
We give some theoretical background on Wasserstein gradient flows extending Section 2. In what follows we refer to (p) == (p, IR") as the set of square integrable measurable functions on I with respect a given measure μ P(R*).
For absolutely continuous curve y, there exists unique minimal norm solution et of the continuity equation (1) in the sense that any solution E of (I) fulfills []ū(;, ()00000) 2 (lo(s,))zza for almost every t. This is the unique solution of (1) such that v(: 1) is contained in the regular tangent space
T Px {A(T- Id): (Id, TT [1, Tw)), X 0} ba(())
An absolutely continuous curve 7: (0, 00) + P2 (IR")) with velocity field i2 (2(())P((R)) is Wasserstein gradient flow with respect to F: (2( (c00, if
12 -F)((()) for a.e. t 1> 0,
where (F()) denotes the reduced Fréchet subdiffential at defined as
/
8F(p)= <ELL)))) F(v) F(v) F(H) 2
inf xErop" (a.r)
(E(x).y-2)) dm +(W2(p, vu P2
Rd
The norm [[:)|La(T(()) of the velocity field of a Wasserstein gradient flow coincides for almost every t with the metric derivative
(aF((M) inf F(H) F(V) セーナ W2(u,v)
For the convergence result from Theorem 2.2, need two more definitions. First, need some convexity assumption. For À E IR, F: (IR2)) IR U }+oo} is called X-convex along geodesics if, for every μr, V E dom F = {1r E P2(RR)) F(u) < jo there exists at least one geodesics 7: [0, - P2 (IR between μ and v such that
F(r(t)) (1 - F(H) +tv()) _/2 t(1 t) (2,
E [0,
To ensure uniqueness and convergence of the JKO scheme, slightly stronger condition, namely being A-convex along generalized geodesics will be in general needed. Based on the set of ttree-plans with base E P2(R)) given by
(o for P2(RR Rd Rd)
(m # a a, (T²) # a m, (Ta) v},
the so-called generalized geodesics y: [0, c] P2 P2(Rd) joining μ and (with base σ)) defined as
(t) = ((1 {) T2 + {TT) a,
t [0,
(8)
where a E VI with (T1,2) E rott (a, (u) and (TI,3) + 1 (o, v). see Definition 9.2.2 in (Ambrosio et al., 2005). The plan may be interpreted as transport from f to via σ. Then a function F: P2(R') - (-o0, oo] is called X-convex along generalized geodesics (see Ambrosio et al., 2005, Definition 9.2.4), if for every a, 14, E dom F, there exists at least one generalized geodesics y: [0, 1] + (2(Rd) related to some a in (8) such that
F(r(t)) < (1 - t) F(1) +tF()) Σे t(1
W (H,
t [0, 1].
where
wa( -= "R"xR<xR4 |yy - 2
Every function being X-convex along generalized geodesics also A-convex along geodesics since generalized geodesics with base 1 actual geodesics. Second, X-convex functional F: (2R R U {+00} } is called coercive, if there exists some 0 such that
/
inf (FF()) 1 P2(RR)) Rd
[|>])dp(()) S > -0,
see (Ambrosio et al., 2005, eq. (11.2.1b)), In particular, any functional which is bounded from below is coercive.
If F is proper, lower semicontinuous, coercive and A-convex along generalized geodesics, one can show that the prox, non-empty and unique for small enough (see Ambrosio et al., 2005, page 295).
14

---

Importance Corrected Neural JKO Sampling
B. Proofs and Examples from Section 3
B.1. Examples fulfilling Assumption 3.1
Assumption 3.1 is fulfiiled for many important divergences and loss functions F. We list some examples below. While straightforward to check that they are proper, lower semicontinuous, coercive and bounded from below, the convexity non-trivial. However, the conditions under which these functionals X-convex along generalized geodesics are well investigated in (Ambrosio et al., 2005, Section 9.3). In the following denote the Lebesgue measure PLeb- For measure μ E pac(RR) we denote by du/dunaa the Lebesgue density of if it exists.
Let v/ E P2 (R") with Lebesgue density and define the forward Kullback-Leibler (KL) loss function
( log
g(zr
dr, if 3dd//du..b and dv/dutaa 9,
F(u) (L,
p(#)
(+00
otherwise.
By (Ambrosio et al., 2005, Proposition 9.3.9), we obtain that fulfills Assumption 3.1.
We can also derive functional, be reversing the arguments in the KI. divergeee Then, arrive the reverse KI. loss function given by
P(x) log
p(x) ((x))
dæ, if 3du//dues and dv/dptcb q,
F(u) KL(H,V)
+00
otherwise.
Given that - -lo(()) is X-convex, we obtain that F fulfills Assumption 3.1, see (Ambrosio et al., 2005, Proposition 9.3.2).
-Finally, we can define F based on the Jensen-Shannon divergence. This results into the funetion
1 [KL (n, { v)) + KL (v. [m)+
F(u) JS(,,
Assume μ and admit Lebesgue densities p and 9 respectively. Then, combining the two previous statements, this fulfills Assumption 3.1 whenever - log(p) and log(g) are X-convex.
All of these functionals are integrals of smooth Lagrangian functional, i.e., there exists some smooth F: Rd x IR x Rd + R such that
F(u) = [ke F(x,,(x), Vp(x))dz, if 3dp/duueb = P. otherwise.
oc
In this case, the limit velocity field from Theorem 2.2 (which appears as a limit in Theorem 3.3) can expressed analytically the gradient of the so-called variational derivative of F, which is given by
6F (y(t) (x) (z)= - Og F2F Vp(x)) + 0g Fz,,p()( vp(z))
where O,F is the derivative of F with respect to the i-th argument and is the Wasserstein gradient flow, see (Ambrosio et al., 2005, Example 11.1.2). For the above divergence functionals, computing these terms lead to a (weighted) difference of the Stein scores of the input measure (t) and the target measure v which is nice link to score-based methods. More precisely, denoting the density of 7tt) by Pt. we obtain the following limiting velocity fields.
For the forward KL loss function we have that F(I,,,2) = q(x) log
(g(x) y1
Thus, we have that & (x()) = 8p p(
g(x)
Hence, the velocity field v(. V 8 is given by
dr((3)
q(x) (p((z)
g(x) Vq(a))
g(z)
u(I t)
(Vlog(p((z)) V Vlog(g(z)))
(x Pr (x(x) (x(x) ((x) pr(r)
15

---

Importance Corrected Neural JKO Sampling
For the reverse KL loss function the Lagrangian is given by F(z, y. 2) = y log
9
Thus, we have that 8 (z)
g(x)
log
p(z) ((x) + 1 = log(p(z)) log(g(z)) + 1. Hence, the velocity field u(., t)
V
5.F
is given by
6y(1)
v(z, t) v(log(p:)(z)) - ((lo(()
For the Jensen- -hannonn divergence the Lagrangian is given by y, = 1 log vtt(+a)
2y
+ q(x))
22(x)
Thus, we have that
6F (x) = 2/ log p(a)+a(m) Hence, the velocity field o(.,t) = 5π(1)
6F
is given by
by
1 (x)) + ((x) Vpc (p) - - px (x) + Vg(x))
(xr t)
2 p((x)
(pe(z)+
1 (pr(z)q(Z) - Pr( V((2))
2 (pr(x) + ((x)()))))) ppt
((x)
Vp((z)
Vg(l)]
(())(z)) + ((z))
pa(a:)
((x)
q(x))
[V(log(p))<))) - v(log())(z)))
2(pe(x) + (()))
Note that computing the score V log(p;) of the current approximation is usually intractable, such that these limits cannot be inserted into a neural ODE directly.
B.2. Proof of Corollary 3.2
Using similar arguments in (Altekrüger et al., 2023; Mokrov et al., 2021; Onken et al., 2021; Xu et al., 2024), let v: Rd x [0, r| + IR such that the solution of
((x, t). 2(x, x,
fulfills +11 1 Since (1,,4s k) is minimizer of (5), obtain that
% //'0 J2d
[
t) + F(st,((' ( +)
1
1// [ (1)(((z,)) ()) F()())))
2
Observing that F(zr,(s, T) +1+) F(e(ss)) F(u ) we obtain that
L / I10,,(Z (z, 1), 1)1 S [- [ [1(s((;
Since chose arbitrary, obtain that Ur the optimal velocity field from the theorem of Benamou-Brentn which now directly implies part (ii) and (ii).
B.3. Proof of Theorem 3.3
In order prove convergence of the velocity fields vr. first introduce some notations. To this end, let T. be the optimal transport maps between ي² and - and define by of (T I)/T the corresponding discrete velocity fields. Then, the velocity fields u, can be expressed as
Uy (x, KT tr) cf(((1 - +)
(9)
Further, denote the piece- wise constant concatenation of the discrete velocity fields by
(x, kT + (r) = o t (0,
16

---

Importance Corrected Neural JKO Sampling
Note that for any т it holds that e, E L2(7,, Rd [0, T)) and ü, E Lg(5,, R [0, T]). To derive limits of these velocity fields, we recall the notion of convergence from (Ambrosio et al., 2005, Definition 5.4.3) allowing that the iterates defined the same space. In this paper, stick to square integrable measurable functions defined on finite dimensional domains, which slightly simplifies the definition.
Definition B.I. Let Ω C IRd be measurable domain, assume that 14k P22(2) converges weakly to Jr P2(22) and let ft L2(Mk, Q)) and f E L2( 92). Then, we say that fx converges weakly to f. if
J (o())
/ (o(t),
as k
for test functions & E 02 (0). We say that fú converges strongly to f if
lim fi (2) & 1lfll k
(10)
Note that (Ambrosio al., 2005, Theorem 5.4.4 (ii)) implies that formula (10) is fulfilled with equality for any strongly convergent sequence fo to f. Moreover is known from the literature that subsequences of the piece- wise constant velocity admit weak limits.
Theorem B.2 (Ambrosio et al., 2005, Theorem 11.1.6.. Suppose that Assumption 3.1 is fulfilled for F. Then, for any n E dom(F) and any sequence (n)) (0, oo), there exists a subsequence (again denoted by (Ti)t) such that
-The piece-wise conslunt curve 3m (1) narrowly converges to some limit curve 7(1) for all & E [0, 0o).
-The velocity field in E L2(5n(() (t), d2d d [O, weakly converges some limit E L2(7(t), Rd [0, TT) according Definition B.1 for any T > O.
The limit i fulfills the continuity equation with respect to %, i.e.,
0t9(() V (0n (s,)((()) 0.
In order to show the desired result, there remain the following questions, which we answer in the rest of this section:
Does Theorem B.2 also hold for the velocity fields vn defined in (9) belonging the geodesic interpolations?
Can we show strong convergence for the whole sequence Ur,
Does the limit é have the norm-minimizing property that (6(,)))lL = (AF)))(())?
To address the first of these questions, we show that weak limits of o, and u, coincide. The proof is a straightforward computation. similar statement slightly different setting was proven in (Santambrogio. 2015. Section 8.3).
Lemma B.3, Suppose that Assumption 3.1 is fulfilled and let (TT) C (0, oo) be sequence with - 0 as + oo such that Un L2(ān, Rd ()) converges weakly to some E L2(9, Rd [0, TT). Then, also Vn E L2(7m, Rd [0,7]) converges weakly to G.
Proof. Let & E C (RR 10, T)). In particular, & is Lipschitz continuous with some Lipschitz constant J. < co. Then, it holds that
Veraaen (0, - (0, o)da
/ / (0,0n)) - (0, 8 ddi
+ +/ (0, on. - (0, u)dd
17

---

Importance Corrected Neural JKO Sampling
Since in converges weakly the second term converges to zero | - oc. Thus in order show that also Un, converges weakly to é, remains show that also the first term converges zero. By using the notations 7i min kn 2 T, & E bz0} and Ki - Zi
we can estimate
Wee ( (0 - Spemoor (0, onddin
J." (olee ())d7, (1)(z) - (o(o), 0r (1,)))),)) (1)x) di (1
1," 1Be [
- / (0(t,, 1), 0z (x,, (0)x)) dt (0(x, (0,,,,,,,cc 8n (1) de
K
( U (), (), (1)
JeT
1=
Ki-1
Σ
(o(z, + tn) u. kn tn)) (KT (x)
7
- (6(2, + tn), in (z. Kn tn))d (kn + tz)(()) dt.
By denoting with T the optimal transport map from p to p we have for E (0,1) that
Vn (KT, + tni) = =(1 t)T + xT wH (n + tn)) = pn:
and the velocity fields satisfy
Br (n + tm) (m) ((11 (m))) (, + tm) = (x).
Then, the above term becomes
Ki-1
| (0(x, kn tm), ((11 t)T (a))) t)T + +T+h)
Σ
k=
+m), f ktm) (x) dt
K1-1
1 (0((1 t)x + +T (m)) KJ + tn) ot (x)) (z)
Σ
km
-f (ola,, (m), of (x) (ar) dt
Ki-1
/ [ [60((1 t)r + (x)) (z)) ln tn) ol2, + tn), v (z)) (x)
E
S n
Using Hölders' inequality and we obtain that this is smaller or equal than
K1-1
Σ '.' (Ne (lo((1 t)r + (x))
(2))), kn + (m) - o(z, km + ann, (x) L loe (2)1 (o)) 1/2
n
18

---

Importance Corrected Neural JKO Sampling
By the Lipschitz continuity of & and inserting the definition of u T this smaller or equal than
K,
{
L2
x1| =/ (a)) f (x) =112 (0)) 1/2 J (T)
7
dt
1.
Ki-1
Σ
(L. ull
(1 =112)(((())
K=
L K1-1
Σ
W2 (1+, M
2 k=0
1.
W3 (u 17 1
S
2 k=0
Finally, we have by the definition of the minimizing movements scheme that
W 1
- F(u F(u
2m
Summing up for i: = 0, 1,.. we finally arrive at the bound
Jenppor (o, u,)d
/
Fl())
JR4x10.7) u)d S 7i F(11))
inf
(EPa(RR)
Since F is bounded from below the upper bound converges to zero as + oc. This concludes the proof.
Finally, we employ the previous results show Theorem 3.3 from the main part of the paper. That is, we show that for any (T)) C (0,oc) with -> the whole Ur, converges strongly v.
Theorem B.4 (Theorem 3.3). Suppose that Assumption 3.1 is fulfilled and let (ri)e S (0, oo) with n 0. Then, (vz, )r
converges strongly the velocity field i E L2(7, IR x [0, of the Wasserstein gradient flow y: (0, 20) - P2(R2) of F starting in po.
Proof. We show that any subsequence of ni admits subsequence converging strongly to i. Using the sub-subsequence criterion this yields the claim.
By Theorem B.2 and Lemma B.3 we know that any subsequence n admits weakly convergent subsequence. an abuse of notations, we denote again by 77 and its limit by ö. Then, we prove that the convergence is indeed strong and that é
Step 1: Bounding lim limsup S ||=,, (;, (,)t)) ( dt from above. Since X-convexity with 2 0 implies X-convexity 1-:00 with À 1, can assume without loss of generality that O. Then, by (Ambrosio al., 2005, Lemma 9.2.7, Theorem 4.0.9), we know that for any 0 it holds
C(T,1)= (1 + 2111 exp (-loy(1 T
log(1 + -
(2()-(1),<() TC(,,t),
v2(1+ AT)
where t, min(kT KT 2 t,k E Not. For simplicity, we use the notations Tmax max 1 E No). K = I and TIaxx T Tmax. Since A < we have that C(T,8) S ((T,Tmax) C(r) for all [0, Tiaa] Moreover, we have that C(T) + X|Tma (FI()))) exp( - X7 as T + 0, such that the sequence (C(T))) is bounded. In particular, there /2
7.
exists C > 0 such that
(t) < nC, for all t [0, Tmax].
Inserting = le, for k O.... Ki gives
W2(n S STC, for all
k = 0, ...,. Ki.
(11)
19

---

Importance Corrected Neural JKO Sampling
Now, we can conclude by the theorem of Benamou--renier that
l. 11 (1.0) (00,,, -
KiTi |0n((.00iiatt,,
Ki-1
(k+1)
(0..0))
km0
km
Ki
W3 1
Now applying the triangular inequality for any k = 1..... K2
W 14 + < (Walue, ((rk)) + Wa(o(mk)) 7(((k (((1( 1)) + Wa(o(n(k ())
and the estimate from (11) yields
L. (1,,,,,)1l2 Re) <
K
(2rC + W2(1(Tk), ((t( +)))
A=
By Jensens' inequality the right hand side can be bounded from above by
Ki
1/2
Ki-1
4 + KiC
Ki (((((())),
+ ((((())) 1)))
Finally, again Renamou-Brenier gives that ([(T((T)))<Tl)) 1))) s Jkn formula is smaller or equal than
(k+1)m
|[0( t) dt such that the above
4K 2T KRVKKC (,,m0 |10( (ll
Kim
1/2
Kiq
+
|16( t) (lle
Since KIT- - and -> 0 as oo this converges to fu [] ()) dl such that can conclude
('" (1 & S/' (1((,))) 10000
lim sup
Step 2: Strong convergence. By (Ambrosio et al., 2005, Theorem 8.3.1, Proposition 8.4.5) we know that for any fulfilling the continuity equation
acy(t) V (v(+)t)(tt)
holds that [lu(. t) 16((; Since ő fulfills the continuity equation by Theorem B.2, this implies that
J'" [15(, 116(00)R l- (16(
lim ( (10,()))) (n 1-+00
In particular, Un E Le(7,, Rd x [0, 7J) converges strongly to such that by (Ambrosio et al., 2005, Theorem 5.4.4 (iii) it holds equality the above equation, i.e.,
L |06 t) (1) = L''
Using again (Ambrosio al., 2005, Theorem 8.3.1, Proposition 8.4.5) this implies that 7
20

---

Importance Corrected Neural JKO Sampling
C. Proofs from Section 4
C.1. Proof of Theorem 4.2
1) Let C. _ Lebesgue-measunee draws of X respectively. By the law of total probability, we have
set. We denote by accept and reject the measurable set of accepted and rejected
PP E A) P E A and X was accepted) + P(x E A and X was rejected).
Since it holds X = X if X is accepted and X X' if X is rejected this can be reformulated
P(X E A) = P(X E A, accept) + P(X' E rejecct))
Now it holds
P(X& A, accept) = / a(x)p(x) dr
([p(o)) dz (1 (-)(a(X))),
P(X' E reject) P(X' A) P(reject)
where we used the fact that X, X' ~ μ are independent. As a consequence, we obtain
P(X A) a(z)p(x)dr + plasate) (1 Ela(X)))
p(x) (a(x) + 1- E[a(X)]) dr.
Thus, f(x) = p(x)(a(()) + -- ([) is the density of X.
(ii) Note, since X ~ we have
p(z)a(z) dx [X) Rd SEL P(y)a(y)dy
p(z)a(z) dr =,
in particular p(a))((x))
[] defines a density. Now let E (2 be the corresponding probability measure. Then, it holds
by part (i) that f [[(((X)] (1 - [()())) We will show that KT. v) S KL,(4, v). Due the convexity of the KL divergence in the linear space of measures this implies the claim.
We denote E[a(X)] Jg. min(p(x), 2(x) dd.. Then holds
[ p(x)a(x)
p(x))(z) dr Zq(xx)
KL(1.v)
log
Jad
Z
min(p(x), ((x))
log (p(x)a(()) dr loo(Ze)
/ min(p(x),
s(x))
8 dz
Rd
Z
g(x)
Rd
min(p(x),
q(xt)
log
min ēp(ar ((x))
11) dr -log( () - -log(e)
Jdd
Z
min (p(x),
((x)
Ōp(x) ((x))
min
Z log
0 dr -log(2) log(2).
Rd
Since log
(p(x)
< 0 if and only if p(x) S q(xr//ā l this can be reformulated as
p(x) Z
Ōp(x)
min
log
. dr log(Z)) log(e)
q(x)
(pp(r)
min
p(x) log
dr log(2) log(e),
(12)
q(x))
21

---

Importance Corrected Neural JKO Sampling
rejection k-th layer
yes
previous layer
accepted?
24
k+1
next layer
CNF/ rejection
CNF/ rejection
no
yields
resample a
n
Figure 2: Illustration of single sample within rejection layer. The independent sample of is either accepted directly else resampled once, and then used to define k+1 as input for the next layer.
where the inequality comes from the fact that Z E(a(X)) E [0, 1]. Moreover, it holds by Jensen's inequality that
2)) log [[)(()))) - - log (/
((xr)
p(a)
min
dæ
Ōp(x:)
gd
g(xx)
p(x)
log
min
dx
ēp(z))
Rd
/ p(x)
g(xr) 0 dr
min
log
Ōp(x))
[ p(x)
(p((r) 0 dx
max
log
gd
q(x))
[ max
(p(ar) <(x)
p(x)) log
0 dz.
Thus, obtain that (12) can be bounded from above by
min
p(x) log
(p(xx)
q(x)) 0) dx +
/ max
p(x)) log
ēp(ar)
. log(e)
Rd
Rd
q(x)
which equals
cp(x) g(x) dr - log(() = =1e p(x) log
p(x) ((x)
p(x) log
dr KL(H,
In summary, we have KL(n, v) S KL(U, v), which implies the assertion.
C.2. Proof of Corollary 4.3
Let Y = B[a(X)] 1/4 ΣN ((x)) denote the random variable representing the error. Since axXw) E [0, 1]. Hoeffding s inequality (Hoeffding, 1963) yields that P(Y > t) < 2exp( Nr²)
Consequently, we have that
B| = / P(Y > t)dt < / exp (- N2
2m
dt
2
VN
D. Algorithms
D.I. Training and Evaluation Algorithms
We summarize the training and evaluation procedures for the neural JKO steps in Algorithm 3 and 4 and the parameter selection for the importance-basee rejection steps Algorithm 5. The evaluation of the importance- based rejection steps is summarized in Algorithm in the main text.
22

---

Importance Corrected Neural JKO Sampling
Algorithm 3 Training of neural JKO steps
Input: Samples æř, ..,,...." VN of u². Minimize the loss function 0 + [()) from (6) using the Adam optimizer. Output: Parameters 0.
Algorithm 4 Sampling and density propagation for neural JKO steps
Samples xf, ...,, srty of μ* Density values p² (xf), .... p (x%).
Input:
for 1,... do
1. Solve the ODE (7) for 2tt
2. Set x k-1 (0) T).
3. Set ph-11 (x. ( (xp((o(x(((7))
p (a
end for
{
Samples uf, .., ory of wht1,
Output:
Density values p h+ (x(tt)),... p x (a))t)).
Algorithm 5 Parameter selection for importance- based rejection steps
Samples #f --,; o of μ² with corresponding densities desired rejection rate r, unnormalized target density g
Input:
Choose e by bisection search such that
1, g(x) cp (x)
1 --r
ag (xf),
ax(x) min
N
Output: Rejection parameter and estimate of Bla x
--r.
Algorithm 6 Density Evaluation of Importance Corrected Neural JKO Models
Input: x E IRdd model
if K then
Return latent density (z).
else if the last step is rejection step then
1. Evaluate n (xx) by applying this algorithm for (Xo, Xx-
2. Return p (x) pk (an(x) + E(ax((x-1)) (X)=1)))
else
last step is a neural JKO step
1. Solve the ODE system (with vi from the neural JKO step)
20(x; 1) ) trrace(Vrg(t((()
v0(2o(x,t).)
(
20(I,T)
(x( ) = (0())
)
2. Set ā = 2 (x, 0).
3. Evaluate pd K 1 (i) by applying this algorithm for (Xo, ..., Xx-1).
4. Return p (x) - exp(fe(x,0)
end Output: Density k(x))
23

---

Importance Corrected Neural JKO Sampling
D.2. Density Evaluation of Importance Corrected Neural JKO Models
Let (Xo, ..., Xx) ) be an importance corrected neural JKO model. We aim to evaluate the density pK of X at some given point R.. Using Theorem 4.2, this can be done recursively by Algorithm 6. Note that the algorithm always terminates since K is reduced by one in each call.
E. Additional Numerical Results and Implementation Details
E.1. Test Distributions
We evaluate our method on the following test distributions.
Mustache: The two-dimensional log- log-density is given log log,V(0, Σ) T with Σ [1,0; 1] and T(21,22)
(I1, (I2 - (xf 1)²)). Note that det( (T)))) = 1 for all æ. In particular, we obtain directly by the transformation formula that the normalization constant is one. Depending on a E [0, 1) close to 1, this probability distribution has very long and narrow tails making hard for classical MCMC methods to sample them. In our experiments we use a = 0.9.
Shifted Modes: wwo-dimensional Gaussian mixture model with 8 equal weighted modes and covariance matrix 10 10-21. The modes are placed in circle with radius 1 and center (-1,0), Due the shifted center classical MCMC methods have difficulties to distribute the mass correctly onto the modes.
Shifted 8 Peaky: This is the same distribution as the shifted 8 Modes with the difference that we reduce the width of the modes to the covariance matrix 5 x 10-$1. Since the modes are disconnected, it becomes harder to sample from them.
Funnel: We consider the (normalized) probability density function given by q(x) ((==10,,f)) 0f cxp(21) where of 9. This example was introduced by (Neal, 2003). Similarly to the mustache example, this distribution has narrow funnel for small values 21 which can be hard to sample.
GMM-d: A d-dimensional Gaussian mixture model with 10 equal weighted modes with covariance matrix x 10-21 and means drawn randomly from a uniform distribution on [-1, 1Jd. This leads to a peaky high-dimensional and multimodal probability distribution which is hard to sample from.
GMM40-50D:
Another Gaussian mixture model, which was used an example in (Blessing al., 2024; Chen et al., 2025) based on (Midgley et al., 2023). It is 50 dimensional GMM with 40 equally weighted modes with the identity as covariance matrix, where the means are drawn uniformly from [-40, 40)d. We use the same seed for drawing the modes as in (Chen et al., 2025).
LGCP: This is high dimensional standard example taken from (Arbel et al., 2021; Matthews et al., 2022: Vargas et al., 2023a). It describes a -og-GGas Cox process on a 40 x 40 grid as arising from spatial statistics (Møller et al., 1998). This leads ((x) x (=]μ,, (4, K) [I exp(x. a cxp(x:)), where μ and K are fixed mean and covariance kemnl,, Yi is some data and a is hyperparameter. For details this example and the specific parameter choice we refer to (Matthews et al., 2022).
1600-dimensional probability distribution with the unnormalized density function
E.2. Error Measures
We use the following error measures.
The energy distance was proposed by (Székely, 2002) and is defined by
1 / nd / ||= vlld(a -i) v)(x)) vv).
= (( D(u,
It is the maximum mean discrepancy with the negative distance kernel K(x, v) = -|, yl (see Sejdinovic et al., 2013) and can be estimated from below and above by the Wasserstein-1 distance (see Hertrich et al., 2024). It is a metric on the space of probability measures. Consequently, a smaller energy distance indicates a higher similarity of the
24

---

Importance Corrected Neural JKO Sampling
Table 3: We report the expected squared empirical W distance (W2) and its standard deviation between generated and ground truth samples of size N = 5000 for the different methods and for all examples where the ground truth model is known. A smaller value of W indicates a better result. Note that the curse of dimensionality present in the sample complexity, might limit the reliability of the results for the nigh-dimensional examples. In particular for the funnel distribution, we observe that the expected empirical Wasserstein-2 distance between two independent sets of ground truth samples is higher than the observed W/2 values.
Sunpler
Distrbbtton
MALA
HMC
DDS
CRAFT
Nearal JKO
Newaal JKO IC (ours)
Sunpling Eroor
Mustache
4.7 10 & 13 10
288 x 10
L2
x 11
5.2 10 2.
10"
5.4 I0
1.2
10
3,x to
1 1.3 10
1.7 x 10 1 6.0 10"
1.2% 10
shitedd Modes
55 10-2
10
47
10
a
56
10-3
8.7 10-2 3.1
10-2
2.4x 10r
2.4
10
5.6
10
1.6
10
x
10-3
10-1
10
shifted Peaky
10
56 10
144
10
10x 10x 10
2.xx 10"1
10
69x 10
10
7.2
113x
5.8 10
Funnel
5. 10+2
10+2
10
10+4
10
10+
7. 10
10F2
94
10-2
85x +
10r2
1,0 10+0
GMM-10
x0
#.4.x.00
+00
88x 10
10
24 10"
10"
10
+14
10
+223%
GMM-21
88 10" +.000
10-1
80x
x
10-1
8&* 10-1
1.0%
10-1
7. HP
1.5
10P
1.2x t/
=2
10-1
x 10-1
10-
3.7%
GMM 50
x 1.00
10P
27
10
298
10-1
36x 10
8.8
10-1
2.6x 10 422.6
108
4 10
= 788
10-1
1.2 x 10"
113x 10
1.2x 107
GMM-110
5 10+1
10P
5.7
10
113
10P
9.9% 10
#29
10P
66x 10
+11
10P
1.1
10-1
+2.
109
2.9 x 100
44.7% 10-1
2.8% 10a
GMM-200
10 2 100
12x 10
228
x 10"
244x 16
4.0
10
11 4:3.0
24 10
#.8. 10"
7.8 x 10" 16.9 10
6,9 10'
input distributions. By discretizing the integrals it can be easily evaluated based on N samples æ = (ax);; NONN
y (yi); vSN as
NV
N Σ
1
D( v)
Σ
1
Ø3
llys 8/
2
i,=1
(,)=
We use. N- 50000 samples in Table 1.
We also evaluate the squared Wasserstein-2 distance which is defined in Section 2. To this end, we use the Python Optimal Transport package (POT, Flamary al., 2021). Note that computing the Wasserstein distance has complexity (n²)) where is the number of points. Hence, compared to the case of other metrics. We use N = 5000 samples in Table 3. In addition, we want to highlight that the expected Wasserstein distance evaluated on empirical measures instead of its continuous counterpart suffers from the curse of dimensionality. In particular, its sample complexity scales as O(n -1/d) (Peyré & Cuturi, 2019, Chapter 8.4.1). Consequently, the sample-based Wasserstein distance in high dimensions can differ significantly from the true Wasserstein distance of the continuous distributions. Indeed, we can see in Table 3 that the sampling error has often the same order of magnitude as the reported errors. Overall we can draw similar conclusions from this evaluation as for the energy distance in Table 1.
evaluate the Wasserstein distance based on fewers samples
We estimate the log normalizing constant (short log(2 ()) estimation) which is used as a benchmark standard in various references (e.g., in Arbel al., 2021; Matthews al., 2022: Phillips et al., 2024; Vargas et al.. 2023a). More precisely, for the generated distribution μ with normalized density p and target measure with density g(x) g we evaluate the term
g p(x))
] = log(Zz) - -Bzwa log ((()) q(x)
Ezzuy
log
log Zg) - KL(H,
Due to the properties of the KL divergence, a higher log(2) (Z) estimate implies a lower KL divergence between μ and u and therefore a higher similarity of generated and target distribution. our experiments we compute the log(2) estimate based on N 50000 samples. The results are given in Table 2.
To quantify how well the mass is distributed on different modes for the mixture model examples (shifted Modes, shifted 8 Peaky, GMM-d), we compute the mode weights. That is, we generate N = 50000 samples and assign each generated samples to the closest mode of the GMM. Afterwards, we compute for each mode the fraction of samples which is assigned each mode. To evaluate this distribution quantitatively, we compute the mean square error (MSE) between the mode weights of the generated samples and the ground truth weights from the GMM. We call this error metric the mode MSE, give the results are in Table 4.
Remark E. (Bias in log(2 (Z) Computation). In the cases, where the importance corrected neural JKO sampling fits the target distribution almost perfectly, we sometimes report in Table log(2) Z)) estimates which are slightly larger than the ground truth. This can be explained by the fact that the density evaluation of the continuous normalizing flows uses the Hutchinson
25

---

Importance Corrected Neural JKO Sampling
Table 4: We report the mode MSEs for the different methods for all examples which can be represented as mixture model. A smaller mode MSE indicates a better result.
Sampler
Distrabution
MALA
HMC
DDS
CRAFT
Neural IKO
Neural IKO KC (ours)
shifsed S Modes
3.2x 10
10
1 10
1. 10
6.8% 10
2.2
10
3.2
x 10
1.1.1 x 10 10
83 10
10 16
1,5 10
4.4
10
Peaky
8. 10
137x
11
18x
10
10
7.9
33
10
84x
167
15x
122
11
10-2
10
15.
GMM-22
6.8 10
122
10
86
10
12
3.6%
10
1.5
10
388
10
19
10
13
10
128
16
24 10
10-6
GMM-50
1.0% 10
+-8
10
98
10
5.6
10-3
9.6
10
4.9
10-8
90
10
6.4
10
12x
10
+.44
16
28
10
6.0 10
GMM 100
1.1 10
1544
10-3
11
2
5.4
10
10
18.
10
9.0
10
48
10
14x
10
7.0
16
15
10
GMM-200
1.8.1
10-3
14
10-3
10-6
Table 5: We report the mode Sinkhorn distance for the different methods for the GMM40-50D example. A smaller mode Sinkhorn distance indicates a better result. The comparison values are taken from (Chen et al., 2025).
Sampler
SMC
DDS
CRAFT
PES
CMCD-IV
SCLD
Nearal IKO IC (ours)
Sunkhoen distance
46370 34 187.79
5816.118 é 172.20
28961L.0 & 354,89
101400.7 75 a 60.41
4208.57 737,
3787. 73 249.75
3025.61 61 418.10
trace estimator for evaluating the divergence and a numerical method for solving the ODE. Therefore, we have a small error in the density propagation of the neural JKO steps. This error amplified by the rejection steps since samples with underestimated density are more likely be rejected than samples with overestimated density.
We would like to emphasize that this effect only appears for examples, where the energy distance between generated and ground truth samples is in the same order of magnitude like the average distance between two different sets of ground truth samples (see Table 1). This means that in the terms of the energy distance the generated and ground truth distribution are indistinguishable. At the same time the bias in the log(Z) estimate appears at the third or fourth relevant digit meaning that it is likely to be negligible.
E.3. Additional Comparison
We additionally compare our importance corrected neural JKO model with several other methods on the example of a 50 dimensional Gaussian mixture model as used (Midgley et al., 2023). We follow the setting of Table 3 in Chen et al. (2025) and evaluate the Sinkhorn distance (Cuturi, 2013) with entropic regularization parameter 10-3 based on 2000 generated and ground truth samples. Moreover average the results over 5 training runs with 25 evaluations from each run. We compare our results with sequential Monte Carlo (SMC, Del Moral et al., 2006), denoising diffusion sampler (DDS, Vargas et al., 2023a), continual repeated annealed flow transport Monte Carloe (CRAFT, Matthews et al., 2022), path integral sampler (PIS, Zhang& & Chen, 2021), controlled Monte Carlo Diffusion (CMCD, Vargas et al., 2023b) and sequential controlled Langevin diffusion (SCLD, Chen et al. (2025)). For CMCD, we consider the log-variance (LV) variant considered in (Chen et al., 2025) since gives better results. The hyperparameters used for our importance corrected neural JKO scheme are given in Table 6. Similar as most comparison methods, we use normal distribution with a larger variance as initial distribution 140 and use a gradient flow on the negative log-density prior the learned CNF and rejection steps, see Remark E.2 for details.
The results are given in Table 5, where the values for the comparison are taken from (Chen et al., 2025). We observe that neural JKO IC significantly outperforms the comparisons on this example.
E.4. Additional Figures and Evaluations
Additionally to the results from the main part of the paper, we provide the following evaluations.
Visualization of the Rejection Steps In Figure 1 and with involved steps in Figure 3, we visualize the different steps of our importance corrected neural JKO model on the shifted 8 Peaky example. Due to the shift of the modes, the modes on the right attract initially more mass than the ones on the left. In the rejection layers, we can see that samples are mainly rejected in oversampled regions such that the mode weights are equalized over time. This can also be seen in Figure 4, where we plot the weights of the different modes over time. We observe, that these weights are quite imbalanced for the latent distribution but are equalized over the rejection steps until they reach the ground truth value.
26

---

Importance Corrected Neural JKO Sampling
(Marginal) Plots of Generated Samples Despite the quantitative comparison of the methods in the Tables 1, 2 and 4, also plot the first coordinates of generated samples for the different test distributions for all methods for a visual comparison.
In Figure 5, we plot samples of the shifted 8 Modes example. We can see that all methods roughly cover the ground truth distribution even though CRAFT and DDS have slight and the uncorrected neural JKO scheme has a severe imbalance in the assigned mass for the different modes.
For the shifted Peaky example in Figure we see that this imbalance increases drastically for CRAFT, HMC, MALA and the uncorrected neural JKO. Also DDS has has slight imbalance, while the importance corrected neural JKO scheme fits the ground truth almost perfectly.
The Funnel distribution in Figure 7 has two difficult parts, namely the narrow but high-density funnel on the left and the wide moderate density fan on the right. We can see that DDS does cover none of them very well. Also MALA, CRAFT and the uncorrected neural JKO do not cover the end of the funnel correctly and have also difficulties to cover the fan. HMC covers the fan well, but not the funnel. Only our importance corrected neural JKO scheme covers both parts in reasonable way.
For the Mustache distribution in Figure 8, the critical parts are the two heavy but narrow tails. We observe that CRAFT and DDS are not able to cover them at all, while HMC and MALA and uncorrected neural JKO only cover them only partially. The importance corrected neural JKO covers them well.
Finally, for the GMM-d example we consider the dimensions d = 10 and d = 200 in the Figures 9 and 10. We can see that CRAFT mode collapses, i.e., for d = 10 it already finds some of the modes and for d 200 it only finds one mode. DDS, HMC, MALA and the uncorrected neural JKO find all modes but do not distribute the mass correctly onto all modes. While this already appears for d 10 it is more severe for d 200. The importance corrected neural JKO sampler finds all modes and distributes the mass accurately.
Note that the distributions generated by MALA and the neural JKO are always similar. This is not surprising since both simulate a Wasserstein gradient flow with respect to the KL. divergence and only differ by the time-discretization and the Metropolis correction in MALA.
Development of Error Measures over the Steps We plot how the quantities of interest decrease over the application of the steps of our model. The results are given in the Figure 1 11. It can be observed that the different steps may optimize different metrics. While the rejection steps improve the log(2) estimate more significantly, the JKO steps focus more on the minimization of the energy distance. Overall, we see that in all figures the errors decrease over time.
E.5. Implementation Details
To build our importance corrected neural JKO model, we first apply E N JKO steps followed by n2 E NN blocks consisting out of a JKO step and three importance- based rejection steps. The velocity fields of the normalizing flows are parameterized by a dense three-layer feed-forward neural network. For the JKO steps, we choose an initial step size To > and then increase the step size exponentially by Tk+1 = 4T₃. The choices of #1, 022 Tu and the number of hidden neurons from the networks is given in Table 6 together with the execution times for training and sampling. For evaluating the density propagation through the CNFs, we use the Hutchinson trace estimator with 5 Rademacher distributed random vectors whenever d > 5 and the exact trace evaluation otherwise. Note that we redraw the Rademacher vectors in every step of the ODE solver such that eтrоrs average out throughout the trajectory. We observed that the estimator of the divergence already has a sufficiently small variance, particularly, since the trajectories are almost straight due to the norm penalty of the velocity field, see also Appendix F.2. For implementing the CNFs, we use the code from Ffjord (Grathwohl et al., 2019) and the torchdi f feq library by (Chen, 2018). We provide the code online at https: / i thub. . com/ jonertooortttto r ch / www.eeeuaaaa..aaa Leeeurrll_KK_ii
For the MALA and HMC we run an independent chain for each generated sample and perform 50000 steps of the algorithm. For HMC use 5 momentum steps and set the step size to 0.1 for 8Modes and Funnel and step size 0.01 for mustache. To stabilize the first iterations of MALA and HMC we run the first iterations with smaller step sizes (0.01 times the final step size in the first 1000 iterations and 0.1 times the final step size for the second 1000 iterations). For MALA we use step size 0.001 for all examples. For DDS and CRAFT we use the official implementations. In particular for the test distributions such as Funnel and LGCP we use the hyperparameters suggested by the original authors. For the other examples we optimized them via grid search.
27

---

Importance Corrected Neural JKO Sampling
Table 6: Parameters, training and sampling times for the different examples. For the sampling time we draw N = 50000
samples once the method is trained. The execution times are measured on a single NVIDIA RTX 4090 GPU with 24 GB memory.
Distribution
Parameter
Mustache
shifted Modes
shifted & Pesky
Punnel
GMM-10
GMM-20
GMM-50
GMM-100
GMM-200
GMM441-SD
LGCP
Dimension
50
1600
Number my of flow steps riz rejection Initial step size to I Latent scaling
2
6
4
3
0.05
0.01
0.01
5
0,0025
0.0025
0,0025
0,0025
0.001
1
5
40
Gradient Flow Initialization
3
Hidden neuroes batch
256
70
150
250
512
256
1024
5000
500
Required GPU memocey
5 GB
5 GB
5 GB
GB
5 GB
5 GB
5 GB
G
6B
6GB
11 GB
Training
Sampling time (sec)
15
79 22
23
72
129
473
530
535
Remark E.2 (Initializations). To deal with target distributions living on different scales, we consider initial distributions No = (0, (dd), where c² is a scaling parameter. The choice of for each target distribution is givein in Table 6 in the row "latent scaling". Moreover, it is beneficial in some cases to choose the velocity field the first CNF not by learning it, but by fixing it to (x)) = v. log(g(=)) = v. log(g(z)). We report the final time T this first gradient flow step Table in the row "gradient flow initialization" with value we do not use gradient flow step. This initialization is closely related to Langevin preconditioning, which is used in most neural sampling methods, (He et al., 2025) for detailed study.
A+1 we do not draw in each training step new batch from pk. Instead, we maintain a dataset During the training of t/ of N = 50000 samples of μa and draw batches from it. After the training is completed, we generate a new dataset of N samples by applying the learned CNF onto new drawn samples from uıg. To avoid overfitting effects, we redraw our dataset after each CNF step.
F. Computational Aspects of Normalizing Flows
In this appendix, we give some more details about the computational aspects of the normalizing flows in our model. First, we discuss the relation between multimodal target distributions, mode collapse and non-convex loss functions of normalizing flows. Afterwards, we discuss some computational aspects of continuous normalizing flows like ODE discretizations and density evaluations with trace estimators. Finally, we run some standard normalizing flow networks on our numerical example distributions and compare them with our importance corrected neural JKO sampling.
F.1 Multimodalities, Mode Collapse and Non-Convex Loss Functions
For the sampling application, normalizing flows are usually trained with the reverse KI. loss function F(w) = KT, v). see, e.g., (Marzouk et al., 2016). More precisely, a normalizing flow aims learn the parameters 0 of family of diffeomorphisms by minimizing (o)) In the case that multimodal it can be observed that this training
To such that 1oo mode collapses. That is. the approximation T# covers not all of the modes of / but instead neglects some of them. Examples of this phenomenon can be seen in the numerical comparison in Appendix F.3. One reason for this effect is that the functional F is convex if and only if v is log-concave and therefore unimodal, see (Ambrosio et al., 2005, Prop. 9.3.2). In particular, for multimodal v. the functional F is non-convex. In the latter case of, then the mode collapses can correspond to the convergence to a local minimum, as the following example shows.
Example F.1. We consider the case d = 1 and the target distribution
v = [,N4(1,00005 + $N(
As latent distribution po choose standard Gaussian. Then, parametrize the normalizing flow To = (1 + 8Ti for 0 E 10, where To is the optimal transporl map between μi and and Ti is the optimal transport map between po and 0.052.)
In particular, To Po perfectly recovers the target distribution for 0 = 0 and produces a mode collapsed version of it for 8 1. Now, we plot the reverse KL loss function L(0) = KL #To 1/) and the densities of the generated distributions To pg in Figure 12. We observe that it has two local minima for and 0 = 1 (for f outside of [0. 11, T no longer invertible), where is the perfectly learned parameter and is the mode collapsed version. At the same
0.052).
28

---

Importance Corrected Neural JKO Sampling
time, we note that the curve 0 + To # Po is a geodesic in the Wasserstein space. In particular, the non-convexity of is direct consequence of the fact that F(u) V)) is geodesically non-convex.
Remark F.2 (Motivation of Wasserstein Regularization). This connection between the non-convexity of the loss function F and mode collapse also motivates the Wasserstein regularization from Section 3. By considering the loss function ((u) F(u) + 27 W (1, ( for > 0 small enough instead of F, we obtain a loss function which is convex in the Wasserstein space, see (Ambrosio et al., 2005, Lem. 9.2.7). Even though this does not imply that the map 0 + g ( convex, we can expect that the training does not get stuck in local minima as long as the architecture of To is expressive enough.
# Mo) is
Theoretically, we need for X-convex FF that T S 1 to ensure that G is geodesically convex. However, if μ² is already close to minimum of F, then it can be sufficient that the functional g is convex locally around μ². In this case, the distribution generated by the CNF will stay in this neighborhood. Note that in practice the constant À is unknown such that we start with small step size To and increase it over time as outlined at the end of Section 4 and in Appendix E.5.
E2. Computational Limitations of (Continuous) Normalizing Flows
Similar to the literature on neural JKO schemes (Vidal et al., 2023; Xu et al., 2024), our implementation of the neural JKO scheme relies on continuous normalizing flows (Chen et al., 2018). This comes with some limitations and challenges, which were extensively discussed in (Chen et al., 2018). Since they are relevant for our method, we give synopsis below.
Derivatives of ODE Solutions For the training phase we need the derivative of the solution of an ODE, For this, we use the torchdi ffeq package (Chen, 2018). In particular, this package does not rely on backpropagation through the steps from the forward solver. Instead, overwrites the backward pass of the ODE by solving an adjoint ODE. This avoids expensive tracing in the automatic differentiation process within the forward pass and keeps the memory consumption low. Moreover, the quadratic regularization of the velocity field leads to straight paths such that the adaptive solvers only require few steps for solving the ODE, see (Onken et al., 2021) for detailed discussion. Indeed, we use in our numerical examples the dopri i5 5 solver from torchdi ffeq. which is an adaptive Runge- Kutta method. However, we still have to solve two ODEs during the training time. This still can be computational costly, in particular when the underlying model is large.
Trace Estimation for Density Computations In order to evaluate the (log-) )density of our model, we have to compute the divergence div Ui = trace( (z t)) of the learned vector field v9 (and integrate it over time), see (7). Computing the trace of the Jacobian of vy becomes computationally costly in high dimensions. As remedy, we consider the Hutchinson trace estimator (Hutchinson, 1989), which states that for any matrix A and a random vector = with mean zero and identity covariance matrix it holds that E[=T A2] (r) Applying this estimator to the divergence, we obtain that the divergence coincides with Vrr (2, f), t)2]. The integrand is now again Jacobian- product which can be computed efficiently. Finally, we estimate the trace by empirically discretizing the expectation by finitely many realizations of z.
In our numerics, we choose = to be Rademacher random vectors, i.e. each entry is with probability 1 either - 1 or 1. For the training of the continuous normalizing flow, an unbiased low- lowpprecision estimator sufficient, such that we discretize the expectation with one realization of z. During evaluation, we require higher precision and use 5 realizations instead.
Initialization In order to find stable initialization of the model, we initialize the last layer of the velocity field vo with zeros such that it holds (2, t) = 0 for all x E Rd, 1 E [0, r). In this case, the solution zg of the ODE ép 1g with initial condition 2g (zx 0)- = x is given by z0 t)= = x for all t. In particular, we have that the transport map To = 20(2, is the identity such that the generated distribution To coincides with the latent distribution for the initial parameters.
Remark F.3 (Other Normalizing Flow Architectures). In general any architecture To of normalizing flows can be considered in the neural JKO scheme. To this end, we can replace the velocity field regularization in (5) by the penalizer fE (TT) xr]] d/p By Brenier's 's theorem (Brenier, 1987) this is again equivalent to minimizing the Wasserstein distance. However, while discrete-ting normalizing flows like coupling- based networks (Dinh et al., 2016; Kingma & Dhariwal, 2018) and autoregressive flows (De Cao et al., 2020; Durkan et al., 2019; Huang et al., 2018; Papamakarios et al., 2017) are often faster than CNFs, this is not generally true in our setting, since their evaluation time does not benefit from the OTregullaiization. Additionally, we observed numerically that the expressiveness of discrete-ting -icre architectures scales much worse to high dimensions and that these architectures are less stable to train. Nevertheless, their evaluation can be cheaper since these architectures do not have to deal with trace estimators within the density evaluation. Residual architectures (Behrmann
29

---

Importance Corrected Neural JKO Sampling
Table 7: Comparison of neural JKO IC with normalizing flows trained with the reverse KL loss function. We evaluate the energy distance (smaller values are better), the log(22 eestimatton (larger values are better) and the expected squared empirical Wasserstein distance (smaller values are better).
Energy distance
log, 2) esstimatioo
squared Wasserstein-2
Distribution
NF
coupling
ncural JKO
coupling NF
NF
coupling
neural JKO
Mustache
2.1 10
4.9 x 10
2.9 x 10-3
-7.1 x 10
-22 10
7.3 x 10
3.8 x 10
3.1 10
1.7 x 10
Modes
4.1
1.2 x 10-5
-14x 10"
4.3
15.1 10-6
1.4
10-
10-a
shifted Peaky
5.7 10
6.8 10
3.4 x 10-5
-2. x 16"
-211 10
-2.1 10
1. x 10"
18x 106
7.2 10-a
Funnel
1.9
2.2 x
14x
-9.1% 10
-288 10
-7.1 10
3.8 10
5.1
8.5x
GMM-10
5.4 10
6.3 10
5.2 x 10
-233 109
-233 10
+3.5
10
3.4 10"
35 10P
14 x 10
GMM-20
10
1.1
1.1 x 10-4
-2.1 10"
-2.3 108
+6.4
10
9.5 x 10"
94 10P
3.7 x 10-1
et al., 2019; Chen et al., 2019) are at least similar to continuous normalizing flows, but they are very expensive to train and evaluate, and additionally rely on trace estimators for computing the density the generated samples.
F.3. Numerical Comparison
Finally, we compare our neural JKO IC scheme with standard normalizing flows trained with the reverse KL loss function as proposed in (Marzouk et al., 2016). In particular, we aim to demonstrate the benefits our neural JKO IC scheme avoid mode collapse.
To this end, train two architectures of normalizing flows for our examples using the reverse KL loss function as proposed in (Marzouk et al., 2016). First, use continuous normalizing flow with the same architecture used in the neural JKO (IC). That is, we parameterize the velocity field eg by dense neural network with three hidden layers and the same hidden dimensions as in Table 6. Second, we compare with a coupling network with Glow-coupling blocks (Kingma & Dhariwal, 2018), where the coupling blocks have two hidden layers with three times the hidden dimension as in Table We found numerically that choosing larger architectures does not bring significant advantages.
We plot some generated distributions of the continuous and coupling normalizing flows well as the neural JKO IC scheme in Figure 13. As have already seen in the previous examples, the neural JKO IC scheme is able to recover multimodal distributions almost perfectly. On the other hand, the normalizing flow architectures always collapse to one small number of modes. We additionally report the eгrоr measures in Table 7. We can see that neural JKO IC performs always better than the normalizing flow architectures (note that for the Funnel distribution all W3 -vlvue are below the sampling error reported in Table 3). For multimodal distributions, the normalizing flow approximations are by several orders of magnitude worse than neural JKO IC, while they work quite well for unimodal distributions (Mustache and Funnel).
30

---

Importance Corrected Neural JKO Sampling
2
2
-2
-
-1
-2
2
2
-2
-2
-1
-2
-2
-2
0
-2
Figure 3: Visualization of the steps of the importance corrected neural JKO model for the shifted 8 Peaky example. We start at the top left with the latent distribution and apply in each image one step from the model. The orange color indicates samples which are rejected in the next step and the blue color marks the resampled points.
Mode weight convergence 8 Peaky
mode
mode
mode
0.25
mode
mode
made
weight mass
0.20
rejection layer
normalizing flow layer
truth 0.125
0.15
0.05
10
14
18
layer
Figure 4: Plot of the mode weights for the 8 Peaky example over the different layers of the importance corrected neural JKO model. observe that the weights are mainly changed by the rejection steps and not by the neural JKO steps.
31

---

Importance Corrected Neural JKO Sampling
HBC
MALA
CRAFT
1.0
0.5
0.5
00
00
0.
-05
-05
-05-
-24
-1.
-1.0
-55
0.0
-20
-15
-1.0
-0.
-25
-15
0.0
Veunt PO
Neyill 8o0 K. Jourst
1.0
1.0
1.2
0.5
-10
Figure 5: Sample generation with various methods for shifted Modes example with ground truth samples and generated samples for each associated method. While most methods recover the distribution well, we can see imbalance the modes for the uncorrected neural JKO, CRAFT and DDS.
MALA
0.5
0.5
0.
0.
0.-
-0.
-05
-04
-1.0
-10
-2.
-15
-1.0
-55
6.
-30
-15
-1.0
-0.5
0.0
-32
-15
-2.0
-0.5
0.0
co5
Meurdl KO
Meunal Jio IC Jounn
1.-
1.0
15-
0.5
0.
0.5
0.2
Figure 6: Sample generation with various methods for shifted 8 Peaky mixtures with ground truth samples and generated samples for each associated method. We can see severe imbalance among the modes for HMC, MALA, CRAFT and neural JKO. Also DDS has slight imbalance between the modes.
32

---

Importance Corrected Neural JKO Sampling
HMC
MALA
CRAFT
20
20
20
15
10
10
5
5
0-
0-
-5
-10
-10
-15
-15
-15
-20
-20
-20
-15
5
155
-10
-5
Neural JKO
JKO (ours)
20
20
20
15
15
15
10
10
5
0-
0-
-5-
-5
-5
-10
-10
-15
-15
-20
-20
-20
-15
-10
10
-15
-10
-5
5
10
-15
-10
-5
0
10
Figure 7: Marginalized sample generation with various methods for the = 10 funnel distribution with ground truth samples and generated samples for each associated method. We observe that only the importance corrected neural JKO covers the thin part the funnel well.
33

---

Importance Corrected Neural JKO Sampling
HMC
NALA
CRAFT
10
10
10
-2
-2
-2
DD5
Neural JKD
Neural JKD IC (ours)
10
8-
8-
6-
Figure 8: Sample generation with various methods for the = 2 mustache distribution with ground truth samples and generated samples for each associated method. We can see that MALA, CRAFT and DDS have difficulties to model the long tails of the distribution properly.
34

---

Importance Corrected Neural JKO Sampling
MBC
MALA
CRART
1.0
0.5
0.5
0.
00
0C
-05
-05
-05-
-14
-05
5.0
0.
1.0
-10
-55
0.0
0.5
-15
-05
1.0
Vount Po
Neyial 8o K. Jourst
1.0
10
15
0.5
Figure 9: Marginalized sample generation with various methods for the GMM-10 distribution with ground truth samples and generated samples for each associated method. We observe that CRAFT mode collapses and that only the importance corrected neural JKO model distributes the mass correctly onto the modes.
MALA
0.5
0.5
05
0.
05-
-0.5
-05
-2.5
-1.0
-10
-3.4
-05
0.0
0.5
1.0
-1.0
-05
0.5
0.5
3.0
-322
-0.
0.0
05
1.0
D05
Yeursl KO
Meural BIo IC Jours
3.5
1.0
1.5
0.5
0.
0.5
0.0
-3.6
Figure 10: Marginalized sample generation with various methods for the GMM-200 distribution with ground truth samples and generated samples for each associated method. We observe that CRAFT mode collapses and that only the importance corrected neural JKO model distributes the mass correctly onto the modes.
35

---

Importance Corrected Neural JKO Sampling
Mustache
rejection layer
10a
rejection layer
10
normalizing layer
normalizing layor
energy stance
10:
log estimation error
10a2
10-14
10a2
10
15 layer
20
25
30
10
1.5
20
25
30
layer
shifted 8 Modes
rejection layer normalizing layer
rejection layer normaliziin layer log Z estimation error
10a1
energy staanc
10
10
10"
10
10-
10
10
12
16
18
10
12
14
16
18
layer
layer
shifted Peaky
10a
layer normalizing ng flow layer
rejection layer
1011
normalizing flow layer
energy stance
10:
log estimation crror
10
10:
10-1-
10
10
10
layer
layer
Funnel
rejection normalizing flow layer energy
rejection
normalizing flow layer
10-1
10-
10
layer
layer
GMM-10
10
layer narmalizing flow layer
10a
rejection layer
normalizing flow layer
energy distance
10.
log estimation error
10 2
10a
10-1
2
10
10
20
2.5
layer
# laayer
Figure 11: We plot the energy distance (left) and log(2) estimate (right) over the steps of our importance corrected neural JKO method for different examples. The error measures decrease the beginning and then saturate at some values.
36

---

Importance Corrected Neural JKO Sampling
Non-convex loss function [ (L (To No. v)
Lonturetee Kineeeeniiiiss Kilhes
09
Density of To Po for different values
2
0
Figure 12: Illustration of the loss function [0, 0 + (() KL (o, v) and the densities of the generated
distributions Po from Example F.1. Both values 0 = 0 and 0 1 correspond to local minima (note that f(1) > 0 (())). In particular 0 corresponds to the case of mode collapse.
05
-12
-10
-12
20
-12
-5
-15
crapiing
20pping
1.5
20
45
1.5
POIC
-14
80
18
45
1.5
25
8 Modes
Peaky
GMM-10
GMM-20
Figure 13: Marginalized sample generation for single normalizing flow compared with neural JKO IC for different example distributions with ground truth samples and generated samples for each associated method. We observe that the standard normalizing flow architectures always collapse to or few modes while neural JKO IC recovers all modes correctly.
37


Reviews:
[{"review_id": "NNJkcOhj6o", "reviewer": "Reviewer_BUrR", "summary": "This paper proposes to sample from an unnormalized probability density via a sequence of interleaved continuous normalizing flows (CNFs) and importance accept/reject steps. The CNFs, which are penalized with a velocity norm regularizer as in OT-Flow (Onken et al. 2021), are interpreted as Wasserstein proximal mappings applied to the reverse KL divergence loss functional by replacing the static form of the $W_2$ distance typically used to define the proximal mapping with its equivalent dynamic formulation. The authors use this interpretation to show convergence of their OT-regularized CNF velocity fields to the velocity field corresponding to the Wasserstein gradient flow of the reverse KL divergence as the proximal mapping step-size $\\tau\\to0$ and that, for nonzero $\\tau$, the velocity fields at each step correspond to the OT velocity fields between the starting and ending measures at each step. The CNF scheme is implemented by representing the velocity fields as neural networks and optimizing the parameters of the networks to minimize the dynamic OT-regularized CNF loss; hence the scheme is termed “neural JKO.”\n\nTo address the issue of slow or incorrect convergence known to plague CNFs when the target density is multimodal, the authors propose to insert “importance-based rejection steps” in between some of the CNF/Wasserstein proximal steps. These rejection steps consist, essentially, of one step of rejection sampling: each particle $X$ from the current ensemble (resulting from previous CNF and rejection steps) is rejected with probability $1 - \\alpha(X)$, where $\\alpha(X) = \\min \\left\\{ 1, \\frac{g(X)}{cf(X)} \\right\\}$, $g(X)$ is the unnormalized density of the target measure, $f(X)$ is the unnormalized density of the current ensemble, and $c > 0$ is a tuning parameter. Thus, each particle has a lower chance of getting rejected if the importance weight $\\frac{g(X)}{f(X)}$ is large. If the particle is rejected, it is replaced by repeating the entirety of the previous CNF/rejection procedure (starting from the reference density) to generate a new sample from the current particle distribution. The authors show that it is possible to write the new density of the particle ensemble after it has been transformed by this one-step rejection procedure and that the rejection sampling step decreases the KL divergence to the target. The density information is carried through the rest of the CNF/rejection sampling procedure, and as such allows the procedure to be used for density estimation as well.\n\nNumerically, the method is exercised on a variety of target distributions in between two and 1600 dimensions and shown to generate samples which result in lower energy distance and higher estimates of the “log normalizing constant” in comparision to MALA, HMC, DDS, CRAFT, and Neural JKO. The experiments in the paper all, in the end, generate an ensemble of 50,000 samples, and the quality metrics for the samples are likewise computed over 50,000 samples.\n\n## update after rebuttal\n\nI thank the authors for their response. I have a positive view of this paper and will thus maintain my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your very detailed and thoughtful review. Please find the answers to your questions and comments below. Additionally, we will correct the typos and grammar errors.\n\n## Questions\n\n1. Our choice of the schedule is based on the following heuristic: Since the rejection layers are (on a long term) more expensive than the CNF layers, we start with using only CNF layers until the progress made by them becomes small. This can be detected without knowing the target distribution, because we have (approximately) access to the Wasserstein distance between input and output of the JKO layer via the norm of the velocity field. Afterwards, we alternate one CNF layer and three rejection steps. Numerically, the CNF layer only does minor local adjustments from this point, while the main work is then done by the rejection steps.\n\n2. If the step size is too large, we might miss some modes of the target distribution (which matches with the theoretical consideration that we loose convexity (with respect to the Wasserstein metric) of the loss in this case). If the step size is small, then the mapping learned by the CNFs is close to the identity and does not have a large effect. While the model is sensitive with respect to the first case, the second case only increases training and sampling time of the model (since more CNF steps are required), but does not effect much the quality of the result. A simple tuning heuristic for the step size can again be established based on the (approximated) Wasserstein distance: We start with a small step size and increase it as long as the Wasserstein distance between input and output of the next CNF step is smaller than a certain threshold. While we have not automatized this procedure so far, we are quite sure that such an adaptive choice of the step size can be established.\n\n3. The curse of dimensionality indeed not specific to normal distribution, but to any proposal which is not close enough to the target. So in fact the use of tailored/CNF proposals can help to reduce or even completely avoid the curse of dimensionality. We will highlight this further. \n\n## Number of Samples\n\nOur main intention behind the large number of samples is the following: Since the samples generated by our model are independent, our main quantity of interest is how much the distribution of a generated samples differs from the target distribution (in MMD, Wasserstein etc.). Of course we cannot measure this quantity directly, but taking a finite number of samples from both distributions. Using a small number of samples introduces an error in this estimator, which is much more related to the properties (sample complexity) of the evaluation metric than with the approximation quality with the model. Therefore, from our viewpoint, the number of samples should be chosen as large as possible.\n\nWhile the $\\log(Z)$ estimation can be viewed in the same way (just with the KL divergence), we agree that this can also viewed differently. If we are not interested in the KL divergence between generated and ground truth distribution (which was our viewpoint so far) but in the $\\log(Z)$ estimate itself, we agree with you that reporting the errors for different numbers of samples could be useful. We will include such a study in the final version. In our first experiments with $5000$ samples the greater picture is about the same as for $50000$ samples.\n\n## Proof Theorem 4.2 (i)\n\nYes, you are right, the expression $P(X=x)$ is meant in a weak sense (where we integrate against all measurable sets $A$) and that this part is written a bit sloppy. We will write it down more formally for the final version. However, neither the proof nor the result change.\n\n## Visualization etc.\n\nThank you for the suggestions. We will add a flow chart-like diagram describing the sampling process to the paper. For the numerical part, we think that box plots often make it harder to assess the details, even though we agree that they give a better first impression. We are happy to add a box plot for visualizing the numbers, but we prefer to keep the tables as well (if the space constraints allow it, we will keep both in the main part of the paper).", "claims_and_evidence": "All claims are sound and substantiated by theoretical and empirical results.", "methods_and_evaluation_criteria": "The proposed method on the whole is sensible in the context of similar, previously successful methods for sampling from unnormalized probability densities (e.g., continuous normalizing flows and OT-regularized variants, annealed flow transport Monte Carlo, etc.). In order to use the method, one must have access to the unnormalized density of the target measure and its score. One potential methodological issue, which is briefly mentioned in remark 4.5, is that the rejection steps rely on the generation of new samples from the current particle distribution. In the context of the proposed method, this generation implies re-sampling from the reference distribution and re-simulating the previous sequence of CNF and rejection layers. The CNF layers shouldn’t pose too much of an issue in themselves, as they involve integration of previously identified ODEs with relatively straight trajectories, but the previous rejection layers have the potential to dramatically increase the time required to regenerate another sample. That is, if a sample is rejected at any layer, we must start over at the reference distribution and re-do all of our previous CNF/rejection layers. If the acceptance probability is $1 - r$, then the probability of successfully passing through $n$ previous rejection layers is $(1 - r)^n$. This feature may result in long runtimes, as we see for the higher dimensional test distributions in Table 5.\n\nIn the exposition of the paper it is not made explicitly clear how this resampling is performed in the rejection steps -- i.e., the rejection procedure presented in section 4 is general and never specified for the method at hand. The reader must, therefore, put the pieces together on their own. I think it would be helpful to explicitly outline how one resamples in the context of this neural JKO-IC method so as to give a clearer picture of what is at stake in the rejection steps. A diagram of an example sample trajectory might even be helpful for this purpose (perhaps with a background similar to that of the plots in Figure 10, but with arrows indicating the procession of one particle through the layers, with at least one rejection step depicted).\n\nThe quality of the samples generated is evaluated using the energy distance (MMD with the negative distance kernel) and by estimation of the log normalizing constant. These evaluation metrics are consistent with those used for similar sampling methods in the literature. In the appendix there are plots of 2D marginals of the samples generated by the various methods, which are also helpful. The method is primarily compared to the Metropolis-Adjusted Langevin Algorithm (MALA), Hamiltonian Monte Carlo (HMC), the Denoising Diffusion Sampler (DDS), Continual Repeated Annealed Flow Transport Monte Carlo (CRAFT), and Neural JKO (equivalent to the proposed method without the rejection steps). MALA and HMC are standard MCMC approaches, while DDS, CRAFT, and Neural JKO are all, to some extent, based on dynamic measure transport. This collection of comparison algorithms is reasonable, and in the appendix there are additional  comparisons to un-regularized continuous normalizing flows (CNFs).", "theoretical_claims": "I checked the proofs of Corollary 3.2, Theorem 3.3, Theorem 4.2, and Corollary 4.3 at a high level. My only quibble is that in the proof of Theorem 4.2, the density $\\tilde p$ is interpreted as a probability, i.e., $\\tilde p(x) = \\mathbb{P}(\\tilde X = x)$,  which isn’t quite right. However, I suspect that if the proof were repeated working with $\\mathbb{P}(\\tilde X \\in A)=\\int_A\\tilde p(x)\\,\\mathrm{d}x$ for some measurable set $A$ the result would still emerge.", "experimental_designs_or_analyses": "The experimental design on the whole seems sound, but I think it would be interesting to compare the proposed Neural JKO-IC method to the competing methods for varying ensemble sizes. As far as I can tell, all experiments performed involved the generation of $N = 50,000$ samples, which is quite large and should render any Monte Carlo estimation error in the losses for the various methods small. Thus, the error metrics reported in the results are likely dominated by the biases of the various methods. What happens if the ensemble size is decreased, for example, to 50, 500, or 5,000 samples? Does Neural JKO-IC also have low variance relative to the competing methods at these smaller ensemble sizes? At the very least, the particle ensemble size should be listed in the main part of the paper (not just the appendix) to give proper context for the results.\n\nThere is also no comparison or discussion of the *computational cost* among the various methods. Training and runtimes are given for Neural JKO-IC are given in Table 5, but no training or runtime information is available for the competing methods. Knowing the relative computational costs of each of the methods is important for contextualizing numerical performance and informing choice of method in application.", "relation_to_broader_scientific_literature": "The connection between OT-Flow (Onken et al. 2021) and Wasserstein proximal mappings was highlighted in (Vidal et al. 2023), and the Neural JKO scheme (absent the rejection layers) which forms the basis for the proposed method was also proposed in Vidal et al. 2023. Annealed Flow Transport Monte Carlo (AFT, Arbel et al. 2021) and Continual Repeated Annealed Flow Transport Monte Carlo (CRAFT, Matthews et al. 2022) are similar in spirit to the proposed method in that they involve the interleaving of normalizing flow layers with importance resampling and mutation steps. One big difference between AFT/CRAFT and the proposed method is that the former methods are designed follow a specific annealing path of densities, while the latter seeks a sequence of densities corresponding to Wasserstein proximal mappings. Moreover, the importance resampling steps used in AFT/CRAFT fundamentally cannot alter the support of the empirical particle distributions, while the rejection sampling approach, while perhaps more costly, generally does modify the support, and evidently to good effect. Importance weights have also been combined with normalizing flows in, e.g. in (Noe et al. 2019), where normalizing flows are used to define a proposal distribution for importance sampling, but this approach is not iterative like AFT/CRAFT and the proposed method.\n\n## References\n\n- Onken, D., Fung, S. W., Li, X., and Ruthotto, L. OTflow: Fast and accurate continuous normalizing flows via optimal transport. In AAAI Conference on Artificial Intelligence, volume 35, pp. 9223–9232, 2021.\n- Vidal, A., Wu Fung, S., Tenorio, L., Osher, S., and Nurbekyan, L. Taming hyperparameter tuning in continuous normalizing flows using the JKO scheme. Scientific Reports, 13(1):4501, 2023.\n- Arbel, M., Matthews, A., and Doucet, A. Annealed flow transport Monte Carlo. In International Conference on Machine Learning, pp. 318–330. PMLR, 2021\n- Matthews, A., Arbel, M., Rezende, D. J., and Doucet, A. Continual repeated annealed flow transport Monte Carlo. In International Conference on Machine Learning, pp. 15196–15219. PMLR, 2022.\n- Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457):eaaw1147, 2019.", "essential_references_not_discussed": "No, all the works that came to mind when I was reading this submission were appropriately cited.", "other_strengths_and_weaknesses": "### Strengths\n\n- I was previously unaware of the connection between OT-regularized CNFs and Wasserstein proximal mappings, and I enjoyed learning about it by reading this paper. While this viewpoint was already advanced in (Vidal et al. 2023), the convergence results and the viewpoint of “piecewise geodesic interpolation” in this work are, to my knowledge, new.\n- While the neural JKO numerical approach already appeared in (Vidal et al. 2023), the addition of the importance-based rejection steps to the neural JKO scheme is new and does seem to improve sample quality significantly.\n- The rejection steps are relatively easy to add if the neural JKO steps have already been implemented.\n\n### Weaknesses\n\n- The numerical examples do not exercise the method across a range of sample sizes.\n- The rejection steps have the potential to cause long runtimes of this method.\n- The numerical results are difficult to quickly interpret due to the tabular format.", "other_comments_or_suggestions": "*(In this section I append the line numbers with “L” to indicate that the text in question is in the left column, and “R” to indicate that it is in the right column.)*\n\nThe preliminary material presented in section 2, while providing a firm mathematical background for the methods presented, is a little bit technical, and I wonder if the level of technicality could be lightened in order to make the material more approachable to a broad machine learning audience. For example, on line 146 (R), the minimal norm velocity field is described as belonging to a “so-called regular tangent space.” While this statement is true, it would be more evocative and approachable to simply state that the velocity must be a gradient. Similarly, the definition of Wasserstein gradient flows could be explained in the main body of the paper without invoking reduced Frechet subdifferentials (velocity is the negative gradient of the first variation of $\\mathcal F$…), with the more technical definition saved for the appendix. Level of technicality is of course a matter of taste, but I wonder if toning down some of the detail in the main body of this paper would help it achieve broader appeal.\n\nPresenting the numerical results in tables makes it difficult to quickly assess the relative performances of the methods, especially given that all of the metrics are in scientific notation. Consider using bar charts, box plots, or another suitable plot instead of tables for the numerical results; these plots would make it easy for the reader to quickly compare the performances of each method without having to squint at nine rows of seven exponents.  \n\nThere are some typographical errors or notational inconsistencies which need to be corrected:\n\n- Lines 152-153 L: the integral in the definition of $\\mathcal{P}_2(\\R^d)$ should be $\\mathrm{d} \\mu(x)$, not $\\mathrm{d} x$\n- Line 157 L and elsewhere: it is confusing to use $\\bf \\pi$ to denote the coupling between $\\mu$ and $\\nu$ and also to use $\\pi_1$ and and $\\pi_2$ to denote projection of $\\bf \\pi$ onto the first and second coordinates. Consider using $\\gamma$ or a different letter for couplings instead ($\\gamma$ matches $\\Gamma(\\mu, \\nu)$).\n- Line 184 R: Here the normalized target density is written $q(x) = Z_gg(x)$, but it should be $q(x) = g(x)/Z_g$ to be consistent with the setup in Section 1 (Line 43 L).\n- Line 190 R: It seems like $\\propto$ is used to indicate equality up to an *additive* constant, but I think this notation is more commonly used for equality up to a multiplicative constant. Perhaps consider alternate notation, or at least specify that the constant is additive.\n- Line 246 R: should the stopping time be $\\tau$ and not $T$ in $\\mu_\\tau^{k+1} = z_\\tau^k(\\cdot, T)_\\sharp \\mu_\\tau^k$?\n- Line 276/Equation (7): the top row of the matrix on the RHS should be $v_\\theta(z_\\theta(x, t), t)$, not $v_\\theta(x, t)$. This error is also in Algorithm 6.\n- Line 298 L: Should this read “this leads to slow *convergence* speeds”?\n- Line 311 L: The base/acceptance rate should be $1-r$, not $1+r$\n\nThere are also spelling and grammatical errors throughout the paper. Below is a list of the ones I noticed, but I would suggest performing a thorough check for others which I may have missed.\n\n- Line 78 L: “the rejection steps *readjust”*\n- Line 83 L: subject-verb agreement needed here: either “our methods *generate*… and *allow*” or “our *method* generates… and allows”\n- Lines 91-92 L: “the velocity fields… *converge*”\n- Lines 97-98 L: “we consider neural JKO schemes *in more detail”*\n- Line 132 L: “directly *following”*\n- Line 148 R: “*absolutely* continuous”\n- Line 322 R: “a constant *ratio”*\n- Line 364 L: “and *which* we can sample from”\n- Line 374 L: I’m not sure that “vice versa” is quite the correct term here. Maybe you are looking for “likewise” or “moreover” instead?\n- Line 380 R: “it is a kernel *metric”*\n- Line 422 L” “distance between *two* sets of…”\n- Line 1052: “Since it holds $\\tilde X = X$ if $X$ *is accepted* and…”\n- Line 1256: “we evaluate the Wasserstein distance based on *fewer* samples”\n- Line 1372: “we run an independent chain for each generated *sample”*\n- Line 1452: “becomes *computationally* costly”\n- Line 1470: “*Brenier’s* theorem”\n- Line 1474: “We observed numerically that the expressiveness of discrete-time architectures *scales…* and *that* *these architectures* are less stable to train”\n- Line 1475: “The evaluation can be cheaper and the density evaluation” -- I’m not sure what this sentence is trying to say. Perhaps that “the sampling and density evaluation can be cheaper”?\n- Line 1476: “Residual architectures *are”*\n- Line 1477: “they are very *expensive*”\n- Line 1505: “On the other side” -- perhaps you mean “on the other *hand*”?", "questions_for_authors": "1. How does one choose the “schedule” of rejection layers and CNF layers? Would a heuristic like ESS, which is used in CRAFT/AFT, be useful here?\n\n2. How sensitive is the method to the initial choice of step-size $\\tau_0$ and step-size schedule?\n\n3. Line 289 R: Is the curse of dimensionality in rejection sampling specific to use of a standard normal proposal distribution? Or would the issue occur with any fixed/non-tailored proposal? If it is the latter, you may want to emphasize this point to further highlight why the use of the tailored/CNF proposals is useful.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hUqJhdkpsj", "reviewer": "Reviewer_JGqE", "summary": "This paper applies the Wasserstein Gradient Flow (WGF) framework to the sampling problem, i.e. sampling from a given target distribution. The proposed approach consists of two key stages:\n\n**Stage 1**: JKO Steps with Continuous Normalizing Flows (CNFs)\n- Given a terminal density, the authors perform Jordan–Kinderlehrer–Otto (JKO) steps, each of which corresponds to an ODE-based control problem.\n- These steps are learned through Continuous Normalizing Flows (CNFs), which parameterize the velocity field.\n- However, a notable drawback of this approach is its slow and suboptimal convergence, primarily due to the complexity of solving high-dimensional JKO problems.\n\n**Stage 2**: Importance Rejection Sampling Enhancement\n- To address the inefficiency of pure JKO-based updates, the authors incorporate an rejection sampling scheme that alternates with JKO steps.\n- Due to the resampling, it shows faster convergence, requiring only few steps till they reach the convergence. Moreover, it converges to more optimal solution. \n- The resulting framework iterates between JKO-based updates and importance sampling, aiming to improve convergence speed and sample quality.\n\nThe proposed method is evaluated on several sampling benchmarks including LGCP and funnel energy functions.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the detailed evaluation of our paper. Please find our comments below.\n\n## Literature\n\nThe process of research is an active field and therefore many interesting contributions are published frequently in particular in this modern field. However, we want to kindly point out, that in accordance to the ICML guidelines, **comparisons to very recent preprints, which are not published yet (like [3,4]) or published a week before the submission deadline (like [2]) or haven't even been available as a preprint until the submission deadline (like [4]) cannot be expected.** We do our best to provide timely comparisons and have now even added a comparison to [2] (details below), but we **can't compare to methods which appeared on arxiv shortly before or even after the ICML submission deadline**. Of course, we are happy to include all papers mentioned by the reviewer in the \"related work\" section (where [1,3] are already cited).\n\n## Computational Cost\n\nWe respectfully disagree with the reviewers opinion that our method would be computational infeasible. Specifically we would like to stress the following points:\n\n- The evaluation of $\\mathrm{trace}(\\nabla v_t^k(x))$ is estimated by a Hutchinson trace estimator ($\\mathrm{trace}(A)=\\mathbb E[z^TAz]$) for $z$ with zero mean and identity covariance matrix. \n- During the training of the CNF a **single sample of $z$ is enough** in the Hutchinson estimator to approximate the trace (since we only need an unbiased estimate of the loss for training the CNF). For the evaluation we take 5 samples of $z$ for each time step. We find that this already leads to a small variance of the resulting densities, while the cost remains comparably low (see also the answer to Question 5 of Reviewer Eq2H).\n- When training $v_\\theta^{k+1}$, we do not draw in each training step \"a fresh batch from $\\mu_k$\". Instead, we maintain during training time a dataset of 50000 samples of $\\mu_k$. During training of $v_\\theta^{k+1}$ we then draw batches from this dataset. After the training of $v_\\theta^{k+1}$ is completed, we generate a dataset of $\\mu_{k+1}$ by applying the CNF onto the samples from $\\mu_k$. We add this detail to the numerical details.\n- During evaluation we **don't need to draw a large batch**. During test time the samples don't interact with each other (see also our reply to the part \"Batch size\" in the paragraph \"Questions\" below)\n- Please note that we **already discussed the computational aspects of CNFs** in Appendix F.2 and list the resulting **training and evaluation times including the required GPU memory in Table 5 in the appendix**. We can see that they remain moderate for all examples considered in the paper.\n\nNevertheless, as already discussed in the limitations paragraph, lowering the computational cost is part of our current and future research. In particular the use of supervised trained generative models as intermediate surrogates for several steps in our model is future work and goes beyond the scope of the paper.\n\n\n## Other Benchmark Problems and Comparisons\n\nWe believe that our test problems are standard among the community and respectfully disagree with the reviewers opinion that \"no sharp benchmarks\" were used. However, upon the reviewers request, **we now ran our method on the GMM40 problem in $d=50$**. For the sake of availability of comparisons, we adapt **the same setting as Table 3 in [2]** and also evaluate the Sinkhorn distance. We can see in the table below that the neural JKO IC outperforms DDS, CRAFT and SCLD on this example.\n\n\n|  | DDS  | CRAFT | SCLD [2] | Neural JKO IC (ours) |\n| :---: | :---: | :---: | :---: | :---: |\n| Sinkhorn distance | $5435.18$ | $28960.70$ | $3787.73$ | $3154.91$ \n\n\n\n## Questions\n\n1. **Stability:** In the case that hyperparameters are chosen inappropriately, the method might miss some modes of the target distribution, but we never observed divergence issues. Once the hyperparameters are chosen appropriately, it consistently produces the same results.\n2. **Ablations:** The paper **already contains an ablation study** with respect to the number of steps in Figure 10, where we plot the error measures over the number of steps. Since our model is trained iteratively, the plots show how the model would behave if we stop the training earlier and we can see that the error measures saturate.\n3. **Batch size:** We would like to clarify that at test time **the sampling procedure is independent of the batch size**. The only parameter in the rejection step depending on several samples, in particular depending on more than one sample is $\\mathbb E[\\alpha(X_k)]$ (it is estimated within the training phase and fixed afterwards). For the evaluation **the samples do not interact with each other** and are consequently independent.", "claims_and_evidence": "The idea of combining importance sampling with a Wasserstein Gradient Flow framework is novel and conceptually interesting. \nTheoretical claims are correct.\n\nHowever, I see major computational challenges in the proposed method:\nA critical aspect of this method is that it heavily relies on accurate trace estimation of the gradient of the velocity field in CNFs, which is required for both the JKO update and importance sampling correction. I have strong concerns regarding the computational feasibility of the proposed approach:\n\n-\tTo obtain a sample from $\\mu_k$, an initial sample $x$ must be passed through $k$ different CNF models $v^i_\\theta$ ($i =1, 2, \\dots k$). Furthermore, it is necessary to estimate the trace of the Jacobian $Tr(\\nabla v^k_\\theta (x))$ along the whole trajectory to perform importance sampling. \n\n- Updating $v^{k+1}_\\theta$ requires drawing a fresh batch of samples from $\\mu_k$ at every iteration, leading to excessive computational overhead. \n\n- The training of CNFs itself is demanding since it involves the integration of the trace $Tr(\\nabla v^k_\\theta (x))$, which further exacerbates the computational burden.\n\n- Even in the evaluation procedure, it still requires to compute the integration of the trace. Moreover, due to the importance sampling, the sample should be drawn by large batch. \n\nDue to this computational burden (in both training/evaluation scheme) and due to the heavy reliance on importance sampling, I guess that JKO-IC (the proposed method) will possess scalability issue.", "methods_and_evaluation_criteria": "My major concern is its computational efficacy as discussed in previous section \"Claims and Evidence\". Accordingly, I don't see the potential application to more realistic dataset. To verify its applicability, I would like to further ask authors to include the following:\n\n- Comparison on the training time, evaluation time, and GPU storage with other benchmark approaches? \n\n- Discussion on the advantages of the proposed method over existing importance weighted based schemes like [1 ,2, 3]?\n\n- I believe there should be a comparison with [1] and [2] for several datasets (e.g. funnel, LGCP). These are also importance sampling based methods. \n\n- I feel that there are no sharp or realistic energy function benchmarks discussed in the paper. Could the authors run on 40-modes GMM implemented in [3] or [4], or more realistic datasets DW4 or LJ13 discussed in [5]? (The implementation of GMM in this paper is different from original 40-mode benchmark.)", "theoretical_claims": "I’ve checked all the theories and proofs. I believe the claims are all correct.", "experimental_designs_or_analyses": "As aforementioned, I would encourage authors to verify the followings:\n\n- Computational Efficacy Study.\n\n- Discussion on related works including [1], [2], and [3].\n\n- Comparison with [1] and [2].\n\n- Comparison on additional benchmark dataset.", "relation_to_broader_scientific_literature": "The sampling problem can be reformulated in various ways. This paper suggests a new sampler by incorporating rejection sampling into the original JKO scheme. I believe this idea is very interesting and worth further investigation to improve its efficacy.", "essential_references_not_discussed": "I believe a discussion and direct comparison with [1] and [2] is necessary to further improve the paper. \n\n**References**\n\n[1] Phillips, Angus, et al. \"Particle Denoising Diffusion Sampler.\" ICML, 2024.\n\n[2] Chen, Junhua, et al. \"Sequential Controlled Langevin Diffusions.\", ICLR, 2025.\n\n[3] Albergo, Michael S., and Eric Vanden-Eijnden. \"Nets: A non-equilibrium transport sampler.\", preprint, 2024.\n\n[4] He, Jiajun, et al. \"No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers.\", preprint 2025.\n\n[5]  Akhound-Sadegh, Tara, et al. \"Iterated Denoising Energy Matching for Sampling from Boltzmann Densities.\" ICML, 2024.", "other_strengths_and_weaknesses": "Strength\n- The paper is well-organized, easy-to-follow.  \n- The idea of integrating importance sampling with JKO updates is a new fresh idea.\n- The paper is well-written and provides a rigorous theoretical foundation.", "other_comments_or_suggestions": ".", "questions_for_authors": "- Did the authors observe any instability or divergence issues when alternating between JKO steps and importance sampling?\n\n- Ablation Studies: What happens if we reduce the batch size (affects importance rejection sampling)? Also, could authors provide ablation studies on number of flow steps (longer flow steps)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1ZHxqnVlbB", "reviewer": "Reviewer_Eq2H", "summary": "This paper contributes a method called \"Importance corrected neural JKO sampling\", based on the well-established Jordan-Kinderlehrer-Otto (JKO) scheme. The method is constituted by a flow-based ordinary differential equation (ODE), which is parameterized by neural networks and learned using standard neural ODE optimization techniques. The authors approximate this solution by using a finite sequence of proximal maps; they show that, under some assumptions on the target distribution, the neural network that minimizes a given loss over the ODE will minimize the reverse KL w.r.t. the target, and produce samples via a flow mapping. In particular, the solution asymptotically approaches the Wasserstein gradient flow and approximates the flow discussed in the well-established theorem from Benamou and Brenier. However, the authors identify that this sampling procedure can be inefficient due to the behavior of the reverse KL divergence, used in variational inference techniques like neural ODEs; to ameliorate these inefficiencies, they adapt this neural JKO scheme by contributing a method hybridizing importance and rejection sampling. This method samples from a proposal distributions and uses a weighted rejection step, whose corrected distribution is consistent with the target and is closer to the target than the proposal in the prescribed reverse KL. Further, due to its construction, the *corrected* distribution's density is known analytically, which allows the user to know the normalized density when retrieving the target samples from this procedure. This procedure gives typical Monte Carlo convergence rates that are nominally dimensionally-independent. Finally, the authors provide numerical results demonstrating competitive results with several typical methods on many benchmark problems spanning low- to high-dimensions.\n\n## Update after rebuttal\nI thank the authors for their response. I maintain my score to 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for the detailed and thoughtful review. Please find our answers below. For the final version, we will additionally correct the typos, extend the literature part (e.g. with SVGD), improve the visual accessibility based on your comments and add the definitions of lsc/coercive/$\\lambda$-convexity in the appendix.\n\n## On the $\\lambda$-Convexity\n\nWe stress that for the theoretical results we only require that the negative log-density of the target is $\\lambda$-convex for some $\\lambda\\in\\mathbb R$ **which explicitly includes negative $\\lambda$**. This assumption is very weak and **automatically fulfilled when $q$ is smooth enough** (plus some asymptotics for $\\\\|x\\\\|\\to\\infty$). In particular, the theoretical results are **also applicable for target densities which are not log-concave**. We will add this explanation after Assumption 3.1 in the paper.\n\n## Questions\n\n1. We stress that in Section 3 we consider the theoretical results (Cor 3.2, Thm 3.3) to be the main novelty of our paper. The actual neural JKO steps are indeed very similar to approaches in the literature, particularly, to the papers of (Vidal et al., 2023; Xu et al., 2024) which also rely on the dynamic formulation of the Wasserstein distance. Since these papers consider the generative modeling setting there are some small technical differences (e.g., we start with the latent distribution while Vidal et al., Xu et al. start with the data distribution; moreover they only consider the Gaussian as target $g$, which makes the whole objective convex regardless of $\\tau$, but is not useful for the sampling application). In the final version, we will clarify the relation to these approaches more detailed in the beginning of Section 3.2, where we derive the neural JKO step.\n2. The statements of Cor 3.2 and Thm 3.3 are not directly related to the rejection steps in Section 4. However, in order to be able to combine the importance-based rejection step in an iterative manner, we need the following requirements on our sampling model: evaluate the density, sample independently, start at an arbitrary latent distribution, avoid mode collapse. Since not many generative models (or sampling methods) fulfill these properties at the same time, we focused on the neural JKO scheme. We will clarify this in the introduction.\n3. It is not surprising that the neural JKO scheme and Langevin sampling produce very similar results, since both approximate the same Wasserstein gradient flow with respect to the KL divergence in the limit. The main difference is that we **can evaluate the density of the distribution generated** by the neural JKO scheme, while we **can't do that for the distribution generated by Langevin sampling**. Thus, we can combine the neural JKO scheme with our importance-based rejection resampling steps, while we can't do that with the Langevin steps. We briefly outlined this relation in the introduction (l 47--50 right side), but will add a reminder to that in the numerical results. Regarding ULA vs. neural JKO: When running our experiments, we also ran a plain ULA sampling (without Metropolis correction). The results are mostly similar to the MALA results despite the Metropolis correction helps a little bit to fill out the narrow tails of the funnel/mustache distribution. Therefore, we omitted ULA in the paper.\n4. The most important hyperparameters are the initial step size $\\tau$ and the size of the network architecture, which require some tuning. If those parameters are chosen inappropriately, the model might miss some modes (which matches the theory, since we loose any convexity of the JKO steps in this case). We outline a possible tuning strategy in the answer to Reviewer BUrR (Question 2). But most hyperparameters from the neural JKO steps are quite standard choices (the velocities are dense 3-layer neural networks, we use the `torchdiffeq` library with rather standard parameters, the training parameters like batch size, learning rate etc. are also quite standard choices for CNFs).\n5. We think that we can resolve this confusion: When computing the Jacobian trace for **5 new Rademacher vectors in each time-discretization step** of the ODE. Since the ODE trajectories are almost straight and therefore the Jacobian only slightly varies over the different time steps. Consequently, the effective number of considered Rademacher vectors is **5 times the number of time discretization steps** which is (depending on the example) in a order of magnitude of 20 to 100. Therefore the errors mostly cancel out over the whole solution of the ODE. Indeed, we observed that the variance of the estimator over the whole solution of the ODE is already very small for 5 Rademacher vectors per time step and taking more does not have much of an effect. Taking less Rademacher vectors causes a bias in the importance weights and lowers the quality of the results. We will add this explanation to the final version of the paper.", "claims_and_evidence": "The claims in this submission are reasonably clear and well-supported, though their consequences are a little under-discussed. A mild criticism is that the provided theoretical results largely hinge on log-densities that are $\\lambda$-concave, with discussion of the non-convexity of the reverse KL and its tendencies to seek modes; however, the results for the importance correction show improvement in this reverse-KL regime. While such a descent property is comforting, it seems undercut by the qualifications the authors lay out regarding the loss the authors choose to descend.", "methods_and_evaluation_criteria": "The methods employed and evaluation metrics seem well-founded. If these results are published with further details, one suggestion would be to report results using MMD with various kernels, rather than relying solely on the energy distance.", "theoretical_claims": "The theoretical proofs were only briefly considered. The simpler proofs, i.e., Corollary 3.2, Theorem 4.2 (i), and Corollary 4.3 were checked, and Theorem 4.2 (ii) was also considered informally. The remaining proofs were not.", "experimental_designs_or_analyses": "The provided experiments largely demonstrate a comprehensive approach to the problem at hand. A slight shortcoming that might be suggested is that, as the importance/rejection step is not tethered to the gradient flow regime, seeing if that performs any different when used in conjunction of Langevin or Hamiltonian dynamics. This would allow the reader to see how much the error decreases due to a good importance correction versus the quality of the JKO. For instance, if the importance corrected Langevin dynamics outperformed MALA substantially due to sample independence, then this is a notable result in and of itself.", "relation_to_broader_scientific_literature": "This reiterates several results and ideas in neural ODEs (e.g., [Chen et al., 2018] in NeurIPS) and JKO sampling literature (e.g., [Jordan et al., 1998] for its introduction, as well as, e.g., [Salim et al., 2020] in NeurIPS for the discussed connection to proximal algorithms). Further, its proposed methods augment typical variational inference techniques, e.g., [Marzouk et al. 2016] in Handbook of UQ, [Lambert, 2022] in NeurIPS. The discussion of functional geometry over the reverse KL divergence echos prior literature, e.g., [Marzouk et al. 2016] and [Grenioux et al., 2023] in ICML, but is clear and reasonably concise in the prior literature's connection to the current methods. Regarding neural JKO details, the authors provide a convincingly comprehensive collection of works reflecting progress towards neural network-based JKO sampling schemes, e.g., [Altekrüger et al., 2023], [Mokrov et al., 2021], [Onken et al., 2021], and [Xu et al., 2024] (publication venues provided in submission). Each of these seems to tackle varying problems of using JKO for sampling unnormalized densities. For instance, [Altekrüger et al., 2023] seems to tackle a flow using MMD-based discrepancies, [Mokrov et al., 2021] seems to minimize the reverse KL divergence (the expression of the functional is in (5), where the constant $\\beta$ seems to be chosen according to (13)), but do so using an explicit construction of transport maps using a discrete sequence input-convex neural networks, [Onken et al., 2021] seems to work very similarly to this paper using neural ODEs but omits the regularization term this submission's authors denote $w_\\theta$ (which corresponds to the Wasserstein proximal mapping). The authors do not seem to articulate the differentiation of their submission's neural JKO step from [Xu et al., 2024] in terms of approaching the JKO from a neural ODE perspective; the reviewer could not distinguish between the two in the time allotted. In particular, (8) from [Xu et al., 2024] seems virtually identical to (7) in this submission.", "essential_references_not_discussed": "The importance-correction rejection steps closely mimic the accept-reject step in typical Metropolis-Hastings algorithms, where the density described in (4.2) seems to be akin to a density derived from the Markov transition kernel, e.g., [Andrieu et al., 2003] in Machine Learning. The divergence seems to be that MH accept/reject steps act locally, whereas this is a globalized accept/reject procedure; critically, the reject step still \"accepts\" a change of position according to the proposal distribution to make such steps global regardless. However, while this superficially seems like a \"Metropolized\" transport algorithm in the vein of, e.g., [Parno and Marzouk, 2018] in SIAM Journal on UQ and [Gabrié et al., 2022] in PNAS, the globalization differentiates itself from such algorithms. It would behoove the authors to investigate parallels in, e.g., \"global\" MH transition kernels (see discussion in, e.g., section 3.3 in [Andrieu et al., 2003]) and differentiate themselves for the benefit of the reader who believes this to be such a Metropolization of Neural JKO. Moreover, the acceleration of Langevin sampling using birth-death processes, proposed in [Lu et al., 2023], was overlooked.\nFinally, there is little to no mention of Stein variational gradient descent (SVGD) [Liu and Wang, 2016] in NeurIPS, despite its popularity as a competing algorithm. SVGD approximates a gradient flow under a kernelized KL-loss ([Liu, 2019] in NeurIPS) or a $\\chi^2$-gradient flow ([Chewi et al., 2020] in NeurIPS), and the recently introduced noisy SVGD variant ([Priser et al., 2024] in ICLR), which addresses some drawbacks of the vanilla SVGD. However, these developments are notably absent from the relevant literature section.", "other_strengths_and_weaknesses": "The major weakness of this article is that it has a very inconsistent writing style and does not seem well-contained. In particular, many of the theoretical background and discussion seemed at-length with little intent; Theorem 2.2 is included from prior literature but only invoked once and hardly discussed for the reader's benefit. Many terms stand undefined (e.g., lsc/lower semi-continuous, coercive, $\\lambda$-convex, etc.). Further, it does lack the discussion of some vital assumptions. In particular, there is some missing context regarding the assumption of $\\lambda-$convexity of the target distribution $q$. Of course, there will always be assumptions for theoretical results, but such an assumption is particularly strong seeing as this is not satisfied by most distributions considered in the numerical results. Perhaps such an assumption is \"softened\" by the rejection steps, or just analytically convenient and stronger than necessary in practice. With careful consideration of the typical ICML reader, the analytical and theoretical discussion could be made significantly more accessible and parochial by focusing on the actual methods and ideas vital to the method. Finally, there is a connected issue that the paper is not well self-contained, where it's unclear what theory, exactly, the authors contribute to the field and what is rehashed from prior literature. In particular, the paper does not differentiate its core neural JKO scheme very well from similar neural JKO approaches.\n\nWith these shortcomings in mind, the paper does invariably provide strong theoretical guarantees, particularly for smooth and log-concave targets. Despite some issues with the presentation of guarantees and terminology, the compact presentation of Bernamou & Brenier, its connection to the algorithm, and the discussion of neural ODEs is reasonably clear and helpful. To add to the theoretical insights, the rejection procedure seems method-agnostic and might be of great benefit to other inference techniques, particularly other variational inference schemes. If anything, the authors undersell the intrigue of such an idea. The numerical results demonstrate an undeniably strong algorithm, with a stark comparison against unadjusted neural-JKO. The independence of the corrected samples (which is not proclaimed until the conclusion) is actually quite remarkable, especially compared with the sequential nature of the correction of Metropolis algorithms. Additionally, Appendix B.1 is well written considering how formal the discussion of the theory assumptions are within the paper.", "other_comments_or_suggestions": "- A discussion of the nominal dimension-free ideas in Corollary 4.3 could be nice:\n\t- The dimension-free guarantee applies to $\\alpha$, so what makes the problem hard? Presumably it has to do with when the target is non-log-concave, which will be exacerbated by dimension, when performing the neural-JKO scheme.\n- If there is sufficient space, it might be interesting to include more explicit insight on how the importance rejection step compares to typical SMC methods as well as, e.g., [Midgley et al., 2023] in ICLR 2023, or the CRAFT algorithm discussed.\n\t- This is partially elaborated in the related work section, but it would be more helpful after the explanation of the algorithm.\n- If space, it might be nice to just remark in the text that 4.3 simply comes from an application of Hoeffding's inequality, as it is unclear why it is a corollary at first glance.\n- While *annealed* importance sampling was largely popularized by [Neal, 2001], the idea of importance sampling dates back significantly further.\n- Similarly, while (Rezende & Mohammed, 2015) certainly introduced the \"CNF\" term, it should be noted that the \"concept\" of coupling complicated target and simple reference dates further back, even within generative modeling communities.\n- Visual accessibility:\n\t- In addition to the highlighting of tabular metrics, the numbers should be bolded for visual accessibility.\n\t- The coloring in Figures 2/3 should probably be adjusted for visual accessibility. Instead of red/green consider using a visually accessible color palette (a little over four percent of people are red/green colorblind).\n- Style/grammar comments:\n\t- The introduction would benefit from more careful attention to its language and style.\n\t- There should be appropriate capitalization in the bibliography.\n\t- The letter $\\lambda$ is used both for Lebesgue measure and convexity.\n\t- The intention of the word \"ration\" on line 322, page 6 is unclear; perhaps the authors meant \"ratio\"?\n\t- In the left column of line 416 on page 8, the phrasing \"We can see, that importance...\" is quite awkward.\n\t- Remark 4.5 is clearly crucial, but difficult to parse. What is \"moderate base of $1+r$\"? The entire remark could benefit from more clear writing.\n\t- The appendices, especially F, should be reviewed for style/grammar (e.g., line 662 \"We give some more backgrounds...\" in appendix A, repeated use of \"Lebesque\" instead of \"Lebesgue\" in appendix B, the sentence in 1303-1305, line 1442 \"...this package does not rely backpropagation...\", line 1476 \"...the evaluation can be cheaper and the density evaluation since these...\", line 1476 \"Residual architectures is at least....\", line 1477 \"...they are very expansive to train...\" etc.).", "questions_for_authors": "1) Does the (uncorrected) neural JKO (N-JKO) scheme described align exactly with prior N-JKO schemes? Is there a difference? What does this paper contribute over, e.g., [Xu et al., 2024]? A different method or just theory on an equivalent method?\n\t- The paper is unclear as to whether such an uncorrected scheme is previously introduced. For example, it cites many papers using JKO-like schemes via W2 proximal mapping iterations, but does not delineate very well from them. Being clearer here would help the paper stand out and better define the contributions to the field, where a reader could clearly see that this is not just adding importance/rejection sampling to a pre-existing neural ODE scheme.\n\n2) Do Cor. 3.2 and Thm 3.3 have anything to do with the rejection steps?\n\t- As suggested in previous comments, it seems like this submission attempts to contribute two different methods that work in tandem. The theoretical extent that the importance correction contributes to _specifically_ the JKO scheme is not well described. It is certainly acceptable that the answer to this theoretical question is \"they are two different items that work well together\", but such a disclaimer would highlight exactly what the authors can say about their work.\n\n3) What are the practical and theoretical parallels with MCMC schemes? Why not just use a \"Metropolis Adjusted Neural JKO\" scheme? Should I regard the result as a CRAFT-like approach to unadjusted Langevin dynamics (ULA)?\n\t- While the guarantees of JKO are great in practice, the uncorrected N-JKO does not seem too different from Langevin. There is unfortunately no comparison to a Metropolis adjustment to neural JKO, nor is there a comparison of the uncorrected N-JKO scheme to ULA. Again, it is imperative that the methods are appropriately benchmarked and indeed that the fact this paper incorporates two different methods is justified.\n\n4) Practically, are there any remarkable features regarding N-JKO's sensitivity to hyperparameters? Time discretizations? Adjoint solves? etc.\n\t- In contrast to MALA or HMC/NUTS, there is a functional optimization loop inside the body of the methods. This is admittedly discussed in part within appendix F, i.e., that the required coupled forward/adjoint can be quite costly, that they use a classic adaptive explicit Runge-Kutta scheme, and that \"choosing larger architectures [than what they have] does not bring significant advantages\". However, in contrast to MALA, HMC, and (unmentioned) SVGD, the proposed methods do not have any asymptotic or mean-field guarantees for a nonasymptotic architecture/function class; therefore, it is certainly worth noting how much effort regarding heuristics and hyperparameter tuning it may take to achieve the results given.\n\n5) It is stated in F.2 that the evaluation of the neural IC only uses a total of five Rademacher vectors for estimating the Jacobian trace. If it's a Monte Carlo estimator, why would this be enough? Would increasing this make the performance better or worse? Could this poison the error metrics in any way?\n\t- While it is entirely acceptable to use just a few evaluations for _training_ the neural networks, it is imperative that the authors are clear that any additional refinement of this discretization works in their favor and not against it. To be more explicit, a case where _higher_ variance in the trace estimator coincidentally improves results over more exact trace estimators would raise questions of method validity.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ELzOoV8DTF", "reviewer": "Reviewer_5Uui", "summary": "This paper presents a method to sample from an probability distribution known through its density, up to an unknown normalizing constant. The method follows the trend of neural parameterizations to solve the proximal steps of the JKO scheme to compute the Wasserstein Gradient Flow of the reverse Kullback-Leibler divergence. The neural parameterization is based on the Benamou-Brenier formulation of optimal transport, which allows to write every step of the JKO scheme using a neural ODE and thus to parameterize the density at a given step via a continuous normalizing flow, that is tuned to minimize the reverse KL divergence. A few theoretical properties of the resulting scheme are presented (mostly in the case of a log concave energy). Besides, to counter the fact that Neural JKO schemes essentially explore the energy landscape locally, the authors propose to introduce rejection steps to enhance sampling. They modify the classical importance sampling scheme using a proposal coming from the Neural JKO scheme (to mitigate the curse of dimensionality) and allowing to propagate the resulting probability density in order to be able to resume subsequent JKO steps. The method is tested on a benchmark of several distributions with closed form densities with various properties (multimodal, with narrow high energy regions, various dimensions) and report improvements over a number of competing methods (classical sampling methods, or more recent JKO, NF or diffusion based). The measures of quality of the generated samples are the energy distance to GT samples or KL divergence values.\n\n## Update after Rebuttal\n\nI thank the authors for their responses and maintain my positive opinion on the paper. I maintain my score to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the detailed and valuable feedback.\n\n## Convexity Assumption\n\nThere appears to be a misunderstanding regarding the assumptions for the theoretical part: We **do not assume that the density is log concave**. Instead, Assumption 3.1 assumes that the functional is $\\lambda$-convex along generalized geodesics. For the KL divergence, this corresponds to the assumption that **the target energy $-\\log(q)$ is $\\lambda$-convex** for some $\\lambda\\in\\mathbb R$. We stress that this explicitly **includes the case of negative $\\lambda$**. This assumption is **much weaker than log-concavity** and is automatically fulfilled, if $-\\log(q)$ is smooth (plus some asymptotics for $\\\\|x\\\\|\\to\\infty$). More intuitively, the condition can be rephrased as: $-\\log(q)$ is $\\lambda$-convex if and only if there exists some (possibly negative) $\\lambda\\in\\mathbb R$ such that $-\\log(q(x))-\\frac{\\lambda}{2}\\\\|x\\\\|^2$ is convex.\n\nGiven that this misconception appeared in more than one review, we will include this discussion in the final version and highlight that $\\lambda$-convexity is usually not an issue as long as the target distribution is smooth.\n\n## Relation to other neural JKO Schemes\n\nWe stress that we consider the theoretical results (Cor 3.2, Thm 3.3) to be the main novelty of Section 3. As outlined in the beginning of Section 3, the implementation (Sec 3.2) is close to previous papers, specifically to (Vidal et al., 2023) and (Xu et al., 2024), who also use the dynamic formulation of the Wasserstein distance. But these papers consider the generative modeling setting instead of sampling. We are not aware of a reference, which adapts the dynamic formulation for sampling, which requires some smaller technical differences (e.g., other $\\mathcal F$, we start with latent distribution, while Vidal et al., Xu et al. start with the data distribution). However, we stress that the main contributions of our paper are:\n\n- The theoretical analysis (Cor 3.2, Thm 3.3) of neural JKO schemes,\n- proposing importance-based rejection/resampling steps which can maintain the density of the generated samples (Section 4) and preserves independence of the generated samples,\n- combining both to an importance corrected neural JKO sampler, which achieves state of the art results.\n\nIn the introduction and in the beginning of Section 3 we already write that similar approximations of the JKO steps exist in the literature. We will add a reminder in the beginning of Section 3.2 (and specifically point to Vidal et al., 2023, Xu et al., 2024 for using the dynamic formulation).\n\n## Literature on Rejection Schemes with Generative Models\n\nIndeed there exist some approaches to combine rejection steps with generative models in the literature. Many of these approaches are based on sequential Monte Carlo techniques (e.g., CRAFT which we used as comparison, but also Arbel et al., Phillips et al.). These methods implement a reweighting step by first approximating importance weights and then sampling from the empirical distribution defined by this weights. However, for SMC-based methods the analytic evaluation of the arising density is usally not possible. Moreover, after SMC reweighting steps, there exist with a high probability several samples at the exact same position. Finally, the generated samples are not exactly independent.\n\nThe paper of (Gabrie et al.) proposes to iteratively train a normalizing flow for sampling by running a Langevin process, adding Metropolis steps with the current normalizing flow as proposal, and retraining the normalizing flow with the updated samples. In contrast to our paper, the rejection steps are not part of the model, but are rather used for training a normalizing flow. In addition, they can't include these steps in the model, because this would require to evaluate the density of these steps, which is not possible for the Metropolis algorithm. Since several papers from the literature [1] found that training normalizing flows to approximate probabilities with disconnected modes (like GMMs) or non-Gaussian tails (like funnel, mustache) are difficult, this might limit the expressiveness of the model.\n\nWe will extend the discussion on these methods in the final version.\n\n[1] https://arxiv.org/abs/1907.04481, https://arxiv.org/abs/2009.02994, https://arxiv.org/abs/2206.14476", "claims_and_evidence": "The paper’s main claim is the capacity of the proposed method to escape local maxima of the target density and prevent model collapse thanks to the proposed rejection scheme. This is indeed supported by the experiments where the methods shows better capacity to sample from multimodal distributions. The proposed rejection scheme combines well with the neural JKO scheme because it allows to propagate densities, which is essential to continue sampling using the JKO scheme.", "methods_and_evaluation_criteria": "The method is tested on several sampling tasks where the target distribution is known and enjoys tractable sampling for comparisons. The quantitative evaluation criteria make sense, i.e. using MMD for two sample hypothesis testing and using estimates of log normalization constants. The method is tested against several types of methods : classical sampling methods (MALA, HMC) that are known to have trouble handling multimodality, and more recent generative model based methods (DDS based on diffusion models, their own neural JKO without resampling) and a last CRAFT method based on sequential MC methods.", "theoretical_claims": "The novel theoretical results are in Corollary 3.2,  Thm 3.3, Thm 4.2 and Corollary 4.3. \n\nThe first two results state results on the convergence of JKO steps, reformulated using a dynamical OT formulation towards tha WFG curve of the reverse KL divergence, in the log-concave case where the KL divergence is geodesically convex (which is not the one that is the objective of the method). This is not a new result per se, but the statement that the functional G (detailed in app F) is geodesically convex for small tau even when F is not is a strong motivation to use the JKO scheme instead of more direct mininmization strategies.\nThe other theorems showcase the good properties of the proposed rejection scheme and how the combine well with the JKO scheme. I only checked the proof of Thm 4.2.", "experimental_designs_or_analyses": "While the combination of JKO scheme and the rejection sampling outperforms all other methods, it would have been interesting to include in the comparison a few methods based on rejection schemes. Naive ones (rejection sampling or importance sampling) are bound to perform poorly in the cases that are tested du to high rejection rates or the curse of dimensionality but surely there must be more recent methods using similar concepts ? Namely the idea of using proposal distributions that are less naive than default choice and are guided by a generative model has been used in some works recently.", "relation_to_broader_scientific_literature": "The paper seems mostly well positioned within the related literature on WGF, sampling algorithms, continuous NF. I still have two small concerns :\n- A number of papers on neural versions of the JKO scheme are cited, but the actual approach taken by the authors is not clearly positioned compared to this work : what is the originality (if any) of the proposed neural parameterization (CNF version, Benamou-Brenier dynamical view on transport) compared to these studies ? Are there already to the authors’ knowldge papers that use JKO schemes specifically for sampling ?\n- As in my comment above, the combination of rejection schemes using generative models as proposal distributions and more classical methods have been explored recently, can the authors say more about the relation between their work and those methods. I think about for instance about Gabrié et al. 2022 (and others cited in the introducion), which the authors cite but do not discuss much.", "essential_references_not_discussed": "See my comment above.", "other_strengths_and_weaknesses": "The theoretical analyis of WGF is interesting and gathers a number of results that are spread in the literature, which is nice, but only applies in the ideal log-concave case, which is precisely the one that the authors aim at extending with their method. This is a minor criticism as the analysis is much harder in the non  log concave case and the authors give convincing arguments about the well behavedness of the JKO scheme in the nonconvex case (app F).\n\nThe rejection scheme seems efficient and importantly is not at odds with the need to propagate densities in the JKO scheme, which is a nice result.", "other_comments_or_suggestions": "There are a number of typos in the paper that I did not list (e.g. Lebesgue, not Lebesque) but should not survive a careful proofreading of the paper.", "questions_for_authors": "See my questions above about\n1) the existence of rejection schemes combined with generative models\n2) the use of JKO for sampling and positioning the present instanciation of neural JKO within the literature.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a novel sampling algorithm, termed Importance-Corrected Neural JKO, which interleaves continuous normalizing flows (CNFs) implementing neural JKO steps with importance-based rejection sampling to improve convergence and robustness when sampling from unnormalized target distributions. 
Strengths:
* Novel combination of methods: While it combines ingredients introduced in recent previous papers (Vidal et al., 2023; Xu et al., 2024) - but for generative modeling instead of sampling-, the neural JKO scheme with an importance-corrected rejection step is new and effective. 
* Empirical performance: The method outperforms classical MCMC (e.g., MALA, HMC) and flow-based baselines (e.g., Neural JKO, DDS, CRAFT) on a range of benchmark tasks, particularly in high dimensions and multimodal settings.
Concerns:
* Computational cost and scalability: Several reviewers (notably JGqE and BUrR) expressed concern over the potential inefficiency of the method due to rejection steps requiring full regeneration of sample paths. However, the authors clarified that costs are manageable and provided implementation details and runtime data in the appendix.
* Positioning within literature: While the authors cite and discuss relevant prior work (e.g., Vidal et al., Xu et al., AFT, CRAFT), reviewers requested clearer articulation of how this work differs from or builds upon those methods, especially regarding- Vidal 2023,Xu et al 2024 - and comparisons with recent importance-weighted or rejection-based sampling schemes which share many similarities with the introduced scheme.
* clarity: one reviewer raised the concern that the authors are not clear regarding the novelty of their theoretical results and their limitations
Final Assessment: All reviewers were generally positive : 4/5, 3/5, 4/5, 3/5  (one reviewer went from 2 to 3 after rebutal). While there are some clarity and scope limitations, the authors’ thorough rebuttal and additional experiments addressed key concerns.
Recommendation: Accept. This paper makes a strong theoretical and practical contribution to the literature on sampling with generative models and optimal transport, and will be of interest to the NeurIPS community.
Decision: Accept (poster)