Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions

Dongze Wu, Yao Xie

Abstract: Sampling from high-dimensional, multi-modal distributions remains a fundamental challenge across domains such as statistical Bayesian inference and physics-based machine learning. In this paper, we propose Annealing Flow (AF), a method built on Continuous Normalizing Flows (CNFs) for sampling from high-dimensional and multi-modal distributions. AF is trained with a dynamic Optimal Transport (OT) objective incorporating Wasserstein regularization, and guided by annealing procedures, facilitating effective exploration of modes in high-dimensional spaces. Compared to recent NF methods, AF significantly improves training efficiency and stability, with minimal reliance on MC assistance. We demonstrate the superior performance of AF compared to state-of-the-art methods through extensive experiments on various challenging distributions and real-world datasets, particularly in high-dimensional and multi-modal settings. We also highlight AF’s potential for sampling the least favorable distributions.

Lay Summary:

Sampling is a key tool in science and engineering—it helps us understand complex systems by generating representative samples from statistical distributions. However, when the distributions are complex and high-dimensional, such as in quantum systems, or Bayesian statistics, standard methods often struggle. Our paper introduces Annealing Flow, a new machine learning method designed to efficiently sample from these complex distributions, especially when they have many distinct patterns (called "modes") or live in high-dimensional spaces.

The method works by gradually transforming samples from a simple distribution into the complex ones of interest, similar to how metals are cooled slowly to reach a stable state—an idea known as annealing. This transformation is learned using neural networks and a mathematical framework called optimal transport, ensuring stability and efficiency.

Annealing Flow outperforms existing methods across a wide range of tests, offering faster, more reliable sampling. It also enables better handling of rare but important scenarios, improving both scientific insight and practical applications.

Link To Code: https://github.com/StatFusion/Annealing-Flow-For-Sampling

Primary Area: Probabilistic Methods->Monte Carlo and Sampling Methods

Keywords: Statistical Sampling, Optimal Transport, Continuous Normalizing Flows

Submission Number: 846