Importance Corrected Neural JKO Sampling

Thank you very much for the review. We carefully revised the paper and inserted the requested changes.

Regarding Weaknesses

1. and 2. During writing the paper, we also experimented with discrete-time architectures for normalizing flows, mainly with coupling flows similar to the Glow model (Kingma and Dhariwal 2018). Compared to the continuous normalizing flows we made the following observations, which also align with the paper of Onken et al. 2021 ("OT-flow: Fast and accurate continuous normalizing flows via optimal transport"):

   • Since the regularization of the velocity field in the continuous normalizing flows (CNFs) leads to geodesic interpolations, the particle trajectories become straight paths. Hence, adaptive solvers only require very few steps to solve the involved ODE. Therefore, training and sampling with OT-regularized CNFs is much faster than with standard CNFs, this is discussed in detail by Onken et al. 2021. In particular, we observed that the training does not speed up significantly with the coupling flows, in high demensions, when large subnetworks are required, discrete-time architectures can be even slower than the ODE formulation.

   • Moreover, we observed in our experiments that the expressiveness of discrete-time architectures scales worse to high dimensions than the expressiveness of continuous normalizing flows.

   • On the other hand, we agree that discrete-time normalizing flows can improve the evaluation time and accuracy due to the exact density evaluations. However, we can also achieve sufficient accuracy for the density evaluation with trace estimators in the CNFs to meet the requirements for the rejection steps.

We added a discussion on other architectures (Remark 19) and point to it at the end of Section 3.2. Moreover, also considering the comments of the other reviewers, we added some more details on the computational cost of OT-regularized CNFs and trace estimators (Appendix F.2) and point to these remarks in the limitations section.

3. We added a comparison with two kinds of normalizing flows with the loss function from Marzouk et al. 2016 in Appendix F.3. First, we use a continuous normalizing flow architecture (since this is what we have used for our method). Second, (also related to point 1. and 2.) we used an explicit coupling architecture based on triangular blocks. We can observe that this works mostly fine for the unimodal examples (good results for funnel and mustache), but often mode collapses for multimodal distributions. We added Figure 12 for an experimental illustration of the prescribed situation.

4. Thanks for the comment. We agree and rewrote the related work section such that we now clearly indicate which references consider sampling and which consider generative modeling.

5. The Wasserstein distance has several drawbacks as evaluation metric. Most important, its sample complexity depends exponentially on the dimension such that the empirical version of the Wasserstein distance differs significantly from continuous version. Despite that it is not independent of our model (which is based on the Wasserstein geometry) and suffers from heavy computation times for large number of samples (cubic scaling). Nevertheless, we added a table evaluating the empirical Wasserstein distance between generated and ground-truth samples in the appendix (Table 3). The conclusions are similar as for the energy distance.

Regarding Questions

1. This was a typo, which we have now corrected. Thanks for spotting it!

2. You are absolutely right. We corrected the error. Many thanks!

3. This comes from the fact that we are adding up a $\lambda$-convex function and $1/\tau$ times the Wasserstein distance (which is $1$-convex). Hence we obtain a function which is $(\lambda+\frac1\tau)$-convex, so in particular strongly convex for $\tau>-\lambda$. A proof of this statement can be found in Lemma 9.2.7 of Ambrosio et al. 2005. We added the statement before Section 3.1. Additionally, we added a Section on the connection between multimodal target distributions, non-convexity of the KL divergence and mode collapse in the Appendix F.1. The new appendix F.3 shows in numerical examples that normalizing flows mode-collapse on the multimodal examples, for which our neural JKO IC works.

I would like to thank the author(s) for the detailed response, which answers most of my previous questions. After reading the revised manuscript, now I have one more question that puzzles me. In the theory, one needs the step size $\tau$ to be sufficiently small, but at the end of Section 4, the author(s) show that an exponentially increasing step size scheduling $\tau_{k+1}=4\tau_k$ is used. Is there any contradiction on this aspect?

Thank you very much for the review. We carefully revised the paper. Please find below our answers to your comments and questions.

Contributions

We want to highlight the following points:

• To the best of our knowledge our importance-based rejection steps are the first resampling/rejection steps which allow it explicitly to access of the density throughout their application (see Proposition 8). This property is crucial for the iteration of several importance steps, in particular their concatination.

