NETS: A Non-equilibrium Transport Sampler
Authors: Michael Samuel Albergo
Eric Vanden-Eijnden
Venue: ICML 2025 poster (2025)
Keywords: jarzynski
sampling
generative models
diffusion
monte carlo
Primary area: probabilistic_methods->monte_carlo_and_sampling_methods
TL;DR: An augmentation of sequential monte carlo and annealed importance sampling with transport, unbiased with IS weights, with connections to non-equilibrium physics.

Abstract: We introduce the Non-Equilibrium Transport Sampler (NETS), an algorithm for sampling from unnormalized probability distributions. NETS builds on non-equilibrium sampling strategies that transport a simple base distribution into the target distribution in finite time, as pioneered in Neal's annealed importance sampling (AIS). In the continuous-time setting, this transport is accomplished by evolving walkers using Langevin dynamics with a time-dependent potential, while simultaneously evolving importance weights to debias their solutions following Jarzynski's equality. The key innovation of NETS is to add to the dynamics a learned drift term that offsets the need for these corrective weights by minimizing their variance through an objective that can be estimated without backpropagation and provably bounds the Kullback-Leibler divergence between the estimated and target distributions. NETS provides unbiased samples and features a tunable diffusion coefficient that can be adjusted after training to maximize the effective sample size. In experiments on standard benchmarks, high-dimensional Gaussian mixtures, and statistical lattice field theory models, NETS shows compelling performances.

NETS: A Non-Equilibrium Transport Sampler
Michael S. Albergo 12
Eric Vanden-Eijnden
Abstract
to Bayesian inference (Neal, 1993), becomes particularly challenging for non-log-concave targets. In such cases, tra- ditional ergodic sampling methods based on Markov Chain Monte Carlo (MCMC) or Langevin dynamics often exhibit prohibitively slow convergence rates.
We introduce the Non-Equilibrium Transport Sampler (NETS), an algorithm for sampling from unnormalized probability distributions. NETS builds on non-equilibrium sampling strategies that transport a simple base distribution into the tar- get distribution in finite time, as pioneered in Neal's annealed importance sampling (AIS). In the continuous time setting, this transport is ac- complished by evolving walkers using Langevin dynamics with time-dependent potential, while simultaneously evolving importance weights to debias their solutions following Jarzynski's 's equal- ity. The key innovation of NETS is to add to the dynamics a learned drift term that offsets the need for these corrective weights by minimizing their variance through an objective that can be estimated without backpropagation and provably bounds the Kullback-Leibler divergence between the estimated and target distributions. NETS pro- vides unbiased samples and features a tunable diffusion coefficient that can be adjusted after training to maximize the effective sample size. In experiments on standard benchmarks, high- dimensional Gaussian mixtures, and statistical lat- tice field theory models, NETS shows compelling performances.
An alternative approach offered by non-equilibrium sam- pling strategies is transport samples on non -oa-saatonnary distributions that evolve from simple hase distribution (e.g., normal distribution) to the target in finite time. Meth- ods like Annealed Importance Sampling (AIS) (Neal, 2001), Sequential Monte Carlo (SMC) (Del Moral, 1997; Doucet et al., 2001), and their continuous-time time variants based on Jarzynski's equality (Hartmann et al., 2018) achieve this transport through dynamical quench, using importance weights correct for the bias arising when the walkers' distribution lags behind their evolving target. However, these methods can fail when the lag is too strong, leading to high high-variance weights.
We address this limitation by modifying the dynamics via some additional transport learned to guide toward the target. In the context of quenches performed via Langevin dynam- ies on an evolving potential, we show that an optimal drift can be added to eliminate the need for importance weights entirely, and can be characterized as the minimizer of objec- tive functions amenable to empirical estimation. This leads to an unbiased sampling strategy where importance weights, while still available for exact correction, have substantially reduced variance due to the improved transport.
1. Introduction
In sum, our work makes the following main contributions:
The aim of this paper is to sample probability distributions on Rdd known only up to normalization. This problem, central to applications from statistical physics (Faulkner & Livingstone, 2023; Wilson, 1974; Hénin et al., 2022)
We introduce the -oullqquli Transport Sampler (NETS), which augments Langevin dynamics on an evolving potential with learned transport. We show that the method remains unbiased through a generalization of Jarzynski's equality, and that can be used in concert with sequential Monte-Carke Carlo (SMC) strategies.
'Society of Fellows, Harvard University, Cambridge, MA, USA NNFF Al Institute for Artificial Intelligence and Fundamental In- teractions, Cambridge, MA, USA 'Machine Learning Lab, Capital Fund Management, Paris. France 'Courant Institute of Mathemati- cal Sciences, New York, NY, USA. Conresoondence to: Michael Albergo maalerggo haw.aad..du.., Eric Vanden- Bijnden sone@@iimsnnyu.com
We demonstrate that the optimal drift for this additional transport minimizes an objective function, based on physics-informin neural networks (PINN), whose gradi- ent can be estimated without backpropagating through the sampling equations.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
We show that the PINN objective has two key properties:

---

NETS: A Non-Equilibrim Transport Sampler
it is off-policy (requiring no samples from the target density) and it bounds the KI--divergence between the sampler and target distributions.
while others reformulate sampling as a stochastic optimal control (SOC) problem (Zhang & Chen, 2021; Behjoo & Chertkov, 2024; Hua et al., 2024). However, these methods face limitation: learning the optimal drift is complicated optimization problem which becomes computationally pro- hibitive in high dimensions, and simplified approaches such as Akhound-Sadegh et al. (2024) introduce bias into their objective function. Alternative perspectives include using denoising oracles to simplify the sampling problem (Bruna & Han, 2024) and adapting graph-based distribution model- ing techniques for sampling, including off-policy training (Malkin et al., 2023; Sendera et al., 2024).
We discuss other possible objectives make connections with Action Matching (AM) as well Controlled Monte Carlo Diffusions (CMCD).
We establish that NETS's performance can be optimized post-training by tuning both the integration time and diffusion coefficient, which we demonstrate through high-dimensional numerical experiments.
1.1. Related work
Vargas et al. (2024) establish Controlled Monte Carlo Dif- fusions (CMCD), another unbiased sampler that, like ours, augments the dynamics with learned transport. The methods differ in how this drift is learned: CMCD derives a gradient- form objective through path integrals and Girsanov's is theo- rem, requiring either backpropagation through the SDE or computation with a numerically unstable reference measure on a fixed grid. Our approach, based on Tokker-Planck equa- tion manipulations, yields new objectives for the additional drift that avoid backpropagation entirely. Moreover, our framework allows for post-training adaptation of both step size and dime-dependent diffusion, providing tunable parameters to enhance performance. One of our objectives is a Physics-Informedd Neural Network (PINN) loss, which has appeared elsewhere for sampling (Mâté & Fleuret, 2023; Tian et al., 2024; Fan et al., 2024). We establish two key results: this objective is valid for an- nealed Langevin dynamies, and it directly controls both the KL divergence and the importance weights arising from Jarzynski's equality.
Dynamical Measure Transport. Modern generative mod- els for continuous data often employ dynamical measure transport, where samples from a base density are trans- formed into samples from a target density by solving ordi- nary stochastic differential equations (ODE/SDH) with learnable drift coefficients. This approach, pioneered in (Chen et al., 2018; Grathwohl et al., 2019), has developed in several directions, notably with score-based diffusion models (Ho et al., 2020; Song et al., 2020) that frame drift estimation as quadratic regression. This formulation has led to more general frameworks (Albergo & Vanden-Eijnde 2022; Lipman et al., 2022; Albergo et al., 2023; Liu al., 2022; De Bortoli et al., 2021; Neklyudov et al., 2023). A key requirement of these methods is access to samples from both base and target distributions. Our work shows that sim- ilar models building on dynamical transport can be learned without prior access to data from the target.
perimize-then-disceeiee
Augmenting sampling with learning. Augmenting MCMC and importance sampling proceures with transport has been active area of research for the past decade. Early work makes use of the independence Metropolis algorithm (Hastings, 1970; Liu, 1996), in which proposals come from a transport map (Parno& & Marzouk, 2018; Noé et al., 2019; Albergo et al., 2019; Gabrié et al., 2022) that are accepted or rejected based off their likelihood ratio with the target. These methods were further improved by combining them with AIS and SMC perspectives, learning incremental maps that connect a sequence of interpolating densities between the base and target (Arbel et al.. 2021: Matthews et al., 2022; Midgley et al., 2023). Similar works in the high- energy physics community posit that interleaving stochastic updates within a sequence of maps can be seen as a form of non-equilibrium sampling (Caselle et al., 2022; Bonanno et al., 2024).
2. Methods
2.1. Setup and Notations
We assume that the target distribution is absolutely contin- uous with respect to the Lebesgue measure on Rd with probability density function (PDF) p1 (x) (t)); here RRd U : _d + R is known energy poten- tial, assumed twice differentiable and bounded below, and Z1 = fR. e -Ui()) dx < 30 is an unknown normalization constant, referred to as the partition function in physics and the evidence in statistics. Our aim is to generate samples from p1 (zx) so as to be able to estimate expectations with respect to this density. Additionally we wish to estimate Z
To this end, we introduce family of time-dependent po- tentials (x)x) that interpolate between a simple initial po- tential Ux(x) (e.g., Uo(x) 1 2) at t = 0 and the tar- get potential U1 (x) at = 1. While linear interpolation, Uz(x)) = (1 - ()))(() + (x)), provides a straightforward choice, more sophisticated interpolation schemes are often
Following the success of generative models based on dy- namical transport, several approaches have emerged to ap- ply these ideas to sampling. Some translate concepts from diffusion models to minimize the KL divergence between model and target (Vargas et al., 2023; Berner et al., 2024),
2

---

NETS: A Non-Equilibriun Transport Sampler
preferable. Our only requirements are that U1=o Uo. U==1 U1 and (x) is twice differentiable in (1) E [0, x Ig We assume that the time-dependent PDF as- sociated with this potential (x( is normalizable for all t te [0, 1] and denote it as
Proposition 2.1 (Jarzynski equality). Let (X Ac) solve
the coupled system of SDE/ODE
dX,- = - (.VUAXX))) + 2e,dW,, dA, = -,,U,
Xo ~ Po: (9)
Ao = 0, (10)
pc(x)) = -U(lx) Zi { e e-Uile) dr < 00. (1)
where é, 2 0 is dependennddntt diffusion coefficient and
<00,
W. E TR is the Wiener process. Then for all t E [0. 1] and any test function h : IR - R, we have
so that P== = Po( and Pt=1(x)= = (2); we also assume that (a) is simple to sample (either directly or via MCMC or Langevin dynamics) and that its partition func- tion Zo is known. To simplify the notations we introduce
J M=))((x)dx= Ble h(x))
EleAr]
(1I)
2,/20= = e = =]
the free energy
F2 log Ze,
where the expectations are taken over the law of (Xt, As)-
and note the useful identity
The proof of this proposition is given in Appendix A and it relies on the identity SRa (x)))((z) fg+s fx a)dadz/ Sga.. fe(a, a) which follows from (8). The second equation in (11) for the free energy F, is what referred to as Jarzynski's equality, and was originally surmised in the context of non-equilibrium thermodynamics (Jarzynski, 1997).
F.F. [ ne
((U)(t)
since- -c log Fox e Us = O,Ue- SRee
2.2. Non-equilibrium Sampling with Importance Weights
Remark 2.2. We stress that it is key to use the weights e in (11) because pr(z) is not the PDF of X, in general. Indeed, if we denote by pe(x)) the PDF of X,, it satisfies
Annealed importance sampling (AIS) uses a finite sequence of MCMC moves that satisfy detailed-balance locally in time but not globally, thereby introducing a bias that can be corrected with weights. Here we use a time-continuous variant of AIS based on Jarzynski equality.
O.p: V (VU,P: + Vipe: pi= =
(12)
This FPE misses the term - -0,U, - 0,Fi)pt at the right
By definition of the PDF in (1), (z) - (U)()))(())
hand-side of (6), and as a result pe(x)) + pr(x)) in general- - intuitively, (xr) lags behind pa (z)) when the potential (x)
and hence, for any & 2 0, we have
evolves and this lag is what the weights in (11) correct for.
E (VUiPt + Vp.)
(4)
It is important to realize that, while (11) exact, estimators based on this relation can be high variance if the lag between the PDF (x) of X, and pt (z) is too pronounced. This issue can be alleviated by using resampling methods as is done in SMC (Doucet et al., 2001). Here we will address it by adding some additional drift in (9) that will reduce the lag and as a result lower the variance of the weights.
Since we also have
a.p: = - - -,F.).
(5)
we can combine these last two equations to deduce that
0ep: = V. (VU,p:+ + Vp.) - (a,U, - a.FFi)p.
(6)
The effect of the last term at the right hand-side of this equation can be accounted for by using weights. To see how, extend the phase space to (Ir a) E Rd+1 and introduce the
2.3. -uu-euuil Sampling with Perfect Transport
We can add a transport term to eliminate the need of the weights. To see how, let (x( E Rd be a velocity field which at all times f [0, 1] satisfies
PDF [x a) solution to the Bokker-Planck equation (FPE)
(7) with initial condition fe=o(1, a) = (())) A direct calculation using (3) (for details Appendix A) shows that
ātft = v (VU,ft + f() + O.U,fff
V (b.p.) = -d.p.
(13)
We stress that this an equation for (xx)) in which (x(xr) is prescribed and given by (I): In Appendix C we show how to express the solution to (13) via Feynman- Kac formula.
Sg e² fi(x pr(z) SE en fi a)dady
If (13) is satisfied, then we can combine this equation with (4) and (5) to arrive at
Therefore we can use the solution to the SDE associated with the FPE (7) in the extended space to estimate expectations with respect to pe (x))::
Dept = V. (VU + Vpt) V -(bept); (14)
3

---

NETS: A Non-Equilibrium Transport Sampler
which is standard FPE. Therefore the solution the SDE associated with (14) allows (without weights). We phrase this result as:
Proposition 2.4 (Nonequilibrium Transport Sampler
sample pa(x)) directly
(NETS)). Let (X), A4) solve the coupled system of SDE/ODE
Proposition 2.3 (Sampling with perfect additional trans- port.). Let bz(z) be a solution to (13) and let X3 satisfy the SDE
dx -- - (U.((a) + (x( + 22ndW,, (19)
dA V (,,((x)) vUi(x,) bi (x
(20)
- -d,Uo(X))dt,
dx; = -e, -. (X)(X be (x.))) V2,,W,, (15)
with X3- ~ Po and AB = O, and where Ct 2 0 is a time- dependent diffusion coefficient. Then for all t E [0, 1] and any test function h. T2 - IR we have
with X8 pu and where E 2 0 is a time- dependent diffu- sion coefficient. Then pe(x) is the PDF of X4,, i.e. for all
t E [0, 1] and, given any test function _2 + R, we have
J h(z) Mlx3
J h())p(()dd = =E((X)),,
(16)
B[e49]
(21)
where the expectation at the right-hand side is taken over the law of (Xt).
Ze/Zo e = Ele
where the expectations are taken over the law of (X). A?).
This proposition is proven in Appendix A and it shows that we can principle get rid of the weights altogether by adding the drift (z)) in the Langevin SDE. This possibil- ity was first noted in Vaikuntanathan & Jarzynski (2008) and is also exploited in Tian et al. (2024) for deterministic dynamics (i.e. setting = 0 in (15)) and in Vargas et al. (2024) using the SDE (15). Of course, in practice we need to estimate (x). and also correct for sampling errors if this drift is imperfectly learned. Let us discuss this second question first, and defer the derivation of objectives to learn (2)) Secs. 2.5 and E. In Appendix C we show how to express the solution to (13) via Feynman Kac formula.
A proof of this proposition using manipulations of the FPE is given in Appendix A, which will allow to write down variety of new loss functions for learning b,; for an al- ternative proof using Girsanov theorem, see Vargas et al. (2024). For completeness, in Appendix B we also give a time-discretized version of Proposition 2.4, and Ap- pendix D we generalize it in two ways: to include inertia and to turn f into vector coordinate for multimarginal sam- pling. We also discuss the connection between NETS and the method of Vargas et al. (2024) in Appendix F.
Notice that, if (x())) = 0, the equations in Proposition 2.4 simply reduce to those in Proposition 2.1, whereas (x( = (xx) solves (13) we can show that
2.4. Non-Equilibrium Transport Sampler
Let us now combine the approaches discussed in Secs. 2.2 and 2.3 to design samplers in which we use added trans- port, possibly imperfect, and importance weights.
4 - Fi Fu.
(22)
i.c. the weights have zero variance and give the free energy difference. Indeed, by expanding both sides of (13) and dividing them by Pc(x) > 0, this equation can equivalently be written as
To this end, suppose that we wish to add an additional transport term- -V (t. ) in (6), where bx(x) E IRd is some given velocity that does not necessarily solve (13). Using the expression in (1) for pc(x), we have the identity
b VU, be 0,U, O.F.
(23)
As a result, when (z(2) = bz)).. (20) reduces to dA = Fi
-V (b.p.) - bept + VU.- bept
(17)
=
(24)
Therefore we can rewrite (4) equivalently as
and the solution to this equation is (22). In practice, achiev- ing zero variance of the weights by estimating b, (z exactly is not generally possible, but having a good approximation of (z) can help reducing this variance dramatically, as we will illustrate below via experiments.
dep: = V (VU, + Vp:) - -V - (bpp)) (V b, VU. b, a,U, F.)pp
(18)
We can now proceed as we did with (4) and extend state space to account for the effect of the terms (V b VU,- b, - a,U,- + Fi) Pt in this equation via weights, while having the term- -V- (zx()) (x)) contribute to some ad- ditional transport. This leads us to a result originally ob- tained in Vaikuntanathan & Jarzynski (2008) and recently de-iieived in Vargas et al. (2024):
2.5. Estimating the Drift b,(x) via a PINN Objective
Equation (23) can be used to derive an objective for both (x()) and F, The reason is that in this equation the un- known O,F, can be viewed as factor that guarantees solv- ability: indeed, integrating both sides of (13) gives 0
4

