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

TL;DR

This work tackles mode collapse in normalizing flows used for Bayesian parameter estimation by introducing an adaptive annealing scheme guided by the effective sample size $n_\text{eff}$. By gradually transforming the target from the prior toward the posterior via ESS-based updates of the annealing parameter $\beta$, the method robustly captures multimodal posteriors without requiring prior mode knowledge. Across a repressilator ODE model, the NF-annealing approach achieves about a 10× speedup over ensemble MCMC, with ESS-based pruning reducing variance in marginal likelihood estimates. The approach is general and scalable, with potential for further improvements in alternative NF architectures and annealing metrics.

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.

Figures (10)

• Figure 1: The annealing-sampling scheme. The parameter $\beta$ increases to interpolate from a simple distribution that is close to the base distribution (top left) to the ultimate target distribution (top right). At each value of $\beta$, samples generated from previously trained NFs are reweighted and used to train a new NF, from which samples are then drawn.
• Figure 2: Repressilator model. (a) Schematic of the system; each circle represents a gene product, and $i\dashv j$ represents repression of $j$ by $i$. (b) Solution used to generate the data for fitting; the parameter values are $X_i(t_0) = 2$ for all $i$, $\alpha_1 = 10$, $\alpha_2 = 15$, $\alpha_3 = 20$, $m=4$, and $\eta = 1$. (c) Time series of the total concentration of gene products (blue line) and the simulated observable (orange dots) produced by adding Gaussian noise with variance 0.25.
• Figure 3: Annealing with a schedule based on the effective sample size (ESS) mitigates mode collapse. (a) Samples from NF models trained with the annealing protocol. Colors correspond to different hyperparameter choices, as specified in (c) and (d). (b) Samples from NF models trained with fixed $\beta=1$. Different colors distinguish samples from three models trained with distinct initial weights and random seeds. Note that the scales in (a) and (b) are different. (c) The change of $\beta_s$ in training from three independent runs with different parameters. (d) The effective sample size (ESS). Translucent lines show the instantaneous ESS values ($n_\mathrm{eff}$) and opaque lines show their exponetial moving average ($\bar{n}_\mathrm{eff}$) with $\lambda= 0.01$.
• Figure 4: Estimating marginal likelihoods. (a, c) Marginal likelihood estimates for different sample sizes. Shaded regions indicate standard deviations from 10 independent sets of samples. In (a), all samples are used for the estimates. In (c), samples are excluded to maximize the effective sample size. (b) Variation of the effective sample size (ESS) as samples with large weights are excluded. Results for three independent batches of samples are shown. (d) Integrand for thermodynamic integration (TI). The red solid line indicates the cutoff below which data points are excluded. The inset provides an expanded view for $\beta_s>0.2$.
• Figure S1: Sampling the parameter space of the repressilator with an NF using a preset schedule for $\beta_s$: $\beta_s =( s/1000 )^4$ for $s = 1, 2, \cdots, 1000$friel2008powerPosteriors. For each $\beta_s$, the NF is trained for $800$ steps with samples updated every $50$ steps ($k=50$). (a) Samples from NF models trained with fixed $\beta=1$. Different colors distinguish samples from three models trained with distinct initial weights and random seeds. (b) The effective sample size (ESS). Translucent lines show the instantaneous ESS values ($n_\mathrm{eff}$) and opaque lines show their exponential moving average ($\bar{n}_\mathrm{eff}$) with $\lambda=0.01$. Colors correspond to the runs in (a). The drop in ESS is consistent with the NFs failing to converge. In particular, the orange run exhibits both a sudden drop and mode collapse.