# Importance Corrected Neural JKO Sampling

arXiv: https://arxiv.org/abs/2407.20444 · HTML: https://mlsanity.com/paper/2407.20444/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 54 of its 93 references that are in the map.
- First posted on arXiv 29 Jul 2024.

## Builds on (most cited first, 30 of 54)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Invertible Residual Networks](https://mlsanity.com/paper/1811.00995/) (2018)
- [Residual Flows for Invertible Generative Modeling](https://mlsanity.com/paper/1906.02735/) (2019)
- [Analyzing Inverse Problems with Invertible Neural Networks](https://mlsanity.com/paper/1808.04730/) (2018)
- [A Conceptual Introduction to Hamiltonian Monte Carlo](https://mlsanity.com/paper/1701.02434/) (2017)
- [Neural Autoregressive Flows](https://mlsanity.com/paper/1804.00779/) (2018)
- [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021)
- [Stochastic Normalizing Flows](https://mlsanity.com/paper/2002.06707/) (2020)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport](https://mlsanity.com/paper/2006.00104/) (2020)
- [Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/1906.04370/) (2019)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (2024)
- [An Introduction to Deep Generative Modeling](https://mlsanity.com/paper/2103.05180/) (2021)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2102.07501/) (2021)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions](https://mlsanity.com/paper/1806.08141/) (2018)

## Built on by (most cited first, 7 of 7)

- [Sequential Controlled Langevin Diffusions](https://mlsanity.com/paper/2412.07081/) (2024)
- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [On the Relation between Rectified Flows and Optimal Transport](https://mlsanity.com/paper/2505.19712/) (2025)
- [Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching](https://mlsanity.com/paper/2603.00530/) (2026)
- [Iterative Importance Fine-tuning of Diffusion Models](https://mlsanity.com/paper/2502.04468/) (2025)
- [Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow](https://mlsanity.com/paper/2601.05583/) (2026)
- [Sampling from Boltzmann densities with physics informed low-rank formats](https://mlsanity.com/paper/2412.07637/) (2024)

## Abstract (arXiv)

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.