NETS: A Non-equilibrium Transport Sampler

We thank the reviewer for their careful reading of our paper and positive feedback. We are glad that you found the work well-written, theoretically sound, and that the numerical experiments demonstrate convincing evidence for our method's effectiveness. Below we address your comments and suggestions:

Experimental Designs Or Analyses: The potential $U_t$ design is a feature we can exploit, similar to FAB or CMCD. In our drive link (https://tinyurl.com/netsicml) we compare with CMCD using identical interpolation, revealing CMCD struggles with this path: Using the same number of sampling steps as us, CMCD achieves only 15% ESS. To perform these experiments we worked with the authors to best benchmark their codebase.

Relation To Broader Scientific Literature and Essential References Not Discussed: Thanks for pointing us to the papers by Shi, Zhekun, et al.; Sun, Jingtong, et al.; and Chen, Junhua, et al. We are happy to add a citation to these works, but we wish to note that they fall within the ICML policy of concurrent works: "https://icml.cc/Conferences/2025/ReviewerInstructions".

We would also like to stress that, while variants of the PINN loss have appeared in the literature, our work makes the following new contributions:

• Novel connection between the PINN loss and the Jarzynski weighting factors and annealed langevin dynamics. This result fits in naturally with physical interpretation of Jarzynski equality -- the PINN loss control on the variance of the Jarzynski weights (the dissipation) in the process connecting $\rho_0$ to $\rho_1$.
• 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.
• Show that the PINN loss directly controls the KL-divergence between the sampled and target distribution.
• Directly characterize the minimizer of the action matching loss, which is also not known from previous work on it.

Our method connects with CMCD, but we derive results through Fokker-Planck manipulations rather than Girsanov theorem.

We'll add references to Richter and Berner's foundational work.

Other comments and suggestions: Could you please specify which equations are stated without context? We will happily try to clarify in editing.

Questions for authors:

-Difference between the proposed PINN objective and other related PINN based sampling methods: See our reply above about the key new insights that our work provides. -Insight of top row, left of Figure 3: The top row (left) are example configurations of samples as we approach the critical point, and the bottom row are samples past the critical point, where they are fully magnetized. They illustrate the juxtaposition. -Benefit of the proposed method compared to CMCD-LV: We have included in the "additional experiments" section below results on CMCD-LV, which we worked with the CMCD authors to set up. Our results show that, unlike our approach, CMCD-LV exhibits instabilities if the number of sampling steps is too small in training: while the CMCD-LV loss is "off-policy" like our PINN loss, it still requires a trajectory on which to perform the optimization, and the generation of such a $n_{step}$ trajectory requires O($n_{step}$ $\times$ network size) memory. If the trajectory generation is performed over too few steps the CMCD LV loss becomes unstable. After discussing with CMCD authors, this is an issue generally when $n_{step} < 256$.

Additional experiments: In the anonymous drive link https://tinyurl.com/netsicml:

• We provide additional comparison between NETS and CMCD on the mean-interpolating time-dependent potential, as asked by another reviewer. We see that NETS performs well regardless of number of discretization steps, but the log-variance loss struggles for small steps and only performs well when $n_{step} = 256$ or more. We worked with the CMCD authors to implement their code and otherwise use their hyperparameters.
• We have also set up a test on the Lennard Jones potential, and have included preliminary results. We are continuing to refine these experiments and will incorporate them in the final version.

We thank the reviewer for the positive feedback. We are glad that you found the contribution novel and the theory sound and thorough. Below we address your comments and suggestions and supply more information on experimental results:

Additional experiments: In the anonymous drive link https://tinyurl.com/netsicml:

• We provide additional comparison between NETS and CMCD on the mean-interpolating time-dependent potential, as asked by another reviewer. We see that NETS performs well regardless of number of discretization steps, but the log-variance loss struggles for small steps and only perform well when $n_{step} = 256$ or more. We worked with the CMCD authors to implement their code and otherwise use their hyperparameters.
• We have also set up a test on the Lennard Jones potential, and have included preliminary results. We are continuing to refine these experiments and will incorporate them in the final version.

We thank the reviewer for their valuable feedback, and we are happy to hear you found the work theoretically sound and novel. Below, we try to address all your questions, and provide some info on new additional experiments. We itemize our response according to the headings that appear in you review:

Methods And Evaluation Criteria:

• estimating Z_1: Yes, it's straightforward to construct an unbiased estimate if we know $Z_0$: $Z_1 = Z_0 \cdot \mathbb{E}[e^{A_1}]$. This follows from equation (21) at $t=1$, which can beestimated empirically over trajectories from the coupled SDE/ODE (19,20).
• complexity of experiments: We have included preliminary results on a Lennard-Jones system in an "Additional Experiments" section discussed below. Note that the $\phi^4$ models studied here are arguably more complex than the LJ-13 system: these $\phi^4$ models are either 256 or 400 dimensional problems, much bigger even than the LJ55 system, and are known as a proxy for one of the hardest sampling problems in research -- studying the strong force in lattice field theory. In particular, the $\phi^4$ models, like these lattice systems, suffer from what is known as critical slowing down in which MCMC algorithmic efficiency dramatically decays as one approaches the critical parameters. It has also been a target for neural samplers for some time [2].
• Dropped Baselines for iDEM, CMCD: We only included baselines where we could either quote authors' results or work with them to verify implementation. We did the latter with CMCD on Funnel and GMM, but couldn't for MoS distribution.
• SMC/HMC: SMC can be incorporated within our approach (this is what we refer to as "resampling"), but on its own isn't efficient enough for our sampling tasks. We tried HMC-type sampling (Inertial NETs, Appendix D.1) but saw little difference.
• CMCD interpolation path: We thank the reviewer for pointing this out. We have included in the "additional experiments" section below results on CMCD on the path interpolating the means, which we worked with the CMCD authors to set up. With this path as well, CMCD-LV exhibits instabilities if the number of sampling steps is too small in training: while the CMCD-LV loss is "off-policy", it still requires a trajectory on which to perform the optimization, and the generation of such a $n_step$ trajectory requires O($n_{step}$ $\times$ network size) memory. If the trajectory generation is performed over too few steps the CMCD LV loss becomes unstable. After discussing with CMCD authors, this is an issue generally when $n_{step} < 256$. This will become apparent in the "additional experiments" section below.

Theoretical Claims:

• eq (6) and (7): SMC can be incorporated within our approach (as "resampling"), but on its own isn't efficient enough for our sampling tasks. We tried HMC-type sampling (Inertial NETs, Appendix D.1) but saw little difference.
• derivation of (7) and (10,11): We'll add derivations in the appendix, as these aren't commonplace in ML literature.
• Why learn $\partial_t F$?: The loss function is a PINN loss which means it is minimized at 0 (i.e., when the physics equation is solved). $\partial_t F_t$ is a free parameter in this equation and must be fit if one wants to learn off policy. We do indeed use that $\hat b_t \rho_t$ vanish at infinity, which is a consequence of the conservation of probability that prohibits the existence of a probability current at infinity.
• uniqueness of PINN minimizer: The minimizer is unique if and only if the velocity field is learned in gradient form $\nabla \hat \phi_t = \hat b_t$.
• integration by parts: This is again a consequence of conservation of probability requiring no probability flux at infinity.

Experimental Designs:

• PIS not included: The PIS authors confirmed to us that they used $\sigma^2 = 3$ (not $\sigma^2=9$ as stated in their paper). This makes the problem much easier, so we omitted this experiment.

Additional experiments:

In the anonymous drive link https://tinyurl.com/netsicml:

• We provide a NETS vs CMCD comparison on mean-interpolating potentials which shows NETS performs well regardless of discretization steps, while log-variance struggles for small steps, improving only when $n_{step} ≥ 256$.
• We've included preliminary Lennard Jones results and will incorporate refinements in the final version.

We'll make your suggested expository edits when allowed to edit the paper. Thank you for your insights - if we've clarified everything, we'd appreciate an increase in your rating.

[1] Alpha Collaboration, "Critical Slowing Down and Error Analysis in Lattice QCD simulations," Nuclear Physics B, 2010.

[2] Albergo, Kanwar, Shanahan, "Flow-based Generative Models for MCMC in lattice field theory," Physical Review D, 2019.

We thank the reviewer for their valuable feedback, and we are happy to hear you found the work theoretically sound and novel. Below, we try to address all your questions, and provide some info on new additional experiments. We itemize our response according to the headings that appear in you review:

References

• Thanks for pointing us to this paper. The SMC that they propose is also something that we propose in this work (when we may choose to do 'resampling'). We are happy to add a citation to this work, but we wish to note that this falls within the ICML policy of concurrent works: "https://icml.cc/Conferences/2025/ReviewerInstructions".

Weaknesses

• Proposition 2.4 in general is not new, though this derivation is. Please note that it is also not new in Vargas et al. The essential machinery of this proposition has been known for 20 years in the statistical physics community (Jarzynski 1999, Vaikuntanathan & Jarzynski 2008), and neither us nor Vargas et al. are claiming to have discovered this equality. However, we both provide different derivations of it that make it more interpretable for sampling. Vargas et al. prove this result through the use of Girsanov, and here we provide a proof of it through simple manipulations of the the Fokker-Planck equation and other PDEs.

Other Comments Or Suggestions:

• We are happy to improve the readability of equation 6 when we can edit the text if the paper is accepted.
• thanks for catching the typo in Eq (40).

Questions:

• backprop thru sde and numerical instability: The first drawback cannot be entirely addressed by the CMCD log-variance loss. It is still a loss function that must use samples along a trajectory as input, and the generation of such a trajectory of lengthn $n_{step}$ requires $O$($n_{step}$ $\times$ network size) memory. By numerical instability, we mean that the CMCD LV loss becomes unstable if the trajectory generation is performed over too few steps. After discussing with CMCD authors, this is an issue generally when $n_{step} < 256$. This will become apparent in the "additional experiments" section below. Let us also stress that, in contrsat, our PINN loss can be evaluated pointwise in space and time.
• Optimize then discretize: Thanks for the observation. We imagine most other diffusion models could do this as well, but the main obstacle would be whether or not their method allows for an adaptable diffusion coefficient post-traing (like our method). If not, increasing the diffusion would require decreasing the time-step for trajectory generation during training.
• PINN vs KL: The benefits of the off-policy nature of the PINN loss manifest themselves in two ways: 1) The PINN loss is valid with respect to any sampling distribution), and 2) if you want to use simulated samples from your model, you do not need to backpropagate through their generation. The KL loss explicitly needs samples from the model distribution, and requires a backprop through the solve. As mentioned earlier, the log-var loss still needs to generate trajectories but has a smaller memory footprint. We did test the log-variance more, too. See: "additional experiments".
• Paper by Chen et al: Thanks again for pointing us to this. We are happy to reference it, though again per ICML policy it is out of the scope of this submission to benchmark it.
• $\epsilon \rightarrow \infty$ limit: Yes, that is exactly the motivation of the additional drift function. For our setup, perfect sampling is achieved in this limit, whether the learned transport is perfect or not. While this limit cannot be reached in practice without transport (as it would require taking astronomically large value of $\epsilon_t$ in general), we show that, with some learned transport added, even moderate values of $\epsilon_t$ can improve the sampling dramatically. This feature can be exploited after training as an explicit knob for tuning performance vs cost. This has not been recognized in other ML sampling literature.

Additional experiments: In the anonymous drive link https://tinyurl.com/netsicml :

• We provide additional comparison between NETS and CMCD on the mean-interpolating time-dependent potential, as asked by another reviewer. We see that NETS performs well regardless of number of discretization steps, but the log-variance struggles for small steps. Note that LV begins to perform better when $n_{step} = 256$. We worked with the CMCD authors to implement their code and otherwise use their hyperparameters. This is the numerical instability we referred to earlier.
• We have also set up a test on the Lennard Jones potential, and have included preliminary results. We are continuing to refine these experiments and will incorporate them in the final version.

Please let us know if you have any other questions, and thanks again for your insights. If we have clarified everything for you, we would greatly appreciate an increase in your rating.