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

TL;DR

This paper introduces an adaptive annealing schedule based on effective sample size to mitigate mode collapse in normalizing flows, significantly improving parameter estimation efficiency in complex models.

Contribution

It proposes a novel adaptive annealing method using ESS to prevent mode collapse in normalizing flows, enhancing their practical utility for parameter estimation.

Findings

Reduces mode collapse in normalizing flows.

Achieves ten-fold faster convergence than ensemble MCMC.

Uses ESS to prune samples and reduce variance.

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.