PaperlayerView exact revision on arXiv

Source-linked AI summary

Flow Annealed Importance Sampling Bootstrap

Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, José Miguel Hernández-Lobato

arXiv:2208.01893v3cs.LGq-bio.QMstat.ML

TL;DR

Flow training can be hampered by mode-seeking objectives, dependence on expensive target samples, or high-variance losses. FAB combines α=2 divergence minimization with AIS to train flows without target-generated samples. On alanine dipeptide, it outperforms MD-sample maximum-likelihood training while using 100 times fewer target evaluations.

  • Problem

    Flow training methods may be mode-seeking, rely on expensive MD-generated target samples, or use mass-covering objectives with high-variance flow samples.

    Source: §1 · pp. 1–2

  • Method

    FAB minimizes the mass-covering α-divergence with α=2 and uses AIS targeted at p²/q to improve flow-training samples.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    For alanine dipeptide, FAB with a replay buffer outperforms maximum-likelihood training on MD samples while using 100 times fewer target-density evaluations; reweighted samples yield a distribution nearly the same as ground truth.

    Source: §4.2 · p. 9

  • Takeaways & Limitations

    FAB produces accurate approximations of complex multimodal targets without samples from those targets, and reweighting its alanine-dipeptide samples gives a distribution nearly the same as ground truth.

    Source: §4.2 · p. 9, §7 · p. 11

    The experiments used relatively few intermediate AIS distributions and minimal hyperparameter search, and further tuning could improve FAB’s performance.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

The introduction frames flow training as a trade-off between mode-seeking objectives, costly target-generated MD samples, and high-variance mass-covering objectives. It introduces FAB, which combines AIS with α=2 divergence minimization and reports strong results on multimodal targets and alanine dipeptide.

1 INTRODUCTION

  • Flow training faces a trade-off: reverse KL can miss modes, while mass-covering objectives such as forward KL can have high-variance flow samples.For molecular Boltzmann distributions, target samples are typically obtained through expensive, long MD simulations that are highly correlated.Source: §1 · pp. 1–2
  • FAB outperforms competing approaches on multimodal targets and trains an alanine-dipeptide Boltzmann generator from the unnormalized density, using over 100 times fewer target evaluations than MD-sample maximum-likelihood training.The reported alanine-dipeptide approximation is better than the MD-trained flow.Source: §1 · p. 2
  • FAB minimizes the mass-covering α-divergence with α=2, using AIS to improve training samples and a prioritized replay buffer to reduce computational cost.The paper states that this objective minimizes importance-weight variance.Source: §1 · p. 2

2 BACKGROUND

The background introduces normalizing flows as tractable density models and contrasts divergence objectives with annealed importance sampling. It explains how AIS moves samples from a flow through intermediate distributions toward a target and supplies weights for expectation estimates.

2 BACKGROUND

  • A normalizing flow transforms a random variable through an invertible map; with target samples it can be trained by maximum likelihood, or with only the density by minimizing reverse KL using flow samples.The background assumes a normalized target for notation but says the methods also apply when it is unnormalized.Source: §2 · p. 2
  • The α-divergence is mode-seeking for α≤0 and mass-covering for α≥1; at α=2, minimizing it is equivalent to minimizing importance-sampling weight variance.The paper motivates this choice for importance sampling at test time.Source: §2 · p. 2
  • AIS starts with a flow sample, transitions through intermediate distributions toward the target, and returns a weighted final sample usable for estimating target expectations.The intermediate log densities interpolate between the initial and target log densities.Source: §2 · pp. 2–3
  • The paper uses “bootstrap” to mean using a flow together with AIS to estimate a loss that improves the same flow.The term is explicitly used in this sense for FAB.Source: §2 · p. 2

3 METHOD

FAB estimates α=2 divergence training signals with AIS targeted at p²/q, while a prioritized replay buffer reuses AIS samples. The method section also analyzes gradient-estimate quality and describes the buffer’s sampling and correction scheme.

3.1 FLOW ANNEALED IMPORTANCE SAMPLING BOOTSTRAP

  • FAB minimizes Dα=2(p∥q) using AIS with q as the initial distribution and p²/q as its target, focusing training on regions where the flow poorly approximates the target.AIS weights support an unbiased loss estimate; the method’s gradient uses stopped gradients through AIS samples and weights, with self-normalized weights improving training stability.Source: §3.1 · pp. 3–4
  • The AIS bootstrapping approach is presented as extensible to other objectives and models, including models that combine flows with stochastic sampling steps.The paper also notes a derivation for arbitrary α values in an appendix.Source: §3.1 · p. 4
  • In a one-dimensional Gaussian setting, FAB with few AIS distributions has a far superior gradient-estimate signal-to-noise ratio to importance sampling from q or p; under stated factorized, perfect-transition assumptions, its variance remains constant as AIS distributions scale with dimension.Under those assumptions, the comparison importance-sampling gradient variance grows exponentially with dimensionality; the authors note that the assumptions are strong.Source: §3.1 · p. 4

3.2 RE-USING SAMPLES THROUGH A REPLAY BUFFER

  • Reusing AIS samples avoids re-evaluating the target density, but updating every buffer weight before sampling is expensive, so FAB instead corrects sampled minibatches.Sampling without replacement worked better in practice but introduces bias into gradient estimates.Source: §3.2 · pp. 4–5
  • The replay buffer reuses AIS samples by prioritizing points with high p and low q, then reducing their weights as the flow fits them and storing q values for weight correction.Corrected weights account for samples generated under an earlier flow parameter setting.Source: §3.2 · p. 5

4 EXPERIMENTS

The experiments test FAB on a multimodal Gaussian mixture and alanine dipeptide. Buffered FAB covers the mixture’s modes and, on alanine dipeptide, outperforms ML trained on MD samples while using fewer target evaluations.

4.1 MIXTURE OF GAUSSIANS IN 2D

  • The Gaussian-mixture test uses 40 bivariate components and deliberately initializes the flow with samples concentrated in a small region.The study compares FAB variants with reverse-KL flows, RBD, SNF, CRAFT, direct α=2 training, and ML using target samples.Source: §4.1 · p. 6
  • On the 40-component Gaussian mixture, both FAB variants and ML fit all target modes, while alternative methods cover only a subset.FAB with a replay buffer performs similarly to ML and leads on ESS, log-likelihood, forward KL, and MAE; FAB without a buffer is next best.Source: §4.1 · pp. 6–7, Figure 2 · p. 6, Table 1 · p. 7

4.2 ALANINE DIPEPTIDE

  • The alanine dipeptide experiment models a 22-atom molecule at 300 K using internal coordinates, with a minimum-energy configuration and typical variable scales replacing MD-derived normalization.Internal coordinates implicitly satisfy rotational and translational invariance; training retains only samples for the L-form.Source: §4.2 · pp. 7–8
  • For alanine dipeptide, evaluation uses parallel-tempering MD samples as ground truth and compares ESS, test log-likelihood, and Ramachandran-plot KL divergence before and after reweighting.The models are evaluated over three random seeds.Source: §4.2 · p. 8, Table 2 · p. 9
  • 100 times fewer target evaluations let buffered FAB outperform ML trained on MD samples, with higher ESS and lower Ramachandran KLD.After reweighting, its distribution is nearly the same as the ground truth.Source: §4.2 · pp. 8–9, Table 2 · p. 9

5 RELATED WORK

Related work connects FAB to flow-based MCMC and SMC methods, and discusses extensions involving AIS gradients, importance-sampling techniques, physical knowledge, and diffusion models.

5 RELATED WORK

  • CRAFT can reduce mode seeking relative to SNFs on the Many Well problem, but fails on the Gaussian mixture, where its loss can favor mode seeking and flow samples provide a poor training signal.SNFs and CRAFT combine flows with MCMC or SMC; unlike ordinary flows, they do not provide likelihoods and can require many target evaluations at inference.Source: §5 · p. 9
  • FAB does not differentiate through AIS or use target gradients in its loss, and the discussion identifies differentiable AIS, control variates, and defensive importance sampling as possible improvements.Target gradients are used by HMC during sampling; force matching or an added KL term are proposed ways to incorporate gradient information into FAB.Source: §5 · p. 9
  • FAB uses the flow as AIS’s base distribution to balance mass across metastable states; AIS can incorporate modes absent from initialization, unlike the cited approach requiring prior knowledge of those states (Gabrié et al., 2022).A single AIS-bootstrap sample from a new mode can prompt a strong flow update.Source: §5 · pp. 9–10
  • Physical knowledge and symmetries can reduce dimensionality and repeated modes; the discussion suggests these may help scale FAB to larger, more complex molecular systems.The cited atomic-solid approach incorporates permutation and translation invariance, periodic boundaries, and metastable states (Wirnsberger et al., 2022).Source: §5 · p. 10
  • An AIS bootstrap similar to FAB is proposed as a potential way to address high-variance training in diffusion models for complex target distributions.The cited diffusion approach estimates a score-matching loss using samples from the model (Jing et al., 2022).Source: §5 · p. 10

6 DISCUSSION

The discussion notes possible extensions to FAB and limits the reported experiments by their AIS settings and hyperparameter search.

6 DISCUSSION

  • Replacing FAB’s AIS with SMC could make resampling produce more useful training points and lower batch importance-weight variance; at test time, flows could remain available for exact densities or serve as SMC bases.Combining FAB with CRAFT is also proposed, but would inherit limitations including no likelihoods and costly target evaluations at inference.Source: §6 · p. 10
  • The experiments use relatively few intermediate AIS distributions and minimal hyperparameter search, so reported performance may improve with tuning.More intermediate distributions trade lower variance for higher compute cost; the authors also suggest reducing their number online as loss variance falls.Source: §6 · p. 10

7 CONCLUSION

FAB combines α=2 divergence minimization with AIS bootstrapping to train flows for multimodal targets. It produces accurate approximations without target samples and outperforms MD-sample maximum-likelihood training on alanine dipeptide with fewer target evaluations.

7 CONCLUSION

  • FAB combines mass-covering α=2 divergence minimization with AIS bootstrapping to focus training on regions where the flow poorly approximates the target.This objective favors coverage of multimodal distributions and minimizes importance-weight variance.Source: §7 · p. 10
  • For alanine dipeptide, FAB outperforms maximum-likelihood training on MD samples while using 100 fewer target evaluations than the MD simulations.The paper also reports that reweighted dihedral-angle histograms are almost identical to ground truth in the abstract, but that claim is not in these passages.Source: §7 · p. 11
  • FAB accurately approximates complex multimodal targets without samples from those targets, while alternative approaches fail in this setting.A prioritized replay buffer reduces FAB’s cost and improves performance.Source: §7 · p. 11

A DERIVATION OF THE LOSS

The section derives an AIS-based gradient estimate for α=2 divergence minimization, then describes broader divergence and model variants. It also gives CRAFT’s annealed flow-transport setup and notes conditions that can make the divergence infinite.

A DERIVATION OF THE LOSS

  • For α=2, FAB uses AIS from qθ toward the minimum-variance target p²/qθ to estimate the divergence gradient without bias.The general construction selects g=f/Zf as the minimum-variance importance-sampling distribution and estimates its expectations through AIS.Source: A DERIVATION OF THE LOSS · p. 14
  • Normalizing AIS weights stabilizes training by changing gradient magnitude without changing its direction, but self-normalization biases the finite-sample surrogate estimate.The method uses qθ to initialize AIS and relatively few intermediate distributions to limit computation.Source: A DERIVATION OF THE LOSS · p. 15
  • Dα=2 can be infinite when q assigns zero density where p has mass or when p’s tails are heavier than q’s.The authors report dropping points with infinite or NaN q density during training; these did not typically occur late in training.Source: A DERIVATION OF THE LOSS · p. 15

