Neural Sampling from Boltzmann Densities: Fisher-Rao Curves in the Wasserstein Geometry

Weakness: Sorry for the high mathematical technicality which are necessary here. We have already tried to put most of the theoretical analysis in the appendix. In order to make the learning part understandable without going into the mathematical details, we started adding pseudocode in Appendix C.

We added the numerical examples.

Questions:

We added the numerical example you proposed in Example 5. Thanks! As expected, the linear interpolation fails, while the learned and the gradient flow interpolation properly approximate the target distribution.

In addition we give a numerical example in 16 dimensions.

For the 2d Gaussian mixture reported in our numerical section (Table 1) the gradient flow method performs better than the learned one.

First order transition: There do seem to be some similarities to the physical phenomenon you described. In our phenomenon described in Figure 1, the distribution $\rho_t$ changes drastically to another distribution $\rho_{t + \delta}$ for late times $t \approx 1$, given that the overlap of the two distributions is very small, and the squared velocity $\|v_t\|_{L_2}^2$ explodes. This might resemble the discontinuity of the first derivative of the energy happening at first-order phase transitions.

I thank the author for his reply, and I would have raised my grade to 7 if that would be possible...

Concerning the novelty of the theoretical results, please see our response to Reviewer m4s7 and their acknowledgement.

Weakness 1: We added certain notation to make it easier to read.

Weakness 2: Thanks for suggesting the three references. We included them in the discussion about related work. The first reference about stochastic interpolants investigates the setting, where samples from the source and target distribution are available, which is not the case in our setting. But it is interesting, since it is investigating different paths in probability space. In the RDMC paper, the score at point $x_t$ is estimated using a unadjusted Langevin algorithm. This score is then used to obtain $x_{t+h}$. They do not learn a vector field $v_t$, but have to rerun the entire algorithm if they want to draw new samples. Similarly, in the paper SLIPS they use MCMC in order to approximate the drift of an SDE and use that approximation for sampling. This is an interesting approach, but again the goal here is not to learn a vector field that can be used for sampling.

Weakness 3: We added numerical results, so far up to dimension 16, and an example proposed by reviewer CXu4.

Minor: We added the assumption $\rho_1 = \rho_d$.

We added the notation in (2) and (5).

We heuristically tried different domains, such that they are not too large which could lead to numerical instability and not too small to create a bias. We did not optimize over this choice.

We added Remark 3 on sampling.

We corrected the typos, thanks for pointing out!

Weakness: We mainly use an adaptive Runge-Kutta-5 stepsize solver for solving the ODE. It is true that solving this ODE is crucial and there are three main obstructions to overcome. First, we can empirically only learn the vector field on a bounded domain and thus outside of this domain the approximation might fail. If we then sample from a Gaussian and obtain a starting point outside this domain, the particle might be pushed in a wrong direction, which sometimes happens in practice.

The second issue concerns the straightness of the ODE paths for the ground truth vector field, which affects how good the ODE solvers work. Both issues are interesting problems to analyze, but they are out of the scope of this paper. The third issue comes up when the vector field becomes irregular meaning, e.g. the $L^2$ norm of the vector field explodes as in the example ''A teleportation issue with the linear interpolation'' at p. 6. Such a vector field is then hard to approximate. We partly solve this problem by choosing the gradient flow interpolation that is in general better behaved as the linear interpolation of the energies. We added numerical results, so far up to dimension 16, and an example proposed by reviewer CXu4.

Question: Concerning the ODE solver, see our answer above. Please note that the Many Well experiment was already conducted in 8 dimensions in the original submission and we scaled it up to 16 dimensions now.

Thanks for acknowledging our nontrivial theoretical results. We think that having analytical evidence that the linear interpolation leads to exploding velocities (instead of just numerical observations) is important.

Weakness 1: Thanks for pointing out the reference [Xu]. They also tackle the task of sampling from an unnormalized target density by approximating a Wasserstein gradient flow of the KL. However, their curve is the gradient flow of $KL(\cdot|\rho_D)$, i.e. the flow from $\rho_Z$ to the data, while we use the gradient flow of $KL(\cdot,\rho_Z)$ starting from $\rho_D$. While they are using one network per gradient flow approximation step, we use one network for the whole flow.

The paper [Boffi] considers a SDE $d x_t=b_t(x_t)dt + \nabla\cdot D_t(x_t)dt + \sqrt{2}\sigma_t dW_t$, where they assume that $b_t$ and $D_t$ are known. This in in contrast to our approach where we do not know $D_t$, but it rather depends on the curve $\rho_t$. Hence their approach is not applicable in our case. We added both reference in our related works section.

Weakness 2: We added numerical results, so far up to dimension 16, and an example proposed by reviewer CXu4.

Questions:

