Importance Corrected Neural JKO Sampling

Authors: Johannes Hertrich, Robert Gruhlke
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: Sampling, Wasserstein Gradient Flows, Normalizing Flows, Rejection Sampling

Abstract:
In order to sample from an unnormalized probability density function, we propose to combine continuous normalizing flows (CNFs) with rejection-resampling steps based on importance weights. We relate the iterative training of CNFs with regularized velocity fields to a JKO scheme and prove convergence of the involved velocity fields to the velocity field of the Wasserstein gradient flow (WGF). The alternation of local flow steps and non-local rejection-resampling steps allows to overcome local minima or slow convergence of the WGF for multimodal distributions. Since the proposal of the rejection step is generated by the model itself, they do not suffer from common drawbacks of classical rejection schemes. The arising model can be trained iteratively, reduces the reverse Kullback-Leibler (KL) loss function in each step, allows to generate iid samples and moreover allows for evaluations of the generated underlying density. Numerical examples show that our method yields accurate results on various test distributions including high-dimensional multimodal targets and outperforms the state of the art in almost all cases significantly.

Decision: reject
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730625916460

thread_type:
official_review

<thread_object>
creator: reviewer_wcjw

<content>
summary:
In this article, the author(s) consider the problem of sampling from unnormalized density functions, and the basic idea is to implicitly construct a diffeomorphism to transport a simple distribution to the target distribution, based on continuous normalizing flows (CNFs). To further improve the quality of sampling, a rejection-resampling step is also used.

strengths:
Sampling is a fundamental task in machine learning, and this article tackles this problem by utilizing some recent advances in deep generative modeling. The use of JKO scheme to iteratively optimizing the model parameters is also interesting.

weaknesses:
1. Overall I like the idea of this article, but one of my majors concerns on this article is the motivation of using a continuous normalizing flow (CNF) rather than the classical discrete-time normalizing flow. If my understanding is correct, CNF uses neural networks to model the velocity field $v_\theta(\cdot)$, and then the transport map $\mathcal{T}\_\theta(\cdot)$ is implicitly defined by $v_\theta(\cdot)$. However, when minimizing the (reverse) KL loss function, it would be easier to work on $\mathcal{T}\_\theta(\cdot)$, which is what classical normalizing flows do. Also, the sampling of classical normalizing flows should be faster than CNFs, since they do not need to solve an ODE.

2. Moreover, for classical normalizing flows, the evaluation of the (log-)density function is *exact*, which is important for rejection sampling. CNF needs to approximate the Jacobian term in evaluating the density.

3. There are some notable existing works on sampling using diffeomorphisms/normalizing flows, for example [1][2][3]. When comparing deep-learning-based sampling algorithms for the experiments, I would suggest at least adding one of them, since they are closer to the method used in this article.

4. In the last paragraph of the introduction section, I think it is helpful to distinguish *modeling* methods (learning from data) with *sampling* methods (learning from a density function). The latter is closer to the topic of this article.

5. For the evaluation metric in the experiments, isn't Wasserstein distance a more natural choice? It has the benefit of avoiding selecting the kernel as a hyperparameter.

[1] Marzouk, Y., Moselhy, T., Parno, M., & Spantini, A. (2016). An introduction to sampling via measure transport. arXiv preprint arXiv:1602.05023.

[2] Hoffman, M., Sountsov, P., Dillon, J. V., Langmore, I., Tran, D., & Vasudevan, S. (2019). Neutra-lizing bad geometry in hamiltonian monte carlo using neural transport. arXiv preprint arXiv:1903.03704.

[3] Qiu, Y., & Wang, X. (2024). Efficient multimodal sampling via tempered distribution flow. Journal of the American Statistical Association.

questions:
Besides the questions in the "Weaknesses" section, I have a few other minor ones that I hope the author(s) can clarify.

1. In Theorem 3, the curve is denoted as $\gamma_\tau$, but I only see $\tilde{\gamma}_\tau$ in equation (3). Are they the same?

2. In Section 4, the "Annealed Importance Sampling" paragraph, I would like to point out that the self-normalized estimator $\sum_{i=1}^N w_i f(x_i)/(\sum_{i=1}^N w_i)$ is in general not unbiased, since the denominator is also a random variable. The numerator itself $\sum_{i=1}^N w_i f(x_i)$ is actually unbiased. The self-normalized version is used when $w_i$ is computed from unnormalized densities.

3. In the introduction section, I see "this regularization converts the objective functional into a convex one". Is there any theoretical guarantee for this?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730649472980

thread_type:
official_review

<thread_object>
creator: reviewer_9ukk

