Importance Corrected Neural JKO Sampling

2024-07-29

arXiv.org

AI Summary

This work addresses efficient sampling from unnormalized probability density functions. Methodologically, it introduces a novel framework integrating continuous normalizing flows (CNFs) with importance-weighted rejection resampling. The CNF training is formulated as an iterative JKO variational scheme, with theoretical guarantees of convergence to the Wasserstein gradient flow velocity field. Crucially, the framework proposes an alternating mechanism between local flow steps and non-local rejection steps, where rejection proposals are generated adaptively by the model itself—eliminating reliance on hand-crafted proposal distributions and mitigating the susceptibility of conventional Wasserstein gradient flows to local minima and slow convergence in multimodal settings. The approach yields i.i.d. samples and supports implicit density evaluation. Experiments demonstrate substantial improvements in accuracy over state-of-the-art methods on high-dimensional multimodal benchmarks, confirming both effectiveness and robustness.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Generative Adversarial Networks (GANs) for Vision

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Sampling from unnormalized densities using CNFs and rejection-resampling
• Overcoming slow WGF convergence for multimodal distributions
• Reducing reverse KL loss while generating iid samples

Innovation

• Combines CNFs with rejection-resampling steps
• Links CNF training to JKO scheme
• Overcomes WGF limitations via alternation