Importance Corrected Neural JKO Sampling
Sampling;Wasserstein Gradient Flows;Normalizing Flows;Rejection Sampling
Yes
No ethics review needed.
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.
The problem studied is an important one, and the algorithm proposed may be useful to address it.
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.
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.
Yes
No ethics review needed.
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?
- 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.
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.
- 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!
Yes
No ethics review needed.
- 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.
- 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.
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.
- 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.
Yes
No ethics review needed.
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?
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.
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.
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.
We accurately sample from unnormalized densities by alternating local flow steps with non-local rejection steps.
@inproceedings{
anonymous2024importance,
title={Importance Corrected Neural {JKO} Sampling},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=eZLckrDOom},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Sampling Wasserstein Gradient Flows Normalizing Flows Rejection Sampling
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/404fcf299e9fbfb9592b39bb2284446840f14cb3.pdf
probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/786515ed7193abeefd7f467871561a56f779ffa7.zip
Importance Corrected Neural JKO Sampling
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission