NETS: A Non-Equilibrium Transport Sampler

## Abstract

Abstract We propose an algorithm, termed the Non-Equilibrium Transport Sampler (NETS), to sample from unnormalized probability distributions. NETS can be viewed as a variant of annealed importance sampling (AIS) based on Jarzynski’s equality, in which the stochastic differential equation used to perform the non-equilibrium sampling is augmented with an additional learned drift term that lowers the impact of the unbiasing weights used in AIS. We show that this drift is the minimizer of a variety of objective functions, which can all be estimated in an unbiased fashion without backpropagating through solutions of the stochastic differential equations governing the sampling. We also prove that some these objectives control the Kullback-Leibler divergence of the estimated distribution from its target. NETS is shown to be unbiased and, in addition, has a tunable diffusion coefficient which can be adjusted post-training to maximize the effective sample size. We demonstrate the efficacy of the method on standard benchmarks, high-dimensional Gaussian mixture distributions, and a model from statistical lattice field theory, for which it surpasses the performances of related work and existing baselines.

## 1 Introduction

The aim of this paper is to sample probability distributions supported on $\mathbb{R}^{d}$ and known only up to a normalization constant. This problem arises in a wide variety in applications, ranging from statistical physics (Faulkner & Livingstone, ; Wilson, ; Hénin et al., ) to Bayesian inference (Neal, ), and it is known to be challenging when the target distribution is not log-concave. In this situation, vanilla methods based on ergodic sampling using Markov Chain Monte Carlo (MCMC) or stochastic differential equations (SDE) such as the Langevin dynamics typically have very slow convergence rates, making them inefficient in practice.

To overcome these difficulties, a variety of more sophisticated methods have been introduced, based e.g. on importance sampling. Here we will be interested in a class of methods of this type which involve non-equilibrium sampling, by which we mean algorithms that attempt to sample from a non-stationary probability distribution. The idea is to first generate samples from a simple base distribution (e.g. a normal distribution) them push them in finite time unto samples from the target. Traditionally, this aim has been achieved by combining some transport using e.g. the Langevin dynamics with reweighing, so as to remove the bias introduced by the non-equilibrium quench. Neal’s Annealed Importance Sampling (AIS) (Neal, ), Sequential Monte Carlo (SMC) methods (Del Moral, ; Doucet et al., ), or continuous-time variants thereof based on Jarzynski’s equality (Hartmann et al., ) are ways to implement this idea in practice. While these methods work better than ergodic sampling in many instances, they too can fail when the variance of the un-biasing weights become too large compared to their mean: this arises when the samples end up being too far from the target distribution at the end of the non-equilibrium quench.

This problem suggests to modify the dynamics of the samples to help them evolves towards the target during the quench. Here we propose a way to achieve this by learning an additional drift to include in the Langevin SDE. As we will show below, there exists a drift that removes the need for the un-biasing weights altogether, and it is the minimizer of a variety of objective functions that are amenable to empirical estimation. In practice, this offers the possibility to estimate this drift using deep learning methods. We exploit this idea here, showing that it results in an unbiased sampling strategy in which importance weights can still be used to correct the samples exactly, but the variance of these weights can be made much smaller due to the additional transport.

To this end, our work makes the following main contributions:

- •
We present a sampling algorithm which combines annealed Langevin dynamics with learnable additional transport. The algorithm, which we call the Non-Equilbrium Transport Sampler (NETS), is shown to be unbiased through a generalization of the Jarzysnki equality.
- •
We show that the drift coefficient contributing this additional transport is the minimizer of two separate objective functions, which can be learned without backpropagating through solutions of the SDE used for sampling.
- •
We show that one of these objectives, a physics informed neural network (PINN) loss, is also an off-policy objective, meaning that it does not need samples from the target density. In addition, this objective controls the KL-divergence between the model and its target.
- •
The resultant samplers can be adapted after training by tuning the integration time-step as well as the diffusivity to improve the sample quality, which we demonstrate on high-dimensional numerical experiments below.

### 1.1 Related work

##### Dynamical Measure Transport.

Most contemporary generative models for continuous data are built upon dynamical measure transport, in which samples from a base density are mapped to samples from a target density by means of solving ordinary or stochastic differential equations (ODE/SDE) whose drift coefficients are estimable. Initial attempts to do this started with (Chen et al., ; Grathwohl et al., ). Recent work built upon these ideas by recasting the challenge of estimating the drift coefficients in these dynamical equations as a problem of quadratic regression, most notably with score-based diffusion models (Ho et al., ; Song et al., ), and since then with more general frameworks (Albergo & Vanden-Eijnden, ; Lipman et al., ; Albergo et al., ; Liu et al., ; De Bortoli et al., ; Neklyudov et al., ). Importantly, these methods work because samples from the base and the target are readily available. Here, we show that analogous equations governing those systems can lead to objective functions that can be minimized with no initial data but still allow us to exploit the expressivity of models built on dynamical transport.

##### Augmenting sampling with learning.

Augmenting MCMC and importance sampling procedures with transport has been an active area of research for the past decade. Early work makes use of the independence Metropolis algorithm (Hastings, ; Liu, ), in which proposals come from a transport map (Parno & Marzouk, ; Noé et al., ; Albergo et al., ; Gabrié et al., ) 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., ; Matthews et al., ; Midgley et al., ). 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-equilbirium sampling (Caselle et al., ; Bonanno et al., ).

| Method | Dynamical<br>Transport | No backprop<br>thru time | Off-Policy<br>Learning | Grid-Free<br>Learning |
| --- | --- | --- | --- | --- |
| FAB | $\times$ | $\checkmark$ | $\times$ | $\times$ |
| DDS | $\checkmark$ | $\times$ | $\times$ | $\times$ |
| PIS | $\checkmark$ | $\times$ | $\times$ | $\checkmark$ |
| CMCD | $\checkmark$ | Partial | $\times$ | $\times$ |
| p/iDEM | $\checkmark$ | $\checkmark$ | Partial | $\checkmark$ |
| NETS-AM (ours) | $\checkmark$ | $\checkmark$ | $\times$ | $\checkmark$ |
| NETS-PINN (ours) | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ |

Following the success of generative models built out of dynamical transport, there has been a surge of interest in applying these perspectives to sampling: Vargas et al. (); Berner et al. () translate ideas from diffusion models to minimize the KL divergence between the model and the target, and Zhang & Chen () as well as Behjoo & Chertkov () reformulate sampling as a stochastic optimal control (SOC) problem. These approaches require backpropagating through the solution of an SDE, which is too costly in high dimensions. In addition, the methods based on SOC must start with samples from a point mass, which may be far from the target. Akhound-Sadegh et al. () avoid the need to backpropagate through an SDE, but in the process introduce a bias into their objective function. An alternative perspective was also recently given in (Bruna & Han, ) by using denoising oracles to turn the original sampling problem into an easier one.

Vargas et al. () establish an unbiased sampler with added transport called Controlled Monte Carlo Diffusions (CMCD) that is similar to ours. The main difference is how we learn the drift. In Vargas et al. () an objective for this drift in gradient form is derived through the use of path integrals and Crook’s identity. This objective either needs backpropagating through the SDE or has to be computed with a reference measure, and is done on a fixed grid. Here, through simple manipulations of Fokker-Planck equations, we propose a variety of new objective functions for the additional drift, none of which require backpropagating through the simulation. In addition, our learning can be done in a optimize-then-discretize fashion so that the sampling can be done with arbitrary step size and time-dependent diffusion after learning the model. This gives us an adaptive knob to increase performance, which we demonstrate below.

## 2 Methods

### 2.1 Setup and notations

We assume that the target distribution is absolutely continuous with respect to the Lebesgue measure on $R^{d}$, with probability density function (PDF) $\rho_{1}(x)=Z_{1}^{-1}e^{-U_{1}(x)}$: here $x\in\mathbb{R}^{d}$, $U:\mathbb{R}^{d}\to\mathbb{R}$ is a known energy potential, assumed twice differentiable and bounded below, and $Z_{1}=\int_{\mathbb{R}^{d}}e^{-U_{1}(x)}dx<\infty$ is an unknown normalization constant, referred to as the partition function is physics and the evidence in statistics. Our aim is to generate samples from $\rho_{1}(x)$ so as to be able to estimate expectations with respect to this density. Additionally we wish to estimate $Z_{1}$.

To this end we will use a series of time-dependent potentials $U_{t}(x)$ which connects some simple $U_{0}(x)$ at $t=0$ (e.g. $U_{0}(x)=\frac{1}{2}|x|^{2}$) to $U_{1}(x)$ at $t=1$. For example we could use linear interpolation:

$$ $U_{t}(x)=(1-t)U_{0}(x)+tU_{1}(x),$ (1) $$

but other choices are possible as long as $U_{t=0}=U_{0}$, $U_{t=1}=U_{1}$, and $U_{t}(x)$ is twice differentiable in $(t,x)\in[0,1]\times\mathbb{R}^{d}$. We assume that the time-dependent PDF associated with this potential $U_{t}(x)$ is normalizable for all $t\in[0,1]$ and denote it as

$$ $\rho_{t}(x)=Z^{-1}_{t}e^{-U_{t}(x)},\qquad Z_{t}=\int_{\mathbb{R}^{d}}e^{-U_{t }(x)}dx<\infty,$ (2) $$

so that $\rho_{t=0}(x)=\rho_{0}(x)$ and $\rho_{t=1}(x)=\rho_{1}(x)$; we also assume that $\rho_{0}(x)$ is simple to sample (either directly or via MCMC or Langevin dynamics) and that its partition function $Z_{0}$ is known.
To simplify the notations we also introduce the free energy

$$ $F_{t}=-\log Z_{t}.$ (3) $$

Since $\partial_{t}F_{t}=-\partial_{t}\log\int_{R^{d}}e^{-U_{t}(x)}dx=\int_{R^{d}}
\partial_{t}U_{t}(x)e^{-U_{t}(x)}dx/\int_{R^{d}}e^{-U_{t}(x)}dx$ we have the useful identity

$$ $\partial_{t}F_{t}=\int_{R^{d}}\partial_{t}U_{t}(x)\rho_{t}(x)dx.$ (4) $$

### 2.2 Nonequilibrium sampling with importance weights

Annealed importance sampling 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 present a time-continuous variant of AIS based on Jarzynski equality that will be more useful for our purpose.

By definition of the PDF in (2), $\nabla\rho_{t}(x)=-\nabla U_{t}(x)\rho_{t}(x)$ and hence, for any $\varepsilon_{t}\geq 0$, we have

$$ $0=\varepsilon_{t}\nabla\cdot(\nabla U_{t}\rho_{t}+\nabla\rho_{t}).$ (5) $$

Since we also have

$$ $\partial_{t}\rho_{t}=-(\partial_{t}U_{t}-\partial_{t}F_{t})\rho_{t},$ (6) $$

we can combine these last two equations to deduce that

$$ $\partial_{t}\rho_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{t}\rho_{t}+\nabla \rho_{t})-(\partial_{t}U_{t}-\partial_{t}F_{t})\rho_{t}.$ (7) $$

The effect of the last term at the right hand-side of this equation can be accounted for by using weights. To see how, notice that if we extend the phase space to $(x,a)\in\mathbb{R}^{d+1}$ and introduce the PDF $f_{t}(x,a)$ solution to the Fokker-Planck equation (FPE)

$$ $\partial_{t}f_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{t}f_{t}+\nabla f_{t})+ \partial_{t}U_{t}\partial_{a}f_{t},\qquad f_{t=0}(x,a)=\delta(a)\rho_{0}(x),$ (8) $$

then a direct calculation using (4) (for details see Appendix 5.1) shows that

$$ $\rho_{t}(x)=\frac{\int_{\mathbb{R}}e^{a}f_{t}(x,a)da}{\int_{\mathbb{R}^{d+1}}e ^{a}f_{t}(y,a)dady}.$ (9) $$

Therefore we can use the solution to SDE associated with the FPE (8) in the extended space to estimate expectations with respect to $\rho_{t}(x)$:

###### Proposition 1 (Jarzynski equality) .

Let $(X_{t},A_{t})$ solve the coupled system of SDE/ODE

$$ $\displaystyle dX_{t}$ $\displaystyle=-\varepsilon_{t}\nabla U_{t}(X_{t})dt+\sqrt{2\varepsilon_{t}}dW_ {t},\qquad$ $\displaystyle X_{0}\sim\rho_{0},$ (10) $\displaystyle dA_{t}$ $\displaystyle=-\partial_{t}U_{t}(X_{t})dt,$ $\displaystyle A_{0}=0,$ (11) $$

where $\varepsilon_{t}\geq 0$ is a time-dependent diffusion coefficient and $W_{t}\in\mathbb{R}^{d}$ is the Wiener process. Then for all $t\in[0,1]$ and any test function $h:\mathbb{R}^{d}\to\mathbb{R}$, we have

$$ $\int_{\mathbb{R}^{d}}h(x)\rho_{t}(x)dx=\frac{\mathbb{E}[e^{A_{t}}h(X_{t})]}{ \mathbb{E}[e^{A_{t}}]},\qquad Z_{t}/Z_{0}=e^{-F_{t}+F_{0}}=\mathbb{E}[e^{A_{t} }],$ (12) $$

where the expectations are taken over the law of $(X_{t},A_{t})$.

The proof of this proposition is given in Appendix 5.1 and it relies on the identity $\int_{\mathbb{R}^{d}}h(x)\rho_{t}(x)=\int_{\mathbb{R}^{+1d}}e^{a}h(x)f_{t}(x,a
)dadx/\int_{\mathbb{R}^{d+1}}f_{t}(x,a)dadx$ which follows from (9). The second equation in (12) for the free energy $F_{t}$ is what is referred to as Jarzynski’s equality, and was originally surmised in the context of non-equilibrium thermodynamics (Jarzynski, ).

###### Remark 1 .

We stress that it is key to use the weights $e^{A_{t}}$ in (12) because $\rho_{t}(x)$ is not the PDF of $X_{t}$ in general. Indeed, if we denote by $\tilde{\rho}_{t}(x)$ the PDF of $X_{t}$, it satisfies the FPE

$$ $\partial_{t}\tilde{\rho}_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{t}\tilde{ \rho}_{t}+\nabla\tilde{\rho}_{t}),\quad\tilde{\rho}_{t=0}=\rho.$ (13) $$

This FPE misses the term $-(\partial_{t}U_{t}-\partial_{t}F_{t})\rho_{t}$ at the right hand-side of (7), and as a result $\tilde{\rho}_{t}(x)\not=\rho_{t}(x)$ in general – intuitively, $\tilde{\rho}_{t}(x)$ lags behind $\rho_{t}(x)$ when the potential $U_{t}(x)$ evolves and this lag is what the weights in (12) correct for.

It is important to realize that, while the relation (12) can be used to compute unbiased estimators of expectations, this estimator on its own can be high variance if the lag between the PDF $\tilde{\rho}_{t}(x)$ of $X_{t}$ and $\rho_{t}(x)$ is too pronounced. This issue can be alleviated by using resampling methods as is done in sequential Monte Carlo (Doucet et al., ). Here we will solve it by adding some additional drift in (10) that will compensate for this lag and reduce the effect of the weights.

###### Proposition 1 (Jarzynski equality) .

###### Remark 1 .

### 2.3 Nonequilibrium sampling with perfect additional transport

To see how we can add a transport term to eliminate the need of the weights, let us introduce a velocity field $b_{t}(x)\in\mathbb{R}^{d}$ which at all times $t\in[0,1]$ satisfies

$$ $\nabla\cdot(b_{t}\rho_{t})=-\partial_{t}\rho_{t}.$ (14) $$

We stress that this is an equation for $b_{t}(x)$ in which $\rho_{t}(x)$ is fixed and given by (2):
In Appendix 5.3 we show how to express the solution to (14) via Feynman-Kac formula.
If (14) is satisfied, then we can combine this equation with (5) and (6) to arrive at

$$ $\partial_{t}\rho_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{t}\rho_{t}+\nabla \rho_{t})-\nabla\cdot(b_{t}\rho_{t}),$ (15) $$

which is a standard FPE. Therefore the solution to the SDE associated with (15) allows us to sample $\rho_{t}(x)$ directly (without weights). We phrase this result as:

###### Proposition 2 (Sampling with perfect additional transport.) .

Let $b_{t}(x)$ be a solution to (14) and let $X^{b}_{t}$ satisfy the SDE

