# Mitigating mode collapse in normalizing flows by annealing with an adaptive schedule: Application to parameter estimation

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 61 references that are in the map.
- First posted on arXiv 6 May 2025.

## Builds on (most cited first, 19 of 19)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Solving Statistical Mechanics Using Variational Autoregressive Networks](https://mlsanity.com/paper/1809.10606/) (2018)
- [Stiff Neural Ordinary Differential Equations](https://mlsanity.com/paper/2103.15341/) (2021)
- [Learning Interpolations between Boltzmann Densities](https://mlsanity.com/paper/2301.07388/) (2023)
- [Pareto Smoothed Importance Sampling](https://mlsanity.com/paper/1507.02646/) (2015)
- [Detecting and Mitigating Mode-Collapse for Flow-based Sampling of Lattice Field Theories](https://mlsanity.com/paper/2302.14082/) (2023)
- [Gradients should stay on Path: Better Estimators of the Reverse- and Forward KL Divergence for Normalizing Flows](https://mlsanity.com/paper/2207.08219/) (2022)
- [Efficient Bayesian Sampling Using Normalizing Flows to Assist Markov Chain Monte Carlo Methods](https://mlsanity.com/paper/2107.08001/) (2021)
- [Deterministic Langevin Monte Carlo with Normalizing Flows for Bayesian Inference](https://mlsanity.com/paper/2205.14240/) (2022)
- [Fixed-kinetic Neural Hamiltonian Flows for enhanced interpretability and reduced complexity](https://mlsanity.com/paper/2302.01955/) (2023)

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

- [Annealing in variational inference mitigates mode collapse: A theoretical study on Gaussian mixtures](https://mlsanity.com/paper/2602.12923/) (2026)
- [Leveraging generative models to assist Monte Carlo sampling](https://mlsanity.com/paper/2608.07648/) (2026)
- [Jeffreys Flow: Robust Boltzmann Generators for Rare Event Sampling via Parallel Tempering Distillation](https://mlsanity.com/paper/2604.05303/) (2026)
- [Generative Nested Sampling of Atomistic Thermodynamic Landscapes](https://mlsanity.com/paper/2609.03193/) (2026)
- [Machine learning kinetics from molecular dynamics data](https://mlsanity.com/paper/2609.17736/) (2026)
- [SPIBER: Reconstructing Free Energy Landscapes from Short, Unconverged Trajectories with Generative Flow Networks](https://mlsanity.com/paper/2609.22663/) (2026)

## Abstract (arXiv)

Normalizing flows (NFs) provide uncorrelated samples from complex distributions, making them an appealing tool for parameter estimation. However, the practical utility of NFs remains limited by their tendency to collapse to a single mode of a multimodal distribution. In this study, we show that annealing with an adaptive schedule based on the effective sample size (ESS) can mitigate mode collapse. We demonstrate that our approach can converge the marginal likelihood for a biochemical oscillator model fit to time-series data in ten-fold less computation time than a widely used ensemble Markov chain Monte Carlo (MCMC) method. We show that the ESS can also be used to reduce variance by pruning the samples. We expect these developments to be of general use for sampling with NFs and discuss potential opportunities for further improvements.