Flow Annealed Importance Sampling Bootstrap (2208.01893v3)

Abstract: Normalizing flows are tractable density models that can approximate complicated target distributions, e.g. Boltzmann distributions of physical systems. However, current methods for training flows either suffer from mode-seeking behavior, use samples from the target generated beforehand by expensive MCMC methods, or use stochastic losses that have high variance. To avoid these problems, we augment flows with annealed importance sampling (AIS) and minimize the mass-covering $\alpha$-divergence with $\alpha=2$, which minimizes importance weight variance. Our method, Flow AIS Bootstrap (FAB), uses AIS to generate samples in regions where the flow is a poor approximation of the target, facilitating the discovery of new modes. We apply FAB to multimodal targets and show that we can approximate them very accurately where previous methods fail. To the best of our knowledge, we are the first to learn the Boltzmann distribution of the alanine dipeptide molecule using only the unnormalized target density, without access to samples generated via Molecular Dynamics (MD) simulations: FAB produces better results than training via maximum likelihood on MD samples while using 100 times fewer target evaluations. After reweighting the samples, we obtain unbiased histograms of dihedral angles that are almost identical to the ground truth.

• The paper introduces FAB, an approach integrating annealed importance sampling with flow training to enhance mode discovery and reduce variance.
• It employs α-divergence minimization (α=2) to optimize sample quality by effectively targeting areas where the flow approximation fails.
• Empirical tests show FAB's superior performance in effective sample size, log-likelihood, and mode coverage on complex multimodal distributions.

Flow Annealed Importance Sampling Bootstrap: An Expert Review

The paper under review presents the Flow Annealed Importance Sampling Bootstrap (FAB), an innovative approach to training normalizing flows to approximate intractable multimodal distributions such as Boltzmann distributions. Traditional methodologies suffer from significant drawbacks including mode-seeking behaviors, reliance on expensive pre-generated Markov chain Monte Carlo (MCMC) samples, or high-variance stochastic losses. The authors introduce FAB to address these limitations, proposing an augmentation of flows with Annealed Importance Sampling (AIS) and advocating for $\alpha$-divergence minimization with $\alpha=2$ to reduce importance weight variance—a strategy that seeks to maintain mass coverage while discovering hidden modes.

Technical Contributions and Methodology

FAB hinges on an adaptation of AIS that targets the region where the flow approximation fails against the target distribution. This is accomplished by incorporating AIS into the flow training process, allowing for the generation of samples from areas of poor flow approximation. Particularly, FAB suggests an objective that minimizes the $\alpha$-divergence with $\alpha=2$. This choice optimizes the balance between exploratory and exploitative sampling phases, thus promoting mode discovery and enhancing sample quality while reducing variance.

Key elements of FAB include:

• Use of AIS: The authors exploit AIS, initiating from the flow and transitioning through MCMC towards a target that minimizes variance—a defined distribution proportional to $p^2/q_\theta$, thereby ensuring focus on critical distribution regions.
• Replay Buffer: To mitigate computational overhead, FAB introduces a prioritized replay buffer. This buffer allows reuse of AIS-generated samples, effectively making training more efficient and less reliant on constantly re-sampling from complex distributions.

Experimental Analysis

The authors rigorously test FAB on a series of challenging multimodal distributions: a 40-component mixture of 2D Gaussians and a 32-dimensional Many Well problem. In both cases, FAB exhibited superior mode coverage compared to traditional KL-divergence minimization approaches. Notably, in the case of the alanine dipeptide—a 22-atom molecule—FAB demonstrated its ability to successfully model the Boltzmann distribution using 100 times fewer target evaluations than models reliant on MD-generated samples.

Key Results

• Effective Samples and Log-Likelihood: The authors report significant improvements in ESS and log-likelihood metrics under FAB compared to other methods, attesting to its robustness in handling distributions with complex energy landscapes and high-dimensionality.
• Importance Sampling Efficiency: FAB shows substantial variance reduction in estimated expectations, corroborated by empirical analysis in gradient performance and variance scaling relative to problem dimension.
• Mode Coverage: Empirical evaluations demonstrate FAB's efficacy in achieving full mode coverage, a critical requirement for molecular simulations and other scientific computing applications.

Implications and Future Prospects

FAB opens new avenues for training density models where sampling from the target distribution is inherently expensive or impractical. Given its emphasis on minimizing variance and emphasizing mode coverage, FAB could significantly enhance the deployment of normalizing flows in high-stakes scientific computations, molecular simulations, or any task requiring robust estimates over complex probability landscapes.

Future research directives could include integrating FAB with more expressive flow architectures such as autoregressive and spline-based models and extending applications to larger biomolecular systems. Furthermore, the potential fusion of FAB and alternative sampling techniques like sequential Monte Carlo presents an exciting opportunity to further push the boundaries of computational feasibility in high-dimensional statistical learning.

In sum, FAB proposes a practical and theoretically grounded contribution to the flow-based modeling repertoire, extending the scope of applications demanding precision in the sampling of complicated energy-based models.