$$ $dX^{b}_{t}=-\varepsilon_{t}\nabla U_{t}(X^{b}_{t})dt+b_{t}(X^{b}_{t})dt+\sqrt{ 2\varepsilon_{t}}dW_{t},\qquad X^{b}_{0}\sim\rho_{0},$ (16) $$

where $\varepsilon_{t}\geq 0$ is a time-dependent diffusion coefficient and $W_{t}\in\mathbb{R}^{d}$ is the Wiener process. Then $\rho_{t}(x)$ is the PDF of $X^{b}_{T}$, i.e. for all $t\in[0,1]$ and, given any test function $h:\mathbb{R}^{d}\to\mathbb{R}$, we have

$$ $\int_{\mathbb{R}^{d}}h(x)\rho_{t}(x)dx=\mathbb{E}[h(X^{b}_{t})],$ (17) $$

where the expectation at the right-hand side is taken over the law of $(X^{b}_{t})$.

This proposition is proven in Appendix 5.1 and it shows that we can in principle get rid of the weights altogether by adding the drift $b_{t}(x)$ in the Langevin SDE. This possibility was first noted in Vaikuntanathan & Jarzynski () and is also exploited in Tian et al. () for deterministic dynamics (i.e. setting $\varepsilon_{t}=0$ in (16)) and in Vargas et al. () using the SDE (16). Of course, in practice we need to estimate this drift, 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 $b_{t}(x)$ to Secs. 2.5 and 2.6. In Appendix 5.3 we show how to express the solution to (14) via Feynman-Kac formula.

###### Proposition 2 (Sampling with perfect additional transport.) .

### 2.4 Non-Equilibrium Transport Sampler

Let us now show that we can combine the approaches discussed in Secs. 2.2 and 2.3 to design samplers in which we use an added transport, possibly imperfect, and importance weights.

To this end, suppose that we wish to add an additional transport term $-\nabla\cdot(\hat{b}_{t}\rho_{t})$ in (7), where $\hat{b}_{t}(x)\in\mathbb{R}^{d}$ is some given velocity that does not necessarily solve (14). Using the expression in (2) for $\rho_{t}(x)$, we have the identity

$$ $-\nabla\cdot(\hat{b}_{t}\rho_{t})=-\nabla\cdot\hat{b}_{t}\rho_{t}+\nabla U_{t} \cdot\hat{b}_{t}\rho_{t}$ (18) $$

Therefore we can rewrite (5) equivalently as

$$ $\partial_{t}\rho_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{t}\rho_{t}+\nabla \rho_{t})-\nabla\cdot(\hat{b}_{t}\rho_{t})+(\nabla\cdot\hat{b}_{t}-\nabla U_{t }-\partial_{t}U_{t}+\partial_{t}F_{t})\rho_{t}.$ (19) $$

We can now proceed as we did with (5) and extend state space to account for the effect of the terms $(\nabla\cdot b_{t}-\nabla U_{t}-\partial_{t}U_{t}+\partial_{t}F_{t})\rho_{t}$ in this equation via weights, while having the term $-\nabla\cdot(\hat{b}_{t}(x)\rho_{t}(x))$ contribute to some additional transport. This leads us to a result originally obtained in Vaikuntanathan & Jarzynski () and recently re-derived in Vargas et al. ()):

###### Proposition 3 (Nonequilibrium Transport Sampler (NETS)) .

Let $(X^{\hat{b}}_{t},A^{\hat{b}}_{t})$ solve the coupled system of SDE/ODE

$$ $\displaystyle dX^{\hat{b}}_{t}$ $\displaystyle=-\varepsilon_{t}\nabla U_{t}(X^{\hat{b}}_{t})dt+\hat{b}_{t}(X^{ \hat{b}}_{t})dt+\sqrt{2\varepsilon_{t}}dW_{t},\qquad$ $\displaystyle X^{\hat{b}}_{0}\sim\rho_{0},$ (20) $\displaystyle dA^{\hat{b}}_{t}$ $\displaystyle=\nabla\cdot\hat{b}_{t}(X^{\hat{b}}_{t})dt-\nabla U_{t}(X^{\hat{b }}_{t})\cdot\hat{b}(X^{\hat{b}}_{t})dt-\partial_{t}U_{t}(X^{\hat{b}}_{t})dt,\qquad$ $\displaystyle A^{\hat{b}}_{0}=0,$ (21) $$

where $\varepsilon_{t}\geq 0$ is a time-dependent diffusion coefficient and $W_{t}\in\mathbb{R}^{d}$ is the Wiener process. Then for all $t\in[0,1]$ and any test function $h:\mathbb{R}^{d}\to\mathbb{R}$, we have

$$ $\int_{\mathbb{R}^{d}}h(x)\rho_{t}(x)dx=\frac{\mathbb{E}[e^{A^{\hat{b}}_{t}}h(X ^{\hat{b}}_{t})]}{\mathbb{E}[e^{A^{\hat{b}}_{t}}]},\qquad Z_{t}/Z_{0}=e^{-F_{t }+F_{0}}=\mathbb{E}[e^{A^{\hat{b}}_{t}}],$ (22) $$

where the expectations are taken over the law of $(X^{\hat{b}}_{t},A^{\hat{b}}_{t})$.

A simple proof of this proposition using simple manipulations of the FPE is given in Appendix 5.1, which will allow us to write down a variety of new loss functions for learning $\hat{b}_{t}$; for an alternative proof using Girsanov theorem, see Vargas et al. ().
For completeness, in Appendix 5.2 we also give a time-discretized version of Proposition 3.

Notice that, if $\hat{b}_{t}(x)=0$, the equations in Proposition 3 simply reduce to those in Proposition 1, whereas if $\hat{b}_{t}(x)=b_{t}(x)$ solves (14) we can show that

$$ $A^{b}_{t}=F_{t}-F_{0},$ (23) $$

i.e. the weights have zero variance and give the free energy difference. Indeed, by expanding both sides of (14) and dividing them by $\rho_{t}(x)>0$, this equation can equivalently be written as

$$ $\nabla\cdot b_{t}-\nabla U_{t}\cdot b_{t}=\partial_{t}U_{t}-\partial_{t}F_{t}.$ (24) $$

As a result, when $\hat{b}_{t}(x)=b_{t}(x)$, (21) reduces to

$$ $dA^{b}_{t}=\partial_{t}F_{t}dt,\qquad A^{b}_{0}=0,$ (25) $$

and the solution to this equation is (23). In practice, achieving zero variance of the weights by estimating $b_{t}(x)$ exactly is not generally possible, but having a good approximation of $b_{t}(x)$ can help reducing this variance dramatically, as we will illustrate below via experiments.

###### Proposition 3 (Nonequilibrium Transport Sampler (NETS)) .

### 2.5 Estimating the drift b t ⁢ ( x ) subscript 𝑏 𝑡 𝑥 b_{t}(x) italic_b start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ( italic_x ) via a PINN objective

Equation (24) can be used to derive an objective for both $b_{t}(x)$ and $F_{t}$. The reason is that in this equation the unknown $\partial_{t}F_{t}$ can be viewed as factor that guarantees solvability: indeed, integrating both sides of (14) gives $0=-\partial_{t}\int_{\mathbb{R}^{d}}U_{t}(x)\rho_{t}(x)dx+\partial_{t}F_{t}$, which, by (4), is satisfied if and only if $F_{t}$ is (up to a constant fixed by $F_{0}=-\log Z_{0}$) the exact free energy (3). This offers the possibility to learn both $b_{t}(x)$ and $F_{t}$ variationally using an objective fitting the framework of physics informed neural networks (PINNs):

###### Proposition 4 (PINN objective) .

Given any $T\in(0,1]$ and any PDF $\hat{\rho}_{t}(x)>0$ consider the objective for $(\hat{b},\hat{F})$ given by:

$$ $L^{T}_{\text{PINN}}[\hat{b},\hat{F}]=\int_{0}^{T}\int_{\mathbb{R}^{d}}\big{|} \nabla\cdot\hat{b}_{t}(x)-\nabla U_{t}(x)\cdot\hat{b}_{t}(x)-\partial_{t}U_{t} (x)+\partial_{t}\hat{F}_{t}\big{|}^{2}\hat{\rho}_{t}(x)dxdt.$ (26) $$

Then $\min_{\hat{b},\hat{F}}L^{T}_{\text{PINN}}[\hat{b},\hat{F}]=0$, and all minimizers are such that $F_{t}$ is the free energy (3) and $b_{t}(x)$ solves (24) for all $t\in[0,T]$.

This result is proven in Appendix 5.1: in practice, we will use $T\in(0,1]$ for annealing but ultimately we are interested in the result when $T=1$. Note that since the expectation over an arbitrary $\hat{\rho}_{t}(x)$ in (26), it is an off-policy objective. It is however natural to use $\hat{\rho}_{t}(x)=\rho_{t}(x)$ since it allows us to put statistical weight in the objective precisely in the regions where we need $b_{t}(x)$ to transport probability mass: we will show below how the expectation over $\rho_{t}(x)$ can be estimated without bias to arrive at an empirical estimator for (26) when $\hat{\rho}_{t}(x)=\rho_{t}(x)$. Note also that, while minimization of the objective (26) gives an estimate $\hat{F}_{t}$ of the free energy, it is not needed at sampling time when solving (20). An objective similar to (26) was also recently posited in Tian et al. () for use with deterministic flows. Here, we devise it in the context of augmenting annealed Langevin dynamics.

One advantage of the PINN objective (26) is that we know that its minimum is zero, and hence we can track its value to monitor convergence when minimizing (26) by gradient descent as we do below. Another advantage of the loss (26) is that it controls the quality of the transport as measured by the Kullback-Leibler divergence:

###### Proposition 5 (KL control) .

Let $\hat{\rho}_{t}$ be the solution to the transport equation

$$ $\partial_{t}\hat{\rho}_{t}=-\nabla\cdot(\hat{b}_{t}\rho_{t}),\qquad\hat{\rho}_ {t=0}=\rho_{0}$ (27) $$

where $\hat{b}_{t}(x)$ is some predefined velocity field. Then, given any estimate $\hat{F}_{t}$ of the exact free energy $F_{t}$, we have

$$ $D_{\text{KL}}(\hat{\rho}_{t=1}||\rho_{1})\leq\sqrt{L^{T=1}_{\text{PINN}}(\hat{ b},\hat{F})}.$ (28) $$

This proposition is proven in Appendix 5.1.

###### Proposition 4 (PINN objective) .

###### Proposition 5 (KL control) .

### 2.6 Estimating the drift b t ⁢ ( x ) = ∇ ϕ t ⁢ ( x ) subscript 𝑏 𝑡 𝑥 ∇ subscript italic-ϕ 𝑡 𝑥 b_{t}(x)=\nabla\phi_{t}(x) italic_b start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ( italic_x ) = ∇ italic_ϕ start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ( italic_x ) via Action Matching (AM)

In general (14) is solved by many $b_{t}(x)$. One way to get a unique (up to a constant in space and time) solution to this equation is to impose that the velocity be in gradient form, i.e. set $b_{t}(x)=\nabla\phi_{t}(x)$ for some scalar-valued potential $\phi_{t}(x)$. If we do so, (14) can be written as $\nabla\cdot(\nabla\phi_{t}(x)\rho_{t})=-\partial_{t}\rho_{t}$, and it is easy to see that at all times $t\in[0,1]$ the solution to this equation minimizes over $\hat{\phi}_{t}$ the objective

$$ $\displaystyle\int_{\mathbb{R}^{d}}\big{[}\tfrac{1}{2}|\nabla\hat{\phi}_{t}(x)| ^{2}\rho_{t}(x)-\hat{\phi}_{t}(x)\partial_{t}\rho_{t}(x)\big{]}dx$ (29) $\displaystyle=\int_{\mathbb{R}^{d}}\big{[}\tfrac{1}{2}|\nabla\hat{\phi}_{t}(x) |^{2}+(\partial_{t}U_{t}(x)-\partial_{t}F_{t})\hat{\phi}_{t}(x)\big{]}\rho_{t} (x)dx.$ $$

If we use (4) to set $\partial_{t}F_{t}=\int_{\mathbb{R}^{d}}\partial_{t}U_{t}(x)\rho_{t}(x)dx$ we can use the objective at the right hand-side of (29) to learn $\phi_{t}(x)$ locally in time (or globally if we integrate this objective on $t\in[0,1]$). Alternatively, we can integrate the objective at the left hand-side of (29) over $t\in[0,T]$ and use integration by parts for the term involving $\partial_{t}\rho_{t}(x)$ to arrive at:

###### Proposition 6 (Action Matching objective) .

Given any $T\in(0,1]$ consider the objective for $\hat{\phi}_{t}(x)$:

$$ $L^{T}_{AM}[\hat{\phi}]=\int_{0}^{T}\int_{\mathbb{R}^{d}}\big{[}\tfrac{1}{2}| \nabla\hat{\phi}_{t}(x)|^{2}+\partial_{t}\hat{\phi}_{t}(x)\big{]}\rho_{t}(x) dxdt+\int_{\mathbb{R}^{d}}\big{[}\hat{\phi}_{0}(x)\rho_{0}(x)-\hat{\phi}_{T}(x )\rho_{T}(x)\big{]}dx.$ (30) $$

Then the minimizer $\phi_{t}(x)$ of (26) is unique (up to a constant) and $b_{t}(x)=\nabla\phi_{t}(x)$ satisfies (14) for all $t\in[0,T]$.

This proposition is proven in Appendix 5.1. This objective is analogous to the loss presented in Neklyudov et al. (), but adapted to the sampling problem. In practice, we will use again $T\in(0,1]$ for annealing, but ultimately we are interested in the resut at $T=1$. Note that, unlike with the PINN objective (26), it is crucial that we use the correct $\rho_{t}(x)$ in the AM objective (30): that is, unlike (26), (30) cannot be turned into an off-policy objective.

###### Remark 2 .

If we use $\hat{b}_{t}(x)=\nabla\hat{\phi}_{t}(x)$ in the SDEs in (20) and (21), we need to compute $\nabla\cdot\hat{b}_{t}(x)=\Delta\hat{\phi}_{t}(x)$, which is computationally costly. Fortunately, when $\varepsilon_{t}>0$, the calculation of this Laplacian can be avoided by using the following alternative equation for $A^{\hat{b}}_{t}$:

$$ $A^{\hat{b}}_{t}=\frac{1}{\varepsilon_{t}}[\hat{\phi}_{t}(X^{\hat{b}}_{t})-\hat {\phi}_{0}(X^{\hat{b}}_{0})]-B_{t},$ (31) $$

where

$$ $dB_{t}=\partial_{t}U_{t}(X^{\hat{b}}_{t})dt+\frac{1}{\varepsilon_{t}}\partial_ {t}\hat{\phi}_{t}(X_{t}^{\hat{b}})+\frac{1}{\varepsilon_{t}}|\nabla\hat{\phi}_ {t}(X^{\hat{b}}_{t})|^{2}dt+\sqrt{\frac{2}{\varepsilon_{t}}}\nabla\hat{\phi}_{ t}(X^{\hat{b}}_{t})\cdot dW_{t}.$ (32) $$

This equation is derived in Appendix 5.1.

###### Proposition 6 (Action Matching objective) .

###### Remark 2 .

## 3 Implementation

The PINN objective (26) can be written as

$$ $L^{T}_{\text{PINN}}[\hat{b},\hat{F}]=\int_{0}^{T}\mathbb{E}\Big{[}\big{|} \nabla\cdot\hat{b}_{t}(\hat{x}_{t})-\nabla U_{t}(\hat{x}_{t})\cdot\hat{b}_{t}( \hat{x}_{t})-\partial_{t}U_{t}(\hat{x}_{t})+\partial_{t}\hat{F}_{t}\big{|}^{2} \Big{]}dt$ (33) $$

where the expectation is taken over $\hat{x}_{t}$ drawn from $\hat{\rho}_{t}(x)$, and it can be estimated empirically over a population of such samples. This is the off-policy variant of the PINN objective, and it is perfectly valid, but may be inefficient for learning $\hat{b}_{t}$ over the support necessary for the problem. If we decide instead to set $\hat{\rho}_{t}(x)=\rho_{t}(x)$, since the SDEs in (20) and (21) can be used with any $\hat{b}_{t}(x)$ to estimate expectation over $\rho_{t}(x)$ via (22), we can write the PINN objective on-policy as

