Importance Corrected Neural JKO Sampling

We would like to thank the reviewer for the detailed and valuable feedback.

Convexity Assumption

There appears to be a misunderstanding regarding the assumptions for the theoretical part: We do not assume that the density is log concave. Instead, Assumption 3.1 assumes that the functional is $\lambda$-convex along generalized geodesics. For the KL divergence, this corresponds to the assumption that the target energy $-\log(q)$ is $\lambda$-convex for some $\lambda\in\mathbb R$. We stress that this explicitly includes the case of negative $\lambda$. This assumption is much weaker than log-concavity and is automatically fulfilled, if $-\log(q)$ is smooth (plus some asymptotics for $\\|x\\|\to\infty$). More intuitively, the condition can be rephrased as: $-\log(q)$ is $\lambda$-convex if and only if there exists some (possibly negative) $\lambda\in\mathbb R$ such that $-\log(q(x))-\frac{\lambda}{2}\\|x\\|^2$ is convex.

Given that this misconception appeared in more than one review, we will include this discussion in the final version and highlight that $\lambda$-convexity is usually not an issue as long as the target distribution is smooth.

Relation to other neural JKO Schemes

We stress that we consider the theoretical results (Cor 3.2, Thm 3.3) to be the main novelty of Section 3. As outlined in the beginning of Section 3, the implementation (Sec 3.2) is close to previous papers, specifically to (Vidal et al., 2023) and (Xu et al., 2024), who also use the dynamic formulation of the Wasserstein distance. But these papers consider the generative modeling setting instead of sampling. We are not aware of a reference, which adapts the dynamic formulation for sampling, which requires some smaller technical differences (e.g., other $\mathcal F$, we start with latent distribution, while Vidal et al., Xu et al. start with the data distribution). However, we stress that the main contributions of our paper are:

• The theoretical analysis (Cor 3.2, Thm 3.3) of neural JKO schemes,
• proposing importance-based rejection/resampling steps which can maintain the density of the generated samples (Section 4) and preserves independence of the generated samples,
• combining both to an importance corrected neural JKO sampler, which achieves state of the art results.

In the introduction and in the beginning of Section 3 we already write that similar approximations of the JKO steps exist in the literature. We will add a reminder in the beginning of Section 3.2 (and specifically point to Vidal et al., 2023, Xu et al., 2024 for using the dynamic formulation).

Literature on Rejection Schemes with Generative Models

Indeed there exist some approaches to combine rejection steps with generative models in the literature. Many of these approaches are based on sequential Monte Carlo techniques (e.g., CRAFT which we used as comparison, but also Arbel et al., Phillips et al.). These methods implement a reweighting step by first approximating importance weights and then sampling from the empirical distribution defined by this weights. However, for SMC-based methods the analytic evaluation of the arising density is usally not possible. Moreover, after SMC reweighting steps, there exist with a high probability several samples at the exact same position. Finally, the generated samples are not exactly independent.

The paper of (Gabrie et al.) proposes to iteratively train a normalizing flow for sampling by running a Langevin process, adding Metropolis steps with the current normalizing flow as proposal, and retraining the normalizing flow with the updated samples. In contrast to our paper, the rejection steps are not part of the model, but are rather used for training a normalizing flow. In addition, they can't include these steps in the model, because this would require to evaluate the density of these steps, which is not possible for the Metropolis algorithm. Since several papers from the literature [1] found that training normalizing flows to approximate probabilities with disconnected modes (like GMMs) or non-Gaussian tails (like funnel, mustache) are difficult, this might limit the expressiveness of the model.

We will extend the discussion on these methods in the final version.

[1] https://arxiv.org/abs/1907.04481, https://arxiv.org/abs/2009.02994, https://arxiv.org/abs/2206.14476

Thank you very much for your very detailed and thoughtful review. Please find the answers to your questions and comments below. Additionally, we will correct the typos and grammar errors.

Questions

1. Our choice of the schedule is based on the following heuristic: Since the rejection layers are (on a long term) more expensive than the CNF layers, we start with using only CNF layers until the progress made by them becomes small. This can be detected without knowing the target distribution, because we have (approximately) access to the Wasserstein distance between input and output of the JKO layer via the norm of the velocity field. Afterwards, we alternate one CNF layer and three rejection steps. Numerically, the CNF layer only does minor local adjustments from this point, while the main work is then done by the rejection steps.

2. If the step size is too large, we might miss some modes of the target distribution (which matches with the theoretical consideration that we loose convexity (with respect to the Wasserstein metric) of the loss in this case). If the step size is small, then the mapping learned by the CNFs is close to the identity and does not have a large effect. While the model is sensitive with respect to the first case, the second case only increases training and sampling time of the model (since more CNF steps are required), but does not effect much the quality of the result. A simple tuning heuristic for the step size can again be established based on the (approximated) Wasserstein distance: We start with a small step size and increase it as long as the Wasserstein distance between input and output of the next CNF step is smaller than a certain threshold. While we have not automatized this procedure so far, we are quite sure that such an adaptive choice of the step size can be established.