B.1 FAB FOR GENERIC α-DIVERGENCE MINIMISATION

  • FAB generalizes to nonnegative divergence integrands; α=2 performs best in the reported analyses, while the construction does not extend to α→0 reverse KL.At α=1, the gradient corresponds to forward-KL estimation using AIS samples targeted at p.Source: B VARIATIONS OF FAB · p. 15, §B.1 · p. 16

B VARIATIONS OF FAB

  • The authors describe combining FAB with CRAFT by replacing its reverse-KL objective with α=2 divergence and using AIS for gradient estimation.They also note that FAB’s gradient-estimation approach can be applied beyond α=2 divergence.Source: B VARIATIONS OF FAB · p. 15, §B.2 · p. 17

B.2.1 CONTINUAL REPEATED ANNEALED FLOW TRANSPORT MONTE CARLO

  • CRAFT transports samples through annealed distributions using interleaved flow steps, resampling, and MCMC, training each flow to match the next distribution.It returns samples with normalized importance weights for target-expectation estimates.Source: §B.2.1 · p. 17

B.2.2 FAB-CRAFT

FAB-CRAFT reuses CRAFT’s post-flow MCMC samples to train flow steps with α=2 divergence and an AIS bootstrap estimate. In a toy Gaussian comparison, AIS bootstrap has the best gradient-estimator SNR; a separate analysis studies dimensional scaling under simplifying assumptions.

B.2.2 FAB-CRAFT

  • FAB-CRAFT trains each flow step with α=2 divergence using the following MCMC samples as an AIS bootstrap estimate, moving resampling to after MCMC.This repurposes CRAFT’s existing transitions with minimal training-pass changes; the original CRAFT inference procedure remains unchanged.Source: §B.2.2 · p. 18
  • The FAB-CRAFT variant keeps the original algorithm similar, but the authors identify prioritized replay and samples across the full chain as potential improvements.They state that prioritized replay could significantly reduce computational requirements.Source: §B.2.2 · p. 19

C.1 GRADIENT ESTIMATION PERFORMANCE

  • AIS bootstrap has the best gradient-estimator SNR in the toy Gaussian comparison, although it may not outperform target-sample importance sampling early when qθ poorly approximates p.Figure 5 plots SNR against sample count and, with 100 samples, against the number of intermediate AIS distributions.Source: §C.1 · pp. 19–20, Figure 5 · p. 20

C.2 SCALING FAB TO HIGHER DIMENSIONS

  • Under factorized p and q assumptions, FAB’s gradient-estimate variance can remain constant as dimension grows if AIS intermediate distributions increase linearly with dimension.Under the same analysis, importance-sampling gradient-estimate variance using samples from q or p grows exponentially with dimension.Source: §C.2 · p. 20

C.2.1 THEORETICAL ANALYSIS ON FACTORIZED p AND q

For factorized p and q, ordinary importance-sampling gradient estimates have variance that grows exponentially with dimension. FAB can keep gradient variance constant by scaling the number of AIS distributions with dimension, at linear computational cost.

C.2.1 THEORETICAL ANALYSIS ON FACTORIZED p AND q

  • The analysis assumes factorized p and q and no parameter sharing across dimensions of q.Source: §C.2.1 · p. 21
  • Importance sampling with samples from q has gradient-estimator variance that increases exponentially with dimension.Each additional dimension multiplies the variance by another factor greater than one.Source: §C.2.1 · p. 21
  • Importance sampling with samples from p also has exponentially increasing gradient-estimator variance, though at a lower rate than sampling from q.Keeping variance constant therefore requires the number of samples to grow exponentially with dimension, even when samples from p are available.Source: §C.2.1 · pp. 21–22
  • Scaling the number of AIS distributions with dimension keeps FAB’s gradient variance constant while increasing estimator cost only linearly.This result assumes independent samples from the intermediate AIS distributions.Source: §C.2.1 · p. 22

C.2.2 EMPIRICAL ANALYSIS ON TOY PROBLEM

In the toy problem, FAB’s log importance-weight variance stays constant and its gradient-estimation SNR stays roughly flat as dimensionality increases. Importance sampling with samples from p instead shows increasing weight variance and declining SNR.