• In contrast to classical MCMC steps (most Metropolis Hastings implementations, Langevin sampling, HMC, etc.), the importance-based rejection steps act non-local, i.e., the outcome of the rejection step does not depend on the input sample but on the whole input distribution. This is a very important property, which enables our algorithm to reweight the mass between disconnected modes, which is impossible for classical MCMC steps relying on local adjustments.

• In addition to the rejection steps, we prove in Theorem 6 that the velocity fields of the JKO scheme converge to the velocity field of the Wasserstein gradient flow. In the context of the approximation with normalizing flows this is an important stability results, because the velocity fields are the object of interest that we aim to approximate.

• Regarding the two papers, you have reference: The reweighting by Noe et al. considers a different setting, where both the (unnormalized) target density and some samples from the target distribution are given. In our setting (without samples from the target distribution), the reweighting is not applicable. In contrast to the paper of Gabrié et al., our neural JKO IC scheme returns a generative model, which can easily generate additional samples and evaluate the density of our approximation. Nevertheless, we agree that, from a high-level viewpoint, these papers consider the same problem as our paper with different tools. Therefore, we added them to the related work section.

Rectified Flow / Flow Matching

Rectified flows and flow matching consider the generative modeling setting. That is, they already assume given samples from the target distribution (which is not the case in our setting). Since the training is based on the construction of explicit couplings between training and latent samples, it is not applicable to sampling (or minimizing our Wasserstein-regularized KL functional). Therefore, we do not see a close relation to our loss function from Section 3.2.

