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

TL;DR

This paper introduces Annealing Flow (AF), a novel method based on Continuous Normalizing Flow, designed to efficiently sample from complex high-dimensional and multi-modal distributions, outperforming existing techniques in stability and effectiveness.

Contribution

The paper proposes AF, a new sampling method combining OT-based training and annealing, improving efficiency and stability over prior normalizing flow approaches.

Findings

AF outperforms state-of-the-art methods on challenging distributions

AF demonstrates superior sampling quality in high-dimensional, multi-modal settings

AF shows potential for sampling difficult, least favorable distributions

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 Flow (CNF) 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 greatly 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 experiments on various challenging distributions and real-world datasets, particularly in high-dimensional and multi-modal settings. We also highlight AF potential for sampling the least favorable distributions.