C.2.2 EMPIRICAL ANALYSIS ON TOY PROBLEM

  • For FAB, the number of AIS distributions increased in proportion to problem dimensionality, giving linearly increasing compute cost.The experiment used factorized Gaussian targets and proposals with unit marginal variances and means −0.5 and 0.5, respectively.Source: §C.2.2 · p. 22, Figure 6 · p. 23
  • The toy experiment used one HMC iteration with five leapfrog steps and step size 0.5, selected by trial and error using the displayed histogram.The histogram compares AIS samples targeting p and p2/q with the PDFs of p, q, and normalized p2/q.Source: §C.2.2 · p. 22, Figure 7 · p. 23
  • FAB keeps log importance-weight variance constant and gradient-estimation SNR roughly flat as dimensions increase, unlike importance sampling with samples from p.The reported FAB SNR stays within a single order of magnitude; p-sample SNR quickly decreases.Source: §C.2.2 · p. 22, Figure 6 · p. 23

D.1 TRAINING SETUP

The experiments specify flow architectures, training procedures, and method-specific settings, then evaluate trained models using sampling and expectation-estimation metrics.

D.1 TRAINING SETUP

  • Models used 15 RealNVP layers with two-layer, 80-unit conditioners, trained with batch size 128 and Adam at learning rate 1 × 10^-4.Training ran for 2 × 10^7 flow evaluations, with three random-seed runs averaged per method.Source: §D.1 · pp. 23–24
  • FAB used batch sizes of 128 for AIS and buffer sampling, one intermediate distribution, and a single Metropolis-Hastings transition per AIS step.Its buffer started with 1,280 samples and had a maximum length of 12,800; updates were skipped when NaN values arose in the correction weights.Source: §D.1 · p. 24

D.2 EVALUATION SETUP

  • Evaluation measured ESS, average log-likelihood on target samples, forward KL, and importance-sampling MAE for expectations, with and without reweighting.ESS used 5 × 10^4 flow samples; MAE used 1,000 samples and was averaged over 100 repetitions.Source: §D.2 · p. 24

D.3 FURTHER RESULTS

On the mixture-of-Gaussians problem, FAB-based methods and maximum-likelihood training cover all target modes, while several alternatives fit only subsets. Further tests find α=2 slightly best without a replay buffer and α=2 or 3 best with one.

D.3 FURTHER RESULTS

  • FAB-based methods and maximum-likelihood training cover all mixture-of-Gaussians modes, while the other methods fit only subsets.Direct Dα=2 training was highly unstable and performed worst.Source: §D.3 · p. 24, Figure 8 · p. 25
  • Without a replay buffer, FAB with α=2 slightly outperforms FAB with other α values.Table 4 reports mixture-of-Gaussians results for FAB without the replay buffer.Source: §D.3 · p. 25, Table 4 · p. 25
  • With a replay buffer, all tested FAB methods with α≥1 fit the target well, with α=2 and α=3 performing best.The authors identify flow expressiveness as the likely limit on further improvement.Source: §D.3 · p. 25
  • For the alanine dipeptide problem, the appendix reports α=2 as best, with larger performance differences across α values.Source: §D.3 · p. 25

E.1 DESCRIPTION AND RESULTS

This section presents supplemental experimental results and descriptions, covering FAB’s performance and mode coverage on the Many Well benchmark alongside computational costs and dihedral-angle plots for MD data and chiral forms.

E.1 DESCRIPTION AND RESULTS

  • The 32-dimensional Many Well target has 65,536 modes; the authors evaluate mode coverage with a test point in each mode and estimate its normalizing constant from exact marginal samples.The independent Double Well marginals permit exact sampling and yield the Many Well normalizing constant; p(x1) uses rejection sampling and p(x2) is standard Gaussian.Source: §E.1 · pp. 26–27, §E.2 · p. 29