As a side note: There is a preprint (appeared less than a month before ICLR submission deadline), which extends flow matching to sampling by iteratively sampling from the current model and creating the couplings this approximation (https://www.arxiv.org/2408.16249). We have now added this preprint to the related works part. However, it is not clear if (and if yes how) this approach could be applied in our setting, i.e. for solving the steps of the neural JKO scheme with Wasserstein regularization.

We added flow matching as an example of a generative model to the related work section and mention its sampling variant.

Adjoint formulation of CNFs

We now discuss some computational aspects of CNFs in Appendix F.2. In particular, we explain that the regularization of the velocity fields leads to straight paths such that the involved ODE can be solved by very few steps, see also the paper by Onken et al. 2021 for a detailed discussion.

Presentation of Section 4

We extended Section 4 with additional explanations. In particular, we added a summary of the importance corrected neural JKO scheme in Algorithm 1, point directly to the summary of the rejection step (Algorithm 5) and discuss more in detail how we build the final model.

Avoiding Mode Collapse

We added an related discussion in Appendix F.3, where we apply standard NFs (one coupling NFs and continuous NF) to our examples. In Figure 12 we plot the corresponding results and compare them with our neural JKO IC results. That standard (C)NFs mode collapse, while our neural JKO IC scheme covers all modes. Additionally, we added in Appendix F.1 a discussion on the relation between non-convex loss functions, local minima and mode collapse for multimodal target distributions. Before Section 3.1 we added a comment that the neural JKO scheme leads to a convex loss objective (as function on the space of probability measures), such that local minima cannot appear.

We would like to thank the reviewer for the evaluation of our paper.

Differences to prior work

We agree with the reviewer that the derivation of our importance-based rejection steps is one key contributions of our paper. To the best of our knowledge our importance-based rejection steps are the first resampling/rejection steps which allow it to access the density of the arising model (see Prop 8). This is fundamental for applying them iteratively. Additionally, we want to highlight the following two points regarding Section 3:

• We derive a convergence result of the velocity fields from the JKO scheme towards the velocity field of the Wasserstein gradient flows, which took quite some effort to prove. In the context of the approximation with normalizing flows this is an important stability results, because the velocity fields are the object which we aim to approximate.

• The paper of Xu et al. 2024 consider a different setting than our paper, even though they derive a similar simulation of the JKO scheme via CNFs. Since they are given samples from the target distribution (as opposed to our setting given its unnormalized density function), their target functional is given by $\mathcal F(\mu)=KL(\mu,\nu)$, where $\nu$ is a standard normal distribution. Due to the strong log-concavity of $\nu$, this is a strongly convex functional in the Wasserstein space simplifying the solution of these subproblems significantly. In particular, Xu et al. 2024 do not have to deal with common problems from the sampling problem like slow convergence of the underlying flow due to multimodal or narrow target distributions leading to poor constants in related Poincaré or log-Sobolev inequalities.

We have rewritten the related work part where we now clearly indicate which models consider the sampling which consider generative modeling. For the exposition provided, we feel that the backgrounds on the Wasserstein distance and its dynamical formulation and time rescaling are required for presenting Thm 6 properly. Therefore, we decided not to relegate more backgrounds to the appendix.

Answer to Questions:

1. We would like to thank the reviewer for this input. We extended the paragraph about Neural JKO Sampling with Importance Correction and added a summary of the proposed scheme in Alg 1. Additional, information about a suggested realization of layer choices is given in the final part of section 4. An algorithm for the application of a rejection step was already included (Alg 5). However, we would like to point out that the design of the rejection steps and their analysis (Prop 8) belong to the main contributions of our paper and cannot be found in the literature (see Rem 7 for the relation to classical rejection sampling). For the introductionary paragraph on annealed importance sampling, we followed your suggestion to shorten it.

2. Please note that there are two different kinds of discretization. Regarding the discretization of the ODEs from the CNFs, we added Appendix F.2 with more details. Regarding the discretization (the $\tau_k$) from the JKO scheme: While, it is very important to choose the $\tau$ small enough in the beginning, the Wasserstein gradient flow slows down fastly, which suggest to increase $\tau$ over time. One possibility for choosing $\tau$ could be to choose $\tau$ adaptively based on the Wasserstein distance between $\mu_k$ and $\mu_{k+1}$ (Xu et al. 2024 have done something like that). However, in our numerics we found that the simpler scheme of choosing $\tau_{k+1}=4\tau_k$ worked fine, such we sticked to it. While the chosen $\tau_k$ were already stated in the implementation details (Appendix E.4), we have now added a comment about that at the end of Sec 4 (and mention the adaptive schedule by Xu et al. 2024).

3. In each rejection step, we have to evaluate (requires no training and no backpropagation!) the whole model again for only a subset of the samples (not all of them!). However, since we do not need to retrain the models and do not require any gradient computations in these resampling steps, the training and evaluation times remain moderate (they are already included in Tab 5). In addition, Rem 11 and the discussion on limitations comment on the runtime effects of the rejection steps.

4. We included the additional references in the related work section of the paper. However, we would like to emphasize that the first one appeared after the submission deadline of ICLR. The second one is already included in the related work and we compared our numerical results with its follow up paper (CRAFT). The last two are conceptually similar to DDS (included in our numerics) and only use importance sampling at final time (not as intermediate steps). We added a comment in the related work section.

5. The detailed numerical setup is specified in Appendix E. If there is some specific point missing, we are happy to add it.

Thank you very much for the review.

Answer to Questions: Relation to SMC

Thank you very much for this question. There are some relations to SMC (alternate the use of importance weights with local adjustments), but there are also some major differnces. In particular we see the following advantages of our importance corrected JKO sampling over SMC:

• After the first step, one cannot compute the density of the distribution of SMC any more. Instead SMC relies on replacing the density by a quotiont of a forward and backward Markov kernel. The quality of this approximation heavily relies on how exactly the backward kernel indeed approximates the inverse of the forward kernel. It is not directly clear how errors in the approximation propagate and add up. In contrast, we can compute the density of our approximation explicitly (up to the Radamacher trace estimation in the CNF).

• Once, the SMC algorithm was running, generating additional samples is not possible in SMC. If additional samples are required, one has to redo the whole SMC procedure. In contrast, we obtain a generative model out of our algorithm, where we can easily generate additional samples after intermediate trainings steps of the intermediate measures.

• Due to the SMC-resampling step the generated samples of SMC are not independent. In constrast, our samples are independent.

We added some of these explanations to the section of related work.

Discussion on mentioned Weaknesses: Computational cost of CNFs

We added in Appendix F.2 a discussion of the computational limitations of CNFs which extends the limitations discussion in the conclusion section 6.

Some synopsis:

For solving the ODEs and differentiate its solutions, we use the package torchdiffeq (Ref.: Chen, 2018). In particular, this package does not backpropagate through the steps from the forward solver. Instead, it overwrites the backward pass of the ODE by sovlving an adjoint ODE. This avoids tracing within the forward pass and keeps the memory consumption low. The largest model (for the 1600 dimensional LGCP example) requires about 11 GB memory on a single GPU while still using a rather large batch size of 500. Hence, memory consumption is not the limiting factor in our model. We added related required GPU memory to table 5 in the implementation details section E.4 (note that we did not optimize our implementation wrt to memory, since this never was an issue throughout our experiments).

In terms of computation time, the OT regularizations improves the ODE solutions significantly, since the minimum of the problems is a geodesic in the Wasserstein space. In particular, the sample trajectories through the ODE are straight paths such that adaptive ODE solvers only require very few steps to solve the ODE. These advantages of OT regularization are discussed in detail in the paper of Onken et al. 2021 ("OT-flow: Fast and accurate continuous normalizing flows via optimal transport"). We added this comment also to Section 2.2.

On the other side, we agree that the evaluation of the divergence of the vector field remains a computational drawback, which is required since we have to evaluate the density of the model. However, this mainly effects the evaluation phase since during training a low-precision estimator (Hutchinson with one slice) is enough, as long as it is unbiased. Nevertheless, for the evaluation, we need a higher precision in order to be able to perform the importance-based rejection steps, which we realize by taking more slices (5 realizations of a Radamacher random vector) in the Hutchinson estimator. But we do not need to differentiate its computation at evaluation time, such that we can run it without gradient tracing.

We added these discussions in Appendix F.2 and point to them in the limitations section.

Thank you for following up in the discussion. After reading your additional comments, we think that we misinterpreted your original comment regarding SMC. We agree that the motivation for our rejection steps is rather (sequential) importance sampling than annealed importance sampling. As already written before, we also agree that, from an abstract viewpoint, the importance-based rejection scheme shares some conceptional motivations with SMC.

In order to address this comment in the manuscript, we rename the "Annealed Importance Sampling" paragraph in the beginning of Section 4 to "Importance Sampling" and reformulated it (see new version below). Moreover, we add a new remark (see below) in the Neural JKO Sampling with Importance Correction paragraph of Section 4 to outline the relation to the resampling procedure in SMC.

We would like to point out, that the resampling procedure from SMC cannot be applied in our setting, since it generates multiple samples at the same position. Since (despite the rejection steps) we solely use deterministig transformations, they would remain at the same position such that the effective number of samples would decrease over time.

Regarding the de-correlation in SMC: Thanks for pointing this out. We wanted to express that in contrast to SMC our update step generates iid samples by construction. We reformulate the last sentence before the contributions paragraph in the introduction accordingly:

However, the corresponding importance sampling step generates non-iid samples, that de-correlate over time by construction.

Reformulated version of the Importance Sampling paragraph:

Importance Sampling    As a remedy, many sampling algorithms from the literature are based on importance weights, see, e.g., sequential Monte Carlo samplers (Del Moral et al., 2006) or annealed importance sampling (Neal, 2001). That is, we assign to each generated sample $x_i$ a weight $w_i=\frac{q(x_i)}{p(x_i)}$, where $p$ is some proposal density and $q$ is the density of the target distribution $\nu$. Then, for any $\nu$-integrable function $f\colon\mathbb{R}^d\to\mathbb{R}$ it holds that $\sum_{i=1}^N w_i f(x_i)$ is an unbiased estimator of $\int_{\mathbb{R}^d} f(x) d\nu(x)$. Note that importance sampling is very sensitive with respect to the proposal $p$ which needs to be designed carefully and problem adapted.

Rejection Steps    Inspired by importance sampling [...].

New remark regarding relation to SMC:

Remark (Relation to SMC).    While the steps of the importance-corrected neural JKO scheme are completely different to those in sequential Monte Carlo (SMC), SMC and neural JKO IC share the motivation of using importance-based resampling throughout the generation process. As a key difference to SMC, our neural JKO IC scheme successively builds a generative model, which allows to efficiently regenerate iid samples at any intermediate step and to evaluate the corresponding density explicitly.