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

arXiv: https://arxiv.org/abs/2409.20547 · HTML: https://mlsanity.com/paper/2409.20547/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 43 of its 85 references that are in the map.
- First posted on arXiv 30 Sep 2024.

## Builds on (most cited first, 30 of 43)

- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions](https://mlsanity.com/paper/2209.11215/) (2022)
- [What Are Bayesian Neural Network Posteriors Really Like?](https://mlsanity.com/paper/2104.14421/) (2021)
- [Stochastic Gradient Descent as Approximate Bayesian Inference](https://mlsanity.com/paper/1704.04289/) (2017)
- [Stochastic Normalizing Flows](https://mlsanity.com/paper/2002.06707/) (2020)
- [Stein Variational Gradient Descent as Gradient Flow](https://mlsanity.com/paper/1704.07520/) (2017)
- [Parallel Sampling of Diffusion Models](https://mlsanity.com/paper/2305.16317/) (2023)
- [Telescoping Density-Ratio Estimation](https://mlsanity.com/paper/2006.12204/) (2020)
- [Accelerating Bayesian Optimization for Biological Sequence Design with Denoising Autoencoders](https://mlsanity.com/paper/2203.12742/) (2022)
- [A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems](https://mlsanity.com/paper/1912.01825/) (2019)
- [Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2102.07501/) (2021)
- [Continual Repeated Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2201.13117/) (2022)
- [Density Ratio Estimation via Infinitesimal Classification](https://mlsanity.com/paper/2111.11010/) (2021)
- [Scaling Hamiltonian Monte Carlo Inference for Bayesian Neural Networks with Symmetric Splitting](https://mlsanity.com/paper/2010.06772/) (2020)
- [Stochastic Particle Gradient Descent for Infinite Ensembles](https://mlsanity.com/paper/1712.05438/) (2017)
- [A Stein variational Newton method](https://mlsanity.com/paper/1806.03085/) (2018)
- [Theoretical Insights for Diffusion Guidance: A Case Study for Gaussian Mixture Models](https://mlsanity.com/paper/2403.01639/) (2024)
- [Learning general Gaussian mixtures with efficient score matching](https://mlsanity.com/paper/2404.18893/) (2024)
- [Normalizing flow neural networks by JKO scheme](https://mlsanity.com/paper/2212.14424/) (2022)
- [Differentiable Annealed Importance Sampling and the Perils of Gradient Noise](https://mlsanity.com/paper/2107.10211/) (2021)
- [Convergence of flow-based generative models via proximal gradient descent in Wasserstein space](https://mlsanity.com/paper/2310.17582/) (2023)
- [Stochastic Normalizing Flows for Inverse Problems: a Markov Chains Viewpoint](https://mlsanity.com/paper/2109.11375/) (2021)
- [Sampling in Unit Time with Kernel Fisher-Rao Flow](https://mlsanity.com/paper/2401.03892/) (2024)
- [Liouville Flow Importance Sampler](https://mlsanity.com/paper/2405.06672/) (2024)
- [Provable Bayesian Inference via Particle Mirror Descent](https://mlsanity.com/paper/1506.03101/) (2015)
- [Sampling with Mollified Interaction Energy Descent](https://mlsanity.com/paper/2210.13400/) (2022)
- [Efficient Bayesian Sampling Using Normalizing Flows to Assist Markov Chain Monte Carlo Methods](https://mlsanity.com/paper/2107.08001/) (2021)
- [Langevin-gradient parallel tempering for Bayesian neural learning](https://mlsanity.com/paper/1811.04343/) (2018)

## Built on by (most cited first, 11 of 11)

- [Density Ratio Estimation with Conditional Probability Paths](https://mlsanity.com/paper/2502.02300/) (2025)
- [Nonlinear Assimilation via Score-based Sequential Langevin Sampling](https://mlsanity.com/paper/2411.13443/) (2024)
- [Energy-Weighted Flow Matching: Unlocking Continuous Normalizing Flows for Efficient and Scalable Boltzmann Sampling](https://mlsanity.com/paper/2509.03726/) (2025)
- [DoFlow: Flow-based Generative Models for Interventional and Counterfactual Forecasting on Time Series](https://mlsanity.com/paper/2511.02137/) (2025)
- [Flow-based Generative Modeling of Potential Outcomes and Counterfactuals](https://mlsanity.com/paper/2505.16051/) (2025)
- [FlowVAT: Normalizing Flow Variational Inference with Affine-Invariant Tempering](https://mlsanity.com/paper/2505.10466/) (2025)
- [Iterative Sampling Methods for Sinkhorn Distributionally Robust Optimization](https://mlsanity.com/paper/2512.12550/) (2025)
- [CoreFlow: Low-Rank Matrix Generative Models](https://mlsanity.com/paper/2604.24959/) (2026)
- [Sampling via Stochastic Interpolants by Langevin-based Velocity and Initialization Estimation in Flow ODEs](https://mlsanity.com/paper/2601.08527/) (2026)
- [Generative Nested Sampling of Atomistic Thermodynamic Landscapes](https://mlsanity.com/paper/2609.03193/) (2026)
- [GazeDiT: Gaze-Accurate Diffusion Image Generation for Eye Tracking via Spatial Conditioning](https://mlsanity.com/paper/2609.17814/) (2026)

## Abstract (arXiv)

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.