$$ $L^{T}_{\text{PINN}}[\hat{b},\hat{F}]=\int_{0}^{T}\frac{1}{\mathbb{E}[e^{A^{ \hat{b}}_{t}}]}\mathbb{E}\Big{[}e^{A^{\hat{b}}_{t}}\big{|}\nabla\cdot\hat{b}_{ t}(X^{\hat{b}}_{t})-\nabla U_{t}(X^{\hat{b}}_{t})\cdot\hat{b}_{t}(X^{\hat{b}}_ {t})-\partial_{t}U_{t}(X^{\hat{b}}_{t})+\partial_{t}\hat{F}_{t}\big{|}^{2}\Big {]}dt$ (34) $$

These expectations can be estimated empirically over a population of solutions to (20) and (21). Crucially, since we can switch from off-policy to on-policy after taking the gradient of the PINN objective, when computing the gradient of (34) over $\hat{b}_{t}(x)$, $(X^{\hat{b}}_{t},A^{\hat{b}}_{t})$ can be considered independent of $\hat{b}_{t}(x)$ and do not need to be differentiated over. In other words, the method is simulation-free even if used on-policy, i.e. even though it uses the current value of $\hat{b}_{t}$ to estimate the loss and its gradient.

Learning $b_{t}(x)$ and $F_{t}$ for $t\in[0,1]$ from the start can be challenging if the initial $\hat{b}_{t}(x)$ is far from exact and the weights gets large variance as $t$ increases. This problem can be alleviated by estimating $b_{t}(x)$ sequentially. In practice, this amounts to annealing $T$ from a small initial value to $T=1$, in such a way that $b_{t}(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 $t\in[0,1]$.

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. (); Bolić et al. ().

Figure: Algorithm 1 Training: Note that for both objectives the resultant set of walkers across time slices $\{x_{k}^{i}\}$ are detached from the computational graph when taking a gradient step (off-policy learning).

We can proceed similarly with the AM loss (30) by rewriting it as

$$ $\displaystyle L^{T}_{\text{AM}}[\hat{\phi}]=\int_{0}^{T}\frac{\mathbb{E}\big{[ }e^{A^{\hat{b}}_{t}}\big{[}\tfrac{1}{2}|\nabla\hat{\phi}_{t}(X_{t}^{\hat{b}})| ^{2}+\partial_{t}\phi_{t}(X_{t}^{\hat{b}})\big{]}\big{]}}{\mathbb{E}[e^{A^{ \hat{b}}_{t}}]}dt+\frac{\mathbb{E}\big{[}e^{A^{\hat{b}}_{0}}\phi_{0}(X_{0}^{ \hat{b}})\big{]}}{\mathbb{E}[e^{A^{\hat{b}}_{0}}]}-\frac{\mathbb{E}\big{[}e^{A ^{\hat{b}}_{T}}\phi_{T}(X_{T}^{\hat{b}})\big{]}}{\mathbb{E}[e^{A^{\hat{b}}_{T} }]}.$ (35) $$

These expectations can be estimated empirically over solutions to (20) and (21) with $\hat{b}_{t}(x)=\nabla\hat{\phi}_{t}(x)$. The above implementation is detailed in Algorithm 1.

## 4 Numerical experiments

In what follows, we test the NETS method, for both the PINN objective (26) and the action matching objective (30), 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). Following 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\times L=20$ lattice).

### 4.1 40-Mode Gaussian Mixture

A common benchmark for machine learning augmented samplers that originally appeared in the paper introducing Flow Annealed Importance Sampling Bootstrap (FAB) Midgley et al. () 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 re-weighting or locally updating MCMC processes.
We choose as a time dependent potential $U_{t}(x)$ the linear interpolation (1) with $U_{0}$ the potential for a standard multivariate Gaussian with standard deviation scale $\sigma=2$.

**Table 1: Performance of NETS in terms of ESS and 2-Wasserstein metrics for 40-mode GMM $(d=2)$ with comparative results quoted from Akhound-Sadegh et al. () for reproducibility.**
| GMM $(d=2)$ |  |  |
| --- | --- | --- |
| Algorithm | ESS $\uparrow$ | $\mathcal{W}_{2}$ $\downarrow$ |
| FAB | $0.653\pm 0.017$ | $12.0\pm 5.73$ |
| PIS | $0.295\pm 0.018$ | $7.64\pm 0.92$ |
| DDS | $0.687\pm 0.208$ | $9.31\pm 0.82$ |
| pDEM | $0.634\pm 0.084$ | $12.20\pm 0.14$ |
| iDEM | $0.734\pm 0.092$ | $7.42\pm 3.44$ |
| NETS-AM (ours) | $0.808\pm 0.031$ | $3.89\pm 0.22$ |
| NETS-PINN $\varepsilon_{t}=0$ (ours) | $\mathbf{0.954\pm 0.003}$ | $\mathbf{3.55\pm 0.57}$ |
| NETS-PINN $\varepsilon_{t}=4$ (ours) | $\mathbf{0.979\pm 0.002}$ | $\mathbf{3.14\pm 0.46}$ |
| NETS-PINN-resample (ours) | $\mathbf{0.993\pm 0.004}$ | $\mathbf{3.27\pm 0.31}$ |

We train a simple feed-forward neural network of width $256$ against both the PINN objective (26), parameterizing $(\hat{b}$, $\hat{F})$, or the action matching objective (30), parameterizing $\hat{\phi}$.
We compare the learned model from both objectives to recent related literature: FAB, Path Integral Sampler (PIS) (Zhang & Chen, ), Denoising Diffusion Sampler (DDS) (Vargas et al., ), and Denoising Energy Matching (pDEM, iDEM) (Akhound-Sadegh et al., ). 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 ($\mathcal{W}_{2}$) distance between the model and the target. As noted in Table 1, all proposed variants of 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 acceptance rate of the same learned PINN model to nearly 100% by using a single resampling step when the ESS of the walkers dropped below $98\%$. NETS uses 100 sampling steps and an $\varepsilon_{t}=0.0,4.0$ in the SDE.

Figure: Figure 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. (). Left: Annealed Langevin run for 250 steps with $\varepsilon_{t}=4.0$, failing to capture the modes with $0\%$ ESS. Center: Learning using the PINN loss and sampling with 100 steps and $\varepsilon_{t}=0$ achieves an ESS of $95\%$. Right: Same learning and now sampling with $\varepsilon_{t}=4.0$ achieves an ESS of $98\%$.
Refer to caption: x1.png

### 4.2 10-d Funnel distribution

**Table 2: Performance of NETS on the Neal’s Funnel distribution, measured in maximum mean discrepancy from the true distribution for reproducibility with the numbers quoted from Blessing et al. ().**
| Funnel $(d=10)$ |  |
| --- | --- |
| Algorithm | MMD $\downarrow$ |
| FAB (Midgley et al., ) | $\mathbf{0.032\pm 0.000}$ |
| GMMVI (Arenz et al., ) | $\mathbf{0.031\pm 0.000}$ |
| PIS (Zhang & Chen, ) | – – |
| DDS (Vargas et al., ) | $0.172\pm 0.031$ |
| AFT (Arbel et al., ) | $0.159\pm 0.010$ |
| CRAFT (Arbel et al., ) | $0.115\pm 0.003$ |
| NETS-AM (ours) | $0.041\pm 0.001$ |
| NETS-PINN (ours) | $\mathbf{0.031\pm 0.002}$ |
| NETS-PINN-resample (ours) | $\mathbf{0.022\pm 0.002}$ |

We next test NETS on Neal’s funnel, a challenging synthetic target distribution which exhibits correlations at different scales across its 10 dimensions. The definitions of the target density and the interpolating potential are given in Appendix 5.4. Heuristically, the first dimension is Gaussian with variance $\sigma^{2}=9$, and the other 9 dimensions are conditionally Gaussian with variance $\text{exp}(x_{0})$, creating the funnel.

We again parameterize $(\hat{b},\hat{F})$ or $\hat{\phi}$ using simple feed forward neural networks, this time of hidden size $512$. We use 200 sampling steps and a mild diffusion coefficient of $\varepsilon_{t}=0.3$ to sample from the model. Following Blessing et al. (), we compute the maximum mean discrepancy (MMD) between 10000 samples from the model and 10000 samples from the target and compare to related methods in Table 2. NETS matches the best performing predecessors, and, when using a single resampling step during generation when the ESS drops below $70\%$, surpasses them.

### 4.3 Scaling on high-dimensional GMMs

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 performance scales. In addition we are curious to understand how the factor in the sampling SDE coming from annealed Langevin dynamics, $\nabla U$, interacts with the learned drift $\hat{b}$ or $\nabla\hat{\phi}$ as we change the diffusivity. We construct $8$-mode target GMMs in $d=36,64,128,200$ dimensions and learn $\hat{b}$ with the PINN loss in each scenario.

Figure: Figure 2: Demonstration of high-dimensional sampling with our method using the PINN loss in (26) and a study of how diffusivity impacts performance, with and without transport. Left: NETS can achieve high ESS through transport alone, and the effect of increased diffusivity has more of a positive effect on performance with sampling than without. AIS cannot achieve ESS above $\approx 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.
Refer to caption: x2.png

We use the same feed forward neural network of width $512$ and depth $4$ to parameterize both $\hat{b}$ and $\hat{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 $\varepsilon_{t}$ and therefore the effect of the Langevin term coming from the gradient of the potential, we note all the 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 $\varepsilon_{t}=80$ data point. The number of sampling steps used to discretize the SDEs in these experiments ranged from $K=200$ for $\varepsilon_{t}=0$ up to $K=2000$ for $\varepsilon_{t}=80$. Nonetheless, it suggests that diffusion can be more helpful when there is already some successful transport than without.

### 4.4 Lattice φ 4 superscript 𝜑 4 \varphi^{4} italic_φ start_POSTSUPERSCRIPT 4 end_POSTSUPERSCRIPT theory

We next apply NETS to the simulation of a 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 $\varphi^{4}$ theory in $D=2$ spacetime dimensions. The random variables in this circumstance are field configurations $\varphi\in\mathbb{R}^{L\times L}$, where $L$ is the extent of space and time. The interpolating energy function under which we seek to sample is defined as:

$$ $\displaystyle U_{t}(\varphi)=\sum_{x}\left[-2\sum_{\mu}\varphi_{x}\varphi_{x+ \mu}\right]+(2D+m^{2}_{t})\varphi_{x}^{2}+\lambda_{t}\varphi_{x}^{4},$ (36) $$

where summation over $x$ indicates summation over the lattices sites, and $m^{2}_{t}$ and $\lambda_{t}$ are time dependent parameters of the theory that define the phase of the lattice (ranging from disordered to ordered, otherwise known as magnetized). A derivation of this energy function is given in Appendix 5.5. Importantly, sampling the lattice configurations becomes challenging when approaching the phase transition between the disordered and ordered phases.

Figure: Figure 3: Comparison of the performance of NETS to AIS on two different settings for the study of $\varphi^{4}$ theory. Top row, left: 10 example generative lattice configurations with parameters $L=20$, $m^{2}=-1.0$, $\lambda=0.9$, which demarcates the phase transition to the antiferromagnetic phase. Top row, right: Performance of AIS (purple curve) vs. NETS (red curve) in terms of effective sample size over time of integration $t$, and a histogram of the average magnetization of $4000$ lattice configurations, sampled with AIS, NETS, and HMC (superposed in 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^{2}=-1.0$, $\lambda=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 to sample the correct distribution, and its weights are too high variance to be used on the histogram.
Refer to caption: x3.png

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 $\hat{\phi}_{t}$ a simple feed forward neural network. We use the free theory $\lambda_{0}=0$ as the base distribution under which we initially draw samples. The definition of the target parameter values $m^{2}_{1},\lambda_{1}$ both at the phase transition and in the ordered phase are given in the Appendix 5.5. 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 in sampling at 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 either all positive or all negative values. Again in this regime, NETS is nearly 2 orders of magnitude more statistically efficient.

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 to use 1500-2000 steps in the integrator to properly resolve the dynamics of the SDE.

## Acknowledgements

We thank Francisco Vargas, Lee Cheuk-Kit, and Hugo Cui for helpful discussions on related work, as well as Ryan Abbott for helpful clarification on the free-field theory. MSA is supported by a Junior Fellowship from the Harvard Society of Fellows and the ONR project under the Vannevar Bush award “Mathematical Foundations and Scientific Applications of Machine Learning”. EVE is supported by the National Science Foundation under Awards DMR1420073, DMS-2012510, and DMS-2134216, by the Simons Collaboration on Wave Turbulence, Grant No. 617006, and by a Vannevar Bush Faculty Fellowship.

## 5 Appendix

### 5.1 Proofs of Sec. 2

Here we provide the proofs of the statements made in Sec. 2 which, for the reader convenience, we recall.

See 1

###### Proof.

Let $f_{t}(x,a)$ with $(x,a)\in\mathbb{R}^{d+1}$ be the PDF of the joint process $(X_{t},A_{t})$ defined by the SDE (10) and (11). This PDF solves the FPE

$$ $\partial_{t}f_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{t}f_{t}+\nabla f_{t})+ \partial_{t}U_{t}\partial_{a}f_{t},\qquad f_{t=0}(x,a)=\delta(a)\rho_{0}(x).$ (37) $$

Define

$$ $g_{t}(x)=\int_{\mathbb{R}}e^{a}f_{t}(x,a)da.$ (38) $$

We can derive an equation for $g_{t}(x)$ by multiplying both sides of the FPE (37) by $e^{a}$ and integrating over $a\in\mathbb{R}$. Using

$$ $\displaystyle\int_{\mathbb{R}}e^{a}\partial_{t}f_{t}(x,a)da$ $\displaystyle=\partial_{t}\int_{\mathbb{R}}e^{a}f_{t}(x,a)da=\partial_{t}g_{t},$ (39) $\displaystyle\int_{\mathbb{R}}e^{a}\varepsilon_{t}\nabla\cdot(\nabla U_{t}f_{t }+\nabla f_{t})da$ $\displaystyle=\varepsilon_{t}\nabla\cdot\Big{(}\nabla U_{t}\int_{\mathbb{R}}e^ {a}f_{t}(x,a)da+\nabla\int_{\mathbb{R}}e^{a}f_{t}(x,a)da\Big{)}$ $\displaystyle=\varepsilon_{t}\nabla\cdot(\nabla U_{t}g_{t}+\nabla g_{t}),$ $\displaystyle\int_{\mathbb{R}}e^{a}\partial_{t}U_{t}\partial_{a}f_{t}da$ $\displaystyle=\partial_{t}U_{t}\int_{\mathbb{R}}e^{a}\partial_{a}f_{t}da$ $\displaystyle=-\partial_{t}U_{t}\int_{\mathbb{R}}e^{a}f_{t}da=-\partial_{t}U_{ t}g_{t},$ $$

where we arrived at the second equality in the third equation by integration by parts, we deduce that

$$ $\partial_{t}g_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{t}g_{t}+\nabla g_{t})- \partial_{t}U_{t}g_{t},\qquad g_{t=0}(x)=\rho_{0}(x)=e^{-U_{0}(x)+F_{0}}.$ (40) $$

The solution to this parabolic PDE is unique and it can be checked by direct substitution that it is given by

$$ $g_{t}(x)=e^{-U_{t}(x)+F_{0}}.$ (41) $$

Note that this solution is not normalized since it contains $F_{0}$ rather than $F_{t}$. In fact it is easy to see that

$$ $\int_{\mathbb{R}^{d}}g_{t}(x)dx=\int_{R^{d+1}}e^{a}f_{t}(x,a)dxda=e^{-F_{t}+F_ {0}},$ (42) $$

where the first equality follows from the definition of $g_{t}$ and the second from its explicit expression and the definition of the free energy that implies $\int_{\mathbb{R}^{d}}e^{-U_{t}(x)}dx=e^{-F_{t}}$. Equation (41) is the second equation in (12). From (41) we also deduce that, given any test function $h:\mathbb{R}^{d}\to\mathbb{R}$, we have

$$ $\displaystyle\frac{\int_{R^{d+1}}e^{a}h(x)f_{t}(x,a)dxda}{\int_{R^{d+1}}e^{a}f _{t}(x,a)dxda}$ $\displaystyle=\frac{\int_{R^{d}}h(x)g_{t}(x)dx}{\int_{R^{d}}g_{t}(x)dx}$ (43) $\displaystyle=\frac{\int_{R^{d}}h(x)e^{-U_{t}(x)+F_{0}}dx}{\int_{R^{d}}e^{-U_{ t}(x)+F_{0}}}$ $\displaystyle=\frac{\int_{R^{d}}h(x)e^{-U_{t}(x)}dx}{\int_{R^{d}}e^{-U_{t}(x)}dx}$ $\displaystyle=e^{F_{t}}\int_{R^{d}}h(x)e^{-U_{t}(x)}dx$ $\displaystyle=\int_{R^{d}}h(x)\rho_{t}(x)dx.$ $$