<content>
summary:
This paper describes a sampling algorithm using proposals from a continuous normalizing flow with an additional reweighting / accept-reject step, which the authors claim enables the use of a data-free variational inference objective that reduces the reverse KL divergence. The methodology is analyzed and studied numerically on simple examples, where it outperforms a number of standard methods and more recent ML-based strategies like CRAFT and DDS.

strengths:
- The paper is technically strong throughout and provides a clear assessment of the algorithm on a suite of relatively standard examples for the MCMC literature.
	- The clear connections to optimal transport provide a nice framework to enable connections to other algorithms.

weaknesses:
- It is somewhat difficult to identify the contribution / novelty of this work. The basic idea of using reweighting and or metropolization has been used heavily with normalizing flows. E.g., https://arxiv.org/abs/2105.12603 or https://arxiv.org/abs/1812.01729. Much of the paper is standard background material. 
- The connection between the objectives written in section 3.2 and the extremely widely used flow-matching / rectified flow objectives is not explored at all. While both of the aforementioned algorithms use (or can be directly formulated with) a Benamou-Brenier OT objective, the authors choose to instead use the more tedious and less efficient adjoint state framework for CNFs. 
- The importance sampling algorithm is not very clearly explained. Because it seems to be the primary novel content in the paper, it would be beneficial to focus more heavily on this. 
- Not at all clear from the paper how this scheme avoids mode collapse.

questions:
- How does the neural JKO approach differ from the rectified flow / flow-matching objective?
- Is it possible to avoid the adjoint formulation, which is widely known to be inefficient?
- Some of the initial motivation focuses on mode collapse. Do any of your examples clearly indicate that this method avoids mode collapse where standard CNFs fail? A simple example in 2d with well-separated modes is typically enough to observe mode collapse when optimizing the reverse KL.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730665048749

thread_type:
official_review

<thread_object>
creator: reviewer_ff9a

<content>
summary:
This paper combines ideas from JKO sampling and annealed importance sampling to established an unbiased algorithm for implementing the Neural JKO protocol. They establish that using a learned velocity field to perform the transport monotonically decreases the KL divergence between the model density at time $t$ to the target density (as we should expect from a transport with an implied interpolating density). As the author understands it, the main contribution of this paper is the establishment of importance weights that allow to correct errors arising from the approximate velocity field in the neural JKO scheme.

strengths:
- The authors exposit their method very nicely, and in addition make useful connections to various facets of the transport problem in the literature that are necessary for understanding the context of their method -- e.g. the motivation for the JKO scheme and its relation to regularized-OT solutions. 

- The experiments are comprehensive as with those in the literature.

- Having a correctable JKO scheme is important. Given the reviewer's current knowledge, this was absent from previously proposed solutions, and for sampling algorithms, being able to show that your method is unbiased (or has controlled bias) is essential.

weaknesses:
- As the reviewer understands it, the proposed method requires backpropagation through the solution of the ODE for training. Could the authors verify that this is true? This has been a perennial issue in using KL-like optimization for training samplers (it does not scale well with the dimensionality of the problem), and it's not clear to the reviewer that this primary obstacle is avoided in the OT regularized version.  

- Moreover, this scales even worse when the ODE solution requires the computation of the trace of the divergence of a vector field. A measure of computational cost compared to the relevant literature / experiments seems important. In high dimensional sampling problems, I imagine the memory footprint of training probably has to be quite large. 

- These issues are left out of the limitation discussion, which seems important, in addition to what they already remark on that the number of discretization steps (and therefore opportunity for rejection sampling) may be large if the flow is not learned well. Adding these reflections would be useful (and does not retract from the contribution of the authors!), and would allow the reviewer to bump their score a bit!

questions:
To the reviewer, the rejection sampling step seems more like sequential monte-carlo than it does annealed importance sampling. These methods share a lot of the same conceptual reasoning, and the reviewer thinks in fact that the resampling procedure in SMC could be interpreted in this context too. Could the authors comment on that? Do they think it would be useful to make this connection?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730719110060

thread_type:
official_review

<thread_object>
creator: reviewer_wgj8

<content>
summary:
The paper considers the problem of sampling a probability distribution known up to a normalization constant, a classical question of interest in statistical mechanics and bayesian inference. It is proposed to do so by combining JKO steps using continuous normalizing flows (CNF) based on neural ODE with importance based rejection steps to correct for the bias introduced by errors in the CNF. More specifically, the JKO scheme is used to build a sequence of piecewise geodesic interpolations between densities performing steepest descent in Wasserstein-2 metric using as objective function the reversed KL divergence of the evolving measure from the target as energy. Since this KL divergence can be expressed as an expectation over the evolving measure of the unnormalized target distribution, it can be written explicitly. The optimization can thereby be performed in steps, which is practice can be done by propagating samples from one measure along the sequence to the next while solving some ODE along the way. A key aspect of the method is claimed to be the addition of an importance sampling step to accept or reject these samples along te way, in order to remove statistical errors that could accumulate as the walkers used in the sampling slowly drift from the sequence of measure generated.