3. The curse of dimensionality indeed not specific to normal distribution, but to any proposal which is not close enough to the target. So in fact the use of tailored/CNF proposals can help to reduce or even completely avoid the curse of dimensionality. We will highlight this further.

Number of Samples

Our main intention behind the large number of samples is the following: Since the samples generated by our model are independent, our main quantity of interest is how much the distribution of a generated samples differs from the target distribution (in MMD, Wasserstein etc.). Of course we cannot measure this quantity directly, but taking a finite number of samples from both distributions. Using a small number of samples introduces an error in this estimator, which is much more related to the properties (sample complexity) of the evaluation metric than with the approximation quality with the model. Therefore, from our viewpoint, the number of samples should be chosen as large as possible.

While the $\log(Z)$ estimation can be viewed in the same way (just with the KL divergence), we agree that this can also viewed differently. If we are not interested in the KL divergence between generated and ground truth distribution (which was our viewpoint so far) but in the $\log(Z)$ estimate itself, we agree with you that reporting the errors for different numbers of samples could be useful. We will include such a study in the final version. In our first experiments with $5000$ samples the greater picture is about the same as for $50000$ samples.

Proof Theorem 4.2 (i)

Yes, you are right, the expression $P(X=x)$ is meant in a weak sense (where we integrate against all measurable sets $A$) and that this part is written a bit sloppy. We will write it down more formally for the final version. However, neither the proof nor the result change.

Visualization etc.

Thank you for the suggestions. We will add a flow chart-like diagram describing the sampling process to the paper. For the numerical part, we think that box plots often make it harder to assess the details, even though we agree that they give a better first impression. We are happy to add a box plot for visualizing the numbers, but we prefer to keep the tables as well (if the space constraints allow it, we will keep both in the main part of the paper).

We would like to thank the reviewer for the detailed evaluation of our paper. Please find our comments below.

Literature

The process of research is an active field and therefore many interesting contributions are published frequently in particular in this modern field. However, we want to kindly point out, that in accordance to the ICML guidelines, comparisons to very recent preprints, which are not published yet (like [3,4]) or published a week before the submission deadline (like [2]) or haven't even been available as a preprint until the submission deadline (like [4]) cannot be expected. We do our best to provide timely comparisons and have now even added a comparison to [2] (details below), but we can't compare to methods which appeared on arxiv shortly before or even after the ICML submission deadline. Of course, we are happy to include all papers mentioned by the reviewer in the "related work" section (where [1,3] are already cited).

Computational Cost

We respectfully disagree with the reviewers opinion that our method would be computational infeasible. Specifically we would like to stress the following points:

• The evaluation of $\mathrm{trace}(\nabla v_t^k(x))$ is estimated by a Hutchinson trace estimator ($\mathrm{trace}(A)=\mathbb E[z^TAz]$) for $z$ with zero mean and identity covariance matrix.
• During the training of the CNF a single sample of $z$ is enough in the Hutchinson estimator to approximate the trace (since we only need an unbiased estimate of the loss for training the CNF). For the evaluation we take 5 samples of $z$ for each time step. We find that this already leads to a small variance of the resulting densities, while the cost remains comparably low (see also the answer to Question 5 of Reviewer Eq2H).
• When training $v_\theta^{k+1}$, we do not draw in each training step "a fresh batch from $\mu_k$". Instead, we maintain during training time a dataset of 50000 samples of $\mu_k$. During training of $v_\theta^{k+1}$ we then draw batches from this dataset. After the training of $v_\theta^{k+1}$ is completed, we generate a dataset of $\mu_{k+1}$ by applying the CNF onto the samples from $\mu_k$. We add this detail to the numerical details.
• During evaluation we don't need to draw a large batch. During test time the samples don't interact with each other (see also our reply to the part "Batch size" in the paragraph "Questions" below)
• Please note that we already discussed the computational aspects of CNFs in Appendix F.2 and list the resulting training and evaluation times including the required GPU memory in Table 5 in the appendix. We can see that they remain moderate for all examples considered in the paper.

Nevertheless, as already discussed in the limitations paragraph, lowering the computational cost is part of our current and future research. In particular the use of supervised trained generative models as intermediate surrogates for several steps in our model is future work and goes beyond the scope of the paper.

Other Benchmark Problems and Comparisons

We believe that our test problems are standard among the community and respectfully disagree with the reviewers opinion that "no sharp benchmarks" were used. However, upon the reviewers request, we now ran our method on the GMM40 problem in $d=50$. For the sake of availability of comparisons, we adapt the same setting as Table 3 in [2] and also evaluate the Sinkhorn distance. We can see in the table below that the neural JKO IC outperforms DDS, CRAFT and SCLD on this example.