Since by definition of $f_{t}(x,a)$ the left hand-side of this equation can be expressed as the ratio of expectations over $(X_{t},A_{t})$ in the first equation in (12) we are done.
∎

See 2

###### Proof.

If $b_{t}$ satisfies (14), then $\rho_{t}$ satisfies the FPE (15). Since (16) is the SDE associated with this FPE, (17) holds.
∎

See 3

###### Proof.

We can follow the same steps as in the proof of Proposition 1 by considering the PDF $f^{\hat{b}}_{t}(x,a)$ of $(X_{t}^{\hat{b}},A^{\hat{b}}_{t})$. This PDF solves the FPE

$$ $\displaystyle\partial_{t}f^{\hat{b}}_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{ t}f^{\hat{b}}_{t}+\nabla f^{\hat{b}}_{t})-\nabla\cdot(\hat{b}_{t}f_{t}^{\hat{b }})-(\nabla\cdot\hat{b}_{t}-\nabla U_{t}-\partial_{t}U_{t})\partial_{a}f^{\hat {b}}_{t},$ (44) $\displaystyle f^{\hat{b}}_{t=0}(x,a)=\delta(a)\rho_{0}(x).$ $$

Define

$$ $g^{\hat{b}}_{t}(x)=\int_{\mathbb{R}}e^{a}f^{\hat{b}}_{t}(x,a)da.$ (45) $$

We can derive an equation for $g^{\hat{b}}_{t}(x)$ by multiplying both sides of the FPE (44) by $e^{a}$ and integrating over $a\in\mathbb{R}$. Using

$$ $\displaystyle\int_{\mathbb{R}}e^{a}\partial_{t}f^{\hat{b}}_{t}da$ $\displaystyle=\partial_{t}\int_{\mathbb{R}}e^{a}f^{\hat{b}}_{t}da=\partial_{t} g^{\hat{b}}_{t},$ (46) $\displaystyle\int_{\mathbb{R}}e^{a}\varepsilon_{t}\nabla\cdot(\nabla U_{t}f^{ \hat{b}}_{t}+\nabla f^{\hat{b}}_{t})da$ $\displaystyle=\varepsilon_{t}\nabla\cdot\Big{(}\nabla U_{t}\int_{\mathbb{R}}e^ {a}f^{\hat{b}}_{t}da+\nabla\int_{\mathbb{R}}e^{a}f^{\hat{b}}_{t}da\Big{)},$ $\displaystyle-\int_{\mathbb{R}}e^{a}\nabla\cdot(\hat{b}_{t}f^{\hat{b}}_{t})da$ $\displaystyle=-\nabla\cdot\Big{(}\hat{b}_{t}\int_{\mathbb{R}}e^{a}f^{\hat{b}}_ {t}da\Big{)}$ $\displaystyle=-\nabla\cdot(\hat{b}_{t}g^{\hat{b}}_{t}),$ $\displaystyle-\int_{\mathbb{R}}e^{a}(\nabla\cdot\hat{b}_{t}-\nabla U_{t}- \partial_{t}U_{t})\partial_{a}f^{\hat{b}}_{t}da$ $\displaystyle=-(\nabla\cdot\hat{b}_{t}-\nabla U_{t}-\partial_{t}U_{t})\int_{ \mathbb{R}}e^{a}\partial_{a}f^{\hat{b}}_{t}da$ $\displaystyle=(\nabla\cdot\hat{b}_{t}-\nabla U_{t}-\partial_{t}U_{t})\int_{ \mathbb{R}}e^{a}f^{\hat{b}}_{t}da$ $\displaystyle=(\nabla\cdot\hat{b}_{t}-\nabla U_{t}-\partial_{t}U_{t})g^{\hat{b }}_{t},$ $$

where we arrived at the second equality in the fourth equation by integration by parts, we deduce that

$$ $\displaystyle\partial_{t}g^{\hat{b}}_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{ t}g^{\hat{b}}_{t}+\nabla g^{\hat{b}}_{t})-\nabla\cdot(\hat{b}_{t}g^{\hat{b}}_{ t})+(\nabla\cdot\hat{b}_{t}-\nabla U_{t}-\partial_{t}U_{t})g^{\hat{b}}_{t},$ (47) $\displaystyle g^{\hat{b}}_{t=0}(x)=\rho_{0}(x)=e^{-U_{0}(x)+F_{0}}.$ $$

The solution to this parabolic PDE is unique and it can be checked by direct substitution that it is given by

$$ $g^{\hat{b}}_{t}(x)=e^{-U_{t}(x)+F_{0}}.$ (48) $$

This solution is not normalized since it contains $F_{0}$ rather than $F_{t}$, and it is easy to see that

$$ $\int_{\mathbb{R}^{d}}g^{\hat{b}}_{t}(x)dx=\int_{R^{d+1}}e^{a}f^{\hat{b}}_{t}(x ,a)dxda=e^{-F_{t}+F_{0}}.$ (49) $$

where the first equality follows from the definition of $g^{\hat{b}}_{t}$ and the second from its explicit expression and the definition of the free energy that implies $\int_{\mathbb{R}^{d}}e^{-U_{t}(x)}dx=e^{-F_{t}}$. Equation (49) is the secind equation in (22). From (48) we also deduce that, given any test function $h:\mathbb{R}^{d}\to\mathbb{R}$, we have

$$ $\displaystyle\frac{\int_{R^{d+1}}e^{a}h(x)f^{\hat{b}}_{t}(x,a)dxda}{\int_{R^{d +1}}e^{a}f^{\hat{b}}_{t}(x,a)dxda}$ $\displaystyle=\frac{\int_{R^{d}}h(x)g^{\hat{b}}_{t}(x)dx}{\int_{R^{d}}g^{\hat{ b}}_{t}(x)dx}$ (50) $\displaystyle=\frac{\int_{R^{d}}h(x)e^{-U_{t}(x)+F_{0}}dx}{\int_{R^{d}}e^{-U_{ t}(x)+F_{0}}}$ $\displaystyle=\frac{\int_{R^{d}}h(x)e^{-U_{t}(x)}dx}{\int_{R^{d}}e^{-U_{t}(x)}dx}$ $\displaystyle=e^{F_{t}}\int_{R^{d}}h(x)e^{-U_{t}(x)}dx$ $\displaystyle=\int_{R^{d}}h(x)\rho_{t}(x)dx.$ $$

Since by definition of $f^{\hat{b}}_{t}(x,a)$ the left hand-side of this equation can be expressed as the ratio of expectations over $(X^{\hat{b}}_{t},A^{\hat{b}}_{t})$ in the first equation in (22) we are done.
∎

See 4

###### Proof.

Clearly the minimum value of (26) is zero and the minimizing pair $(\hat{b},\hat{F})$ must satisfy

$$ $\nabla\cdot\hat{b}_{t}-\nabla U_{t}\cdot\hat{b}_{t}-\partial_{t}U_{t}+\partial _{t}\hat{F}_{t}=0$ (51) $$

By multiplying both sides of this equation by $\rho_{t}$ is can be written as

$$ $\nabla\cdot(\hat{b}_{t}\rho_{t})-\partial_{t}U_{t}\rho_{t}+\partial_{t}\hat{F} _{t}\rho_{t}=0$ (52) $$

This equation requires a solvability condition obtained by integrating it over $\mathbb{R}^{d}$. This gives

$$ $-\int_{\mathbb{R}^{d}}\partial_{t}U_{t}(x)\rho_{t}(x)dx+\partial_{t}\hat{F}_{t }=0,$ (53) $$

which, by (4), implies that $\partial_{t}\hat{F}_{t}=\partial_{t}F_{t}$. In turn, this implies that
(52) is equivalent to (14), i.e. $\hat{b}_{t}$ solves (14).
∎

See 5

###### Proof.

Consider

$$ $D_{\text{KL}}(\hat{\rho}_{t}||\rho_{t})=\int_{\mathbb{R}^{d}}\log\left(\frac{ \hat{\rho}_{t}(x)}{\rho_{t}(x)}\right)\hat{\rho}_{t}(x)dx$ (54) $$

where $\hat{\rho}_{t}$ satisfies (27). Taking the time-derivative of this expression we deduce that (using (27), $\rho_{t}(x)=e^{-U_{t}(x)+F_{t}}$, and multiple integrations by parts)

$$ $\displaystyle\frac{d}{dt}D_{\text{KL}}(\hat{\rho}_{t}||\rho_{t})$ $\displaystyle=\int_{\mathbb{R}^{d}}\left[\log\left(\frac{\hat{\rho}_{t}(x)}{ \rho_{t}(x)}\right)\partial_{t}\hat{\rho}_{t}(x)-\frac{\partial_{t}\rho_{t}(x) }{\rho_{t}(x)}\hat{\rho}_{t}(x)\right]dx$ (55) $\displaystyle=\int_{\mathbb{R}^{d}}\left[-\log\left(\frac{\hat{\rho}_{t}(x)}{ \rho_{t}(x)}\right)\nabla\cdot(\hat{b}_{t}(x)\hat{\rho}_{t}(x))+(\partial_{t}U _{t}(x)-\partial_{t}F_{t})\hat{\rho}_{t}(x)\right]dx$ $\displaystyle=\int_{\mathbb{R}^{d}}\left[\hat{b}_{t}(x)\cdot\nabla\log\left( \frac{\hat{\rho}_{t}(x)}{\rho_{t}(x)}\right)+\partial_{t}U_{t}-\partial_{t}F_{ t}\right]\hat{\rho}_{t}(x)dx$ $\displaystyle=\int_{\mathbb{R}^{d}}\left[\hat{b}_{t}(x)\cdot\nabla\hat{\rho}_{ t}(x)+\left(\hat{b}_{t}(x)\cdot\nabla U_{t}(x)+\partial_{t}U_{t}-\partial_{t}F _{t}\right)\rho_{t}(x)\right]dx$ $\displaystyle=\int_{\mathbb{R}^{d}}\left[-\nabla\cdot\hat{b}_{t}(x)+\hat{b}_{t }(x)\cdot\nabla U_{t}(x)+\partial_{t}U_{t}-\partial_{t}F_{t}\right]\rho_{t}(x)dx$ $$

Therefore

$$ $\displaystyle D_{\text{KL}}(\hat{\rho}_{t=1}||\rho_{1})$ $\displaystyle=\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[-\nabla\cdot\hat{b}_{t}(x )+\hat{b}_{t}(x)\cdot\nabla U_{t}(x)+\partial_{t}U_{t}-\partial_{t}F_{t}\right ]\rho_{t}(x)dxdt$ (56) $\displaystyle\leq\left[\int_{0}^{1}\int_{\mathbb{R}^{d}}\left|-\nabla\cdot\hat {b}_{t}(x)+\hat{b}_{t}(x)\cdot\nabla U_{t}(x)+\partial_{t}U_{t}-\partial_{t}F_ {t}\right|^{2}\rho_{t}(x)dxdt\right]^{1/2}$ $\displaystyle=\sqrt{L^{T=1}_{\text{PINN}}(\hat{b},F)}$ $\displaystyle\leq\sqrt{L^{T=1}_{\text{PINN}}(\hat{b},\hat{F})}$ $$

where the last equality holds for any $\hat{F}$ since $F$ is the minimizer.
∎

See 6

###### Proof.

By integrating by parts in time the term involving $\partial_{t}\phi_{t}$ in the AM objective (57), we can express is as

$$ $L^{T}_{AM}[\hat{\phi}]=\int_{0}^{T}\int_{\mathbb{R}^{d}}\big{[}\tfrac{1}{2}| \nabla\hat{\phi}_{t}(x)|^{2}\rho_{t}(x)-\phi_{t}(x)\partial_{t}\rho_{t}(x)\big {]}dxdt.$ (57) $$

This is a convex objective in $\hat{\phi}$ whose minimizers satisfy

$$ $\nabla\cdot(\nabla\hat{\phi}_{t}\rho_{t})=-\partial_{t}\rho_{t}.$ (58) $$

This is (14) written in terms of $b_{t}(x)=\nabla\phi_{t}(x)$. The solution of this equation is unique up to a constant by the Fredholm alternative since its right hand-side satisfies the solvability condition $\int_{\mathbb{R}^{d}}\partial_{t}\rho_{t}(x)dx=0$.
∎

###### Proof.

###### Proof.

###### Proof.

###### Proof.

###### Proof.

###### Proof.

##### Derivation of ( 31 ).

If $\hat{b}_{t}(x)=\nabla\hat{\phi}_{t}(x)$, the SDEs (20) and (21) reduce to

$$ $\displaystyle dX^{\hat{b}}_{t}$ $\displaystyle=-\varepsilon_{t}\nabla U_{t}(X^{\hat{b}}_{t})dt+\hat{\nabla}\phi _{t}(X^{\hat{b}}_{t})dt+\sqrt{2\varepsilon_{t}}dW_{t},\qquad$ $\displaystyle\hat{X}^{\hat{b}}_{0}\sim\rho_{0},$ (59) $\displaystyle dA^{\hat{b}}_{t}$ $\displaystyle=\Delta\hat{\phi}_{t}(X^{\hat{b}}_{t})dt-\nabla U_{t}(X^{\hat{b}} _{t})\cdot\nabla\hat{\phi}(X^{\hat{b}}_{t})dt-\partial_{t}U_{t}(X^{\hat{b}}_{t })dt,\qquad$ $\displaystyle A^{\hat{b}}_{0}=0,$ (60) $$

Since by Itô formula we have

$$ $\displaystyle d\hat{\phi}_{t}(X^{\hat{b}}_{t})$ $\displaystyle=\partial_{t}\hat{\phi}_{t}(X^{\hat{b}}_{t})dt-\varepsilon_{t} \nabla\hat{\phi}_{t}(X^{\hat{b}}_{t})\cdot\nabla U_{t}(X^{\hat{b}}_{t})dt+| \nabla\hat{\phi}_{t}(X^{\hat{b}}_{t})|^{2}dt$ (61) $\displaystyle+\sqrt{2\varepsilon_{t}}\nabla\hat{\phi}_{t}(X^{\hat{b}}_{t}) \cdot dW_{t}+\varepsilon_{t}\Delta\hat{\phi}_{t}(X^{\hat{b}}_{t})dt,$ $$

we can express

$$ $\displaystyle\Delta\hat{\phi}_{t}(X^{\hat{b}}_{t})dt$ $\displaystyle=\frac{1}{\varepsilon_{t}}d\hat{\phi}_{t}(X^{\hat{b}}_{t})dt- \frac{1}{\varepsilon_{t}}\partial_{t}\hat{\phi}_{t}(X^{\hat{b}}_{t})dt+\nabla \hat{\phi}_{t}(X^{\hat{b}}_{t})\cdot\nabla U_{t}(X^{\hat{b}}_{t})dt$ (62) $\displaystyle-\frac{1}{\varepsilon_{t}}|\nabla\hat{\phi}_{t}(X^{\hat{b}}_{t})| ^{2}dt-\sqrt{\frac{2}{\varepsilon_{t}}}\nabla\hat{\phi}_{t}(X^{\hat{b}}_{t}) \cdot dW_{t}.$ $$

If we insert this expression in the SDE (60), we can write it as

$$ $dA^{\hat{b}}_{t}=\frac{1}{\varepsilon_{t}}d\hat{\phi}_{t}(X^{\hat{b}}_{t})dt+ dB_{t}.$ (63) $$

where $dB_{t}$ is given by (32). Integrating (63) gives (31).

### 5.2 Time-discretized version of Proposition 3

Here we show how to generalize the result in Proposition 3 if we time discretize the SDE in (20) using Euler-Marayuma scheme and use some suitable time-discretized version of the ODE (21).

###### Proposition 7 .

Let $0=t_{0}<t_{1}<\cdots<t_{K}=1$ be a time grid on $[0,1]$, denote $\mathit{\Delta}t_{k}=t_{k+1}-t_{k}$ for $k=0,\ldots,K-1$, set $\tilde{X}_{0}^{\hat{b}}\sim\rho_{0}$ and $\tilde{A}^{\hat{b}}_{0}=0$, and for $k=0,\ldots,K-1$ define $\tilde{X}_{t_{k+1}}^{\hat{b}},\tilde{A}^{\hat{b}}_{t_{k+1}}$ recursively via

