Research map / Probabilistic and scientific modeling / Variational and generative sampling
Importance Corrected Neural JKO Sampling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Combines CNF-based neural JKO steps with importance-weighted rejection-resampling steps to sample unnormalized densities, with proven convergence of neural JKO velocity fields to the Wasserstein gradient flow.
Problem
Sampling from unnormalized densities with generative models trained via reverse KL, where the non-convex objective leads to mode collapse, wrong mode weights, and slow convergence of Wasserstein gradient flows on multimodal targets.
Approach
Alternate local neural JKO steps (CNFs with OT/kinetic regularization, trained to approximate Wasserstein proximal steps) with non-local importance-based rejection-resampling steps whose proposal is the model itself, and propagate density values so that the final model yields iid samples and density evaluations.
Key results
- On GMM-10 to GMM-200 the energy distance of Neural JKO IC is around 1e-4 to 6e-4 while MALA, HMC, DDS, CRAFT and Neural JKO are typically 1e-2 to 1e+1.
- On shifted 8 Modes Neural JKO IC reports energy distance 1.2e-5 versus 1.3e-1 for Neural JKO, close to the sampling error 2.6e-5.
- On LGCP Neural JKO IC reports log Z estimate 508.2 vs 507.6 for CRAFT and 503.0 for DDS.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add importance-based rejection steps to neural JKO (ablation) | Neural JKO sampling · Mustache | energy distance: 1.8e-2 -> 2.9e-3 | Helped |
| Add importance-based rejection steps to neural JKO (ablation) | Neural JKO sampling · shifted 8 Modes | energy distance: 1.3e-1 -> 1.2e-5 | Helped |
| Add importance-based rejection steps to neural JKO (ablation) | Neural JKO sampling · shifted 8 Peaky | energy distance: 1.3e-1 -> 3.4e-5 | Helped |
| Add importance-based rejection steps to neural JKO (ablation) | Neural JKO sampling · Funnel | energy distance: 4.6e-2 -> 1.4e-2 | Hurt |
| Add importance-based rejection steps to neural JKO (ablation) | Neural JKO sampling · GMM-10 | energy distance: 1.1e-2 -> 5.3e-5 | Helped |
| Add importance-based rejection steps to neural JKO (ablation) | Neural JKO sampling · GMM-20 | energy distance: 1.0e-2 -> 1.1e-4 | Helped |
| Add importance-based rejection steps to neural JKO (ablation) | Neural JKO sampling · GMM-50 | energy distance: 2.7e-2 -> 1.0e-4 | Helped |
| Add importance-based rejection steps to neural JKO (ablation) | Neural JKO sampling · GMM-100 | energy distance: 4.7e-2 -> 6.0e-4 | Helped |
| Add importance-based rejection steps to neural JKO (ablation) | Neural JKO sampling · GMM-200 | energy distance: 8.9e-2 -> 3.3e-3 | Helped |
| Use neural JKO IC instead of HMC | Sampling · Funnel | energy distance: 3.1e-3 -> 1.4e-2 | Hurt |
| Use neural JKO IC instead of HMC | Sampling · Mustache | energy distance: 1.7e-2 -> 2.9e-3 | Helped |
| Use neural JKO IC instead of DDS | Sampling · shifted 8 Peaky | energy distance: 1.1e-2 -> 3.4e-5 | Helped |
Surprising findings
- Neural JKO without rejection steps can perform worse than baselines on several targets, e.g. shifted 8 Modes log Z of -0.34 versus about 0 for ground truth, showing the importance of the rejection steps.
Limitations
- Sample generation time grows exponentially with the number of rejection steps.
- CNF training is computationally expensive, as discussed in the appendix.
- The rejection steps rely on the model's own proposal, and the scheme depends on hyperparameters such as the acceptance ratio r and step size growth.
Experimental setup
Datasets: unnormalized target densities (Mustache, shifted 8 modes, shifted 8 peaky, Funnel, GMM-10 to GMM-200, LGCP)
Benchmarks: energy distance on synthetic test distributions log normalizing constant estimation on synthetic test distributions and LGCP
Metrics: energy distance estimated log Z
Baselines: MALA HMC DDS CRAFT Neural JKO
Builds on: OT-Flow JKO scheme continuous normalizing flows annealed importance sampling sequential Monte Carlo
Closest papers by meaning
- Liouville Flow Importance Sampler (ICML 2024 · similarity 0.88)
- Discrete Neural Flow Samplers with Locally Equivariant Transformer (NeurIPS 2025 · similarity 0.86)
- Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions (ICML 2025 · similarity 0.86)
- Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport (ICML 2024 · similarity 0.85)
- Reflected Flow Matching (ICML 2024 · similarity 0.85)
- Reparameterized Importance Sampling for Robust Variational Bayesian Neural Networks (ICML 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.