---

NETS: A Non-Equilibrim Transport Sampler
-à, Sus ( + orF, which, by (3), is satisfied if and only if F. is (up to constant fixed by Fo log %o) the exact free energy (2). This offers the possibility to learn both by (z) and F' variationally using an objective fitting the framework of physics informed neural networks (PINNs):
Fokker-Plance equation
dept + V. (bepe) - V. (VU... + Vit),
(27)
with pt-o - pa and where by (x) is some predefined velocity field and E 2 0. Then, we have
Proposition 2.5 (PINN objective). Given any T E (0. and any PDF (e(z)) > 0 consider the objective for (b, F)
DRr Dxr(pe-1llp
v LTTNM ( F).
(28)
given by:
where Ft is the free energy. In addition, given any estimate Fs such that Jo |0, P. 0F.l2dtt S 5 for some 8 2 0, we
L, F1 f
//" E [[))())))
(25)
have
where 3t ~ Pr and we denote
DR((i=all S 2L F) + 28.
(29)
ge = V-b(z)) -
(26)
This proposition is proven in Appendix A. Notice that the bound (28) can be estimated by using O.Fi BlcA) in the PINN loss (25).
Then ming.p L' F1 = O, and all minimizers (b, F) are such that and bx(x) solves (23) and F; is the free energy (2) for all t E [0, T.
2.7. Implementation
This result is proven in Appendix A: in practice, we will use T E (0, 1] for annealing but ultimately we are interested in the result when T = 1. Note that since the expectation over an arbitrary (x) in (25), it can be used as an off-policy objective. It is however natural to use (x) = (a()r) since allows us to put statistical weight in the objective precisely in the regions where need by (z:) to transport probabil- ity mass. In either case, there is no need to backpropagate through simulation of the SDE used to produce data. We show in Section 2.7 below how the expectation over (r( can be estimated without bias to arrive at an empirical esti- mator for (25) when (x(x=) = Pe (z)) Note also that, while minimization of the objective (25) gives an estimate F; of the free energy, it is not needed sampling time when solving (19). An objective similar to (25) was also recently posited in Máté & Fleuret (2023); Tian et al. (2024) for use with deterministic flows. Here, we devise it in the context of augmenting annealed Langevin dynamics.
If we minimize (25) off-policy, i.e. with samples xy from some å: + Pt. this is perfectly valid, but may be inefficient for learning by over the support necessary for the problem. If we decide instead to set (x( = (x()), since the SDEs in (19) and (20) can be used with any b.(x)) to estimate expectation over pr(z) via (21), we can write the PINN objective on- on-policy as
B la (x})) TE [[eA4)]
6 F1
dt
(30)
These expectations can be estimated empirically over a pop- ulation of solutions to (19) and (20). Crucially, since we can switch from off- off-policy to on-policy after taking the gra- dient of the PINN objective, when computing the gradient of (30) over (z()). (Xf. Ab) can be considered indepen- dent of (z)) and do not need to be differentiated over. In other words, the method does not require backpropagation through the simulation even if used on-policy, i.c. even though it uses the current value of b, to estimate the loss and its gradient. Finally note that we can use the ODE (20) for A to write (30) as
In addition to the above results, in Appendix D we supply extensions to the setup to the case where there are mul- tiple marginals and where the stochastic dynamics have inertia. We also discuss an alternative objective, relying the action matching formalism (Neklyudov et al., 2023), in Appendix E.
T
1.A9 +
Or
Fel
LTINN F1
dt
(31)
18 E[EAt]
2.6. Control of the Kullback-Leible Divergence
Since Elc-4 0z 0,441/B(c4 B, log clearly shows that this loss controls the variance of 0, Ab. which directly connects the Jarzynski weights to the PINN objective.
E 41 = - a.F, (31)
One advantage of the PINN objective (25) is that we know that its minimum is zero, and hence we can track its value to monitor convergence when minimizing (25) by gradient descent, as we do below. Another advantage of the loss (25) that it controls the quality of the transport as measured by the Kullback-Leein divergence:
The computation of the divergence V (x( in the PINN objective given in (25) can be avoided by using Hutchinson's trace estimator, see Appendix G.
Proposition 2.6 (KI. control). Let Pe be the solution to the
5

---

NETS: A Non-Equilibrium Transport Sampler
Algorithm 1 ITraiinng: Note that the resultant set of walkers across time slices { are detached from the computational graph when taking a gradient step oof-pplicy learning). Initialize: walkers, zi ~ No. Ag O, K time steps, model parameters for {b,, F.). diffusion coefficient E learning rate 7)
where the coefficients on Sr satisfy og = 31 1, 01
Be 0 and on < 0, Br > O. The PDF Pr of I, satisfies (13)
with a drift b, given by
1:
by (x) Elárvo + Bur | 1. - sr),
(33)
where the expectation taken over the law of zo. 21 condi- tional on 4 x. When Pu == (0, an application of Stein's 's lemma indicates that the gradient of the potential Uz associated with p: is given by
2: repeat
3: Randomize time grid: to: lo..1,,..,tt .kk ~ Uniform(O, T), sort such that to < 11 < ... tk
4: for k = O..... K do
5: At& = tk+1 tk.
VUr(2) ar Lpzo 4 = =x].
(34)
6: for each walker i = 1, ................................ do
7:
2k411 4/ + 1b,, (1rt) 6t VU,, (xf)
While theoretically appealing, these relations are not im- mediately useful for two reasons. First, this velocity field can only be regressed when samples from P1 are readily available, which is not the case in our setting. Second, the potential U, associated with this pt is not analytically known. However, by combining (33) with (34) and using x = = I, = z²), we notice that b, can be written as
+v2E,, (W?, 1x+1 - )
8:
1/+1 = A + VV 8 (at)) 0,U,, (xf)
-b., (1tt (x,,,,)))
9:
Estimate (30) by replacing the expectation by an em- pirical average over the n walkers and the time inte- gral by an empirical average over tu...., fx- Take gradient descent step to update the model pa- rameters.
10:
(x()) (ára, o. + A,8, (35)
11:
until converged
which we can exploit to directly learn the U, associated with the interpolant and, in the process, solve the transport problem. A convenient choice is to set, for example, o V1 1² and Se = f, so that (35) reduces to
Learning (x) and Fr for & E [0, 1] from the start can be challenging if the initial (x(x) is far from exact and the weights gets large variance increases. This problem can be alleviated by estimating (Ir) sequentially. In practice, this amounts to annealing T from a small initial value to T = 1, in such a way that by (x) is learned sufficiently accurately so that variance of the weights remains small. This variance can be estimated on the fly, which also give us an estimate of the effective sample size (ESS) of the population at all times E [0, 1].
tbe(2) = VU(()).
(36)
If we parameterize the potential as U, = 0/p with
- (1- (U)()) + (1111111))) - f))(()), (37)
where fe(x): [0, 1] x Rd + R is a neural network, (36) gives us an expression for the drift also in terms of fe:
Note that we can also employ resampling strategies of the type used in SMC to keep the variance of the weights low (Doucet et al., 2001; Bolić et al., 2004).
(38)
= = vU) - (1-)))()))
Details for the numerical implementation of the minimiza- tion of the objective (25) is summarized in Algorithm
This allows us to write the PINN objective (25) as an objec- tive for (f, F):
2.8. Learning the U. of Stochastic Interpolants
L.' TE [lla (sn)) di,
TINNN (f. F1 [oow[.:F)/// F
The choice of the potential U, used in the annealing can have a significant impact on both the learnability of b, and the numerical stability of solving (19). A desirable char- acteristic is that U, gives a density P: that is geometrically smooth in its evolution between Po and (r, so that transport via vector fields is simple. One approach toward achieving this is to use the drift associated with a stochastic interpolant (Albergo& & Vanden- Fijjden,, 2022; Lipman et al., 2022), i.e. the stochastic process defined as
(39)
where of (ax) is obtained from (26) by replacing Un with (37) and by with (38):
- -i)
(40)
When minimized to zero, the objective (39) yields both drift 6/ and potential U/ such that the resulting density Pc = e- + matches the PDF of the interpolant 1, and possesses favorable transport properties.
1 0,30 + B Zo ~ Po. I1 Pl,
(32)
6

---

NETS: A Non- Bon-lquili Transport Sampler
GMM (d
NETS outperform existing methods. In addition, because our method can be turned into an SMC method by including resampling during the generation, we can push the accep- tance rate of the same learned PINN model to nearly 100% by using a single resampling step when the ESS of the walk- ers dropped below 98%. NETS uses 100 sampling steps and an é, = 0.0, 4.0 in the SDE. Note that NETS and CMCD as published use different interpolating potentials on this example based code conventions, but they could be the same.
Algorithm
ESST
W: 4
FAB
0.653 :: 0.017
12.0 2: 5.73
PIS
0.295 0.018
7.64 + 0.92
DDS
0.687 : 0.208
9.31 0.82
PDEM
0.634 0.084
12.20 0.14
IDEM
0.734
7.42
CMCD-KL
0.288 :: 0.069
9.32 + 0.71
CMCD-LV
0.055 4: 0.023
4.01 + 0.25
NETS-AM e, (ouns)
3.89
NETS-PINN < 0 (ours)
0.954 4: 0.003
3.55 : 0.57
NETS-PINN (ours)
0.979 :: 0.002
3.14 + 0.46
NETS- resample (ours)
0.993 0.004
3.27 1 0.31
Table 1. Performance of NETS in terms of ESS and W22 metrics for 40-mode GMM (d = 2) with comparative results quoted from Akhound-Sadegh al. (2024) for reproducibility.
3.2. Funnel and Student-t Mixture
We next test NETS Neal's funnel, a challenging syn- thetic target distribution which exhibits correlations at dif- ferent scales across its 10 dimensions, as well as the 50- dimensional Mixture of Student-T (MoS) distribution used in (Blessing et al., 2024). The definitions of the target densi- ties and the interpolating potentials are given in Appendix H. Heuristically, the first dimension is Gaussian with variance の² 9, and the other 9 dimensions are conditionally Gaus- sian with variance exp(2e), creating the funnel.
3. Numerical Experiments
In what follows, we test the NETS method, for both the PINN objective (25) and the action matching objective (119), on standard challenging sampling benchmarks. We then study how the method scales in comparison to baselines, particularly AIS on its own, by testing it on an increasingly high dimensional Gaussian Mixture Models (GMM). Fol- lowing that, we show that it has orders of magnitude better statistical efficiency as compared to AIS on its own when applied to the study of lattice field theories, even past the phase transition of these theories and in 400 dimensions (an L= 20 x L. 20 lattice).
We again parameterize (6, F) or & using simple feed for- ward neural networks, this time of hidden size 512. We use 100 sampling steps for both, with diffusion coefficients given in the caption of Table 2. Following (Blessing et al., 2024), we compute the maximum mean discrepancy (MMD) and W/2 distance between 2000 samples from the model and 2000 samples from the target and compare to related meth- ods in Table 2. NETS outperforms other methods with both losses on the high dimensional MoS target in both metrics. In addition this can be improved using SMC-style resam- pling in the interpolation when the ESS drops below 70%. NETS matches the best performance in MMD for the Funnel distribution, but it is slightly worse in W2.
3.1. 40-Mode Gaussian Mixture
A common benchmark for machine learning augmented samplers that originally appeared in the paper introduc- ing Flow Annealed Importance Sampling Bootstrap (FAB) (Midgley et al., 2023) is a 40-mode GMM in 2-dimensions for which the means of the mixture components span from -40 to 40. The high variance and many wells make this problem challenging for c--eeghting or locally updating MCMC processes. We choose as time dependent potential (x) that linearly interpolates the means of the GMM with (z(z) the potential for a standard multivariate Gaussian with standard deviation scale a = 2.
Algenities
Faxera (f 6
Mac 50)
NMD
NMD
FOR MMligyylyyye al 100
030 1:0000
115 993991
0.955 5 0.015
1934
1477
DADIV1 (Aees aia., 2013)
0.031 2. 000
105. 02.0000-2 8 3.473
0.135 2: D.017
1235-216 22 215.0
bits riapeet al 20031
0.172 8-0 0:4-551
142 9m0 0 8.555
0.131 1 0-001
2154 1114
2: (1466)
Junel at el. 2001)
8
8 6.003
0.315 8 0.006
2543 410 -61.3
CRAFT
3021)
0.115 8.0.00
234.383
We train a simple feed-forward neural network of width 256 against both the PINN objective (25), parameterizing (b. F). or the action matching objective (119), parame- terizing G. We compare the learned model from both objectives to recent related literature: FAB, Path Integral Sampler (PIS) (Zhang& & Chen, 2021), Denoising Diffu- sion Sampler (DDS) (Vargas al., 2023), and Denoising Energy Matching (pDEM, iDEM) (Akhound-Sadee) et al., 2024). For reproducibility with the benchmarks provided in the latter method, we compute the effective sample size (ESS) estimated from 2000 generated samples as well as the 2- Wasserstein (W2) distance between the model and the target. As noted in Table 1, all proposed variants of
COR ki /Wegg elidl, 3000
01055 2-0.000
513 9091
21 199.4
NETS JINN (oar))
010533 0000
566
161
0.022 : 00.003
669 2912 114..
NEIS-FFSS-oorrrte teans
0.007 0.0000
$41.76 25
0.000 2
400.07H 01-31
Table 2. Performance of NETS on Ncal's Funnel and Mixture of Student-T distributions, measured in MMD and Wi distances from the true distribution. Benchmarking is in accordance with the setup of (Blessing al., 2024). Diffusion coefficient e = 5, was used for NETS-AM on the Funnel and MoS, respectively. Equivalently, 6 5,5 were used by NETS-PINN. Bold numbers are within standard deviation the best performing. Note that NETS still has perfect sample in the e, + oc limit, but would require finer time discretization than the 100 sampling steps used here (see Figure 4).
7