$$ $\displaystyle\tilde{X}^{\hat{b}}_{t_{k+1}}$ $\displaystyle=\tilde{X}^{\hat{b}}_{t_{k}}-\varepsilon_{t_{k}}\nabla U_{t_{k}}( \tilde{X}^{\hat{b}}_{t_{k}})\Delta t_{k}+\hat{b}_{t_{k}}(\tilde{X}^{\hat{b}}_{ t_{k}})\Delta t_{k}+\sqrt{2\varepsilon_{t_{k}}}(W_{t_{k+1}}-W_{t_{k}}),$ (64) $\displaystyle\tilde{A}^{\hat{b}}_{t_{k+1}}$ $\displaystyle=\tilde{A}^{\hat{b}}_{t_{k}}+U_{t_{k}}(\hat{X}^{\hat{b}}_{t_{k}}) -U_{t_{k+1}}(\hat{X}^{\hat{b}}_{t_{k+1}})+R^{+}_{k}(\tilde{X}^{\hat{b}}_{t_{k} },\tilde{X}^{\hat{b}}_{t_{k+1}})-R^{-}_{k}(\tilde{X}^{\hat{b}}_{t_{k+1}}, \tilde{X}^{\hat{b}}_{t_{k}}),$ (65) $$

where we defined

$$ $R^{\pm}_{k}(x,y)=\frac{1}{4\varepsilon_{t_{k}}\mathit{\Delta}t_{k}}\big{|}y-x+ \mathit{\Delta}t_{k}(\varepsilon_{t_{k}}\nabla U_{t_{k}}(x)\mp b_{t_{k}}(x)) \big{|}^{2}$ (66) $$

Then for all $k=0,\ldots,K$ and any test function $h:\mathbb{R}^{d}\to\mathbb{R}$, we have

$$ $\int_{\mathbb{R}^{d}}h(x)\rho_{t_{k}}(x)dx=\frac{\mathbb{E}[e^{\tilde{A}^{\hat {b}}_{t_{k}}}h(\tilde{X}^{\hat{b}}_{t_{k}})]}{\mathbb{E}[e^{\tilde{A}^{\hat{b} }_{t_{k}}}]},\qquad Z_{t_{k}}=e^{-F_{t_{k}}}=\mathbb{E}[e^{\tilde{A}^{\hat{b}} _{t_{k}}}],$ (67) $$

where the expectations are taken over the law of $(\tilde{X}^{\hat{b}}_{t_{k}},\tilde{A}^{\hat{b}}_{t_{k}})$

Note that the weights in (67) correct for the bias coming for both the time evolution of $U_{t}(x)$ 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 time-discretization of the ODE (21).

###### Proof.

For simplicity of notations we will prove (67) for $k=K$: the argument for all the other $k=1,\ldots,K-1$ is similar.
The update rule in (65) implies that

$$ $\displaystyle\tilde{A}^{\hat{b}}_{t_{K}}$ $\displaystyle=\sum_{k=0}^{K-1}\Big{(}U_{t_{k}}(\hat{X}^{\hat{b}}_{t_{k}})-U_{t _{k+1}}(\hat{X}^{\hat{b}}_{t_{k+1}})+R^{+}_{k}(\tilde{X}^{\hat{b}}_{t_{k}}, \tilde{X}^{\hat{b}}_{t_{k+1}})-R^{-}_{k}(\tilde{X}^{\hat{b}}_{t_{k+1}},\tilde{ X}^{\hat{b}}_{t_{k}})\Big{)}$ (68) $\displaystyle=U_{0}(\tilde{X}^{\hat{b}}_{t_{0}})-U_{K}(\tilde{X}^{\hat{b}}_{t_ {K}})+\sum_{k=0}^{K-1}\Big{(}R^{+}_{k}(\tilde{X}^{\hat{b}}_{t_{k}},\tilde{X}^{ \hat{b}}_{t_{k+1}})-R^{-}_{k}(\tilde{X}^{\hat{b}}_{t_{k+1}},\tilde{X}^{\hat{b} }_{t_{k}})\Big{)},$ $$

Now, given the test function $h:\mathbb{R}^{d}\to\mathbb{R}$, consider

$$ $\displaystyle I[h]$ $\displaystyle\equiv\mathbb{E}\big{[}e^{\tilde{A}^{\hat{b}}_{t_{K}}}h(\tilde{X} ^{\hat{b}}_{t_{K}})\big{]}$ (69) $$

Since the transition probability density function of the Euler-Maruyama update in (64) reads

$$ $\rho^{+}_{t_{k}}(x_{k+1}|x_{k})=(4\pi\varepsilon_{t_{k}}{\mathit{\Delta}}t_{k} )^{-d/2}\exp\big{(}-R^{+}_{k}(x_{k},x_{k+1})\big{)},$ (70) $$

the joint probability density function of the path $(\tilde{X}^{\hat{b}}_{t_{0}},\tilde{X}^{\hat{b}}_{t_{1}},\ldots,\tilde{X}^{
\hat{b}}_{t_{K}})$ is given by

$$ $\displaystyle\rho(x_{0},\ldots,x_{K})$ $\displaystyle=\exp\left(-U_{0}(x_{0})+F_{0}\right)\prod_{k=0}^{K-1}\rho^{+}_{t _{k},\mathit{\Delta}t_{k}}(x_{k+1}|x_{k})$ (71) $\displaystyle=C\exp\left(-U_{0}(x_{0})+F_{0}-\sum_{k=0}^{K-1}R^{+}_{k}(x_{k},x _{k+1})\right)$ $$

where $C=\prod_{k=0}^{K-1}(4\pi\varepsilon_{t_{k}}{\mathit{\Delta}}t_{k})^{-d/2}$.
We can use this density along with the explicit expression for $\tilde{A}^{\hat{b}}_{T_{K}}$ in (68) to express the expectation (69) as an integral over $\rho(x_{0},x_{1},\ldots,x_{K})$