Questions

1. Stability: In the case that hyperparameters are chosen inappropriately, the method might miss some modes of the target distribution, but we never observed divergence issues. Once the hyperparameters are chosen appropriately, it consistently produces the same results.
2. Ablations: The paper already contains an ablation study with respect to the number of steps in Figure 10, where we plot the error measures over the number of steps. Since our model is trained iteratively, the plots show how the model would behave if we stop the training earlier and we can see that the error measures saturate.
3. Batch size: We would like to clarify that at test time the sampling procedure is independent of the batch size. The only parameter in the rejection step depending on several samples, in particular depending on more than one sample is $\mathbb E[\alpha(X_k)]$ (it is estimated within the training phase and fixed afterwards). For the evaluation the samples do not interact with each other and are consequently independent.

Many thanks for the detailed and thoughtful review. Please find our answers below. For the final version, we will additionally correct the typos, extend the literature part (e.g. with SVGD), improve the visual accessibility based on your comments and add the definitions of lsc/coercive/$\lambda$-convexity in the appendix.

On the $\lambda$-Convexity

We stress that for the theoretical results we only require that the negative log-density of the target is $\lambda$-convex for some $\lambda\in\mathbb R$ which explicitly includes negative $\lambda$. This assumption is very weak and automatically fulfilled when $q$ is smooth enough (plus some asymptotics for $\\|x\\|\to\infty$). In particular, the theoretical results are also applicable for target densities which are not log-concave. We will add this explanation after Assumption 3.1 in the paper.

Questions

1. We stress that in Section 3 we consider the theoretical results (Cor 3.2, Thm 3.3) to be the main novelty of our paper. The actual neural JKO steps are indeed very similar to approaches in the literature, particularly, to the papers of (Vidal et al., 2023; Xu et al., 2024) which also rely on the dynamic formulation of the Wasserstein distance. Since these papers consider the generative modeling setting there are some small technical differences (e.g., we start with the latent distribution while Vidal et al., Xu et al. start with the data distribution; moreover they only consider the Gaussian as target $g$, which makes the whole objective convex regardless of $\tau$, but is not useful for the sampling application). In the final version, we will clarify the relation to these approaches more detailed in the beginning of Section 3.2, where we derive the neural JKO step.
2. The statements of Cor 3.2 and Thm 3.3 are not directly related to the rejection steps in Section 4. However, in order to be able to combine the importance-based rejection step in an iterative manner, we need the following requirements on our sampling model: evaluate the density, sample independently, start at an arbitrary latent distribution, avoid mode collapse. Since not many generative models (or sampling methods) fulfill these properties at the same time, we focused on the neural JKO scheme. We will clarify this in the introduction.
3. It is not surprising that the neural JKO scheme and Langevin sampling produce very similar results, since both approximate the same Wasserstein gradient flow with respect to the KL divergence in the limit. The main difference is that we can evaluate the density of the distribution generated by the neural JKO scheme, while we can't do that for the distribution generated by Langevin sampling. Thus, we can combine the neural JKO scheme with our importance-based rejection resampling steps, while we can't do that with the Langevin steps. We briefly outlined this relation in the introduction (l 47--50 right side), but will add a reminder to that in the numerical results. Regarding ULA vs. neural JKO: When running our experiments, we also ran a plain ULA sampling (without Metropolis correction). The results are mostly similar to the MALA results despite the Metropolis correction helps a little bit to fill out the narrow tails of the funnel/mustache distribution. Therefore, we omitted ULA in the paper.
4. The most important hyperparameters are the initial step size $\tau$ and the size of the network architecture, which require some tuning. If those parameters are chosen inappropriately, the model might miss some modes (which matches the theory, since we loose any convexity of the JKO steps in this case). We outline a possible tuning strategy in the answer to Reviewer BUrR (Question 2). But most hyperparameters from the neural JKO steps are quite standard choices (the velocities are dense 3-layer neural networks, we use the torchdiffeq library with rather standard parameters, the training parameters like batch size, learning rate etc. are also quite standard choices for CNFs).
5. We think that we can resolve this confusion: When computing the Jacobian trace for 5 new Rademacher vectors in each time-discretization step of the ODE. Since the ODE trajectories are almost straight and therefore the Jacobian only slightly varies over the different time steps. Consequently, the effective number of considered Rademacher vectors is 5 times the number of time discretization steps which is (depending on the example) in a order of magnitude of 20 to 100. Therefore the errors mostly cancel out over the whole solution of the ODE. Indeed, we observed that the variance of the estimator over the whole solution of the ODE is already very small for 5 Rademacher vectors per time step and taking more does not have much of an effect. Taking less Rademacher vectors causes a bias in the importance weights and lowers the quality of the results. We will add this explanation to the final version of the paper.