NETS: A Non-Equilibrium Transport Sampler

We thank the reviewer for their positive assessment of our results and for their questions that have helped us improve our presentation.

Weakness 1: Eq. (25) vs Eq. (46): Thank you for pointing this out. It is an unfortunate misprint: Proposition 4 with Eq. (25) is the same as Proposition 7 with Eq. (46). To help the reader, we re-state the propositions in the appendix, but somehow we made a mistake in restating Proposition 4, and this error percolated back through the text where references to Eq. (46) should have been references to Eq. (15). We have corrected this.

Weakness 2: computational cost for evaluating the objective: The PINN objective must eventually be evaluated all the way through $T=1$ (recall that we use $T$ for annealing only). There is indeed a cost in generating the data and computing the time integral, but this is similar to what one also needs to do in every competing methods based on estimating an additional drift, since this drift must be learned for all $t\in[0,1]$. In the revised vesion we will include the wall-clock times of all the approaches we used.

Weakness 2: Scalability with dimension: We have obtained new numerical results and more thorough benchmarks of the method. In particular we have added a harder test with benchmarks to compare to using the 50-dimensional Mixture of Student-T distribution studied in the Beyond Elbos paper. We have also verified the performance of the PINN on the $\phi^4$ model (in addition to the existing action matching demonstration), that demonstrate the scalability of our approach as the dimensionality increases (this is 400 dimensions).

We hope you find these clarifications and revisions useful. Thanks again for the valuable feedback.

We thank the reviewer for their questions and suggestions, which we have tried to address next. If you find these clarifications and revisions useful, we would appreciate it if you improved your score.

Weakness 1: comparison with methods based on SOC: The problem of sampling a target distribution can indeed be tackled via the solution of a stochastic optimal control (SOC) problem -- this is what is proposed e.g. in Ref. [1] (which we cite) where the authors introduce a method to estimate a Föllmer process, i.e. a Schrödinger bridge between a point mass distribution and the target. The main difference is that, in the SOC formulation, the control = drift in the SDE must be learned from scratch, which is typically challenging since the reference process (i.e. the SDE without the control) has solutions whose law at final time is far from the target distribution in general. In contrast, our approach can be viewed as a guided search in which we predefine the path in distribution space between the base and the target. The resulting process may not be optimal in the sense of SOC, but the learning of the drift is facilitated in our formulation. In addition it allows us to choose any suitable base suitable distribution, whereas the SOC formulation used in Ref. [1] leads to a Föllmer process whose base distribution must be a point mass. This too limits the flexibility/scalability of this approach compared to ours.

[1] Zhang, Q. & Chen, Y. (2022). Path Integral Sampler: a stochastic control approach for sampling. arxiv preprint arXiv:2111.15141.

Weakness 2: choice of hyperparameters: the main metaparameters to choose in the training procedure are: (i) the annealing schedule controlled by $T$ in the loss, and the choice of number of discretization points. Both can be adjusted on-the-fly by monitoring the ESS, and ensuring that it does not deteriorates as the annealing proceeds. This is what we did in all of our experiments: if the annealing is too agressive, the training has a hard time to converge. We explain this point better in the revised version.

Regarding the diffusion coefficient $\epsilon_t$, it is important to stress again that it can be adjusted post-training. Theoretically, the $\epsilon_t \rightarrow \infty$ limit is perfect sampling. We demonstrate in Figure 2 that this limit can be more practically realized when you include transport than without. We also discuss this point in more details in the revised version.

Weakness 3: calculation of expectations: The main purpose of NETS is to calculate expectations as well as estimate the partition function/Bayes factor $Z_1 = \int e^{-U_1(x)} dx$ (which is an important quantity in practice). In the paper we focus on $Z_1$ (as is also done in many of the sampling papers cited), mostly because this factor allows us to benchmark the accuracy/efficiency of our method. For the $\phi^4$ example we compare expectations via analysis of the magnetization given in the figure.

Weakness 4: computational cost: Training times are drastically influenced by how good a coder you are :), so we do not think they are a great scientific metric. But for comparison it took about 23 minutes for the 40-mode GMM to train. Compared to the ESS and $W_2$ metrics quoted from iDEM, this is fast. But again, they may have just coded things differently, and we do not feel that it is fair for us to claim something here. Please see the discussion in the appendix on the Hutchinson trick, e.g. for dealing with the cost of the divergence.

