Neural method improves sampling from complex probability distributions
Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme
Machine Learning
Summary
Sampling from complex probability distributions is a common problem in many fields like physics and machine learning. The authors propose a new neural network method to do this more effectively by combining advanced geometry and probabilistic modeling. They prove their approach reliably converges to the desired distribution even in difficult cases, and they demonstrate its ability to handle complicated, multimodal distributions. This method uses a novel way to represent and transform data distributions with neural networks called reweighted normalizing flows.
What this means in practice
- •For machine learning engineers: Generate complex probability samples efficiently for training and inference in probabilistic models with multimodal distributions.
- •For computational physicists: Model physical systems by sampling unnormalized Boltzmann distributions without restrictive assumptions on their shape.
Authors
Chenguang Duan, Johannes Hertrich, Gabriele Steidl
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.