---

NETS: A Non-Equilibrium Transport Sampler
FO transport.
tranepert.
and trarapori,
Pigure 1. Comparison of the performance of annealed Langevin dynamics alone, transport alone, and annealed Langevin coupled with transport when sampling the 40- mode GMM from (Midgley et al., 2023). Left: Annealed Langevin run for 250 steps with 6, = 4.0, failing to capture the modes with 0% ESS. Center: Learinng using the PINN loss and sampling with 100 steps and é, = achieves an ESS of 95%. Right: Same learning and now sampling with C1 = 4.0 achieves an ESS of 98%.
K = 100 for e, = 0 up to K = 2000 for c = S0. Nonethe- less, it suggests that diffusion can be more helpful when there is already some successful transport than without.
Tage
3.4. Lattice xpt Theory
Mode
We next apply NETS to the simulation of statistical lattice field theory at and past the phase transition from which the lattice goes from disordered, to semi-ordered, to fully ordered (neighboring sites are highly correlated to be of the same sign and magnitude). We study the lattice ee theory in D = 2 spacetime dimensions (Vierhaus, 2010; Albergo et al., 2019). The random variables in this circumstance are field configurations yp E RRXX,, where the extent of space and time. The interpolating energy function under which we seek to sample is defined as:
Figure 2. Demonstration of nigh-dimennionns sampling with our method using the PINN loss in (25) and study of how diffusivity impacts performance, with and without transport. Left: NETS can achieve high ESS through transport alone, and the effect of increased diffusiivty has more of positive effect on performance with sampling than without. AIS cannot achieve ESS above = 0 in high dimension. Right: Kernel density estimates of 2-d cross sections of the high-dimensional, multimodal distribution arising from the model and ground truth.
Σ lp - -Pyl + Σ m2p² + sop!) (41) x~y
U2(4)
3.3. Scaling on High-dimensional GMMs
where the sums are taken over the lattices sites x, or adjacent sites æ ~ y, and a are time-dependent parameters of the theory that define the phase of the lattice (ranging from disordered ordered, otherwise known as magnetized), A derivation of this energy function is given in Appendix I. Importantly, sampling the lattice configurations becomes challenging when approaching the phase transition between the disordered and ordered phases. As an example, we identify the phase transition on L = 16 (d = 256) and L = 20 (d = 400) lattices and run NETS with the action matching loss, with Φe simple feed forward neural network. We use the free theory An 0 as the base distribution under which we initially draw samples. The definition of the target parameter values m. A both at the phase transition and in the ordered phase are given in the Appendix I. In Figure 3, the top row shows samples from NETS for L = 20 at the phase transition, where correlations begin to appear in the lattice configurations. NETS is almost 2 orders of magnitude more statistically efficient than AIS (the same setup without the transport) in sampling the critical point, as seen in the plot showing ESS over time. Note also that NETS can produce unbiased estimates of the magnetization as compared to a Hybrid Monte Carlo (HMC) ground truth. The bottom row shows samples past the phase transition and into the ordered phase, where the lattices begin to take on
In order to demonstrate that the method generalizes to high dimension, we study sampling from multimodal GMMs in higher and higher dimensions and observe how the perfor- mance scales. In addition, we are curious understand how the factor in the sampling SDE coming from annealed Langevin dynamics, VU, interacts with the learned drift or VáJ as we change the diffusivity. We construct 8-mode target GMMs in d = 36, 64, 128, 200 dimensions and learn b with the PINN loss in each scenario. We use the same feed forward neural network of width 512 and depth 4 to param- eterize both i and F for all dimensions tested and train for 4000 training iterations. Figure 2 summarizes the results. On the left plot, we note that AIS on its own cannot produce any effective samples, while even in 200 dimensions, NETS works with transport alone with 60% ESS. As we increase the diffusivity é, and therefore the effect of the Langevin term coming from the gradient of the potential, we note that all methods converge to nearly independent sampling, and the discrepancy in performance across dimensions is diminished. Note that the caveat to achieve this is that the step size in the SDE integrator must be taken smaller to accommodate the increased diffusivity, especially for the é = 80 data point. The number of sampling steps used to discretize the SDEs these experiments ranged from
8

---

NETS: A Non-Equilibrium Transport Sampler
ESS over Time
M(ole) -, slar)
Oi~iU U% AIS
NETS
Reweighted NETS
HMC
Average Magnetization
Figure 3. Comparison of the performance of NETS to AIS on two different settings for the study of 4p theory. Top row, left: 10 example generative lattice configurations with parameters L = 20, m² = - 1.00 À 0.9, which demarcates the phase transition to the antiferromagnetic phase. Top row, right: Performance of AIS (purple curve) vs. NETS (red curve) terms of effective sample size over time of integration f, and a histogram of the average magnetization of 1000 lattice configurations, sampled with AIS, NETS, and HMC (superposed this order). Note that NETS is closer to the HMC target and re- weights correctly. Re-weighted AIS was not plotted because the weights were too high variance. Bottom row: Equivalent setup for L 16, m² 1.0, À 0.8, past the phase transition and into the ordered phase. Note that the field configurations generated by NETS are either all positive across lattice sites or all negative. AIS fails sample the correct distribution, and its weights are too high variance be used on the histogram.
either all positive or all negative values. Again in this regime, NETS is nearly 2 orders of magnitude more statistically efficient.
Machine Learning Research, pp. 760-786. PMLR, 21- 27 Jul 2024. URL https: htttt::://proceedingg.. .lr.. press / /22//223555/// Satess/v35//aaaa / html.
While NETS performs significantly better than conventional annealed samplers on the challenging field theory problem, algorithms built out of dynamical transport still experience slowdowns near phase transitions because of the difficulty of resolving the dynamics of the integrators near these critical points. As such, we need use 1500-2000 steps in the integrator to properly resolve the dynamics of the SDE.
Albergo, M. S. and Vanden-Eiindens E. Building normaliz- ing flows with stochastic interpolants. In The Fleventh International Conference on Learning Representations, 2022. Albergo, M. S., Kanwar, G., and Shanahan, P. E. Flow- based generative models for markov chain monte carlo in lattice field theory. Phys. Rev. D, 100: 034515, Aug 2019. doi: 10.1103/PhysevvD.10 034515. URL https: http:::://likkaappp/ wwtp///////ooooieooooooeooo doi/10. 1103/Phys 100. 034515.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
A Albergo, M. S., Boff,, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.
References
Arbel, M., Matthews, A. G. D. G., and Doucet, A. An-
nealed flow transport monte carlo. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Machine Learning Research, 18-24 Jul 1021.
Akhound-Sadegh, T. Rector-Brooks, J., Bose, J., Mittal,
S., Lemos, P., Liu, C.-H., Sendera, M., Ravanbakhsh, S., Gidel, G., Bengio, Y.. Malkin, N., and Tong, A. Iterated denoising energy matching for sampling from boltzmann densities. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, (eds.), Proceedings of the 41st International Conference Machine Learning, volume 235 of Proceedings of
Arenz, O., Dahlinger, P., Ye, Z., Volpp, M., and Neu- mann, G. A unified perspective on natural gradient varia- tional inference with gaussian mixture models. Transac- tions on Machine Learning Research, 2023. ISSN 2835-

---

NETS: A Non-Equilibrium Transport Sampler
8856. URL htt..uttee//pppeeeeeeee..eeeee// id=tLBjsx4tjjs.
De Bortoli, V., Thornton, J., Heng, J., and Doucet, A. Diffu- sion schrödinger bridge with applications score-based generative modeling. In Advances in Neural Information Processing Systems, volume 34, pp. 17695-17709, 2021.
Behjoo, H. and Chertkov, M. Harmonic path integral diffu- sion, 2024. URL hwwpst http:s///aa//////eeeeeaeeeeaaavr http://///www/./ba//a/eeseeessse 15166.
Del Moral, Nonlinear filtering: Interacting particle resolu- tion. Comptes Rendus de l'Académie des Sciences Series 1- Mathemafics, 325(6):6533655,, 1997. ISSN 0764-4442. doi: http://bbboog//ooooo666664444442 URL scceence//rrtaaeee//////76122222
Berner, J., Richter, L., and Ullrich, K. An optimal control perspective on diffusion- based generative modeling, 2024. URL https:ttp///... / /wt.ivivivv.///ee//eeeeeaaaeeeve http://wwtv....aaaaa.a//.a abs/ 2211 01364.
Blessing, D., Jia, X., Esslinger, J., Vargas, F., and Neu- mann, G. Beyond ELBOs: A large-scale evaluation of variational methods for sampling. In Salakhutdi- nov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceed- ings of the 41st International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 2205- PMLR, 21-27 Jul 2024. httpstss: httptttp://pproeeeeee http///ooeeeens..ceeo////o mlr. mal..press/ v235 9235/blessing244 html.
Doucet, A., de Freitas, N., and Gordon, N. (eds.). Se- quential Monte Carlo Methods in Practice. Statis- ties for Engineering and Information Science. Springer, 2001. ISBN 978-1-4419-28895
doi: 10.1007/
978-1-4757-3437-.. URL https: / /tt./.i.oggg//////
Fan, M. Zhou, R., Tian, C., and Qian, X. Path-guided particle-based sampling. In Forty-first International Con- ference Machine Learning, 2024. URL https: / //openreview, net./ forum? 4
Bolić, M. Djurić, P. M., and Hong. S. Resampling algorithms for particle filters: A computational com- plexity perspective. EURASIP Journal on Advances in Signal Processing, 2004(15) 2004. doi: 10. 10.1155/S1110865770008 org 10. 1155 htt.:1111/1111111111
Faulkner, M. F. and Livingstone, S. Sampling algorithms in statistical physics: guide for statistics and machine learning, 2023. URL https: http:///xxvvvvorr /bb/ 2208. 04751.
URL https::
Gabrié, M., Rotskoff, G. M., and Vanden-Eijnden, Adaptive monte carlo augmented with normalizing flows. Proceedings of the National Academy of Sci- ences, 119(10):e2194420111 2022. doi: 10.1073/pnas. 2109420119. URL https:/ /ww.. www..ppaas...mmm.m org doi/ abs/10 0/sp/10.1073/ppeee /wt..//as...... 21 9420119.
Bonanno, C., Nada, A., and Vadacchino, D. Mitigating topological freezing using out-of-equilbrrium simula- tions. Journal of High Energy Physics, 2024(4):126, 2024. doi: 1010077/HHEP4440441226 URL https: //doi worii.org/100
(2024) 126.
Bruna, J. and Han, J. Posterior sampling with denois- ing oracles via tilted transport, 2024. URL https: //arxiv. org/ abs/ httg/abb/2400..00 00745.
Grathwohl, W., Chen, R. T. Q., Bettencourt, J., and Du- venaud, D. Scalable reversible generative models with free-form continuous dynamics. In International Confer- ence Learning Representations, 2019. URL https: / //openeeview.com hwtp::nntt/oootuueeeeeeeee
Caselle, M., Cellini, E., Nada, A., and Panero, M. Stochas- tic normalizing flows as non-equilibrium transforma- tions. Journal of High Energy Physics, 2022(7):15, 2022. doi: 00.1007//HEP07/202200555 URL https: //doi. waiil.org11..0 wtt.:///g/./..////.....
Hartmann, C., Schütte, C., and Zhang, W. Jarzyn- ski's equality, fluctuation theorems, and variance re- duction: Mathematical analysis and numerical algo- rithms. Journal of Statistical Physics, 175:1214-1261, 2018. URL https: //tpp/appa/////// www.ntisschoooooo.om...o htts//oopsII::53949100
(2022) 015.
Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equa- tions. In Bengio, S., Wallach. H. Larochelle, H., Grauman, K., Cesa-Bianchi. N., and Garnett, R. (eds.), Advances in Neural Information Process- ing Systems, volume 31. Curran Associates, Inc., 2018. URL http:://proeeddinss.eee..e ce sc//appe_ffilee/ppper/////////// 6 (aB.Fe000.eccee22222243333399999 pdf.
Hastings, K. Monte carlo sampling methods using markov chains and their applications. Biometrika, 57 (1):97-109, 1970. ISSN 00063444, 14643510. URL http:/ www. jstorr org/ wtttaalee2334944
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. In Advances in neural information pro- cessing systems, volume 33, pp. 6840-6851, 2020.
10

---