Weakness 5: typos: Thank you for pointing these out. We have corrected them and have gone through a careful read of our paper to correct a few more.

Question: divergence calculation: In practice this divergence can be efficiently calculated using Hutchinson's trace estimator with antithetic variates (similar to what is used to estimate the divergence of the score in score based diffusion models). More specifically, we can use

where $0<\delta \ll1$ is an adjustable parameter and $\eta \sim N(0,\text{Id})$. This estimator can be made unbiased by using two independent copies of $\eta$ for the terms involving the square of $\nabla \cdot b_t(x)$ in the loss. We tested it on some of the examples and didn’t seem to make much of a difference in final performance as compared direct calculation. If the reviewer would like we could include some of these numbers in a camera-ready version.

We thank the reviewer for their questions and suggestions. Our answers follow - please let us know if there is any other question we could address in order for you to raise your score.

Weakness 1: need to compute the divergence in the PINN loss: In practice this divergence can be efficiently calculated using Hutchinson's trace estimator with antithetic variates (similar to what is used to estimate the divergence of the score in score based diffusion models). More specifically, we can use

where $0<\delta \ll1$ is an adjustable parameter and $\eta \sim N(0,\text{Id})$. This estimator can be made unbiased by using two independent copies of $\eta$ for the terms involving the square of $\nabla \cdot b_t(x)$ in the loss. We have added a discussion of this in the appendix in the implementation section. We have tested that this approach works in practice and it seems to work fine. If the reviewer likes, we can add experimental infor for this to a potential camera-ready version.

Weakness 2: more detailed discussion of related works: The results in the paper by Tian et al. are indeed very relevant to ours, as already mentioned in our original submission. With respect to this work, our results bring several important novelties:

On the theoretical side, we introduce several losses to estimate the drift, and show that the PINN loss controls the variance of the variance of the weights as well as the KL divergence. We also show that the method can be applied to Langevin equation with time-dependent diffusion coefficient $\epsilon_t$, even though the additional drift is the same for all $\epsilon_t \ge 0$ - in contrast the paper by Tian et al. focus on the probability flow ODE obtained when $\epsilon_t=0$). Finally, in the revised version, we show that the method can be generalized to the overdamped non-equilibrium dynamics (i.e, adding inertia) and can also be used to sample distributions that vary across more parameters than the single scalar $t$.

On the numerical side, we show that adjusting the diffusion coefficient (and making it non-zero in general) is key to improve performance, both at learning and sampling stages. For the latter, this adjustment can be done post-training. We also demonstrate the performance of our method in a larger class of examples.

The method in the paper by Fan et al. also aims at constructing a probability flow ODE ($\epsilon_t=0$). As far as we could tell, the results in this paper are not fundamentally different from those in the paper by Tian et al.

We will add a more detailed discussion of these works in the revised version, to highlight the difference with our proposed approach.

Weakness 3: sensitivity of the PINN loss to the choice of $\hat{\rho}_t$. Indeed, in practice the choice of this reference density matters, and that is why we always aim to use $\hat{\rho}_t = \rho_t$. The key point, however, is that, since the PINN objective can be used off-policy, its performance are less affected by the errors we make in the sampling (i.e. by the variance of the weights used to sample wrt $\rho_t$) than, say, those of the AM loss (which must be used on-policy). We will add a more careful discussion in the revised version to stress this point.

We thank the reviewer for the detailed and itemized review that has helped us greatly to contextualize our contributions and improve our presentation. Below we give an itemized response to each of your theoretical and experimental remarks. We hope that this answer will allow us to make our contributions clear to you and that you will consider raising your score as a result.

Related works

Controlled Langevin dynamics (Vargas et al. 2024):

We have corresponded extensively with the authors of this work, both to properly benchmark their approach as well as highlight its differences with our method. Please note that the essential machinery of Proposition 3.3 in Vargas et al. has been known for 20 years in the statistical physics community (Jarzynski 1999, Vaikuntanathan & Jarzynski 2008), and neither us nor Vargas et al. are claiming to have discovered this equality. However, we both provide different derivations of it that make it more interpretable for sampling. Vargas et al. prove this result through the use of Girsanov, and here we provide a proof of it through simple manipulations of the the Fokker-Planck equation and other PDEs. This uniquely allows us to:

• Directly relate a PINN loss to the Jarzynski weighting factors.
• Show that we can better approach $\epsilon_t \to \infty$ limit in practice. Note that for our setup, perfect sampling is achieved in this limit, whether the learned transport is perfect or not. While this limit cannot be reached in practice without transport (as it would require taking astronomically large value of $\epsilon_t$ in general), we show that, with some learned transport added, even moderate values of $\epsilon_t$ can improve the sampling dramatically. This feature can be exploited after training as an explicit knob for tuning performance vs cost. This has not been recognized in other ML sampling literature nor in Vargas et al.
• Incorporate SMC resampling into the generation process which allows for increased performance (see both tables), which is also not realized in Vargas et al.
• Directly control the KL-divergence with the PINN loss.
• Directly characterize the minimizer of the action matching loss, which is also not known from previous work on it.

We would like to stress that in our exchanges with authors of Vargas et al. they have explicitly differentiated our work from theirs with these points and suggested we make these delineations clear in this rebuttal. We are glad that you can discern a connection to our KL-bound through more analysis to a proposition in their work, but would like to stress that it is a non-trivial result in general, and a wider audience who is less familiar with this topic would not intuit this result as it is not proven anywhere except in this work.

PINN objective: The PINN objective in Section 2.5 is already stated in...

We appreciate your help in finding these related works and citations, and we have now incorporated them directly into the related works of the text. Please let us know if you think we can improve these citations further. One of these works seems to be essentially co-incident as it only appeared on the arxiv 2 months before this submission.

We would like to stress that this PINN loss was not known to be intrinsically connected to the Jaryznski equation and annealed langevin dynamics, which we elucidate here. This fits in naturally with physical interpretation of this equality -- it is a direct control on the variance of the Jarzynski weights (the dissipation) in the process connecting $\rho_0$ to $\rho_1$. We have added an additional statement to this in the text.

Thus, I would not say that NETS-PINN is "simulation-free even if used on-policy," as stated in line 344

This is a misalignment of the meaning of simulation-free on our part, which we say to mean "not having to backpropagate through the solution of the SDE" as would be necessary, e.g. with a KL-type loss. We have changed this accordingly in the text.

Moreover, many of the shortcomings mentioned in related works ...

We have entirely removed the table because we realize some of the column headings drifted from our original meaning. For example, when we said unbiased, we meant very specifically that the importance weights along the dynamical transport even correct discretization errors in the SDE. We have included an appendix deriving this result (Sec 5.2). Thanks for pointing us to a few other citations, which we have now included in the related works.

We thank the reviewer for their reply. Below we provide answers to your questions and address your follow-ups. Please note that we cannot edit the PDF at this point until the potential camera-ready version, but we reference simple changes we would make to it in such a case:

Question 1 ($\epsilon_t \to \infty$ limit): The limit as $\epsilon\to\infty$ is well-understood: in this limit, the walkers evolve infinitely fast compared to the potential, and as a result are always in quasi-equilibrium with it, so that the weights become unnecessary -- this observation can be made precise using standard limit theorems for diffusions evolving on multiple time scales (see e.g. the book by Pavliotis & Stuart referenced as [1] below), and it is at the basis of thermodynamic integration strategies. Of course, the integration of the SDE must be done on the fast time scale $\tau = \epsilon t$, i.e. the computational cost increases with $\epsilon$. The question therefore becomes whether we can reach the limit $\epsilon\to\infty$ in practice (i.e. whether we can work with values of $\epsilon$ large enough that we are effectively in this limit). Without transport, this is not possible in general -- the limit is only attained for values of $\epsilon$ that are astronomically large (e.g. because of mestastability effects). However our numerical results show that the situation can change dramatically with some transport added, in which case increasing $\epsilon$ do increase the ESS significantly (see Fig. 2 for illustration).

Question 2 (ODE weights): The weight computation that we use is a generalization of the one obtained with the probability flow ODE (when $\epsilon_t=0$). This can be seen as follows:

Suppose that we evolve the walkers using the probability flow ODE

with an imperfect drift $\hat b_t(x)$. In this case the importance weights to use are

where $\hat \rho_t(x)$ is the solution to the PDE

This equation can be solved by the method of characteristics, which shows that

i.e.

Therefore the weights are

On the other hand, if $\epsilon_t =0$, we have (using $\dot X_t = \hat b_t(X_t)$)

so that the equation for the weights reduces to

and we deduce that

Since $\mathbb{E}[e^{A_t}] = e^{-F_t+F_0}$, this shows that

[1] Pavliotis & Stuart, Multiscale Methods, 2008.