$$ $\displaystyle I[h]$ $\displaystyle=C\int_{\mathbb{R}^{d(K+1}}dx_{0}\cdots dx_{K}\exp\left(-U_{0}(x_ {0})+F_{0}-\sum_{k=0}^{K-1}R^{+}_{k}(x_{k},x_{k+1})\right)$ (72) $\displaystyle\times\exp\left(U_{0}(x_{0})-U_{K}(x_{K})+\sum_{k=0}^{K-1}\Big{(} R^{+}_{k}(x_{k},x_{k+1})-R^{-}_{k}(x_{k+1},x_{k})\Big{)}\right)h(x_{K})$ $$

where the second exponential comes from the factor $e^{\tilde{A}^{\hat{b}}_{K}}$. (72) simplifies into

$$ $\displaystyle I[h]$ $\displaystyle=C\int_{\mathbb{R}^{d(K+1}}dx_{0}\cdots dx_{K}\exp\left(-U_{K}(x_ {K})+F_{0}-\sum_{k=0}^{K-1}R^{-}_{k}(x_{k+1},x_{k})\Big{)}\right)h(x_{K})$ (73) $$

In this expression we recognize a product of factors involving

$$ $\rho^{-}_{t_{k}}(x_{k}|x_{k+1})=(4\pi\varepsilon_{t_{k}})^{-d/2}\exp\big{(}-R^ {-}_{k}(x_{k+1},x_{k})\big{)},$ (74) $$

which is the transition probability density function of the time-reversed update

$$ $\displaystyle\tilde{Y}^{\hat{b}}_{t_{k}}$ $\displaystyle=\tilde{Y}^{\hat{b}}_{t_{k+1}}-\varepsilon_{t_{k}}\nabla U_{t_{k} }(\tilde{Y}^{\hat{b}}_{t_{k+1}})\Delta t_{k}-\hat{b}_{t_{k}}(\tilde{Y}^{\hat{b }}_{t_{k+1}})\Delta t_{k}+\sqrt{2\varepsilon_{t_{k}}}(W_{t_{k+1}}-W_{t_{k}}).$ (75) $$

This implies in particular that we can perform the integrals in (73) sequentially over $x_{0}$, $x_{1}$, .., $x_{K-1}$ to be left with

$$ $\displaystyle I[h]$ $\displaystyle=\int_{\mathbb{R}^{d}}\exp\left(-U_{K}(x_{K})+F_{0}\right)h(x_{K} )dx_{K}$ (76) $$

Therefore

$$ $\displaystyle I[1]$ $\displaystyle=\int_{\mathbb{R}^{d}}\exp\left(-U_{K}(x_{K})+F_{0}\right)dx_{K}= e^{-F_{K}+F_{0}},$ (77) $$

which is the second equation in (67), and

$$ $\displaystyle\frac{I[h]}{I[1]}$ $\displaystyle=e^{F_{K}-F_{0}}\int_{\mathbb{R}^{d}}\exp\left(-U_{K}(x_{K})+F_{0 }\right)h(x_{K})dx_{K}=\int_{\mathbb{R}^{d}}h(x)\rho_{t_{K}}(x)dx$ (78) $$

which is the first equation in (67).
∎

###### Proposition 7 .

###### Proof.

### 5.3 Solving for the optimal drift via Feynman-Kac formula

Without loss of generality, we can always look for a solution to (24) in the form of $b_{t}(x)=\nabla\phi_{t}(x)$, so that this equation becomes the Poisson equation

$$ $\Delta\phi_{t}-\nabla U_{t}\cdot\nabla\phi_{t}=\partial_{t}U_{t}-\partial_{t}F _{t}.$ (79) $$

The solution to this equation can be expressed via Feynman-Kac formula:

###### Proposition 8 .

Let $X^{t,x}_{\tau}$ satisfy the following SDE

$$ $dX^{t,x}_{\tau}=-\nabla U_{t}(X^{t,x}_{\tau})d\tau+\sqrt{2}dW_{\tau},\qquad X^ {t,x}_{\tau=0}=x$ (80) $$

where $U_{t}$ is evaluated fixed at $t\in[0,1]$ 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 $\rho_{t}(x)$. Then for all $(t,x)\in[0,1]\times\mathbb{R}^{d}$ we have

$$ $\phi_{t}(x)=\int_{0}^{\infty}\mathbb{E}\big{[}\partial_{t}F_{t}-\partial_{t}U_ {t}(X^{t,x}_{\tau})\big{]}d\tau$ (81) $$

where the expectation is taken over the law of $X^{t,x}_{\tau}$.

###### Proof.

By Ito formula,

$$ $\displaystyle d\phi_{t}(X^{t,x}_{\tau})$ $\displaystyle=\big{(}\Delta\phi_{t}(X^{t,x}_{\tau})-\nabla U_{t}(X^{t,x}_{\tau })\cdot\nabla\phi_{t}(X^{t,x}_{\tau})\big{)}d\tau+\sqrt{2}\nabla\phi_{t}(X^{t, x}_{\tau})\cdot dW_{\tau}$ (82) $\displaystyle=\left(\partial_{t}U_{t}(X^{t,x}_{\tau})-\partial_{t}F_{t}\right) d\tau+\sqrt{2}\nabla\phi_{t}(X^{t,x}_{\tau})\cdot dW_{\tau}$ $$

where the differential is taken with respect to $\tau$ at $t$ fixed, and we used (79) to get the second equality. If we integrate this relation on $\tau\in[0,T]$ and take expectation, we deduce that

$$ $\displaystyle\mathbb{E}\big{[}\phi_{t}(X^{t,x}_{T})\big{]}-\phi_{t}(x)$ $\displaystyle=\int_{0}^{T}\mathbb{E}\big{[}\partial_{t}U_{t}(X^{t,x}_{\tau})- \partial_{t}F_{t}\big{]}d\tau$ (83) $$

where we use Ito isometry to zero the expectation of the martingale term involving $\sqrt{2}\nabla\phi_{t}(X^{t,x}_{\tau})\cdot dW_{\tau}$.
If we let $T\to\infty$, by ergodicty the first term at the left hand side converges towards a constant independent of $(t,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 $\mathbb{E}\big{[}\partial_{t}U_{t}(X^{t,x}_{\tau})\big{]}\to\partial_{t}F_{t}$ exponentially fast as $\tau\to\infty$ by assumption of geometric ergodicity.
∎

###### Proposition 8 .

###### Proof.

##### Example: moving Gaussian distribution.

Let us consider the case where

$$ $U_{t}(x)=\tfrac{1}{2}(x-b_{t})^{T}A_{t}(x-b_{t}),$ (84) $$

where $b_{t}\in\mathbb{R}^{d}$ is a time-dependent vector field and $A_{t}=A_{t}^{T}\in\mathbb{R}^{d}\times\mathbb{R}^{d}$ is a time-dependent positive-definite matrix: we assume that both $b_{t}$ and $A_{t}$ are $C^{1}$ in time, and also that $\dot{A}_{t}A_{t}=A_{t}\dot{A}_{t}$. The free energy in this example is

$$ $F_{t}=-\log Z_{t},\qquad Z_{t}=(2\pi)^{d/2}|\det A_{t}|^{-1/2},$ (85) $$

so that

$$ $\partial_{t}U_{t}(x)=-\dot{b}_{t}^{T}A_{t}(x-b_{t})+\tfrac{1}{2}(x-b_{t})^{T} \dot{A}_{t}(x-b_{t}),\qquad\partial_{t}F_{t}=\tfrac{1}{2}\text{tr}(A_{t}^{-1} \dot{A}_{t}).$ (86) $$

In this case, the SDE (80) reads

$$ $dX^{t,x}_{\tau}=-A_{t}(X^{t,x}_{\tau}-b_{t})d\tau+\sqrt{2}dW_{\tau},\qquad X^{ t,x}_{\tau=0}=x,$ (87) $$

and its solution is

$$ $X^{t,x}_{\tau}=e^{-A_{t}\tau}x+\big{(}1-e^{-A_{t}\tau}\big{)}b_{t}+\sqrt{2} \int_{0}^{\tau}e^{-A_{t}(\tau-\tau^{\prime})}dW_{\tau^{\prime}}.$ (88) $$

This implies that (using Ito isometry)

$$ $\displaystyle\mathbb{E}\big{[}\partial_{t}U_{t}(X^{t,x}_{\tau})\big{]}$ $\displaystyle=-\dot{b}_{t}^{T}A_{t}e^{-A_{t}\tau}(x-b_{t})+\tfrac{1}{2}(x-b_{t })^{T}e^{-A_{t}\tau}\dot{A}_{t}e^{-A_{t}\tau}(x-b_{t})$ (89) $\displaystyle\quad+\int_{0}^{\tau}\text{tr}\big{(}e^{-A_{t}\tau}\dot{A}_{t}e^{ -A_{t}\tau}\big{)}d\tau$ $\displaystyle=-\dot{b}_{t}^{T}A_{t}e^{-A_{t}\tau}(x-b_{t})+\tfrac{1}{2}(x-b_{t })^{T}e^{-A_{t}\tau}\dot{A}_{t}e^{-A_{t}\tau}(x-b_{t})$ $\displaystyle\quad+\tfrac{1}{2}\text{tr}\big{(}A_{t}^{-1}\dot{A}_{t})-\tfrac{1 }{2}\text{tr}(A_{t}^{-1}\dot{A}_{t}e^{-2A_{t}\tau}).$ $$

Therefore, from (81), we have (using also (85))

$$ $\displaystyle\phi_{t}(x)$ $\displaystyle=\int_{0}^{\infty}\big{(}\dot{b}_{t}^{T}A_{t}e^{-A_{t}\tau}(x-b_{ t})-\tfrac{1}{2}(x-b_{t})^{T}e^{-A_{t}\tau}\dot{A}_{t}e^{-A_{t}\tau}(x-b_{t})$ (90) $\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad+\tfrac{1}{2}\text{tr}(A_{t}^ {-1}\dot{A}_{t}e^{-2A_{t}\tau})\big{)}d\tau$ $\displaystyle=\dot{b}_{t}\cdot(x-b_{t})-\tfrac{1}{4}(x-b_{t})^{T}\dot{A}_{t}A_ {t}^{-1}(x-b_{t})+\tfrac{1}{4}\text{tr}(A_{t}^{-1}\dot{A}_{t}A_{t}^{-1}).$ $$

This solution checks out since it implies that

$$ $-\nabla U_{t}(x)\cdot\nabla\phi_{t}(x)+\Delta\phi_{t}(x)=-\dot{b}_{t}^{T}A_{t} (x-b_{t})+\tfrac{1}{2}(x-b_{t})^{T}\dot{A}_{t}(x-b_{t})-\tfrac{1}{2}\text{tr}( A^{-1}_{t}\dot{A}_{t}),$ (91) $$

which is $\partial_{t}U_{t}(x)-\partial_{t}F_{t}$ as it should.

### 5.4 Details on Numerical Experiments

In the following we include details for reproducing the experiments presented in Section 4. An overview of the training procedure is given in Algorithm 1. Note that the SDE for the weights can replaced with (31) when learning with $\hat{\phi}_{t}$, as one would do with the action matching loss (30).

#### 5.4.1 Performance metrics

##### Effective sample size.

We can compute the self-normalized ESS as

$$ $\displaystyle\operatorname{ESS}_{t}=\frac{\left(N^{-1}\sum_{i=1}^{N}\exp\left( A_{t}^{i}\right)\right)^{2}}{N^{-1}\sum_{i=1}^{N}\exp\left(2A_{t}^{i}\right)}$ (92) $$

at time $t$ along the SDE trajectory. We can use the ESS both as a 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., ; Bolić et al., ). 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. These

##### 2-Wasserstein distance.

The 2-Wasserstein distance reported in Table 1 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 following Blessing et al. () to benchmark the performance of NETS on Neal’s funnel. We use the definition of the MMD as

$$ $\operatorname{MMD}^{2}\left(\hat{\rho},\rho\right)\approx\frac{1}{n(n-1)}\sum_ {i,j}^{n}k\left(\hat{x}_{i},\hat{x}_{j}\right)+\frac{1}{m(m-1)}\sum_{i,j}^{m}k \left(x_{i},x_{j}\right)-\frac{2}{nm}\sum_{i}^{n}\sum_{j}^{m}k\left(\hat{x}_{i },x_{j}\right)$ (93) $$

where $\hat{x}\sim\hat{\rho}$ is from the model distribution and $x\sim\rho$ is from the target and $k:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow R$ is chosen to be the radial basis kernel with unit bandwidth.

#### 5.4.2 40-mode GMM

The 40-mode GMM is defined with the mean vectors given as:

$$ $\displaystyle\mu_{1}$ $\displaystyle=\left(-0.2995,\;21.4577\right),$ $\displaystyle\mu_{2}$ $\displaystyle=\left(-32.9218,\;-29.4376\right),$ $\displaystyle\mu_{3}$ $\displaystyle=\left(-15.4062,\;10.7263\right),$ $\displaystyle\mu_{4}$ $\displaystyle=\left(-0.7925,\;31.7156\right),$ $\displaystyle\mu_{5}$ $\displaystyle=\left(-3.5498,\;10.5845\right),$ $\displaystyle\mu_{6}$ $\displaystyle=\left(-12.0885,\;-7.8626\right),$ $\displaystyle\mu_{7}$ $\displaystyle=\left(-38.2139,\;-26.4913\right),$ $\displaystyle\mu_{8}$ $\displaystyle=\left(-16.4889,\;1.4817\right),$ $\displaystyle\mu_{9}$ $\displaystyle=\left(15.8134,\;24.0009\right),$ $\displaystyle\mu_{10}$ $\displaystyle=\left(-27.1176,\;-17.4185\right),$ $\displaystyle\mu_{11}$ $\displaystyle=\left(14.5287,\;33.2155\right),$ $\displaystyle\mu_{12}$ $\displaystyle=\left(-8.2320,\;29.9325\right),$ $\displaystyle\mu_{13}$ $\displaystyle=\left(-6.4473,\;4.2326\right),$ $\displaystyle\mu_{14}$ $\displaystyle=\left(36.2190,\;-37.1068\right),$ $\displaystyle\mu_{15}$ $\displaystyle=\left(-25.1815,\;-10.1266\right),$ $\displaystyle\mu_{16}$ $\displaystyle=\left(-15.5920,\;34.5600\right),$ $\displaystyle\mu_{17}$ $\displaystyle=\left(-25.9272,\;-18.4133\right),$ $\displaystyle\mu_{18}$ $\displaystyle=\left(-27.9456,\;-37.4624\right),$ $\displaystyle\mu_{19}$ $\displaystyle=\left(-23.3496,\;34.3839\right),$ $\displaystyle\mu_{20}$ $\displaystyle=\left(17.8487,\;19.3869\right),$ $\displaystyle\mu_{21}$ $\displaystyle=\left(2.1037,\;-20.5073\right),$ $\displaystyle\mu_{22}$ $\displaystyle=\left(6.7674,\;-37.3478\right),$ $\displaystyle\mu_{23}$ $\displaystyle=\left(-28.9026,\;-20.6212\right),$ $\displaystyle\mu_{24}$ $\displaystyle=\left(25.2375,\;23.4529\right),$ $\displaystyle\mu_{25}$ $\displaystyle=\left(-17.7398,\;-1.4433\right),$ $\displaystyle\mu_{26}$ $\displaystyle=\left(25.5824,\;39.7653\right),$ $\displaystyle\mu_{27}$ $\displaystyle=\left(15.8753,\;5.4037\right),$ $\displaystyle\mu_{28}$ $\displaystyle=\left(26.8195,\;-23.5521\right),$ $\displaystyle\mu_{29}$ $\displaystyle=\left(7.4538,\;-31.0122\right),$ $\displaystyle\mu_{30}$ $\displaystyle=\left(-27.7234,\;-20.6633\right),$ $\displaystyle\mu_{31}$ $\displaystyle=\left(18.0989,\;16.0864\right),$ $\displaystyle\mu_{32}$ $\displaystyle=\left(-23.6941,\;12.0843\right),$ $\displaystyle\mu_{33}$ $\displaystyle=\left(21.9589,\;-5.0487\right),$ $\displaystyle\mu_{34}$ $\displaystyle=\left(1.5273,\;9.2682\right),$ $\displaystyle\mu_{35}$ $\displaystyle=\left(24.8151,\;38.4078\right),$ $\displaystyle\mu_{36}$ $\displaystyle=\left(-30.8249,\;-14.6588\right),$ $\displaystyle\mu_{37}$ $\displaystyle=\left(15.7204,\;33.1420\right),$ $\displaystyle\mu_{38}$ $\displaystyle=\left(34.8083,\;35.2943\right),$ $\displaystyle\mu_{39}$ $\displaystyle=\left(7.9606,\;-34.7833\right),$ $\displaystyle\mu_{40}$ $\displaystyle=\left(3.6797,\;-25.0242\right)$ $$

These means follow the definition given in the FAB (Midgley et al., ) code base that has been subsequently used in recent papers.

#### 5.4.3 Neal’s 10- d 𝑑 d italic_d funnel

The Neal’s Funnel distribution is a 10-$d$ probability distribution defined as

$$ $x_{0}\sim\mathsf{N}(0,\sigma^{2}),\quad x_{1:9}\sim\mathsf{N}(0,e^{x_{0}})$ (94) $$

where $\sigma=3$ and we use subscripts here as a dimensional index and not as a time index like in the rest of the paper. Following this, we use as a definition of the interpolating potential:

$$ $U_{t}(x)=\frac{1}{2}x_{0}^{2}(1-t+\frac{t}{\sigma^{2}})+\frac{1}{2}\sum_{i=1}^ {d-1}e^{-tx_{0}}x_{i}^{2}+(d-1)tx_{0}$ (95) $$

so that at $t=0$, we have $U_{0}(x)=\frac{1}{2}x_{0}^{2}+\frac{1}{2}\sum_{i=1}^{d-1}x_{i}^{2}$ and at time $t=1$ we have the funnel potential given as $U_{1}(x)=\frac{1}{2\sigma^{2}}x_{0}+\frac{1}{2}\sum_{i=1}^{d-1}e^{-x_{0}}x_{i}
^{2}+(d-1)x_{0}$.

### 5.5 Details of the φ 4 superscript 𝜑 4 \varphi^{4} italic_φ start_POSTSUPERSCRIPT 4 end_POSTSUPERSCRIPT model

We consider the Euclidean scalar $\phi^{4}$ theory given by the action

$$ $\displaystyle S_{\text{Euc}}[\varphi]=\int\,\big{[}\partial_{\mu}\varphi(x) \partial^{\mu}\varphi(x)+m^{2}\varphi^{2}(x)+\lambda\varphi^{4}(x)\big{]}\,d^{ D}x$ (96) $$

where we use Einstein summation to denote the dot product with respect to the Euclidean metric and $D$ is the spacetime dimension. We are interested in acquiring a variant of this expression that provides a fast computational realization when put onto the lattice. Using Green’s identity (integrating by parts) we note that

$$ $\displaystyle\int(\partial_{\mu}\varphi(x)\partial^{\mu}\varphi(x)d^{d}x=\int \,\partial_{\mu}\varphi\cdot\partial_{\mu}\varphi\,d^{d}x=-\int\varphi(x) \partial_{\mu}\partial^{\mu}\varphi(x)d^{d}x\,+\text{vanishing surface term}$ (97) $$

so that

$$ $\displaystyle S_{\text{Euc}}[\varphi]=\int-\varphi(x)\partial_{\mu}\partial^{ \mu}\varphi(x)+m^{2}\varphi^{2}(x)+\lambda\varphi^{4}(x)\,d^{d}x.$ (98) $$

Discretizing $S_{\text{Euc}}$ onto the lattice

$$ $\Lambda=\{a(n_{0},\dots,n_{d-1})\mid n_{i}\in\{0,1,2,\dots,L\},\,i=0,1,\dots,d ,a\in\mathbb{R}_{+}\},$ $$

where $a$ is the lattice spacing used to define the physical point $x=an$, we use the forward difference operator to define

$$ $\displaystyle\partial_{\mu}\varphi(x)\rightarrow\tfrac{1}{a}[\varphi(x+\mu)- \varphi(x)]\qquad\partial_{\mu}\partial^{\mu}\varphi(x)\rightarrow\tfrac{1}{a^ {2}}[\varphi(x+\mu)-2\varphi(x)+\varphi(x-\mu)].$ (99) $$

Using these expressions, we write the discretized lattice action as

$$ $\displaystyle S_{\text{Lat}}$ $\displaystyle=\sum_{x\in\Lambda}a^{D}\left[\sum_{\mu=1}^{D}-\tfrac{1}{a^{2}} \big{[}\varphi_{x+\mu}\varphi_{x}-2\varphi_{x}^{2}+\varphi_{x-\mu}\varphi_{x}] +m^{2}\varphi_{x}^{2}+\lambda\varphi_{x}^{4}\right]$ (100) $\displaystyle=\sum_{x}a^{D}\left[2Da^{-2}\varphi_{x}^{2}-a^{-2}\sum_{\mu}\big{ [}\varphi_{x+\mu}\varphi_{x}+\varphi_{x-\mu}\varphi_{x}]+m^{2}\varphi_{x}^{2}+ \lambda\varphi_{x}^{4}\right]$ (101) $\displaystyle=\sum_{x}a^{d}\left[2Da^{-2}\varphi_{x}^{2}-2a^{-2}\sum_{\mu}[ \varphi_{x}\varphi_{x+\mu}]+m^{2}\varphi_{x}^{2}+\lambda\varphi_{x}^{4}\right]$ (102) $\displaystyle=\sum_{x}a^{D}\left[-2a^{-2}\sum_{\mu}\varphi_{x}\varphi_{x+\mu}+ (2a^{-2}D+m^{2})\varphi_{x}^{2}+\lambda\varphi_{x}^{4}\right]$ (103) $$

where we have used the fact that on the lattice $\sum_{x}\varphi_{x}\varphi_{x+\hat{\mu}}=\sum_{x}\varphi_{x-\hat{\mu}}\varphi_
{x}$ to get the third equality.
It is useful to put the action in a form that is independent of the lattice spacing $a$. To do so, we introduce the re-scaled lattice field as

$$ $\displaystyle\varphi_{x}\rightarrow a^{D/2-1}\varphi_{x},\quad m^{2} \rightarrow a^{2}m^{2},\quad\text{and}\quad\lambda\rightarrow a^{4-D}\lambda.$ (104) $$

Plugging these rescalings into (103) gives us the final expression

$$ $\displaystyle S_{\text{Lat}}=\sum_{x}\left[-2\sum_{\mu}\varphi_{x}\varphi_{x+ \mu}\right]+(2D+m^{2})\varphi_{x}^{2}+\lambda\varphi_{x}^{4},$ (105) $$

which we are to use in simulation.

#### 5.5.1 Free theory λ = 0 𝜆 0 \lambda=0 italic_λ = 0

Turning off the interaction makes it possible to analytically solve the theory. To do this, introduce the discrete Fourier transform relations

$$ $\displaystyle\varphi_{k}$ $\displaystyle=\frac{1}{\sqrt{L^{D}}}\sum_{x}\varphi_{x}e^{-ik\cdot x}$ (106) $\displaystyle\varphi_{x}$ $\displaystyle=\frac{1}{\sqrt{L^{D}}}\sum_{k}\varphi_{k}e^{ik\cdot x}$ (107) $$

for discrete wavenumbers $k=\frac{2l\pi}{L}$ with $l=0,\cdots,L-1$. Plugging in (107) into the first part of (103), we get the expanded sum

$$ $\displaystyle\sum_{x}\left[-2\sum_{\mu}\hat{\varphi}_{x}\hat{\varphi}_{x+\mu}\right]$ $\displaystyle\rightarrow-\frac{2}{L^{d}}\sum_{x}\sum_{\mu}\sum_{k}\sum_{k^{ \prime}}\varphi_{k}\varphi_{k^{\prime}}e^{i(k+k^{\prime})\cdot x}e^{ik^{\prime }\cdot\mu}$ (108) $\displaystyle=-2\sum_{\mu}\sum_{k}\sum_{k^{\prime}}\delta_{k,-k^{\prime}} \varphi_{k}\varphi_{k^{\prime}}e^{ik^{\prime}\cdot\mu}$ (109) $\displaystyle=-2\sum_{\mu}\sum_{k}\varphi_{k}\varphi_{-k}e^{-ik_{\mu}}$ (110) $\displaystyle=-2\sum_{\mu}\sum_{k}\varphi_{k}\varphi_{k^{*}}e^{-ik_{\mu}}$ (111) $\displaystyle=-2\sum_{\mu}\sum_{k}|\varphi_{k}|^{2}[\cos{k_{\mu}}+\cancel{i \sin{k_{\mu}}}]=-\sum_{\mu}\sum_{k}|\varphi_{k}|^{2}\cos{k_{\mu}}$ (112) $$

where $\phi^{*}$ 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 $\varphi^{2}$ gives us the expression

$$ $\displaystyle S_{k}=\sum_{k}\left[m^{2}+2D-2\sum_{\mu}\cos{k_{\mu}}\right]| \varphi|^{2}$ (113) $$

The above equation can be written in quadratic form to highlight that the field may be sampled analytically

$$ $\displaystyle S_{k}$ $\displaystyle=\frac{1}{L^{d}}\sum_{k}\varphi_{k}M_{k,-k}\varphi_{-k}$ (114) $\displaystyle\text{where}\quad M_{k,-k}=\left[m^{2}+2D-2\sum_{\mu}\cos{k_{\mu} }\right]\delta_{k,-k}$ (115) $$

Note that this free theory can be sampled for any $m^{2}>0$.

#### 5.5.2 φ 4 superscript 𝜑 4 \varphi^{4} italic_φ start_POSTSUPERSCRIPT 4 end_POSTSUPERSCRIPT numerical details

We numerically realize the above lattice theory in D=2 spacetime dimensions. We use an interpolating potential with time dependent $m^{2}_{t}=(1-t)m^{2}_{0}+tm^{2}_{1}$, $\lambda_{t}=(1-t)\lambda_{0}+t\lambda_{1}$ where $\lambda_{0}$ is always chosen to be 0 (though we note that you could run this sampler for any $U_{0}$ that you could sample from easily, not just analytically but also with existing MCMC methods). For the $L=20$ (d = $L\times 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 $M[\varphi^{i}(x)]=\sum_{x}\varphi^{i}(x)$, where summation is taken over all lattice sites on the $i^{th}$ lattice configuration. We identify this at $m^{2}_{1}=-1.0$, $\lambda_{1}=0.9$ and use these as the target theory parameters on which to perform the sampling. For the $L=16$ test ($d=256$), we go past this phase transition into the ordered phase of the theory, which we identify via HMC simluations at $m^{2}_{1}=-1.0$, $\lambda_{1}=0.8$.

### 5.6 Link with Vargas et al. ( 2024 )

Consider the process $Y_{t}^{\hat{b}}$ solution to the SDE

$$ $dY_{t}^{\hat{b}}=\varepsilon_{t}\nabla U_{t}(Y_{t}^{\hat{b}})dt+2\varepsilon_{ t}\nabla\log\rho_{t}^{\hat{b}}(Y_{t}^{\hat{b}})dt+\hat{b}_{t}(Y_{t}^{\hat{b}}) dt+\sqrt{2\varepsilon_{t}}dW_{t},\quad Y_{0}^{\hat{b}}\sim\rho_{0}.$ (116) $$

where $\rho_{t}^{\hat{b}}$ denotes the PDF of the process $X^{\hat{b}}_{t}$ defined by the SDE (20), i.e. the solution to the FPE

$$ $\partial_{t}\rho^{\hat{b}}_{t}=\varepsilon_{t}\nabla\cdot(\nabla U_{t}\rho^{ \hat{b}}_{t}+\nabla\rho^{\hat{b}}_{t})-\nabla\cdot(\hat{b}_{t}\rho^{\hat{b}}_{ t}),\qquad\rho^{\hat{b}}_{0}=\rho_{0}$ (117) $$

The process $Y^{\hat{b}}_{t}$ has a simple interpretation: it is the time-reversed of the process run using the time-reversed potential $U_{1-t}$ and $-\hat{b}_{1-t}$: that is, if the additional drift $\hat{b}_{t}$ was the perfect one solution to (24), the law of $X^{\hat{b}}=(X_{t}^{\hat{b}})_{t\in[0,1]}$ and $Y^{b}=(Y_{t}^{\hat{b}})_{t\in[0,1]}$ should coincide. This suggests to learn $b$ using as objective a divergence of the path measure of $X^{\hat{b}}$ from that of $Y^{\hat{b}}$. This is precisely what is suggested in Vargas et al. (), 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 mesaure of $X^{\hat{b}}$ from that of $Y^{\hat{b}}$ reads

$$ $\displaystyle\text{KL}(X^{\hat{b}}\|Y^{\hat{b}})$ $\displaystyle=\frac{1}{4}\varepsilon_{t}\int_{0}^{1}\mathbb{E}\big{[}|\nabla U _{t}(X^{\hat{b}}_{t})+\nabla\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})|^{2}\big{] }dt$ (118) $$

This objective is akin to the one used in score-based diffusion modeling (SBDM) and simply says that one way to adjust $\hat{b}$ is by matching the score of $\rho^{\hat{b}}_{t}$ to that of $\rho_{t}$. As written (118) is not explicit since we do not know
$\nabla\log\rho^{\hat{b}}_{t}$. 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

$$ $\displaystyle d\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})$ $\displaystyle=\nabla\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})\cdot(-\varepsilon_ {t}\nabla U_{t}(X^{\hat{b}}_{t})+\hat{b}_{t}(X^{\hat{b}}_{t}))dt+\varepsilon_{ t}\Delta\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})dt$ (119) $\displaystyle\quad+\sqrt{2\varepsilon_{t}}\nabla\log\rho^{\hat{b}}_{t}(X^{\hat {b}}_{t})\cdot dW_{t}$ $$