• 1. $C_{c}^{\infty}$ denotes the set of infinitely-often differentiable functions with compact support. We added a corresponding definition in the updated version. You are also right that in line 118, we should already introduce the notions of $s_t, \alpha_t$ etc. We will update this in our revised version.
• 2. Yes, as mentioned the linear and the learned interpolation methods are known. Also, the approach leading to lines 261-267 for solving the PDE (16) was already proposed (see line 257) by [Mate&Fleuret]. As you stated, [Mate&Fleuret] sample $x$ from the trajectory of the learned vector field which we explained in lines 309-315. We moved this comment further up right after Equation (17) to make the lines 261-267 clearer. In Section 4, we experimented using both, the uniform and trajectory based sampling strategy, for learning. More details are given at the beginning of Section 4.
• 3. Given a Boltzmann target distribution and the dynamics governed by the Fokker-Planck equation it is not a-priori clear that the intermediate densities $\rho_t$ themselves satisfy $\rho_t \propto e^{-f_t}$ for some $f_t$, in particular it is not clear that $\rho_t$ is strictly $>0$. This is however necessary, since $f_t$ is the quantity we aim to approximate. We added further explanation to make this point clearer.

Thanks, we corrected the two misprints!

Many thanks for your careful and detailed review, in particular for pointing to additional references! We will incorporate them in the revised version. Below we address your concerns.

In our work we fundamentally consider two different situations. First, we assume that we are given a prescribed path $\rho_t$ with tractable $\log \rho_t$ (up to a constant). Then we consider learning a vector field $v_t$ which governs the evolution of $\rho_t$ via the continuity equation. This is also the focus of the works [1,3] and leads naturally to the theoretical question for which curves of densities $\rho_t$ a vector field $v_t$ with suitable properties exists. More precisely this leads to the question of existence of a solution to the Poisson equation $\nabla\cdot(\rho_t\nabla \phi_t) =\rho_t(\alpha(t,\cdot)-\mathbb E_{\rho_t}[\alpha])$. We interpret the Poisson equation as combining the Wasserstein geometry on the left and the Fisher Rao geometry on the right. There are three very important differences between the treatment of [1] and our theory.

• They assume that $\rho_t\in C^\infty([0,1],\mathbb R^n)$ and $\alpha\in C^\infty([0,1]\times \mathbb R^n)$ and a Poincare inequality for each time $\rho_t$. We need much weaker conditions: we assume that $t\mapsto\rho_t$ is weakly continuous, $\rho=\int \rho_t dt$ fulfills a (partial in $x$) Poincare inequality and $\alpha\in L^2(\rho)$. In particular, we allow single time points $t$, where the Poincare inequality might fail for $\rho_t$.
• Our solution $\phi$ is an element of $H^1(\boldsymbol{\rho})\subset L^2([0,1]\times \mathbb{R}^d; \boldsymbol{\rho})$ and thus automatically measurable on the whole time domain $t \in [0,1]$. This makes it possible to show that $\nabla \phi_t(x)$ is a Borel vector field with regularity in time which implies absolutely continuity of the flow $\rho_t$ in the Wasserstein geometry. In contrast, [1] only considers a solution only for every single $t\in[0,1]$ without any regularity in $t$. This is a result already shown in [laug] Theorem 2.2 (cited by us) and motivated our investigation into a framework across the time $t$.
• Furthermore, we give precise mild assumption under which the partial Poincare inequality is satisfied. In particular, we cover the case of the linear interpolation. As a result we obtain that these curves are absolutely continuous (AC) in the Wasserstein geometry.

Hence our work gives a rigorous theoretical foundation which can be applied to the setting in [1]. For linear interpolation of energies which is of interest for many authors, we do not only show the existence of a suitable vector field, but also provide a nontrivial analytical example, where the corresponding vector field becomes exponentially irregular, see line 268 and Appendix A4.

Second, motivated by our above findings, we consider approaches to learn more regular curves without tractable $\log \rho_t$. The method proposed by [5] adds a learned regularizer $t(1-t) \psi_t$ to the linear interpolation. This is exactly what we call learned interpolation and used for comparisons (see your question). We instead consider the case of learning a regular AC curve directly by considering a Wasserstein gradient flow in the Fisher Rao geometry. When finishing this paper, we became aware of [2], we thank the reviewer for pointing out [4] which we were not aware of. In [2,4] the authors make similar considerations but coming from an SDE based stochastic optimal control perspective without considering the relation to Fisher Rao and Wasserstein geometry. While we simulate the corresponding ODE instead of the SDE and use a different sampling strategy for the PDE collocation points, eventually, up to time reparametrization, we end up with the same loss (see your question).

The advantage of using the gradient flow interpolation is that we know the resulting curve whose favorable properties were extensively studied in the literature. This opens the door for further analysis. In contrast, when learning the curve with learned regularizer [5], there is no knowledge of the resulting curve.

[laug]: R. S. Laugesen, P. G. Mehta, S. P. Meyn, and M. Raginsky. Poisson’s equation in nonlinear filtering. arXiv preprint arXiv:1412.5845, 2014.