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

2025-05-06

AI Summary

Normalizing flows (NFs) suffer from mode collapse and incomplete coverage of multimodal posterior distributions in parameter estimation tasks. To address this, we propose an Effective Sample Size (ESS)-driven adaptive annealing strategy: (i) ESS is introduced as the primary dynamic metric for controlling annealing progression, jointly optimizing both mode coverage and sampling quality; and (ii) ESS-guided sample pruning is incorporated to reduce estimation variance. Unlike conventional approaches, our method eliminates the need for manual annealing schedule design. Evaluated on parameter estimation for a biochemical oscillator model, it achieves a tenfold speedup in convergence over state-of-the-art ensemble MCMC methods, while rigorously preserving the integrity and accuracy of the multimodal posterior distribution.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Security and Privacy: Large-scale security measurements

Abstract

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.

Problem

• Mitigating mode collapse in normalizing flows for parameter estimation
• Improving convergence speed in complex distribution sampling
• Reducing variance by pruning samples using effective sample size

Innovation

• Annealing with adaptive ESS schedule
• Reducing variance by pruning samples
• Faster convergence than ensemble MCMC