which implies that

$$ $\displaystyle\varepsilon_{t}\nabla\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})\cdot \nabla U_{t}((X^{\hat{b}}_{t})dt$ $\displaystyle=-d\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})+\nabla\log\rho^{\hat{b }}_{t}(X^{\hat{b}}_{t})\cdot\hat{b}_{t}(X^{\hat{b}}_{t})dt$ (120) $\displaystyle\quad+\varepsilon_{t}\Delta\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t} )dt+\sqrt{2\varepsilon_{t}}\nabla\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})\cdot dW _{t}$ $$

Inserting this expression in (118) after expanding the square, and noticing that the martingale term involving $dW_{t}$ disappears by Ito isometry and that the term $d\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})$ can be integrated in time we arrive at

$$ $\displaystyle\text{KL}(X^{\hat{b}}\|Y^{\hat{b}})$ $\displaystyle=\frac{1}{4}\varepsilon_{t}\int_{0}^{1}\mathbb{E}\big{[}|\nabla U _{t}(X^{\hat{b}}_{t})|^{2}+|\nabla\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})|^{2} +2\nabla U_{t}(X^{\hat{b}}_{t})\cdot\nabla\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{ t})\big{]}dt$ (121) $\displaystyle=\frac{1}{4}\int_{0}^{1}\mathbb{E}\big{[}\varepsilon_{t}|\nabla U _{t}(X^{\hat{b}}_{t})|^{2}+\varepsilon_{t}|\nabla\log\rho^{\hat{b}}_{t}(X^{ \hat{b}}_{t})|^{2}+\varepsilon_{t}\nabla U_{t}(X^{\hat{b}}_{t})\cdot\nabla\log \rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})\big{]}dt$ $\displaystyle\quad+\frac{1}{4}\int_{0}^{1}\mathbb{E}\big{[}\nabla\log\rho^{ \hat{b}}_{t}(X^{\hat{b}}_{t})\cdot\hat{b}_{t}(X^{\hat{b}}_{t})+\varepsilon_{t} \Delta\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})\big{]}dt$ $\displaystyle\quad+\frac{1}{4}\mathbb{E}[\log\rho_{0}(X_{0})]-\frac{1}{4} \mathbb{E}[\log\rho^{\hat{b}}_{1}(X_{1})]$ $$

where we used $\rho^{\hat{b}}_{0}=\rho_{0}$. We can now use the following identities, each obtained using $\rho^{\hat{b}}_{t}\nabla\log\rho^{\hat{b}}_{t}=\nabla\rho^{\hat{b}}_{t}$ and one integration by parts:

$$ $\displaystyle\mathbb{E}\big{[}|\nabla\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t})|^ {2}\big{]}$ $\displaystyle=\int_{\mathbb{R}^{d}}|\nabla\log\rho^{\hat{b}}_{t}(x)|^{2}\rho_{ t}^{\hat{b}}(x)dx$ (122) $\displaystyle=\int_{\mathbb{R}^{d}}\nabla\log\rho^{\hat{b}}_{t}(x)\cdot\nabla \rho_{t}^{\hat{b}}(x)dx$ $\displaystyle=-\int_{\mathbb{R}^{d}}\Delta\log\rho^{\hat{b}}_{t}(x)\rho_{t}^{ \hat{b}}(x)dx$ $\displaystyle=-\mathbb{E}\big{[}\Delta\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t}) \big{]},$ $$

$$ $\displaystyle\mathbb{E}\big{[}\nabla U_{t}(X^{\hat{b}}_{t})\cdot\nabla\log\rho ^{\hat{b}}_{t}(X^{\hat{b}}_{t})\big{]}$ $\displaystyle=\int_{\mathbb{R}^{d}}\nabla U_{t}(x)\cdot\nabla\log\rho^{\hat{b} }_{t}(x)\rho_{t}^{\hat{b}}(x)dx$ (123) $\displaystyle=\int_{\mathbb{R}^{d}}\nabla U_{t}(x)\cdot\nabla\rho_{t}^{\hat{b} }(x)dx$ $\displaystyle=-\int_{\mathbb{R}^{d}}\Delta U_{t}(x)\rho_{t}^{\hat{b}}(x)dx$ $\displaystyle=-\mathbb{E}\big{[}\Delta U_{t}(X^{\hat{b}}_{t})\big{]}$ $$

and

$$ $\displaystyle\mathbb{E}\big{[}\nabla\log\rho^{\hat{b}}_{t}(X^{\hat{b}}_{t}) \cdot\hat{b}_{t}(X^{\hat{b}}_{t})\big{]}$ $\displaystyle=\int_{\mathbb{R}^{d}}\nabla\log\rho^{\hat{b}}_{t}(x)\cdot\hat{b} _{t}(x)\rho_{t}^{\hat{b}}(x)dx$ (124) $\displaystyle=-\int_{\mathbb{R}^{d}}\nabla\cdot\hat{b}_{t}(x)\rho_{t}^{\hat{b} }(x)dx$ $\displaystyle=-\mathbb{E}\big{[}\nabla\cdot\hat{b}_{t}(X^{\hat{b}}_{t})\big{]}$ $$

Inserting these expressions in (121), it reduces to

$$ $\displaystyle\text{KL}(X^{\hat{b}}\|Y^{\hat{b}})$ $\displaystyle=\frac{1}{4}\int_{0}^{1}\mathbb{E}\big{[}\varepsilon_{t}|\nabla U _{t}(X^{\hat{b}}_{t})|^{2}-\varepsilon_{t}\Delta U_{t}(X^{\hat{b}}_{t})-\nabla \cdot b_{t}(X^{\hat{b}}_{t})\big{]}dt$ (125) $\displaystyle\quad+\frac{1}{4}\mathbb{E}[\log\rho_{0}(X_{0})]-\frac{1}{4} \mathbb{E}[\log\rho^{\hat{b}}_{1}(X_{1})]$ $$

This is Equation (24) in Vargas et al. () in which we set $\nabla\hat{\phi}_{t}(x)=\hat{b}_{t}(x)$ and we used that, for any $c_{t}:\mathbb{R}^{d}\to\mathbb{R}^{d}$, we have

$$ $\mathbb{E}\int_{0}^{1}c_{t}(X^{\hat{b}}_{t})\cdot\overleftarrow{d}W_{t}=\sqrt{ 2\varepsilon_{t}}\int_{0}^{1}\mathbb{E}\big{[}\nabla\cdot c_{t}(X^{\hat{b}}_{t })\big{]}dt.$ (126) $$

Note that we can neglect the term $\frac{1}{4}\mathbb{E}[\log\rho_{0}(X_{0})]$ in (125) since it does not depends on $\hat{b}$, so that the minimization of (125) can be cast into the minimization of (after multiplication by 4)

$$ $\displaystyle\int_{0}^{1}\mathbb{E}\big{[}\varepsilon_{t}|\nabla U_{t}(X^{\hat {b}}_{t})|^{2}-\varepsilon_{t}\Delta U_{t}(X^{\hat{b}}_{t})-\nabla\cdot b_{t}( X^{\hat{b}}_{t})\big{]}dt-\mathbb{E}[\log\rho^{\hat{b}}_{1}(X_{1})]]$ (127) $\displaystyle=\int_{0}^{1}\int_{\mathbb{R}^{d}}\big{[}\varepsilon_{t}|\nabla U _{t}(x)|^{2}-\varepsilon_{t}\Delta U_{t}(x)-\nabla\cdot\hat{b}_{t}(x)\big{]} \rho^{\hat{b}}_{t}(x)dxdt-\int_{\mathbb{R}^{d}}\log[\rho^{\hat{b}}_{1}(x)]\rho ^{\hat{b}}_{1}(x)dx$ $$

where $\rho^{\hat{b}}_{t}$ solves (117).

Let us check that the minimizer of (127) is $\hat{b}_{t}=b_{t}$, the solution to (24), so that we also have $\rho^{\hat{b}}_{t}=\rho^{b}_{t}=\rho_{t}$. To this end, notice that the minimization of (127)
can be performed with the method Lagrange multiplier, using the extended objective

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\big{[}\varepsilon_{t}|\nabla U_ {t}(x)|^{2}-\varepsilon_{t}\Delta U_{t}(x)-\nabla\cdot\hat{b}_{t}(x)\big{]} \rho^{\hat{b}}_{t}(x)dxdt-\int_{\mathbb{R}^{d}}\log[\rho^{\hat{b}}_{1}(x)]\rho ^{\hat{b}}_{1}(x)dx$ (128) $\displaystyle+\int_{0}^{1}\int_{\mathbb{R}^{d}}\lambda_{t}(x)\big{(}\partial_{ t}\rho^{\hat{b}}_{t}-\varepsilon_{t}\nabla\cdot(\nabla U_{t}\rho^{\hat{b}}_{t} +\nabla\rho^{\hat{b}}_{t})+\nabla\cdot(\hat{b}_{t}\rho^{\hat{b}}_{t})\big{)}dxdt$ $$

where $\lambda_{t}(x)$ is the Lagrange multiplier to be determined. Taking the first variation of this objective over $\lambda$, $\rho^{\hat{b}}$, and $\hat{b}$, we arrive at the Euler-Lagrange equations

$$ $\displaystyle 0$ $\displaystyle=\partial_{t}\rho^{\hat{b}}_{t}-\varepsilon_{t}\nabla\cdot(\nabla U _{t}\rho^{\hat{b}}_{t}+\nabla\rho^{\hat{b}}_{t})+\nabla\cdot(\hat{b}_{t}\rho^{ \hat{b}}_{t}),\qquad$ $\displaystyle\rho^{\hat{b}}_{0}=\rho_{0},$ (129) $\displaystyle 0$ $\displaystyle=\varepsilon_{t}|\nabla U_{t}|^{2}-\varepsilon_{t}\Delta U_{t}- \nabla\cdot\hat{b}_{t}$ $\displaystyle\quad-\partial_{t}\lambda_{t}+\varepsilon_{t}\nabla U_{t}\cdot \nabla\lambda_{t}-\varepsilon_{t}\Delta\lambda_{t}-\hat{b}_{t}\cdot\nabla \lambda_{t}$ $\displaystyle\lambda_{1}=\log\rho^{\hat{b}}_{1}$ $\displaystyle 0$ $\displaystyle=\nabla\rho^{\hat{b}}_{t}-\rho^{\hat{b}}_{t}\nabla\lambda_{t}$ $$

We can check that $\hat{b}_{t}(x)=b_{t}(x)$, $\rho^{\hat{b}}_{t}(x)=\rho_{t}(x)=e^{-U_{t}(x)+F_{t}}$, and $\lambda_{t}(x)=-U_{t}(x)+F_{t}$ is a solution: indeed this solves the first and the last equations in (129) and reduces the second to

$$ $\displaystyle 0$ $\displaystyle=\big{[}\varepsilon_{t}|\nabla U_{t}|^{2}-\varepsilon_{t}\Delta U _{t}-\nabla\cdot b_{t}\big{]}$ (130) $\displaystyle\quad+\partial_{t}U_{t}-\partial_{t}F_{t}-\varepsilon_{t}|\nabla U _{t}|^{2}+\Delta U_{t}+\nabla b_{t}\cdot\nabla U_{t}$ $\displaystyle=-\nabla\cdot b_{t}+\partial_{t}U_{t}-\partial_{t}F_{t}+\nabla b_ {t}\cdot\nabla U_{t}$ $$

which is satisfied since $b_{t}$ solves (24).

##### A modified, more practical objective.

The objective (125) is not practical because it involves $\log\rho_{1}^{\hat{b}}$, which is unknown. There is however a simple way to fix this, by adding a term in the Kullback-Leibler divergence (118)

$$ $\displaystyle\text{KL}^{\prime}(X^{\hat{b}}\|Y^{\hat{b}})$ $\displaystyle=\text{KL}(X^{\hat{b}}\|Y^{\hat{b}})+\frac{1}{4}\mathbb{E}[\log( \rho_{1}^{\hat{b}}(X^{\hat{b}}_{1})/\rho_{1}(X^{\hat{b}}_{1})])$ (131) $$

This additional term is proportional to the Kullback-Leibler divergence of $\rho_{1}^{\hat{b}}$ from the target PDF $\rho_{1}$. Proceding as before, and using $\rho_{1}(x)=e^{-U_{1}(x)+F_{1}}$, it is easy to see that we can express (131) as

$$ $\displaystyle\text{KL}^{\prime}(X^{\hat{b}}\|Y^{\hat{b}})$ $\displaystyle=\frac{1}{4}\int_{0}^{1}\mathbb{E}\big{[}\varepsilon_{t}|\nabla U _{t}(X^{\hat{b}}_{t})|^{2}-\varepsilon_{t}\Delta U_{t}(X^{\hat{b}}_{t})-\nabla \cdot b_{t}(X^{\hat{b}}_{t})\big{]}dt$ (132) $\displaystyle\quad+\frac{1}{4}\mathbb{E}[\log\rho_{0}(X_{0})]+\frac{1}{4} \mathbb{E}[U_{1}(X_{1})]-\frac{1}{4}F_{1}.$ $$

Neglecting the constants in $\hat{b}$ in this objective and multiplying it by 4 we arrive at

$$ $\displaystyle\int_{0}^{1}\mathbb{E}\big{[}\varepsilon_{t}|\nabla U_{t}(X^{\hat {b}}_{t})|^{2}-\varepsilon_{t}\Delta U_{t}(X^{\hat{b}}_{t})-\nabla\cdot b_{t}( X^{\hat{b}}_{t})\big{]}+\mathbb{E}[U_{1}(X_{1})]$ (133) $\displaystyle=\int_{0}^{1}\int_{\mathbb{R}^{d}}\big{[}\varepsilon_{t}|\nabla U _{t}(x)|^{2}-\varepsilon_{t}\Delta U_{t}(x)-\nabla\cdot b_{t}(x))\big{]}\rho^{ \hat{b}}_{t}(x)dx+\int_{\mathbb{R}^{d}}U_{1}(x)\rho_{1}^{\hat{b}}(x)dx$ $$