strengths:
The problem studied is an important one, and the algorithm proposed may be useful to address it.

weaknesses:
The idea of building JKO schemes using is not new, see in particular the paper by Xu, Cheng, and Tie (cited here),. It would therefore be helpful to better emphasizes what is new here (which seems to be the importance sampling step, and demonstrate via example that it solves the slow convergence and instability issues that otherwise could arise. In this respect, the author could summarize (or relegate to the appendix) the background material on Wasserstein spaces, CNF and OT,  the JKO flow, and AIS as well as Metropolis-Hastings: while this material is explained well here, this exposition takes a lot of space and as a result the details on the actual scheme being proposed are rather sparse (in particular in the way the importance sampling bit of the scheme is implemented), and so are the numerical experiments reported.

questions:
1. Please give details about the importance sampling scheme: the introduction in Sec. 4 could be shortened so that the **Neural JKO Sampling with Importance Correction** is expanded and the scene proposed clarified. 

2. The paper blends CNF to perform steepest descent in continuous-time via the solution of ODE, with the JKO scheme which is a discrete-time proximal scheme. Doing so therefore requires discretizing the ODE, which of course is a step that always needs to be done, but the details of this discretization step should be discussed more carefully. In particular, how is the parameter $\tau$ in the JKO scheme chosen: is the same at every step?  Is the discretization step used to solve the CNF (and related ODE) taken to be $\tau$, or are multiple integration steps performed for every $\tau$? Can the time step by adapted along the way?

3. The method seems to require restarting the integration of the CNF from the beginning each time a rejection step is performed. Is this indeed the case? This seems quite costly in practice. 

4. Several papers have recently appeared that propose similar schemes, with an importance sampling step also added, see in particular:

https://arxiv.org/abs/2410.02711

https://arxiv.org/abs/2102.07501

https://arxiv.org/abs/2111.15141

https://arxiv.org/abs/2307.01050

These papers should be discussed here, by giving a careful comparison between the methods and their numerical results.

5. More details about the numerics should be given.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731854233776

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to the Reviewers

comment:
We would like to thank all reviewers for the thorough evaluation of our paper. We carefully revised our paper. Changes and additions are indicated in blue. We respond to each reviewer seperately with detailed answers how we addressed the raised comments and questions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a novel method for sampling from unnormalized distributions by combining continuous normalizing flows (CNFs) with rejection sampling based on importance weights to correct biases. The CNF framework enables fast sampling and density evaluation using neural networks, while the biases introduced by CNFs are appropriately corrected through rejection sampling.

All reviewers acknowledged the significant contribution of the JKO scheme in this study. However, concerns were raised regarding the lack of discussion on importance sampling, computational costs, and parameter tuning. As the paper is borderline, I conducted a detailed review and reached a similar conclusion. While the mathematical contributions are substantial, as an algorithmic proposal, the paper requires major revisions, as also highlighted by Reviewer wGJ8.

Specifically, although the authors added a pseudo code (Algorithm 1) in response to reviewer comments, the algorithm calls several additional sub-algorithms (Algorithms 2–4 in the Appendix), which significantly impacts readability. Moreover, the critical aspect of how the JKO scheme is combined with importance sampling is not explicitly detailed in the pseudo code. For instance, the description, "In practice, we first utilize CNF layers only, then use several blocks consisting of a single CNF layer composed with 3 rejection steps", remains vague. While Appendix E.2 provides details about the combination, this should be clearly presented in the main text as an explicit algorithm.

Additionally, there is no discussion on the stability of the proposed method with respect to how the JKO scheme and importance sampling are combined. The pseudo code also lacks details about key hyperparameters, such as step sizes, which are critical for the algorithm’s implementation. Multiple reviewers raised concerns about hyperparameters, and while these were partially addressed during the discussion phase, such discussions should also be included in the main text due to their importance.

Another significant aspect of the study is the ability to calculate densities, which is an important advantage of the proposed method. However, the specific method for this calculation is only described in the pseudo code in the Appendix. This too should be presented in the main text for clarity and completeness.

Addressing these issues requires substantial revisions to ensure the algorithmic aspects are adequately described. While the mathematical contributions of this study are significant, there is a notable discrepancy in quality between the theoretical and algorithmic components. Given the inadequacies in the latter, I recommend rejection at this time, with the hope that the paper will be resubmitted after major revisions.