NETS: A Non-Equilibrium Transport Sampler
Hua, M., Laurière, M., and Vanden-Eijnden, E. An efficient on-policy deep learning framework for stochastic optimal control. arXiv preprint arXiv:2410.05166 2024.
Neal, R. M. Annealed importance sampling. Statistics and Computing, 1 11(2):125-33, 2001. doi: 10.1023/ A:1008923215028. URL //tp:///oooooo/...//// 1023/A:100892321000
Hénin, J., Lelièvre, T., Shirts, M. R., Valsson, O., and Delemotte, L. Enhanced sampling methods for
Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. Action matching: Learning stochastic dynamics from samples. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 25858-25889, PMLR, 23-29 Jul 2023. httptttss:/pppeeeeeeeeeeeee v202 / nek yudov233. html.
molecular dynamics simulations [article v1.0]. Liv- ing Journal of Computational Molecular Science, 4(1): 1583, Dec. 2022. doi: /0.330011//veeeom444111111
URL htttpe:://iveecomsjooraal..ogoo http:/Lvecomm/artttclllee/ee////
Jarzynski, C. Nonequilibrium equality for free en- ergy differences. Phys. Rev. Letr. 78:2690-266
mlr press/
Apr 1997. doi: 10. 00.113PPhysReeeLett..2..00 URL /ttps:///likk...// www.n..ap..rg//oooooooo..o.ooo PhysRevLett, . 78.2690.
Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltz- mann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457)xcaaw 147, 2019. doi: 10.1 10.1126/science. aaw 147. URL https:
Lipman, Y., Chen, R. T., Ben- Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2022.
org/
doi/abs/ / /10.
Parno, M. D. and Marzouk, Y. M. Transport map accel- erated markov chain monte carlo. SIAM/ASA Journal on Uncertainty Quantification, 6(2):645-682, 2018. doi: 10.1137/17M1114444 URL https: org/10. 1137/17M113 640.
Liu, J. S. Metropolized independent sampling with com- parisons to rejection sampling and importance sampling. Statistics and Computing, 6(2):113-119) 1996. doi: 10. 10.1007/BFF0162522 1007/BF00 62521.
http:: //oii www.oiiorgg/1...
Sendera, M., Kim, M., Mittal, S., Lemos, P., Scimeca, L., Rector-Brooks, J., Adam, A., Bengio, Y., and Malkin, N. Improved off pooccy training of diffusion samplers, 2024. URL https: //axxiv org/ abs / /20.2/4222222 05098.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference Learning Representations, 2022.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poolc, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011. 2020.
Malkin, N., Lahlou, S., Deleu, T., Ji, X., Hu, E., Everett, K., Zhang, D., and Bengio, Y. Gflownets and variational inference, 2023. URL https: /tttpp://arxxv.orrt abs/ 2210. 00580.
Tian, Y., Panda, N., and Lin, Y. T. Liouville flow impor- tance sampler. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp. F. (eds.), Proceedings of the 41st International Conference Machine Learning, volume 235 of Proceedings of Ma- chine Learning Research, pp. 48 0866-88210 PMLR, 21- 27 Jul 2024. URL https: / wwwppooccediooo iigss .mlr. gress/v23355 https//223//tian2c..hete
Matthews, A. G. D. G., Arbel, M., Rezende, D. J., and Doucet, A. Continual repeated annealed flow transport monte carlo. In Proceedings of the 39th International Conference on Machine Learning, Proceedings of Ma- chine Learning Research, Jul 2022.
Midgley, L. I., Stimper, V., Simm, G. N. C., Schölkopf, B., and Hernández-Lobato, J. M. Flow annealed importance sampling bootstrap. The Eleventh International Confer- ence on Learning Representations, 2023. URL https: //openrviee net htwp:e/wwwet/fttnet//eeee -ww
Vaikuntanathan, S. and Jarzynski, C. Escorted free energy simulations: Improving convergence by re- ducing dissipation. Phys. Rev. Lett., 100:190601, May 2008. doi: 10.1003PhysReeeeett1000.0 19060000.. URL https://link: aps PhysRevLett 1 100..........000
Máté, B. and Fleuret, F. Learning interpolations between boltzmann densities, 2023. URL https://arxiv/
atg/abs/2333 07388.
Neal, R. M. Probabilistic inference using markov chain monte carlo methods. Technical Report CRG-TR-93-1, Department of Computer Science, University of Toronto, September 1993.
Vargas, F., Grathwohl, W. S., and Doucet, A. Denoising diffusion samplers. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: / //opnneviiew net forum? d=BpvnfTAbulf. FTAbulf.

---

NETS: A Non-Equilibrium Transport Sampler
Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. Trans- port meets variational inference: Controlled monte carlo diffusions, In The Twelfth International Conference on Learning Representations, 2024. URL https:// d=PP1 rudnxiw.
openreview.nen http:/ww...eo...eeeeeeeeeee
Vierhaus, Simulation of Φ² theory in the strong coupling expansion beyond the Ising Limit. PhD thesis, Humboldt University of Berlin, 07 2010.
Wilson, K. G. Confinement of quarks. Phys. Rev. D,
10:2445-2459, Oct 1974. doi: 10.1103/Physeev.. 10.2445. URL https:/ /tt:://1ia /wwnkiin.....// aps ops.org/ooi/1. 1103/PhysRevD. 2445.
Zhang, Q. and Chen, Y. Path integral sampler: stochastic control approach for sampling. In International Confer- on Learning Representations, 2021.
Zhang, Q. and Chen, Y. Path integral sampler: stochastic control approach for sampling. In International Confer- ence on Learning Representations, 2022. URL https: / /twpeeeeeeee..eeeee/eeeeeeeeeeee forum -d== uCb2ynRu7Y. 1
12

---

