Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme

2026-10-07

AI Summary

This study addresses the challenge of efficient sampling from unnormalized Boltzmann distributions. It proposes a neural sampling algorithm based on the Jordan-Kinderlehrer-Otto (JKO) scheme under the Wasserstein-Fisher-Rao (WFR) geometry, employing reweighted normalizing flows to realize the transport and reaction components within an end-to-end deep learning framework. Theoretically, this work establishes exponential convergence for exact iterations with arbitrary step sizes, without requiring structural assumptions such as log-concavity. Empirically, the method is validated on highly challenging, complex multimodal target distributions, demonstrating superior sampling performance.

Abstract

We propose a neural algorithm for sampling from distributions specified by unnormalized Boltzmann densities. Our approach is based on the Jordan--Kinderlehrer--Otto scheme for the Kullback--Leibler divergence in the Wasserstein--Fisher--Rao geometry (WFR JKO scheme). Our contributions are twofold. First, we prove that, for any fixed step size, the exact WFR JKO iterates converge exponentially fast to the target as the number of iterations tends to infinity. Notably, this result requires no structural assumptions on the target, such as log-concavity or a logarithmic Sobolev inequality. Second, we develop a neural implementation of the WFR JKO scheme that parametrizes its transport and reaction components using reweighted normalizing flows. Numerical experiments on challenging multimodal targets demonstrate the promising performance of the proposed method.

Problem

• Neural Sampling
• Unnormalized Boltzmann Densities
• Normalizing Flows
• Multimodal Distributions

Innovation

• Wasserstein-Fisher-Rao geometry
• JKO scheme
• normalizing flows
• neural sampling
• Boltzmann densities