E.3 SETUP

  • FAB uses four linearly spaced AIS intermediates with one HMC transition of five leapfrog steps, while prioritized-buffer FAB uses L = 8 updates per AIS sampling step.The Many Well experiments use 10 RealNVP layers and train most methods for 10^10 flow evaluations; CRAFT uses two configurations, with the spline-flow configuration using more temperatures and HMC steps.Source: §E.1 · pp. 27–28, §E.3 · p. 31

E.4 FURTHER RESULTS

  • FAB with a replay buffer and both CRAFT configurations sample across the Many Well marginal modes, whereas the KL(q∥p) baseline misses several modes.Across further method plots, FAB-based methods and maximum-likelihood training place samples at each displayed mode, while other methods fail to do so.Source: §E.1 · p. 28, Figure 11 · p. 29, Figure 12 · p. 30

E.1 DESCRIPTION AND RESULTS

  • FAB with a replay buffer and maximum-likelihood training were the two best Many Well methods, leading on ESS and average log-likelihood while minimizing forward KL and normalizing-constant MAE.Combining trained FAB flows with AIS targeting p gave ESS values of 89.9% with a buffer and 13.6% without; log-likelihoods on target samples and mode test points were -27.4 and -20.9, close to FAB with a buffer.Source: §E.1 · p. 28, Table 6 · p. 28

F.1 SETUP

  • For alanine dipeptide, the setup transforms Cartesian coordinates into translation- and rotation-invariant internal coordinates and normalizes them using the minimum-energy position rather than MD samples.Bond lengths, bond angles, and dihedral angles use fixed scales; freely rotating dihedrals are treated as periodic coordinates.Source: §F.1 · p. 35
  • MD dataset generation used 100 nodes for roughly 15.7 hours each, totaling about 113 kCPUh.The simulations ran on Intel Xeon IceLake-SP 8360Y servers with 72 cores and 256 GB RAM.Source: §F.1 · p. 37
  • Flow experiments used around 1.02 kGPUh; although ML training was faster than FAB, generating its MD training data required an additional 9.4 kCPUh.That data-generation cost would take 131 hours on one server; training runtimes are listed in Table 9, but the supplied passage gives no individual values.Source: §F.1 · p. 37, Table 9 · p. 37
  • Figure 17 shows selected models’ φ marginal distributions on a normal density scale, using the same data as Figure 3b’s log-scale plot.Source: Figure 17 · p. 38
  • For a model trained on both alanine dipeptide chiral forms, the dihedral-angle distributions of the L- and D-forms are mirror images.Figure 18 shows the combined density and the samples separated into the two forms.Source: Figure 18 · p. 38

F.2 FURTHER RESULTS

Further results show that FAB with a replay buffer outperforms baselines when AIS is applied to models that already approximate the target fairly well. The experiments also find that α = 2 performs best across nearly all metrics, and that modeling both chiral forms yields a mixture of their densities.

F.2 FURTHER RESULTS

  • Training on both chiral forms yields equal shares of L- and D-form samples, whose dihedral-angle marginals are mirror images.The resulting test-set log-likelihood is 210.80, 0.74 below the L-form-only model, close to log(2) ≈ 0.69.Source: §F.2 · p. 38
  • AIS improves models that already approximate the target fairly well, and FAB with a replay buffer outperforms the baselines.Table 10 reports ESS and KL divergence for Ramachandran plots, with and without reweighting.Source: §F.2 · pp. 38–39, Table 10 · p. 39
  • With or without a replay buffer, α = 2 outperforms other tested α values in almost all performance metrics.Tables 11 and 12 report ESS, test-set log-likelihood, and Ramachandran-plot KL divergence, averaged over three runs.Source: §F.2 · p. 39, Table 11 · p. 39, Table 12 · p. 39
  • Ramachandran plots are shown for the test set and all trained models in the first run, including samples drawn via AIS.The model plots include flows trained with α = 2, KL, maximum likelihood, and FAB with or without a replay buffer, as well as an SNF.Source: §F.2 · p. 39, Figure 19 · p. 40, Figure 20 · p. 40

Related papers

Loading 2208.01893v3…