NETS: A Non-Equilibrium Transport Sampler
A. Proofs of Section 2
Here we provide the proofs of the statements made in Sec. 2 which, for the reader convenience, we recall.
Proof of Proposition 2.1. Let ft a) with (z, a) E [Rd+1 be the PDF the joint process (Xe; defined by the SDE (9)
and (10). This PDF solves the FPE
deJ. = (VU.J + S(=o(z, =((())o))(())(((1) 6(a)
(42)
Define
gc(z) / fe(x,a)da
(43)
We can derive an equation for (x) by multiplying both sides of the FPE (42) by el and integrating over a E R. Using
/ e @.f((lz, = à e f((2 deg.
| V (VU.J++ Viiida =v. =VV..VV, /
(' V / e (e(ee
JR
=4V (VUi+++ +
(44)
J e @,U,,Oafiaa = 0,U, /
R
= - /
e -Ü,0g9t
where we arrived the second equality in the third equation by integration by parts, deduce that
0.g. V (Vug + Vge)= DeUige.
(t=o(z)= Po(x) Fo(
(45)
The solution this parabolic PDE is unique and can be checked by direct substitution that is given by
ge (x) Fo
(46)
Note that this solution is not normalized since it contains Fo rather than Fg. In fact is easy to see that
J s((x)dr=
ga++ Jra.. fe a)drda Po+Fo
(47)
where the first equality follows from the definition of 9i and the second from its explicit expression and the definition of the free energy that implies Sg. Ue(()) dx = FeF. Equation (46) is the second equation in (11). From (46) we also deduce that, given any test function h : Idd R, we have
Jgan (((xfee( a)dada Snee e fe a)drda
JR. (()))ee))da
Sna ge (e(x))
SR h(x)e Fo
Snu U.(x)+Fo -nle-Ut())t
Ste M(x)e Uc(s)
(48)
Seu e -Uil))d
=eFi / Jp h((x)e
(x))
Rd
Since by definition of f( a) the left hand-side of this equation can be expressed as the ratio expectations over (X,, A.) the first equation (11) we are done.
13

---

NETS: A Non-Equilibrium Transport Sampler
Proof of Proposition 2.3. If b, satisfies (13), then Pj satisfies the FPE (14). Since (15) is the SDE associated with this FPE, (16) holds.
Proof of Proposition 2.4. We can follow the same steps as the proof of Proposition 2.1 by considering the PDF s (x. a)
of (Xg, A.). This PDF solves the FPE
a (VU,)) + +VS.) v v. -(b..i)) (V. b, 7U2-0,U,)0,ff
(49)
() (()))
Define
<<(x) = 1000[f00 st
(50)
We derive an equation for (x) by multiplying both sides of the FPE (49) by e and integrating over E R. Using
[ @a 01
/ Ea ff = a.g.
JR
J 6" CVV. (VU.. + vs = d < VVV. / (tdāa + V @sdda),
/ e (b.ff)da = - (b, / (t
= (brn)))
(51)
e e<(V b, - VU, - fa = - (V by - VUe - -
JR
/
=V 42 VU: -,U+))
J S/
(V b VU. -.,.))
where arrived the second equality in the fourth equation by integration by parts, deduce that
d.fk V. (VUis) + Vggi v (brg.) (V b, VU, a,U.)9
(52)
L (z) Po(x) e +%
The solution to this parabolic PDE is unique and it can be checked by direct substitution that it is given by
(e(() = (o
(53)
This solution is not normalized since it contains Fo rather than F,, and it is easy to see that
JoL. ( -f
fpo. e² fh a)dxda = Fot
(54)
where the first equality follows from the definition of and the second from its explicit expression and the definition the free energy that implies fga e- -e(()) e- F., Equation (54) the second equation in (21). From (53) we also deduce
14

---

NETS: A Non-Equilibrium Transport Sampler
that, given any test function h R I, we have
Jaat e( (z. a)drda Sn e [6 a addada
Jg (z)
Su st (c)de
Jn ((lu))-)
JEL u
Sn l1(x)ee Ul(e))d
(55)
Sa e v-vil
[ Ma)e
=eft
JRd h(z)
Since by definition of (2r, a) the left hand-side of this equation can be expressed the ratio of expectations over (X). A4) in the first equation in (21) we are done.
Proof of Proposition 2.5. Clearly the minimum value of (25) is and the minimizing pair (6, F) must satisfy
V b -
+ a, =
(56)
By multiplying both sides of this equation by P: is can be written
v -6.p.)) -.U.p
(57)
This equation requires solvability condition obtained by integrating it over Rdd. This gives
(.))()))())dz + a.F, = 0,
(58)
which, by (3), implies that O.F, F.F. In turn, this implies that (57) is equivalent to (13), i.e. by solves (13).
Proof of Proposition 2.6. Consider
/ log
(e(xr)
(pille))
(e(xx)dr
(59)
(((x))
gd
where de satisfies (27). Taking the time-derivative of this expression we deduce that (using (27), pc(x) e -U.())+FF the
15

---

NETS: A Non-Equilibrium Transport Sampler
identity 6,
(VUi Vice) V. (ArV Pt)): and multiple integrations by parts)
(.Dk((Pellpt) = / log
pc(x)) Oxpx(x) @epr(xt) pt(x) Px(xz) dx
(c(x))
(-v (t)()))))))) + (VU)()))L())+ + va((z))) da (x(TT)
log
+
(0.U,(1) - -r.F.)
be(l)) (ooo)-Voog (pe(x)
+ -.U,--80 F.F: (r(x)do
p((x))
(((x))
V (P((x)
(D((1)
c7
Vlog
Pr(x)
P((x)) P((x)dr
(60)
See (1(x)) Vil(a)) + (bc(z) (U(()) +UU -d(Fi) po(a)) dz
V Vlog (M(la) (p(x)) (x(x)x
[- -V. -x)( bo(x)) VUi(z) + 0,U, 0.Fi p((x)dz
Vlog (a(xa)
(x(x))
(((x))
Therefore
1/- Jzd
(RL(Pi=allP)) =
v be(x)) + b((x). VU((2) + 0,U, - 0.Fi px(x)dodt
/ E1
(Pc(z)
V loo
(((x)drdt
(c(x))
Rd
- (x(x) VU(()) + 0,U, - - Pe pc(x)dadt
(61)
RR
1/2
v (x(x) + b.(x)) VUi(x) + a.U, - -O.F. (e(())ddtt
Rd
LINN (6, F)
which gives (28). To establish (29) observe that
[6 F)
1/' Rd
bL( (2()) - (c(x)dddt
S:¹ [
V bx(x)) [..b(())
- 1x(x) (U.(x) O,U, + a 1 + 0 - a. 22
(62)
-2 (6,F) f)+2 / |0,F; 2
Therefore, if fo f. d. S 6, we have
L.INNN (b, F) < 2L F) + 25
(63)
Combining this bound with (28) gives (29).
B. Time-discretized Version of Proposition 2.4
Here we show how to generalize the result in Proposition 2.4 if we time discretize the SDE in (19) using Euler-Marayuma scheme and use some suitable time-discretized version of the ODE (20).
16

---

NETS: A Non-Equilibrium Transport Sampler
Proposition B.1. Let 0 to 1
kk 1 be a time grid [0, 1.. denote Alx 4--1 La sor h: ....., K
set X po and Ab = 0, and for k - 0,..., ,....,KK K - define X l4+1" b2+1 recursively via
X l++1 = Xx, -EN VUre (x,, (xx,)ttt Atx t + bee (x,, At + 2 (Wea4s - Wen),
(64)
1b dle + Ula (8 )= Utats (X (x+1) R (X tx
le41
x xhh Rg (X reS1 X bx
(65)
where we defined
R (Ir. v) = 4cr, Ats ly Als(ee VUl (x) F bb, (2:)
(66)
Then for all k = O,..., K and any test function h Rdd + R, we have
Ele Ai x(X2)1
J(c (z)d
Zie = =c Fix
Ble As
(67)
Ai.]
where the expectations are taken over the law of (X) A
Note that the weights in (67) correct for the bias coming for both the time evolution of (z) and the fact that the Euler- Maruyama update in (64) does not satisfy the detailed-balance condition locally. It cab be checked by direct calculation that (65) is a consistent fime-dicreeizzation the ODE (20).
Proof. For simplicity of notations we will prove (67) for lo = K: the argument for all the other 1,..., K 1 is similar. The update rule (65) implies that
A = Σ (8h))
Uee (X
(Kt
Rt
(X: xb (br11)
RI (X x%
k=0
(68)
K-1
= Uo((Ht) - UK (Xxx) ) + Σ
R (Xe,
xh x(en)) - Rs (8fe,,; x2h))
k=0
Now, given the test function h: Rdd - IR, consider
21 [h = El ( h ((x)x
(69)
Since the transition probability density function of the Muler-Maruyame update in (64) reads pt. ( = (4TE,, Ata) Atx) exp d/2
R (Ix zx+1));
(70)
the joint probability density function of the path (X:) 1o:
xb ....., x is given by
K
((oo....,KK)) = exp (-Uo(zo) + Fo)
II
Pt. (Ix+1l0))
k=0
(71)
K-1
CCexp (Uo((0)) + Fo
Σ k=0
R (1x,
where C - -I (4TE,, A44)-4/2, We can use this density along with the explicit expression for Ab in (68) to express the expectation (69) as an integral over p(ro. II....,
K-1
[b)] c JEATRI day dux exp
-Uo(20) + Fo
Σ k=0
R (1KK,T44))
(72)
K-1
x exp -
- UK((K) +
R (IK. - RT (2 M(*K)
Rm
17

---

NETS: A Non- Non-Equilibrium Transport Sampler
where the second exponential comes from the factor eAk (72) simplifies into
K-1
[b] C Se dao- dæg exp
-UK((E)++ + Fo
Σ Rs (2))L) h(zk)
(73)
In this expression we recognize product of factors involving
Pri (aa| (ATTE,,)) d/2 exp Ri (x))),
(74)
which is the transition probability density function of the time-reversed update
6- 11
C22 VU, (Yth ) ( tki A + 22 (W - Ww..
(75)
This implies particular that we can perform the integrals in (73) sequentially over 20, I1,.. 2K-1 to be left with
[|1] =
exp (-Ux(IK) + Fb) h( )d
(76)
JRd
Therefore
[[1] exp (-p--U((X))+ + Fo) dx c- FF
(77)
Rd
which is the second equation in (67), and
e Fk F Fo
/ cxp exp(-Ux(EX) + Fo) Mxx)dax [ h(x) (x)dr
(78)
[[]]
which is the first equation in (67).
C. Solving for the Optimal Drift via Feynman-Kac Formula
Without loss of generality, we can always look for solution to (23) in the form of bx(x) Vor(x), so that this equation becomes the Poisson equation
AΦ. VU, VΦ. 0,U, - O.F..
(79)
The solution this equation can be expressed via Feynman- Kac formula:
Proposition C.1. Let Xt. satisfy the following SDE
dx VU.,(X; + V2dW,,
x =x
(80)
where Ug evaluated fixed at [0, fixed. Assume geometric ergodicity of the semi-group associated with (80), i.e. the probability distribution of the solutions to this SDE converges exponentially fast towards their unique equilibrium distribution with density pr(x). Then for all (1, x) E [0, x Rd we have
(x
B (xU(({("))
(81)
where the expectation is taken over the law of Xt.
Proof. By Ito formula,
= x (A.)(
Vol(X;")) (2V0.(X(.) dW, (x, 0.Fi) dr + V2V dW,
(82)
where the differential is taken with respect & fixed, and we used (79) get the second equality. If we integrate this relation on 5 E [o,T] T] and take expectation, we deduce that
dx(x)) / B(OUX(() a.
E [xpr(x
(83)
where we use Ito isometry to zero the expectation of the martingale term involving V (Xt.*)) dW,, If we let T - co, by ergodicty the first term at the left hand side converges towards a constant independent of (f, x) which we can neglect this fixes the gauge of the solution to (79) which is unique only up to a constant. What remains in this limit is the expression (81).
Note that the integral in this expression converges since E [(a,U((Xt - F exponentially fast as T - oo by assumption of geometric ergodicity.
18

---

NETS: A Non-Equilibrium Transport Sampler
Example: moving Gaussian distribution.
Let us consider the case where
(x) -b,) A, bt), - where b, Rd is time-dependent vector field and A, A/ R x IR is time-dependent positive-defiite matrix: assume that both b, and A, are C' in time, and also that A, A.A.. The free energy in this example is
(84)
Fr - -log Zz,
4/2 Ze = (2/) | det Arl-1/2
(85)
so that
0,U,(r) = - Ar(æ - bb) + { b1)² Az( - -b),
0,F, - jurA, At).
A.))
(86)
In this case, the SDE (80) reads
dxt== = -A - -b,)dT + V2dW
x =
(87)
and its solution is This implies that (using Ito isometry)
X = = AcT x (1 c Ac5) + v2 {" c-A((r-))
(88)
E E(0,U,(X?:))
b7 AIC AcT (I (b) 2xx by) e Are Age-AAr
AcT (x be)
+ L ure Air A1fe-Ar5 A,e
A.T A.e-A.1) ) dr
(89)
= 67 A/e ACT - b1)) (1)+ + (=(xxbb) -b) Te -Ai
A:5 A1T - bi)
+ tlt A, AA) - =)rr(A, -A, AAe--2A1T)
Therefore, from (81), we have (using also (85))
dx(2) = [ Aee A,e Ar? Ae7 (x (6-bb) b -1 b,)² Te e AcT A.T- b1)
ArT
+ ( Ale -2ALT) dr
(90)
b. (x b.) [x - bb)² Á. A (x b.) + t Ar ArA
This solution checks out since it implies that
-VU((x) (z,(x) + Ap,(x) = -8 Ar b.) + (x - b,)² A,
b,)- - -
(91)
which is (z), O,F, as it should.
D. Extensions and Generalizations
D.1. Inertial NETS
It is straightforward to generalize Proposition D.1 so that the stochastic dynamics involves some memory/ineria:: Proposition D.1. Ler (xf:", R². Ah:1)) solve the coupled system of SDE/ODE axh b (x.xi1)) + Rh.'d.. dt.
x ~ Pu:
(92)
d.R." VU (xf"")tt )dt - peç - "R dt + 2G, dW.. 2e,, d V bx )dt VU.(x!!)) b,(xf. (
R N(O.
(93)
( ytt, 0,U,
Abw = 0.
(94)
where 6 is time-dependent diffusion coefficient, μ 2 0 is mobility coefficient, and W, Rdd is the Wiener process. Then for all t E [0, 1] and any test function h: Rd + R. we have
[ (x)pe(x)dd
Ele h(Xh")) EleA?")
7 e FFi 11
(95)
where the expectations are taken over the law of (X (X/1 Ab.H..
19

---

NETS: A Non-Equilibrium Transport Sampler
The proof of this proposition be found at the end of this subsection. Note that when 8 b, the solution (23), (94) is simply
A = F2 - Fo.
(96)
Lee. the weights are again deterministic with zero variance. In general, will not be the optimal one, in which case using the SDE (92)-(94) gives us the extra parameter μ to play with post-training to improve the ESS. Below we show that (92)-(94) reduce to (19)-(20) the limit as μ + oo. It is also easy to see that, if we set μ = in (92)-(94), we simply get that R 0 and hence (97) reduces to the ODE dx, = b, (X."")dt. Finally it is worth noting that (92)-(93) be cast into Langevin equations with some extra forces. Indeed, if we introduce the velocity vhl = (x(xx.))) ( + R2 (92)-(93) can be written as
dxp# V..ldt dt avBi bi UXUU,(X pei - (x, (x.x")d (xb Vbe(xbm)) fon pei "uiivV v.l dddd uV 2e, dW,, 2
x po.
(97)
" vol ((b((x uld)
(98)
In these equations, the terms pe, 0g can interpreted as non-conservative forces added to VVU and the term Vb.V extra friction term added to -ue, V.S.
Proof of Proposition D.I. Denote by pi. "x, r, a) the joint PDF of (xLe RI Abit). This PDF satisfies the FPE
x 0.f - (lb,
+) uVU, V,f uei uei ', (rfb + (V
- -(V (V. b VU. b, 0..
(99)
p (z, r. = (oll 1
Let
9: 6 (x,
[ e (x, a)da.
(100)
bi (x) by multiplying both sides of the FPE (99) by e" and integrating over E R. Using
We can derive an equation for 9t equations similar to (51), we arrive at
0.g. b. = -v (l6 + rlgg b:11 + uV V.f + -uu (rg. + +V.gi biy)
+ (V b VU, b, - 0.U.) b.7. ba (x, r) (2(x) (2mTu) -d/2 1r 1r| /(2))
(101)
Since Po(x)= = e Uo(s))+ Fo, it can be checked by direct substitution that the solution to this equation is
91 (x,()r)) - Fu (2(T)))))(22)m)))))))))))) 4/2
412 -11
(102)
Therefore
J gr. (x, r)dzdr = f é²
= e FitFo
(103)
where the first equality follows from the definition of g. b. and the second from its explicit expression and the definition of the free energy that implies Sga e- -U.(()) e F., Equation (103) is the second equation in (95). From (102) we also
20

---

NETS: A Non-Equilibrium Transport Sampler
deduce that, given any test function h Rd R. we have
SEzan eh(x)S( (s, r, a)dxdrda Jaann e fl, dxddra
Jaa Mo)gl
(x))gi" (a, r)d
Saeu di r)dddr
Ja h(x)e Uo
SRs (nee-u())+
Ja h(x)e (olx))--ull)d U.(o)dr
(104)
Sud e -vi(x)d
e [ h(z)ee dr
(x)dd.
Since by definition of "x, r, a) the left hand-side of this equation be expressed as the ratio of expectations (xbr Abdt in the first equation in (95) we are done.
To see what happens when - oo, let us assume that = (time-independent) and integrate (93) using Duhamel principle as
R2""" = = -He & ( e -e (t - / (t-)) dW,,
(105)
Letting -> 20, we see that the first term at the right hand side of (105) tends to zero, whereas the second one gives
lim S "- "
(106)
10000
Finally, the third term the right hand side of (105) is Gaussian process with covariance
min(t,
Ch 2
e (1-)) (((-)))) 2p (e-ae-" - -he
(107)
As result, given any test function di, have
6.C% = 2e f oidt
lim 100000
(108)
10.112
which indicates that Cf converges weakly towards the Dirac distribution eō(l - 1'). Putting these results together shows that in the limit μ 0o, R2 dt converges weakly towards (xVUi(x /2edW, which, if inserted in (92), reduces this equation to (19). The case where e, depends on time can be treated similarly.
D.2. Multimarginal NETS
Let U(a 2) be potential depending D C. RN with N e N well Rd. and assumed be continuously differentiable in both arguments. Assume that e ((a(aa)) is integrable in for all E D, and define the family of PDF
e(0,2) () = -((,))
F(a) - log
Ullaa)) dx.
(109)
gd
Finally, define the family of matrix-valued B(a, x): x Rd + R Rd. assumed to be continuously differentiable in both arguments. These quantities allow us to give a generalization of Proposition 2.4 in which we can sample the PDF g(a,2) æ) along any differential path 04 E
21

---

NETS: A Non-Equilibrium Transport Sampler
Proposition D.2. Let a [0,1] D be differentiable path in D and define the vector field b: [0, 1] x IR IRd as
(xx) di ((
(110)
as well as
(x) = U(0,,
F(or) = ==) c
(111)
Let (X) Abaa solve the coupled system of SDE/ODE
axba = 8 (x. dt c VU, (X: dt + totht-22cce 22,WW
xb ~ Po:
(112)
dA h = V Bx dt (X (x 0g (x("") - 0.U (X, ba ddt,
Ab
0,
(113)
where 6 > 0 is a time-dependent diffusion coefficient and We E RRd is the Wiener process. Then for all E [0, 1] and any test function h Rd - R. we have
J (x)d ElE ((x,1))
FF+
BBe
(114)
Ele 12
where the expectations are taken over the law of xhe shrn))
We will omit to give the proof of this proposition since is simple consequence of Proposition 2.4. The interest in formulating the problem in this new way is that is it easy to see that the right hand side of (113) (with F added for
convenience) can be written as
(8)(()) (U)(x) (3(x))
(115)
a (.. B(ar, x) Bl x.) V(l(,,, vall(an, VaF(a)))
Therefore we zero this term for all (a, x) D x Rd by picking the right B(a, x) we will obtain that (113) reduces to dA - -F dt, i.e. A/.m- = F Fp. Finding this optimal B( x) can be obtained using the following result:
Proposition D.3 (Multimarginal PINN objective). Consider the objective for (B,F) given by:
L [B, F1
=fo ( Rd
V= B(a,x) =(a,x)) Vall(a + F(a) 0( x)f(a)drda
(116)
where x(a,x) 0 is PDF in a for all D, and f(a) a PDF in a. Then ming L F 0. and all minimizers (BFF) are such that and B(a,, x) solves
V(a,z) x) E Rd
0=V2. (a.B(,,)) ((( Vallax
+ VaF(a),
(117)
and F(a)) is the free energy (109) for all D.
We will omit to give the proof of this proposition since it is a simple generalization of the proof of Proposition 2.5.
E. Estimating the Drift (x()) Vor(z) via Action Matching (AM)
In general (13) solved by many (x)). One way to get unique (up to constant in space) solution this equation is to impose that the velocity be in gradient form, i.e. set (x()) (z)(x) for some scalar-valued potential (x).. If do so, (13) be written V (V))())))) -dpt and it easy to see that at all times & [0, 1] the solution to this equation minimizes over òe the objective
(o(l()))P(()) - (.()))
(118)
( (z)
22

---

NETS: A Non-Equilibrium Transport Sampler
If (3) to set O,F, Sea (z) (2r)dz we can use the objective the right hand-side of (118) learn (1(x) locally in time (or globally if we integrate this objective E [0, 1). Alternativeyy, we can integrate the objective the left hand-side of (118) over & E [o and use integration by parts for the term involving (:pr(x) to arrive at:
Proposition E.1 (Action Matching objective). Given any E (0, 1] consider the objective for do(z):
T. [5) = S." [=VV)(())))) (.0.(()))
(119)
+ S
Then the minimizer (x) of (25) is unique (up to a constant) and by (x() = Vol(1) satisfies (13) all E [0, T)
This objective is analogous the loss presented in Neklyudov al. (2023), but adapted to the sampling problem. In practice, we will use again T (0, 1] for annealing, but ultimately we interested in the result at T 1. Note that, unlike with the PINN objective (25), it is crucial that be turned into an off- off-policy objective.
the correct (z( in the AM objective (119): that is, unlike (25), (119) cannot
Proof. By integrating by parts in time the term involving 0: in the AM objective (120), we can express is as
1/"
(120)
Rd
This is convex objective in I whose minimizers satisfy
(V = -Ppt--
(121)
This is (13) written in terms of bx()) = Vér (z). The solution of this equation is unique up to a constant by the Fredholm alternative since its right hand-side satisfies the solvability condition Jzu O.P: = 0.
If we use (x)) (x) in the SDEs in (19) and (20), we need to compute V bx()) = a which is computationally costly. Fortunately, when 6 > 0, the calculation of this Laplacian can be avoided by using the following alternative equation for AB:
1 (x+)) 6o(x60) - Br: A
(122)
6t
where
dB. = O,Ui + - + - Et
Voi(x
(123)
+ V 2 Vx((x) dW.
To derive these equations, notice (z) = (2)( the SDEs (19) and (20) reduce to
axp - e, VU + Vo. + 2e,,W,
XB po:
(124)
d -U vo(xf)d (,U,(x)
AB 0,
(125)
Since by Itô formula we have
dd)( ( (x)( vu + 22a v2aV aw, a Adi
(126)
we can express
1
1 a VV) (u)(aa) d (xp)
4
(127)
1 VVi(()))
2dt - 2 vv dW..
61
23

---

NETS: A Non-Equilibrium Transport Sampler
If we insert this expression in the SDE (125), can write it
Et d.A. dx dBr.
(128)
where dB, is given by (123). Integrating (128) gives (122).
In terms of implementation, we can write the AM loss (119) as
(x,(1 - 0.or dt + IE ([cAi
E B oJ (X+)]
AM [
(129)
B[e4
E[c4à]
B(c=49)
These expectations can be estimated empirically over solutions to (19) and (20) with (xx)) = Vol(a)) The above implementation is detailed in Algorithm 1.
F. Link with CMCD (Vargas et al., 2024)
Consider the process y solution to the SDE
ay e (V, dt 26, V log pi (Y:)) dt b. + 22,WW, Yd Po-
(130)
where & denotes the PDF of the process X defined by the SDE (19), i.e the solution the FPE
app = c V (VU.p: + vp!) V (bip)),
p = Po
(131)
The process y has a simple interpretation: it is the time-reversed of the process run using the time-reversed potential U1-1 and -6 that is, if the additional drift was the perfect one solution to (23), the law of X (e.))e and y (Y")eepo.)) should coincide. This suggests to learn b using objective divergence of the path measure of X from that of yrh This is essentially what is suggested in (Vargas et al., 2024), and for the reader convenience let us re-derive some of their results in our notations.
The Kullback-Leibler divergence (or relative entropy) of the path measure of x from that of y reads
1 XL(X"11r))= = 4 J. [[]VUi(x)) Vlog (xh)
(132)
This objective is akin the one used in score-based diffusion modeling (SBDM) and simply says that one way to adjust ô is by matching the score of o to that of pr. As written (132) is not explicit since we do not know V log p We can however make it explicit after a few manipulations similar to those used in SBDM. To this end, notice first that, by Ito formula, we have
d dlogo V log) (-xVV,,(})) + (x())) + log pf (x))dt
(133)
+ 22, log (x{)) aW.
which implies that
e, v log Aolog&((X) vU.(()))) -d ld di (xP) + Vog (x!)) b, (x G log / (x))dt 22, log p (x!) dW.
(134)
Inserting this expression in (132) after expanding the square, and noticing that the martingale term involving dW, disappears by Ito isometry and that the term d log p (x²) can be integrated in time we arrive at
<L(X*I1T) = 4
E (V
( Vlog ((x)) dt
1
[(ee|VU,
L
Vlog (x())
(135)
1
E E (xh) (x/)) + e log pi (x+)) dt
+
4
Jo
1 Eloog (Xx))
+
Ellog (X
4
24

---

NETS: A Non-Equilibrium Transport Sampler
where we used p = Po- We can now use the following identities, each obtained using V log P Vpp. and one integration by parts:
E V log pl
(x
v log (()))
2d
log (()) V
Rd
(136)
A log (2)
Rd
A log p (xx)),
E
(U)) [loggi(x}))
VU((2) v log (z) (2)))
/ VU)()) (
Rd
(137)
(U)(t))(())de
Rd
BE [AAU(X)))]
and
E (x)) b.,))
/ v Vlog (og)(()) b. (x)
(z)
(138)
Rd
E V bx())))
Inserting these expressions in (135), it reduces to
KL(X'1IXX) 4
/' BlclVUi((X)) - (xU,( - V V (x((x dt
(139)
Ellog po (Xu)) Eloog pi (X1)
This objective is still not practical because involves log pf. which is unknown. There is however simple way to fix this, by adding term in the Kullback-Leibler divergence (132)
KL' XL'(X'IXX*) (
(140)
This additional term is proportional to the Kullback-Leible divergence of p from the target PDF f1. Using (139) as well as pi(x) = e Fi we can now express (140) as
J B ((E]|VU.(X))) - -,AU, (x(x))]
KL'(X"1X*")
(141)
(Xo)) [ F.
+
This is Equation (24) in (Vargas al., 2024) which we set Vo.(=) (x)) and we used that, for any e R - Rd, we have
E %
Tw, = 22a [ (xx)
E[V a (x)))
(142)
Note that we can neglect the term Elog Po (X,) in (139) since it does not depend on b, so that the minimization of (139) can be cast into the minimization of (after multiplication by 4)
J E B(elVU, - V bx(x)) + ((UI(()) Fil
(143)
[./ / ( - v -(((=))) / (Ui(z) FF))
25

---

NETS: A Non-Equilibrium Transport Sampler
where & solves (131).
Let us check that the minimizer of (143) is b, = br the solution to (23), so that we also have p - p Pr. To this end, notice that the minimization of (143) can performed with the method Lagrange multiplier, using the extended objective
L. [<IVU.(=)) - -(A()) v P) (o) + / (UI(2) (o)d
(144)
L. (0.) - V. (VU,p) + V + (bep)))
where )x( is the Lagrange multiplier to be determined. Taking the first variation of this objective over X, p.: and b, we arrive at the LugranLagrange equations
= a.p. e V (VUip) + vp.) + V. (bip:),
O Po.
0 2 - GAU - v.b.
(145)
- -cee +G Vd aAde b. Ve
A1 = U1 + Fi
= - -
We can check that (x) = (z)), (xx) = (x)) e FU))+ and de(x)) = -U(x) + F& is solution: indeed this solves the first and the last equations in (145) and reduces the second to
0 [ VU1 CAU V be]
+ 0,U, O.F. c V + AU, + Vb, VU,
(146)
- b. 0,U 0,1i Vb, VU,
which is satisfied since b, solves (23).
G. Hutchinson's Trace Estimator for the Evaluation of (x(2))
It is well-known that, if (z) is bounded,
V (x) 25 B[n (fe(x + 5n) - be(a 5n))] + 0(0²),
(147)
where <6 < is an adjustable parameter and ~ N Id). Indeed we have
28" (be 8n) (z 8n)) (/Vb((2)+0(6))
(148)
which implies (147) after taking the expectation over We can use this formula estimate the PINN loss via
T.8 16, Fl 1-" E B. (I ( 1) r/)
(149)
where the expectation is now taken independent over It. ~ pr, n N(0, Id), and v/' ~ N(0, Id), and we defined
R (z. n) = 25" (b.(1 + 8m)) - b.( 8n)) (U) (x(x) - a,U,(x) + a.
(150)
The expectation in (149) is unbiased since 1 .. n/,, and its accuracy be controlled by lowering 6.
H. Details on Numerical Experiments
In the following we include details for reproducing the experiments presented in Section 3. An overview of the training procedure is given in Algorithm 1. Note that the SDE for the weights can replaced with (122) when learning with &. as one would do with the action matching loss (1 (11)..
26

---

NETS: A Non-Equilibrium Transport Sampler
H.1. Performance metrics
Effective sample size.
We can compute the self-normalized ESS
N i= exp (4:)
ESS,
(151)
N N N x==1 exp (2At)
at time along the SDE trajectory. We can use the ESS both as quality metric and as a trigger for when to perform resampling of the walkers based on the weights, using, e.g. systematic resampling (Doucet et al., 2001; Bolić et al., 2004). Systematic resampling is one of many resampling techniques from particle filtering wherein some walkers are killed and some are duplicated based on their importance weights.
2- Wasserstein distance. The 2- Wasserstein distance reported in Table were computed with 2000 samples from the model and the target density using the Python Optimal Transport library.
Maximum Mean Discrepancy (MMD). We use the MMD code from (Blessing et al., 2024) benchmark the performance of NETS Neal's funnel. We use the definition of the MMD
MMD? (p.p) n(n 1)
Σ le (2., R2)) + 2
1
be (Zi, x,)
2
ΣD Σ (z),
(152)
m(m 1)
2
where å ~ ý is from the model distribution and x ~ p is from the target and h: Rdd x Rdd - R is chosen to be the radial basis kernel with unit bandwidth.
H.2. 40-mode GMM
The 40-mode GMM is defined with the mean vectors given as:
μ41 = ( (-0.299) 0.9999, 21.4577)
#22 = (-32.9218, -29.4376)
μ4 = (- 15.4062, 10.7263),
14 = ( (-0.7925, 31.7156).
145 3..5988, 10.5845),
(46 12.0885. -7.8626).
J77 = (-38.2139, 26.4913).
A48 (- 16.4889, 1.4817),
μ/ = (15.8134, 24.0009).
P10 = (-27.1176, -17.4185).
J11 = (14.5287, 33.2155)
P12 = (-8.2320, 29.9325).
J41 = ( (-6.4473, 4.2326).
A11 = (36.2190, -37.1068),
N15 (-25.1815, 10.1266).
P16 (-55.5920, 34.5600),
N17 (-25.9272, - -1844133.
P18 (-27.9456, -37.4624).
119 = (-23.3496, 34.3839),
1220 (17.8487, 19.3869),
= (2.1037, 20.5073)
(6.7674, -37.3478)
M23 = ( 28.99026, 20.6212).
P24 (25.2375, 23.4529),
μ25 = (-17.7398, 1..433),
P226 (25.5824, 39.7653),
#27 (15.8753, 5.4037),
P28 (26.8195, -23.5521),
H229 (7.4538, -31.0122),
Р30 (-27.7234, -20.6633).
/331 = (18.0989, 16.0864).
[A32 = (-23.6941, 12.0843),
= (21.9589, -5.0487,,
AA34 (1.5273, 9.2682),
μ45 = (24.8151, 38.4078).
Р36 = (-30.8249, -14.6588),
M43 (15.7204, 33.1420),
Pa8 (34.8083, 35.2943),
M399 (7.9606, -34.7833),
P40 (3.6797, -25.0242)
27

---

NETS: A Non-Equilibrium Transport Sampler
These means follow the definition given in the FAB (Midgley et al., 2023) code base that has been subsequently used in recent papers. The time-dependent potential (1) is given by the interpolation of means. To directly compare CMCD and NETS on this task, i.e compare the log-variance path loss versus the PINN, we test the same mean-interpolation potential for CMCD using their code. The results are summarized in Table 3.
GMM-40
ESS
Reep
100
napp
Depp
naep
NETS-PINN
0.979 + 0.002
0.988 0.001
3.14 + 0.46
3.01 0.49
CMCD-LV
+ 0.030
0.51
9.13 + 1.17
Table Comparison NETS and CMCD with the log-variance loss on GMM-40 using Effective Sample Size (ESS) and Wasserstein-2 distance for the mecn-interpolation potential U..
H.3. Neal's 10-d funnel
The Neal's Funnel distribution a 10-d probability distribution defined as
2g
N(0, o²)
F1.9 N(0,e*0)
(153)
where and we use subscripts here as a dimensional index and not time index like in the rest of the paper. Following this, definition of the interpolating potential:
U,(x) = - - 0z
c
1z0 x + (d 1)tz0
(154)
1=1
so that at t = 0, we have Uo(x)= = 1 + Σ 23
and at time t = 1 we have the funnel potential given as (1())
- x (d - 1)
20 To
H.4. 50-d Mixture of Student-T distributions
Following (Blessing et al., 2024), we use their mixture of 10 student-T T distributions 50 dimensions. We construct U, via interpolation of means from a single standard student-T distribution (mean 0). We use the same neural network as used in the GMM experiments.
To further drive home the fact that our annealed Langevin dynamics with transport be taken post-training to the co limit to approach perfect sampling, we provide the following ablation from our model learned with the action matching loss given in Figure 4.
I. Details of the x² model
We consider the Euclidean scalar 6² theory given by the action
Spuc [pc / p(1)) + m²) + (x) d=)))))>> dDx
(155)
where we use Einstein summation to denote the dot product with respect to the Fuclidean metric and D is the spacetime dimension. We are interested in acquiring variant of this expression that provides fast computational realization when put onto the lattice. Using Green's 's identity (integrating by parts) we note that
((uplt))' ((x))'' ʃ Oup Oup d² x===d==---- =
/ 0(()d_ + vanishing surface term
(156)
so that
SEaclyp = -xxxx)00,00((1) + (x) + (x)
(157)
28

---

NETS: A Non-Equilibrimm Transport Sampler
MoS (d 50) Target. NETS-AM
360
W2
290
220
Diffusion coefficient
Figure 4. Reduction in W22 distance from taking the limit in sampling with NETS. Note that the resolution of the SDE integration must increase to accommodate the higher stochasticity. Average taken over 3 sampling runs of 2000 walkers each.
Discretizing Spne onto the lattice
A = {a(no........ nd-1) n, E
; i 0,,,...,dd, R+}.
where a is the lattice spacing used to define the physical point æ = an, we use the forward difference operator to define
0-f(x)= l μ4) sxxx)]
0,0" (p(x p.) 2x(xx) pp(x 14).
(158)
Using these expressions, we write the discretized lattice action as
SLa =
Σ D
Σ
10x+u02 -242 + 6x [(x-wpz] + mm?pp 2 + Xpp
(159)
a
zEA
Σ a
2Da 22p - -² 22 Σ' [PottWA + Pa--ula + m² + Xu2
(160)
Σ d
2Da 2 44 2a " [ + m + Xup
(161)
Σ
D
-2a 2 Σ
PuPett + (2a-2 + m²) Xxp
(162)
where we have used the fact that on the lattice En Po = Ex Px-pP2 to get the third equality. useful put the action in form that independent of the lattice spacing a. so, we introduce the re-scaled lattice field as
6x + D/2-1 1/2 Puy
m² + a²m² and x 4-
(163)
Plugging these rescalings into (162) gives the final expression
SLat
Σ 2 2 PoPatH
+ (2D m²) + Xxp
(164)
which we are use in simulation.
29

---

NETS: A Non-Equilibrium Transport Sampler
I.I. Free theory
Turning off the interaction makes it possible analytically solve the theory. To do this, introduce the discrete Fourier transform relations
6Pk
Pae
ik.r
(165)
VLD
1
Pkethe
9x V
(166)
for discrete wavenumbers
21 with 1 = 0,... L 1. Plugging in (166) into the first part of (162), get the expanded
sum
2 P2PItH
2
ΣΣΣΣ
e
(167)
La
-22 Σ
(168)
Σ Σ
PK-P-K i
(169)
-
Σ
Σ
(170)
2 Σ>
Σ
1px|2(con k + isintpl = -E 1spk cos kp &
(171)
where o² indicates conjugation, and we got the first equality by the orthogonality of the Fourier modes, the second by the Kronecker delta, and the third by the reality of the scalar field. Proceeding similarly for the terms proportional to y² gives us the expression
S
Σ m + 2D - > cos kip
(172)
The above equation can be written in quadratic form to highlight that the field may be sampled analytically
S I
1
Σ sok -
(173)
where
Mr.-k
m² +DD -0-22 Σ
cos ki 5k.-k
(174)
P
Note that this free theory can be sampled for any m² > 0.
1.2. 5t numerical details
We numerically realize the above lattice theory in D-2 spacetime dimensions. We use interpolating potential with time dependent m = (1 1))m + Lmf, A (1 - i) + LYi where Å is always chosen to be 0 (though we note that you could run this sampler for any Up that you could sample from easily, not just analytically but also with existing MCMC methods), For the - 20 (d = I. x L = 400 dimensional) experiments, we identify the critical point of the theory (where the lattices go from ordered to disordered) using HMC by studying the distribution of the magnetization of the field configurations as (x)) Σn (ar), where summation is taken over all lattice sites on the ęth lattice configuration. We identify this at mf = -L0, A 0.9 and use these as the target theory parameters on which perform the sampling. For the L 16 test (d 256), go past this phase transition into the ordered phase of the theory, which we identify via HMC simulations at
mf = 1100, \ = 0.8.
30


Reviews:
[{"review_id": "khcKwKSB64", "reviewer": "Reviewer_VhJb", "summary": "This paper investigates sampling from a target distribution within the annealed importance sampling (AIS) framework. Inspired by Jarzynski equality, a continuous-time version of AIS can be formulated using an SDE for samples and an ODE for weights. Building on this, the paper proposes NETS by introducing an additional drift function into AIS, and propose learning this drift function using either a PINN loss or an action matching loss. Experimental results demonstrate the effectiveness of the proposed algorithm.", "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 thank the reviewer for their valuable feedback, and we are happy to hear you found the work theoretically sound and novel. Below,  we try to address all your questions, and provide some info on new additional experiments. We itemize our response according to the headings that appear in you review:\n\n**References**\n- Thanks for pointing us to this paper. The SMC that they propose is also something that we propose in this work (when we may choose to do 'resampling'). We are happy to add a citation to this work, but we wish to note that this falls within the ICML policy of concurrent works:  \"https://icml.cc/Conferences/2025/ReviewerInstructions\". \n\n\n**Weaknesses**\n- Proposition 2.4 in general is not new, though this derivation is. Please note that it is also not new in Vargas et al. The essential machinery of this proposition has been known for 20 years in the statistical physics community (Jarzynski 1999, Vaikuntanathan & Jarzynski 2008), and neither us nor Vargas et al. are claiming to have discovered this equality. However, we both provide different derivations of it that make it more interpretable for sampling. Vargas et al. prove this result through the use of Girsanov, and here we provide a proof of it through simple manipulations of the the Fokker-Planck equation and other PDEs.\n\n\n**Other Comments Or Suggestions:**\n- We are happy to improve the readability of equation 6 when we can edit the text if the paper is accepted.\n- thanks for catching the typo in Eq (40).\n\n**Questions:**\n- *backprop thru sde and numerical instability:* The first drawback cannot be entirely addressed by the CMCD log-variance loss. It is still a loss function that must use samples along a *trajectory* as input, and the generation of such a trajectory of lengthn $n_{step}$ requires $O$($n_{step}$ $\\times$ network size) memory. By numerical instability, we mean that the CMCD LV loss becomes unstable if the trajectory generation is performed over too few steps. After discussing with CMCD authors, this is an issue generally when $n_{step} < 256$. This will become apparent in the \"additional experiments\" section below. Let us also stress that, in contrsat, our PINN loss can be evaluated pointwise in space and time. \n- *Optimize then discretize:* Thanks for the observation. We imagine most other diffusion models could do this as well, but the main obstacle would be whether or not their method allows for an adaptable diffusion coefficient *post-traing* (like our method). If not, increasing the diffusion would require decreasing the time-step for trajectory generation during training. \n- *PINN vs KL*: The benefits of the off-policy nature of the PINN loss manifest themselves in two ways: 1) The PINN loss is valid with respect to any sampling distribution), and 2) if you want to use simulated samples from your model, you do not need to backpropagate through their generation. The KL loss explicitly needs samples from the model distribution, and requires a backprop through the solve. As mentioned earlier, the log-var loss still needs to generate trajectories but has a smaller memory footprint. We did test the log-variance more, too. See: \"additional experiments\".\n- *Paper by Chen et al:* Thanks again for pointing us to this. We are happy to reference it, though again per ICML policy it is out of the scope of this submission to benchmark it. \n- *$\\epsilon \\rightarrow \\infty$ limit:*   Yes, that is exactly the motivation of the additional drift function. For our setup, perfect sampling is achieved in this limit, whether the learned transport is perfect or not. While this limit cannot be reached in practice without transport (as it would require taking astronomically large value of $\\epsilon_t$ in general), we show that, with some learned transport added, even moderate values of $\\epsilon_t$ can improve the sampling dramatically. This feature can be exploited after training as an explicit knob for tuning performance vs cost. This has not been recognized in other ML sampling literature.\n\n\n**Additional experiments**:\nIn the anonymous drive link https://tinyurl.com/netsicml :\n- We provide additional comparison between NETS and CMCD on the mean-interpolating time-dependent potential, as asked by another reviewer. We see that NETS performs well regardless of number of discretization steps, but the log-variance struggles for small steps. Note that LV begins to perform better when $n_{step} = 256$. We worked with the CMCD authors to implement their code and otherwise use their hyperparameters. This is the numerical instability we referred to earlier.\n- We have also set up a test on the Lennard Jones potential, and have included preliminary results. We are continuing to refine these experiments and will incorporate them in the final version.\n\nPlease let us know if you have any other questions, and thanks again for your insights. If we have clarified everything for you, we would greatly appreciate an increase in your rating.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "This paper introduces a novel sampling algorithm inspired by concepts from nonequilibrium physics, aimed at enabling efficient sampling from unnormalized distributions.", "essential_references_not_discussed": "Yes. One related reference is missing: Junhua Chen et al. Sequential Controlled Langevin Diffusions. ICLR, 2025.", "other_strengths_and_weaknesses": "Strengths:\n\n- The paper is generally well-structured and theoretically sound.\n\n- A novel objective based on PINN or action matching is proposed to learn an additional drift function in the context of AIS.\n\nWeaknesses:\n\n- The main proposition of the paper (Proposition 2.4) has already been presented in Vargas et al. (2024), although the current derivations are arguably simpler to some extent.\n\n- Please see the questions below for the authors.\n\n---\n\nReference:\n\nFrancisco Vargas et al. Transport meets Variational Inference: Controlled Monte Carlo Diffusions. ICLR, 2024.\n\nJingtong Sun et al. Dynamical Measure Transport and Neural PDE Solvers for Sampling. Arxiv, 2024.", "other_comments_or_suggestions": "- 2.2. Non-equilibrium sampling with importance weights: I didn't initially see how the right-hand side of Eq. (6) could be related to weights until the introduction of Jarzynski equality. The readability could be improved.\n\n- Line 669: the term $\\hat{b}_t$ is missing in Eq.(40). Should it be $\\nabla U_t \\cdot \\hat{b}_t$? It seems that this omission also occurs in other equations in the appendix. Please check it.", "questions_for_authors": "- I have the following questions when comparing NETS with CMCD:\n\n  - Regarding the sentence 'requiring either backpropagation through the SDE or computation with a numerically unstable reference measure on a fixed grid' (Line 73), I guess that the first drawback can be addressed by off-policy divergences, e.g., log-variance. Could you clarify what is meant by 'computation with a numerically unstable reference measure on a fixed grid'? Should the grids be tuned to avoid instability?\n\n  - Regarding the sentence 'Moreover, our optimize-then-discretize framework allows for post-training adaptation of both step size and time-dependent diffusion, providing tunable parameters to enhance performance.'  (Line 79), I am curious if this approach can be applied to other diffusion-based samplers, or what might prevent other samplers from doing the same?\n\n  - Aside from the fact that the PINN loss can be trained in an off-policy manner, are there any other advantages over KL divergence optimization? Additionally, have you considered experimenting with CMCD-LV, which can also be trained in an off-policy manner?\n\n  - Missing reference: Junhua Chen et al. Sequential Controlled Langevin Diffusions. ICLR, 2025, where CMCD combined with SMC (AIS with resampling) is proposed. I am courious about the comparision in terms of both concept and experiment.\n\n- Figure 4: Taking $\\epsilon_t \\rightarrow \\infty$ shows improved performance, although with more discretization steps. This should be the motivation for introducing an additional drift function - it may be imperfect for most of the time but still helpful, thus leading to fewer discretization steps?\n\n- Line 614 - 628: Is the definition of $g_t$ sepicial? Usually, we don't need $e^a$ to define the marginal distribution. The motivation here is to link it to the weight function? If so, does Eq.(36) describe the evolution of a *weighted* distribution?\n\n\n---\n\nUpdate after rebuttal: I have raised my score to 4.\n\n---", "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": "0BjlYRu7Md", "reviewer": "Reviewer_kVWz", "summary": "The authors propose NETS, a Non-Equilibrium Transport Sampler that interpolates between two unnormalized densities $\\rho_0$ and $\\rho_1$ based on a user-defined choice of interpolant. A key contribution of the proposed approach is to introduce learning in the dynamics of continuous time annealed importance sampling by leveraging a learned drift which aims to minimize the variance of the importance weights, which can be combined with Sequential Monte Carlo by re-sampling during the simulation procedure based on the evolved importance weights. The paper introduces a physics informed neural networks (PINN) based loss to learn the drift function, and show that it additionally provides bounds on the KL divergence. Experiments are conducted on multiple different unnormalized densities, ranging from low and high dimensional Gaussian mixture models, Neal's funnel and mixture of Student-t distributions, as well as $\\phi^4$ lattice theory and highlights the superiority of NETS over some of the established baselines.", "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 thank the reviewer for their valuable feedback, and we are happy to hear you found the work theoretically sound and novel. Below, we try to address all your questions, and provide some info on new additional experiments. We itemize our response according to the headings that appear in you review:\n\n**Methods And Evaluation Criteria:**\n- *estimating Z_1:* Yes, it's straightforward to construct an unbiased estimate if we know $Z_0$: $Z_1 = Z_0 \\cdot \\mathbb{E}[e^{A_1}]$. This follows from equation (21) at $t=1$,  which can beestimated empirically over trajectories from the coupled SDE/ODE (19,20).\n- *complexity of experiments:* We have included preliminary results on a Lennard-Jones system in an \"Additional Experiments\" section discussed below. Note that the $\\phi^4$ models studied here are arguably more complex than the LJ-13 system: these $\\phi^4$ models are either 256 or 400 dimensional problems, much bigger even than the LJ55 system, and are known as a proxy for one of the hardest sampling problems in research -- studying the strong force in lattice field theory. In particular, the $\\phi^4$ models, like these lattice systems, suffer from what is known as critical slowing down in which MCMC algorithmic efficiency dramatically decays as one approaches the critical parameters. It has also been a target for neural samplers for some time [2].\n- *Dropped Baselines for iDEM, CMCD:* We only included baselines where we could either quote authors' results or work with them to verify implementation. We did the latter with CMCD on Funnel and GMM, but couldn't for MoS distribution.\n- *SMC/HMC*: SMC can be incorporated within our approach (this is what we refer to as \"resampling\"), but on its own isn't efficient enough for our sampling tasks. We tried HMC-type sampling (Inertial NETs, Appendix D.1) but saw little difference. \n- *CMCD interpolation path:* We thank the reviewer for pointing this out. We have included in the \"additional experiments\" section below results on CMCD on the path interpolating the means, which we worked with the CMCD authors to set up. With this path as well, CMCD-LV exhibits instabilities if the number of sampling steps is too small in training: while the CMCD-LV loss is \"off-policy\", it still requires a trajectory on which to perform the optimization, and the generation of such a $n_step$ trajectory requires O($n_{step}$ $\\times$ network size) memory. If the trajectory generation is performed over too few steps the CMCD LV loss becomes unstable. After discussing with CMCD authors, this is an issue generally when $n_{step} < 256$. This will become apparent in the \"additional experiments\" section below.\n\n**Theoretical Claims:**\n- *eq (6) and (7):* SMC can be incorporated within our approach (as \"resampling\"), but on its own isn't efficient enough for our sampling tasks. We tried HMC-type sampling (Inertial NETs, Appendix D.1) but saw little difference.\n- *derivation of (7) and (10,11):* We'll add derivations in the appendix, as these aren't commonplace in ML literature.\n- *Why learn $\\partial_t F$?:* The loss function is a PINN loss which means it is minimized at 0 (i.e., when the physics equation is solved). $\\partial_t F_t$ is a free parameter in this equation and must be fit if one wants to learn off policy. We do indeed use that $\\hat b_t \\rho_t$ vanish at infinity, which is a consequence of the conservation of probability that prohibits the existence of a probability current at infinity. \n- *uniqueness of PINN minimizer:* The minimizer is unique *if and only if* the velocity field is learned in gradient form $\\nabla \\hat \\phi_t = \\hat b_t$.\n- *integration by parts:* This is again a consequence of conservation of probability requiring no probability flux at infinity. \n\n**Experimental Designs:**\n- *PIS not included:* The PIS authors confirmed to us that they used $\\sigma^2 = 3$ (not $\\sigma^2=9$ as stated in their paper). This makes the problem much easier, so we omitted this experiment.\n\n\n**Additional experiments**:\n\nIn the anonymous drive link https://tinyurl.com/netsicml:\n- We provide a NETS vs CMCD comparison on mean-interpolating potentials which shows NETS performs well regardless of discretization steps, while log-variance struggles for small steps, improving only when $n_{step} ≥ 256$. \n- We've included preliminary Lennard Jones results and will incorporate refinements in the final version.\n\n\nWe'll make your suggested expository edits when allowed to edit the paper. Thank you for your insights - if we've clarified everything, we'd appreciate an increase in your rating.\n\n[1] Alpha Collaboration, \"Critical Slowing Down and Error Analysis in Lattice QCD simulations,\" Nuclear Physics B, 2010. \n\n[2] Albergo, Kanwar, Shanahan, \"Flow-based Generative Models for MCMC in lattice field theory,\" Physical Review D, 2019.", "claims_and_evidence": "The fundamental claims made by the work is proposing a framework for sampling according to an unnormalized density which uses learned methods and can be trained without back-propagating through the dynamics. The loss considered for this training is well motivated and the claims are supported by mostly clear theoretical evidence, with some concerns that I have raised in sections below. Broadly, I think the authors have provided good evidence (mostly theoretical) regarding the claims that they make and I outline more specific questions and concerns regarding the evidence provided in the later sections.", "methods_and_evaluation_criteria": "Yes, the methods and evaluation criteria considered in this work make sense. I do, however, have a few questions regarding the evaluation criteria and the benchmarks considered.\n\n- In the proposed approach, is it possible to approximate the density $\\hat{\\rho}_1(x)$ corresponding to any sample $x$? In particular, I am wondering if it would be possible to additionally evaluate the proposed method to approximate the partition function $Z_1$?\n- Is there a reason why the authors authors did not evaluate their models on the Lennard-Jones potential or more complex problems to better benchmark against related work?\n- Why were some of the baselines (eg. iDEM, CMCD-LV) dropped for Table 2?\n- Is there a reason why the authors do not consider SMC or HMC as one of the baselines in Tables 1,2?\n- The authors should also consider evaluating against CMCD with the same interpolation of potentials as it is imperative to perform this fair comparison when evaluating the benefits of NETS.", "theoretical_claims": "I went over some of the theoretical claims made in the paper as well as some of the proofs outlined in Appendix A. While I enjoyed reading the theory outlined in the paper, there were some steps that were not obvious to me\n\n- It was not clear as to how the authors jump from equation (6) to equation (7). In the latter, when talking about gradients and divergences, is it w.r.t the augmented state or the original state? It would be beneficial if the authors could actually sketch out the proof underlying this step.\n- Assuming we have arrived at equation (7), how do the authors convert the partial differential equation to the coupled system of differential equation considered in Proposition 2.1? While the subsequent details and proof regarding expectation w.r.t the measure $\\rho_t$ is clear, where and how does the couple dynamics come from the partial differential equation itself was not clear. I would request the authors to also put more focus and background into this step.\n- Why is learning $\\partial_t F_t$ a good idea? Is it because equation (49) implies that the objective is minimized if and only if the learned $F_t$ perfectly approximates the true free energy? But how does equation (48) lead to equation (49)? Are there any regularity conditions needed that $\\rho_t$ vanishes faster than $\\hat{b}_t$?\n- In Section 2.5, can the authors talk about uniqueness of the solution that can be obtained? Why can there be more than one minimizer $(b, F)$?\n- In Appendix A, within equation (35) when the authors use integration by parts, how do they ensure that the term at the boundary vanishes?", "experimental_designs_or_analyses": "The experimental design and analyses conducted in this work is sound, except for one potential concern. The authors highlight that PIS did not converge on the Funnel task, while [1] shows that it does converge for this task. Could the authors clarify on what was the issue they came across?\n\n[1] Sendera, Marcin, et al. \"Improved off-policy training of diffusion samplers.\" Advances in Neural Information Processing Systems 37 (2024): 81016-81045.", "relation_to_broader_scientific_literature": "In my opinion, this is an interesting work and has key contributions to the scientific literature. In particular, it combines ideas from diffusion models and annealed importance sampling / sequential monte carlo and pushes the frontier of learned samplers.", "essential_references_not_discussed": "The authors cover most of the relevant literature off the top of my head.", "other_strengths_and_weaknesses": "**Strengths**\n\n- The work is well motivated and tackles a challenging and very relevant problem. The combination of annealed importance sampling and diffusion style transport methods is novel and a good contribution to the field.\n- The experiments highlight that the method seems to work well when compared to other learned samplers on the suite of tasks considered.\n- It also provides some nice properties in the sense that one can compute expectations w.r.t the target measure through the use of importance weights.\n\n**Weaknesses**\n\n- I think the writing needs a bit more work. The way I understood it is that the authors consider defining an interpolation $\\rho_t$ and then assume a drift $b$ which is trained with a loss directly derived from the continuity equation, i.e. $b$ is learned to satisfy the continuity equation for the marginals $\\rho_t$. This learned drift is then artificially added to the Fokker Planck equation and then relevant terms are grouped together to define the drift, diffusion and importance sampling weight updates part of the coupled dynamics. In my opinion, this story could have been more clearly explained so that it makes it easier for a broader community to understand the work.\n- The authors should consider a few more complex experimental settings, in particular the Lennard-Jones potential would be a good candidate.\n- The authors mention that they can estimate the drift off-policy as well but they do not provide any experimentation with it. I think this would be a good ablation to add, and can be shown on perhaps a relatively easier task for some off-policy path.", "other_comments_or_suggestions": "- Did the authors forget to include a concluding section or is this by design? Currently, the end of the paper feels quite abrupt.", "questions_for_authors": "I just have one additional question, apart from the ones raised already. \n\n- Within the related work section on augmenting sampling with learning, the authors talk about minimizing the KL divergence between model and target as well as stochastic optimal control. Aren't these two approaches the same, where the latter minimizes KL divergence between the considered path and a reference path measure through application of Girsanov's theorem?\n\nApart from this, I think the paper is well motivated and well positioned, so I will be happy to raise my score if the authors could clarify the questions that I have raised.", "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": "0LlkzXTMW0", "reviewer": "Reviewer_5UR2", "summary": "This paper introduces an algorithm for sampling from unnormalised probability distributions, through non-equilibrium sampling approaches.   When computing expectations with respect to the final-time marginal distribution, classical approaches to this would leverage AIS or equivalently Jarzynski / Crooks.      This approach, instead introduces a correction term, which corrects the discrepancy between the marginal distribution of the non-stationary 'Langevin' dynamics and the true distribution.    This discrepancy can be expressed  as a solution to a PDE, which is solved using a PINN-type loss.\n\nThe authors provide advice on tuning, suggestions for reducing cost of computations, and demonstrate the method on a Gaussian mixture model, Neal’s Funnel and Mixture of Student-T distributions, and statistical lattice field theory models.", "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 thank the reviewer for the positive feedback. We are glad that you found the contribution novel and the theory sound and thorough. Below we address your comments and suggestions and supply more information on experimental results:\n\n\n**Additional experiments**:\nIn the anonymous drive link https://tinyurl.com/netsicml:\n- We provide additional comparison between NETS and CMCD on the mean-interpolating time-dependent potential, as asked by another reviewer. We see that NETS performs well regardless of number of discretization steps, but the log-variance loss struggles for small steps and only perform well when $n_{step} = 256$ or more. We worked with the CMCD authors to implement their code and otherwise use their hyperparameters.\n- We have also set up a test on the Lennard Jones potential, and have included preliminary results. We are continuing to refine these experiments and will incorporate them in the final version.", "claims_and_evidence": "The theoretical claims are supported by proofs in the Supplementary material.   The numerical experiments demonstrate the accuracy of the samples with respect to MMD and W2 distance, compared against a number of relevant comparable ML approaches to sampling from unnormalised distributions, which demonstrate convincing evidence that this method is achieves its claimed objective.", "methods_and_evaluation_criteria": "The proposed approach is sensible, and scales well (compared to competing approaches) in high dimensions.   The evaluation criteria is both sensible and comprehensive.", "theoretical_claims": "There are several theoretical claims made in the paper:\n1. The Jarzynski inequality.\n2. The fact that the corrected SDE (15) has exact marginals (prop 2.3)\n3.  A weighted version, recovering [ Vaikuntanathan & Jarzynski (2008) ] prop 2.4\n4. Derivation of the PINN objective (2.5)\n5. KL controlled by PINN loss (2.6).\n6. Discretised version of  Vaikuntanathan & Jarzynski (2008) (B1)\n7. Connection with Feynman Kac (C1) for a specific form of b (gradient form).\n8. D onwards - generalisations\n\nAll the proofs are sensible,", "experimental_designs_or_analyses": "There are no experimental designs in this paper.", "relation_to_broader_scientific_literature": "This contribution sits in a wider body of literature which seek to correct dynamics by introducing KL-optimal control, building on ideas such as Follmer drift (e.g. [Huang et al, Convergence Analysis of Schrodinger-Follmer Sampler without Convexity, 2021], [Tzen and Raginski, 2019], [Vargas et al, Bayesian Learning via Neural Schrodinger-Follmer Flows, 2021]), and more recently [Vargas et al, Transport meets Variational Inference: Controlled Monte Carlo Diffusions, 2024].      The task of learning the vector field to correct the drift, as is done in this paper,  has been studied approached in various manners, [Reich, A dynamical systems framework for intermittent data assimilation, 2011], [Heng, Jeremy, Arnaud Doucet, and Yvo Pokern. Gibbs flow for approximate transport with applications to Bayesian computation. 2021] and earlier, [Vaikuntanathan & Jarzynski (2008)].    \n\nThe approach has strong connections with CMCD [Vargas et al, 2024], which leverages a similar PINN loss (for an associated HJB equation) to this paper.   This is studied in detail in the supplementary information.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Strengths:  The paper is well written, and the main arguments of the per are straightforward to follow.\nWeakenesses:  While this is a good paper, the numerical experiments remain a bit lacking.  It would have been nice to see a wider range of potnential interactions.", "other_comments_or_suggestions": "Prop 2.3 Then ρt(x) is the PDF of X_t not X_T?", "questions_for_authors": "No further questions", "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": "DLJay0jwVZ", "reviewer": "Reviewer_z8Lc", "summary": "The authors propose a method for sampling from unnormalized probability distributions. The method builds on diffusion-based sampling, where a learnable drift is added to the stochastic differential equation. The authors propose a PINN objective which allows for off-policy optimization and does not require differentiating through the simulations. Moreover, the authors addtionally propose an objective that is based on action matching. The methods are tested on a variety of sampling problems.", "questions": "", "limitations": "", "rating": 2, "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 thank the reviewer for their careful reading of our paper and positive feedback. We are glad that you found the work well-written, theoretically sound, and that the numerical experiments demonstrate convincing evidence for our method's effectiveness. Below we address your comments and suggestions:\n\n**Experimental Designs Or Analyses:** The potential $U_t$ design is a feature we can exploit, similar to FAB or CMCD. In our drive link (https://tinyurl.com/netsicml) we compare with CMCD using identical interpolation, revealing CMCD struggles with this path: Using the same number of sampling steps as us, CMCD achieves only 15% ESS. To perform these experiments we worked with the authors to best benchmark their codebase.\n\n**Relation To Broader Scientific Literature** and **Essential References Not Discussed:** Thanks for pointing us to the papers by Shi, Zhekun, et al.; Sun, Jingtong, et al.; and Chen, Junhua, et al. We are happy to add a citation to these works, but we wish to note that they fall within the ICML policy of concurrent works: \"https://icml.cc/Conferences/2025/ReviewerInstructions\".\n\nWe would also like to stress that, while variants of the PINN loss have appeared in the literature, our work makes the following new contributions:\n\n- Novel connection between the PINN loss and the Jarzynski weighting factors and annealed langevin dynamics. This result fits in naturally with physical interpretation of Jarzynski equality -- the PINN loss control on the variance of the Jarzynski weights (the dissipation) in the process connecting $\\rho_0$ to $\\rho_1$.\n- Show that we can better approach $\\epsilon_t \\to \\infty$ limit in practice. Note that for our setup, perfect sampling is achieved in this limit, whether the learned transport is perfect or not. While this limit cannot be reached in practice without transport (as it would require taking astronomically large value of $\\epsilon_t$ in general), we show that, with some learned transport added, even moderate values of $\\epsilon_t$ can improve the sampling dramatically. This feature can be exploited after training as an explicit knob for tuning performance vs cost. This has not been recognized in other ML sampling literature nor in Vargas et al.\n- Show that the PINN loss directly controls the KL-divergence between the sampled and target distribution.\n- Directly characterize the minimizer of the action matching loss, which is also not known from previous work on it.\n\nOur method connects with CMCD, but we derive results through Fokker-Planck manipulations rather than Girsanov theorem. \n\nWe'll add references to Richter and Berner's foundational work.\n\n**Other comments and suggestions:** Could you please specify which equations are stated without context? We will happily try to clarify in editing.\n\n**Questions for authors:**\n\n-*Difference between the proposed PINN objective and other related PINN based sampling methods:* See our reply above about the key new insights that our work provides.\n-*Insight of top row, left of Figure 3:* The top row (left) are example configurations of samples as we approach the critical point, and the bottom row are samples past the critical point, where they are fully magnetized. They illustrate the juxtaposition.\n-*Benefit of the proposed method compared to CMCD-LV:* We have included in the \"additional experiments\" section below results on CMCD-LV, which we worked with the CMCD authors to set up. Our results show that, unlike our approach, CMCD-LV exhibits instabilities if the number of sampling steps is too small in training: while the CMCD-LV loss is \"off-policy\" like our PINN loss, it still requires a trajectory on which to perform the optimization, and the generation of such a $n_{step}$ trajectory requires O($n_{step}$ $\\times$ network size) memory. If the trajectory generation is performed over too few steps the CMCD LV loss becomes unstable. After discussing with CMCD authors, this is an issue generally when $n_{step} < 256$.\n\n**Additional experiments:**\nIn the anonymous drive link https://tinyurl.com/netsicml:\n\n- We provide additional comparison between NETS and CMCD on the mean-interpolating time-dependent potential, as asked by another reviewer. We see that NETS performs well regardless of number of discretization steps, but the log-variance loss struggles for small steps and only performs well when $n_{step} = 256$ or more. We worked with the CMCD authors to implement their code and otherwise use their hyperparameters.\n- We have also set up a test on the Lennard Jones potential, and have included preliminary results. We are continuing to refine these experiments and will incorporate them in the final version.", "claims_and_evidence": "Claims are supported by theory. The reviewer has no concerns.", "methods_and_evaluation_criteria": "The paper uses a variety of evaluation critera such as ESS, Wasserstein distance or MMD which fit for evaluating sampling methods. Moreover, the paper uses a variety of baselines, most of which are quite recent, which is also good.", "theoretical_claims": "The reviewer skimmed the Proofs in Appendix A which appear to be correct.", "experimental_designs_or_analyses": "The comparision between the different methods seems highly unfair. The authors use knowledge of the target density that significantly reduces the complexity of the problems. For instance, the method linearly interpolates the means of the 40 mode GMM with the prior potential, which effectively removes the difficuly of exploration. Moreover, similar handcrafted potentials are used for other targets as well. \n Are the authors using the same interpolation scheme for the baselines like FAB, or CMCD?", "relation_to_broader_scientific_literature": "It is not quite clear to the reviewer what the key contribution/novelty of this paper is. Sampling with PINNs was already done in [1,2]. The connection with Jarzynski’s equality was recently shown in [4] as noted by the authors. Combing diffusion-based sampling with SMC was also recently proposed in [4]. Moreover, off-policy learning with diffusion-sampler is possible using the log-variance loss, see [5].\n\n\n[1] Shi, Zhekun, et al. \"Diffusion-PINN Sampler.\" arXiv preprint arXiv:2410.15336 (2024).\n[2] Sun, Jingtong, et al. \"Dynamical measure transport and neural PDE solvers for sampling.\" arXiv preprint arXiv:2407.07873 (2024).\n[3] Vargas, Francisco, et al. \"Transport meets variational inference: Controlled monte carlo diffusions.\" arXiv preprint arXiv:2307.01050 (2023).\n[4] Chen, Junhua, et al. \"Sequential controlled langevin diffusions.\" ICLR 25\n[5] Richter, Lorenz, and Julius Berner. \"Improved sampling via learned diffusions.\" arXiv preprint arXiv:2307.01198 (2023).", "essential_references_not_discussed": "The authors did not cite related work that uses PINNs to sample from unnormalized densities, see [1,2]. Moreover, one of the seminal papers [3] on diffusion-based sampling was not cited.\n\nAnother paper that might be of interest for the authors is [4] which also builds on dynamic measure transport with resampling schemes\n\n[1] Shi, Zhekun, et al. \"Diffusion-PINN Sampler.\" arXiv preprint arXiv:2410.15336 (2024).\n[2] Sun, Jingtong, et al. \"Dynamical measure transport and neural PDE solvers for sampling.\" arXiv preprint arXiv:2407.07873 (2024).\n[3] Richter, Lorenz, and Julius Berner. \"Improved sampling via learned diffusions.\" arXiv preprint arXiv:2307.01198 (2023).\n[4] Chen, Junhua, et al. \"Sequential controlled langevin diffusions.\" ICLR 25", "other_strengths_and_weaknesses": "Strenghts:\n\n- The paper proposes an off-policy objective which has potential to increase sample efficiency\n- The methods avoids having to backpropagate though the simulation \n\nWeaknesses:\n\n- See Relation To Broader Scientific Literature and Experimental Designs Or Analyses", "other_comments_or_suggestions": "A lot of the equations are stated but no context is provided which makes it difficult to read the paper.", "questions_for_authors": "- What is the difference between the proposed PINN objective and other related PINN based sampling methods, see e.g. [1,2]. The latter in particular also considers a prescribed density evolution.\n- What insights does top row, left of Figure 3. give?\n- The authors use CMCD-LV and CMCD-KL in Table 1. I assume this refers to the log-variance loss and the KL loss (the former is not mentioned cited). What is the benefit of the proposed method compared to CMCD-LV? \n\n[1] Shi, Zhekun, et al. \"Diffusion-PINN Sampler.\" arXiv preprint arXiv:2410.15336 (2024).\n[2] Sun, Jingtong, et al. \"Dynamical measure transport and neural PDE solvers for sampling.\" arXiv preprint arXiv:2407.07873 (2024).", "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 introduces NETS, a non-equilibrium transport sampler that extends annealed importance sampling (AIS) by learning an adaptive drift term using physics-informed neural networks (PINNs) or action-matching losses. Inspired by Jarzynski’s equality, the approach aims to reduce the variance of importance weights. The method builds connections to recent advances in controlled Langevin dynamics and is evaluated on several benchmark distributions, including high-dimensional Gaussian mixtures and lattice field theory models.

The work is grounded in theoretical foundations and addresses a timely challenge in probabilistic inference. Reviewers particularly appreciated the novel connection between the PINN loss and the Jarzynski weighting factors. Experimental results on synthetic and structured benchmarks show promising improvements over established baselines.

One concern raised by reviewers is the overlap with existing work on diffusion-based samplers and PINN-based training. The authors provided a detailed and thoughtful rebuttal clarifying their contributions and distinctions.

Overall, this paper offers technically sound contributions and proposes a compelling new direction for improving sampling efficiency. The connection between the PINN loss and Jarzynski’s equality is particularly insightful. Most reviewers are positive, and the remaining concerns appear addressable in revision.
Decision: Accept (poster)