All the terms in this objective are now explicit.

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: 1. It is unclear why the paper writes that off-policy objectives do not need samples from the "target density" (line 63) since we do not seem to have samples from the target density for both on- and off-policy losses. In this context, "off-policy" typically means that the trajectories do not need to follow the ones from the *learned* model. 

2. What exactly is meant by "grid-free learning," i.e., why are PIS and NETS grid-free but DDS and CMCD are not? All these methods are based on simulating SDEs on some time grid. In particular, for off-policy methods, the time grid can basically be chosen arbitrarily.

3. It is unclear why DDS is only "partially" unbiased. For diffusion-based samplers, one can perform importance weighting in path-space, which is unbiased (up to bias introduced from self-normalization as also used in NETS).

4. What prevents other samplers based on diffusion models (such as DDS and DIS) from using "arbitrary step size and time-dependent diffusion after learning the model"? As already shown in https://arxiv.org/abs/2106.02808, there is a family of equivalent SDE that can be used for inference and which only requires knowledge of the score. Naturally, the approaches might overfit to a fixed time discretization. However, the time steps can be randomized during training, as is done in NETS.

5. How does the last step in (52) follow (since $\hat{b}$ is not at the optimum)?

---

**Suggestions:** While the largely algebraic calculations might yield easier proofs of existing results, the paper would benefit from providing further explanations and context. Two examples:

- The sentence "The effect of the last term at the right hand-side of this equation can be accounted for by using weights." only becomes clear after reading Remark 1. It might generally be good to add some more intuition to these algebraic calculations, e.g., (5) is just showing that $\rho_t$ is the invariant measure of the Langevin dynamics (with noise $\sqrt{2\varepsilon_t}$).

- One could also mention that Eq. (13) is just the continuity equation with predefined density evolution, and Eq. (14) are the FPEs of the corresponding family of SDEs with the same marginals (by standard manipulations of FPEs, e.g., Appendix G in https://arxiv.org/pdf/2106.02808). This also directly implies the statement of Prop. 2.

See weaknesses above for additional suggestions and questions.

---

**Typos:**

- $R^d$ on line 122

- $\partial_t \log F_t$ on line 140

- $R^{d+1}$ in line 160

- $U$ (without time index) in (10) and line 123

- 9 in line 183 should probably be (8)

- 46 in lines 272 and 275 should probably be (25)

- THe in line 852

- These in line 971

- There seem to be $\hat{cdot}$ missing over the $\phi$ in Section 2.6 and the corresponding proofs.

- closing $)$ missing in (18).

- it should be integration against $\hat{\rho}_t$ at several places in the proof of Prop. 5.
soundness: 2
strengths: The paper is generally well-structured and theoretically sound. In view of existing work (see weaknesses below), the paper provides the following contributions:
1. New (and arguably simpler) derivation of existing results on (controlled) annealed Langevin dynamics,
2. Novel objective based on action matching (AM),
3. Evaluation of additional tricks, such as usage of importance weights for the PINN objective, curriculum learning for the time interval, and combination with resampling.
summary: This work presents an approach to sample from densities specified up to a normalizing constant using controlled annealed Langevin dynamics. The approach is first motivated by Jarzynski's equality, which can be interpreted as a time-continuous variant of annealed importance sampling (AIS). Then, the paper extends this framework to *controlled* Langevin dynamics and proposes to learn the drift either by physics-informed neural networks (PINNs) or a version of action matching. The resulting method, called Non-Equilbrium Transport Sampler (NETS), is numerically evaluated on Gaussian mixtures, the Funnel distribution, and the simulation of a lattice field theory.
weaknesses: **Related works:** Large parts of the paper are already included in existing work:

- *Controlled Langevin dynamics:* While Vargas et al. (2024) are mentioned in the related works, it is important to note that the *main* proposition of the present paper, Proposition 3 (Nonequilibrium Transport Sampler (NETS)), is already presented in Vargas et al. (2024, Proposition 3.3). Using the data processing inequality, the latter result also seems to yield the bound on the KL divergence in Proposition 5. 

- *PINN objective:* The PINN objective in Section 2.5 is already stated in https://arxiv.org/abs/2301.07388 (Section 3, "The continuity loss"), https://arxiv.org/pdf/2405.06672 (Section 2.2), and https://arxiv.org/pdf/2407.07873 (Section 3.2, "Annealing"), where also on- ("uniform") and off-policy ("along trajectories") methods for $\hat{p}_t$ are explored. While both sampling variants do not require backpropagating through the SDE integrator, the latter *does* need to simulate the SDE (at least periodically if using a buffer). Thus, I would not say that NETS-PINN is "simulation-free even if used on-policy," as stated in line 344. 

Because of these observations, the present work only provides little *novel* contributions (see strengths above). In particular, the PINN objective seems to outperform the AM objective in the first two experiments, and curriculum learning is a relatively standard trick for time-dependent PINNs.

Moreover, many of the shortcomings mentioned in related works are either not completely accurate or have already been tackled. In particular, there seem to be several methods having checkmarks for all columns of the table on page 2. Some details (see also questions below):

1. There exist several off-policy losses for samplers based on dynamical transport, which, in particular, do not require backpropagating through the SDE integrator and lead to unbiased sampling methods; see (sub-)trajectory/detailed-balance losses for GFlowNets (e.g., https://arxiv.org/abs/2210.00580, https://arxiv.org/abs/2310.02679, https://arxiv.org/pdf/2402.05098) and log-variance (LV) losses (https://arxiv.org/abs/2307.01198; also used in Vargas et al. (2024)). In particular, these losses also fall into the category of "optimize-then-discretize".

2. I assume the LV loss is meant in the second part of the sentence "this objective either needs backpropagating through the SDE or has to be computed with a reference measure which may make it high variance". However, for the typical choice of using the on-policy drift, the reference measure is actually *reducing* variance (as compared to the KL divergence) (https://arxiv.org/abs/2307.01198, https://arxiv.org/abs/2005.05409).

3. The fact that the objectives of DDS (and DIS, see Berner et al., 2024) can also be viewed as SOC problems shows that methods based on SOC do *not* need to start with samples from a point mass.

---

**Experiments:** Important baselines and ablations seem to be missing:

1. CMCD: As argued above, CMCD has the same theoretical framework, and it can also be viewed as a BSDE (instead of PINN) loss, as shown in https://arxiv.org/pdf/2407.07873 (Prop. 3.1). However, NETS is currently not compared against CMCD. 

2. ODE importance weights: The advantage of using the SDE instead of the ODE during generation should be validated. Specifically, one could also use the learned vector field $\hat{b}$ to simulate the ODE (aka continuous-time normalizing flow) and compute the importance weights using the change-of-variables formula (as is done in some of the works mentioned above that also use the PINN objective).

3. Sections 4.3 and 4.4 do not provide any baselines except for AIS and only evaluate one of the proposed NETS versions. Moreover, no details on the AIS settings seem to be provided.

4. Since NETS seems to require a very high number of discretization steps (all experiments use >= 200 and some up to 1500-2000 steps), it would be interesting to see how well NETS is working for fewer steps (which is relevant in settings where target evaluations are expensive). Is the same number of steps also used during training?

5. More benchmarks (e.g., from Blessing et al. (2024)) are required to judge the performance of NETS against other state-of-the-art samplers, and it would be good to see an ablation of some components of NETS (e.g., the curriculum learning or the importance weights in the PINN objective).

When quoting baseline results from other works, it is crucial to follow the same evaluation protocol, in particular, the number of integrator steps and evaluation samples:

1. *Funnel:* NETS: 200 steps and 10000 samples. Baseline: 128 steps and 2000 samples. 
2. *Gaussian Mixture:* NETS: 250 steps and 2000 samples. Baseline: 100 steps (at least for PIS; it might also be that ESS is computed using an "exact likelihood from a continuous normalizing flow (CNF) model [trained] using optimal transport flow matching on samples from each model," which "keeps the number of function evaluations per integration around 100 in practice") and 1000 samples.

Please reevaluate your results using the same protocol. Finally, it would be interesting to also compare training costs, e.g., time and number of target (& gradient) evals, etc.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: see weaknesses.

Minor comments: typo in Line 140-142: $\partial_t \log F_t$ should be $\partial_t F_t$.
soundness: 3
strengths: - The idea of combining annealed importance sampling and additional transport is interesting and novel.
- The theoretical result of the KL control for the PINN objective is interesting.
- Considering both $\hat{b}_t$ and $\hat{F}_t$ simultaneously in the PINN objective is interesting and novel.
- This method contains correction for sampling errors if the additional transport is imperfectly learned.
summary: This paper proposes a sampling algorithm using annealed Langevin dynamics and additional learnable transport, which can be viewed as a variant of annealed importance sampling. The additional learnable transport can be used to reduce the high variance appeared in original annealed importance sampling. The paper proposes to learn the additional transport by minimizing two objective functions. Theoretically, they show the PINN objective controls the KL divergence between the model and its target. Last, they provide numerical examples outperforming several baseline state-of-the-art techniques and testify the scalability of the method.
weaknesses: -  In practice, using PINN objective need compute the divergence of a velocity field in $\mathbb{R}^d$, which can be extremely inefficient in high-dimensional case.
-  Several studies have investigated sampling algorithms by solving the velocity field through partial differential equations ([1], [2]). Therefore, a more detailed discussion of these related works would enhance the author’s contribution by clarifying connections and advancements in this area.
-  Although the PINN objective is off-policy, making it more computationally efficient, the performance of the algoirthm can be sensitive to the choice of $\hat{\rho}_t$. For example, it may lead to poor exploration when $\hat{\rho}_t$ is significantly far away from $\rho_t$.  A more detailed discussion on this trade-off would strengthen and complete the author’s analysis in this context.

References

1. Tian Y, Panda N, Lin Y T. Liouville Flow Importance Sampler[J]. arXiv preprint arXiv:2405.06672, 2024.
2. Fan M, Zhou R, Tian C, et al. Path-Guided Particle-based Sampling[C]//Forty-first International Conference on Machine Learning.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: - For the PINN loss in equation (25), we need to compute the divergence of $\hat{b}_t(x)$ to obtain the gradient. Then how do you efficiently compute it, especially in the high-dimensional case?
soundness: 3
strengths: - To the best of my knowledge, the proposed method is novel and firstly tackles the biasedness problem of AIS with an additional transport function. The weight-augmented SDE formulation is also inspiring for follow-up works.
- The authors derive two practical training methods for the transport function $b$. Especially, for the action matching objective, computing the Hessian matrix of $\phi_t(x)$ can be avoided by a smart trick on $A_t^b$.
- The experiments are illustrative and well-designed to demonstrate the superiority of NETS. I especially favor the illustration in Figure 1 that shows the effectiveness of the transport.
summary: This paper proposed NETS which targets sampling from unnormalized probability distributions.
NETS originates from annealed importance sampling, which is widely known to result in biased samples and relies on importance sampling for correction.
NETS views AIS as an SDE augmented with important weights and incorporates a learnable transport drift function $b(x)$ to correct this biasedness.
This idea is similar to the stochastic control.
The core ingredient of NETS, $b(x)$, can be learned by PINN or action matching, and in the first case, the learning is off-policy and does not require backpropagating the SDE.
The alternative perspective on AIS is very interesting and insightful, and improving it with an additional transport map seems novel and effective in practical applications.
weaknesses: - NETS is somewhat similar to stochastic control for me. I suggest the authors clarify the similarities and differences between them.
- There are some critical hyperparameters that can affect the performance of NETS, e.g., the $\varepsilon_t$ parameter and the stepsize of discretization. I would like to see how these parameters affect the convergence of the proposed method. Besides, it would be better for the authors to give some instructions on how to select these hyperparameters in practice.
- Although the authors target sampling tasks, an important application of NETS is to calculate expectations based on equation (21). Exploring or discussing this application would be interesting.
- For methods that require simulating the trajectories or computing divergence/hessian, computational cost is always a major concern. The authors should report the training time and memory consumption of NETS and other baselines, to fairly compare different methods.
- There are many typos in mathematical equations that can cause misunderstanding. I strongly suggest the authors carefully check their manuscript.
    - The expression of $\rho_t(x)$ in equation (9).
    - The reference to equation (9) in line 193 seems incorrect.
    - The expression of the RHS in equation (18).
    - The expression of $g_t$ in equation (40).
    - 'Them' in line 38; 'resut' in line 312; and so on.
- The authors give a guarantee of convergence in KL divergence. However, no KL divergence results have been compared or reported in the experiments.

---

confidence: 2
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. In Line 342, could you clarify how $(X^{\hat{b}}_t, A^{\hat{b}}_t)$ can be considered independent of $\hat{b}_t(x)$? A detailed explanation would be helpful.
1. Could you provide insight into how the method achieves multi-modal sampling without simulation? Without simulation, it seems challenging to capture information on unexplored modes.
soundness: 3
strengths: The paper presents a diffusion-based, simulation-free method for multi-modal sampling, which demonstrates significant potential in sampling efficiency and effectiveness. The method provides a powerful alternative to traditional sampling techniques and is applicable to complex, multi-modal distributions.
summary: The paper proposes a novel sampler for unnormalized probability distributions by augmenting the Stochastic Differential Equations (SDE) framework of Annealed Importance Sampling (AIS). The method introduces an additional drift term in the SDE, which enhances the effective sample size of the collected samples, thus reducing bias. The drift term can be estimated through either a Physics-Informed Neural Network (PINN) or Action matching objectives, which encourages the SDE to be nearly unbiased.
weaknesses: 1. Equation 46, frequently referenced in Section 2.5, would be more accessible if directly included in that section or if Equation 25 were referenced instead.
1. The computational cost for evaluating the objective from $t=0$ to $T$ seems a lot. An analysis of wall-clock time or time complexity compared to baseline methods would be beneficial.
1. The paper's comparisons with diffusion-based baseline samplers focus on low-dimensional examples (e.g., 2D Gaussian Mixture Model and 10D funnel distribution). However, high-dimensional comparisons are only performed as part of an ablation study. To show a practical scalability, comparisons in a high-dimensional setting would strengthen the argument for the method’s applicability beyond small-scale examples. I think the section 4.3 is showing a possibility to scale, not a powerful performance even in high dimensions.

Meta-review: This study proposes a novel sampling method inspired by insights from nonequilibrium physics to enable efficient sampling from unnormalized distributions. The method employs stochastic differential equations (SDEs) and introduces auxiliary dimensions to reformulate the process within a new framework of annealed importance sampling (AIS). Ultimately, an additional drift term is incorporated to mitigate the impact of weight variance in AIS, resulting in a new sampling approach. Considering the quality of the reviews, I (AC) had carefully examined the manuscript. While the research direction is interesting and important, significant concerns remain regarding the presentation, mathematical rigor, and the numerous typographical errors, which cast doubts on the manuscript's accuracy. For these reasons, I recommend rejection at this time. Below are the specific points of concern:
- The authors frequently rely on the relationship between SDEs and the PDEs (FP equations) without providing proofs or referencing existing literature. While the kinetic FP equation seems to be corresponding their approach, many of the PDEs appear to be proposed newly in this paper and thus, a formal discussion about the  derivation is necessary. Additionally, there is no discussion of the functional space in which these PDEs are defined. For instance, it is unclear whether Equation (36) does not diverge. Since these equations are used to ensure the mathematical rigor of the sampling process, such discussions are essential.

- In the proof of Proposition 1 in Section 4.1, the authors use integration by parts (e.g., at the end of Equation (37)). Since this holds in under moderate assumptions, the authors should include a discussion about it, such as the support of the distribution and the behavior of the density function at its boundaries to justify this step.

- The paper contains many typos, as pointed out by all reviewers, significantly impacting readability. Even after the discussion, in Equation (19), the left-hand side's $\nabla U_t$ should be $\nabla U_t \cdot \hat{b}_t$. Similar copy-paste errors are found in multiple other parts of the manuscript.
- In Proposition 2, $T$ is not defined.
- The proof of Equation (9) is claimed to be in Appendix 4.1, but no such proof exists. While Equation (41) in the appendix appears to be related, it does not constitute a complete proof.

Decision: Reject