Bootstrap Your Flow
Abstract
Normalizing flows are flexible, parameterized distributions that can be used to approximate expectations from intractable distributions via importance sampling.
However, current flow-based approaches are limited on challenging targets where they either suffer from mode seeking behaviour or high variance in the training loss, or rely on samples from the target distribution, which may not be available.
To address these challenges, we combine flows with annealed importance sampling (AIS), while using the α -divergence as our objective, in a novel training procedure, FAB (Flow AIS Bootstrap).
Thereby, the flow and AIS improve each other in a bootstrapping manner.
We demonstrate that FAB can be used to produce accurate approximations to complex target distributions, including Boltzmann distributions, in problems where previous flow-based methods fail. We provide code for our experiments at https://github.com/lollcat/FAB-2021 .
1. Introduction
Estimating expectations with respect to target distributions that cannot be sampled from is a challenging task with many real world applications, such as estimating equilibrium properties of physical systems governed by the Boltzmann distribution (Leli`evre et al., 2010).
Boltzmann generators (No´e et al., 2019), which use normalizing flows to approximate the Boltzmann distribution, are a recent approach with growing interest (Dibak et al., 2020;
K¨ohler et al., 2021).
For challenging problems, current approaches to training Boltzmann generators rely partly on samples from the target for training by the maximum likelihood Wu et al.
Such samples are obtained through expensive Molecular Dynamics simulations (Leimkuhler and Matthews, 2015).
Although flows can be trained without samples from the target, current methods for this suffer from either being mode seeking or high variance in the loss, which leads to inferior performance on challenging problems (Stimper et al., 2021).
To address these challenges, we propose using the α -divergence with α = 2 as the training objective, which is mass covering, and employ annealed importance sampling (AIS) to bring the samples from the flow model closer to the target, reducing variance in the objective.
In our experiments, we apply our method, F low A IS B ootstrap ( FAB ), to a 2D Gaussian mixture distribution as well as to the “Many Well” problem and show that it outperforms competing learning algorithms.
Although we focus on a toy problem and the Boltzmann distribution, our approach for training flows to approximate intractable distributions is not specific to these problems, and can be applied to any target distribution for which an un-normalized probability density function is defined.
One key application of interest that we leave for future work is approximate inference over Bayesian posterior distributions (for example the posterior of a
Bayesian Neural Network).
In the context of Bayesian inference, in terms of computational complexity, FAB lies between Markov chain Monte Carlo, the computationally expensive gold-standard for Bayesian inference, and Bayes-by-backprop (Blundell et al., 2015) which provides a computationally cheaper alternative, but suffers from worse performance.
2. Background Normalizing flows
Given a random variable z with distribution q ( z ) a normalizing flow (Rezende and Mohamed, 2015; Tabak and Vanden-Eijnden, 2010) uses an invertible map
F : R d → R d to transform z yielding the random variable x = F ( z ) with the distribution
[Equation]
If we parameterize F we can use this as a model to approximate a target distribution p .
If the target density p ( x ) is available, the flow is usually trained by minimizing the KL divergence, which is estimated via Monte Carlo with samples from the flow model.
Alternatively, we could use the α -divergence (Zhu and Rohwer, 1995), which is defined by
[Equation]
as an objective (Campbell et al., 2021; M¨uller et al., 2019).
If α = 2, minimizing it corresponds to minimizing the variance of the importance weights w IS ( x ) = p ( x ) q ( x ) .
In contrast to
KL divergence, which is mode seeking, D α =2 ( p ∥ q ) is mass covering which is more desirable when approximating multimodal targets.
In this case, the α -divergence can be rewritten as
[Equation]
This objective can be estimated either with samples from p ( x ) or q ( x ).
Since the integral is dominated by regions with high p ( x ) and low q ( x ), the estimate will exhibit higher variance if we sample from q than if we sample from p .
Annealed importance sampling
AIS begins by sampling from an initial proposal distribution x 0 ∼ p
= q , being the flow in our case, and then transitions via MCMC through a sequence of intermediate distributions, p 1 to p N − 1 , to produce samples x N −closer to the target distribution p N = p (Neal, 2001).
Each transition T j is a Markov chain that leaves the intermediate distribution p j invariant.
AIS conventionally returns the importance weights for the samples, which are in the form
[Equation]
where x ( i ) =
x N − 1 and we indicate that the probability density functions may be unnormalized with ˜ p .
They exhibit variance reduction compared to their counterparts w IS ( x ).
Hamiltonian Monte Carlo (HMC) provides a suitable choice of transition operator for challenging problems (Neal, 1995).
3. Normalizing Flow Annealed Importance Sampling Bootstrap
FAB, defined in Algorithm 1, uses log( D α =2 ( p ∥ q )) as a training objective.
Furthermore, we introduce AIS into the training loop, improving the gradient estimator for minimising
D α =2 ( p ∥ q ) by writing the loss function to train the flow as an expectation over p ( x ), and estimating it with the samples and importance weights generated by AIS with the flow as the initial distribution.
If we plug in (4) in (3), compute the expectation over p ( x ) through
AIS and use Jensen’s inequality, we obtain our loss
[Equation]
where θ are the flow’s parameter and ¯ w AIS and ¯ x AIS have blocked gradients 1 .
This method obtains the benefit of bootstrapping, where AIS is used to improve the flow by improved estimation of the loss’s gradient, which is used to update the flow.
This in turn improves
AIS by improving its initial distribution.
The effective sample size of the trained flow model, which may be limited e.g. through expressiveness constraints of the specific flow, can also be improved by using AIS after training.
Set target p
Initialise proposal q parameterized by θ for iteration = 1, M do
Sample batch x (1: L ) q , log q ( x (1: L ) q ) from q
Generate batch x (1: L ) AIS , log w (1: L ) AIS from AIS seeded with x (1: L ) q , log q ( x (1: L ) q )
Calculate FAB loss: L ( θ )
Perform gradient descent on L ( θ ) end
4. Experiments 4.1. Mixture of Gaussians Problem
We begin with a simple two dimensional mixture of Gaussians target distribution.
To estimate expectations E p ( x ) [ f ( x )] with our proposal distribution, we set f ( x ) to be the toy quadratic function f ( x ) = a T ( x − 2 b ) + 2 ( x − 2 b ) T C ( x − 2 b ) , where the elements of a , b and C are sampled from a unit Gaussian and then fixed for the problem.
This allows us to inspect the bias and variance of estimates of the expectation of this toy function, as a further measure of performance.
We compare FAB 2 to flows trained through minimisation of KLD 3 , as well as Stochastic Normalizing Flows (SNFs)
(Wu et al., 2020).
Like FAB,
SNFs also combine flows with stochastic sampling methods such as MCMC, but instead focus on improving the flow’s expressive power.
The SNF is trained with the KL divergence as well.
For all models we choose real NVP (Dinh et al., 2017) as the flow architecture.
In Figure 1 we see that FAB allows us to train a flow that captures the shape of the target distribution well, while the flow and SNF trained with KLD both fail, capturing only a subset of the modes.
Table 1 4 shows that with FAB the trained flow may be used for accurate computation of expectations with respect to the target, while the alternative approaches yield highly biased estimates.
[Table]
4.2. The Many Well Problem
For a more challenging problem, we test FAB against a flow trained by KLD on a 16 dimensional “Many Well” problem, which we create by repeating the Double Well Boltzmann distribution from (No´e et al., 2019) 8 times.
We create a hand-crafted test-set for this problem where we place a point on each of the 256 modes of the target.
In Table 2 we see that FAB allows us to train a model that has a far superior test set log-likelihood and ESS than training a flow with KLD.
In Figure 2 where we visualise a subset of the marginal distributions for pairs of dimensions, we see that the flow trained with FAB has captured the shape of the target well, while the flow trained with KLD fits only a subset of the modes.
[Table]
5. Discussion
In SNFs, MCMC is combined with normalizing flows by introducing sampling layers between the standard flow layers to improve the expressiveness of the flow (Nielsen et al., 2020; Wu et al., 2020).
They are usually trained with samples from the target and perform poorly when trained exclusively with samples from the flow, see section 4.1.
Instead, FAB focuses on improving the training procedure.
However, the two contributions are orthogonal, so SNFs could be trained with FAB as well.
Similarly, FAB works with other flow architectures (Chen et al., 2019; Grci´c et al., 2021) and base distributions (Papamakarios et al., 2017; Stimper et al., 2021) than those we used.
In Ding and Freedman (2019) AIS is applied in the context of variational inference to improve estimation of the marginal likelihood gradient for varitational autoencoders’ decoder.
This approach may be extended by FAB to also improve the encoder training, through minimising the FAB loss with respect to the parameters of the encoder, with the latent posterior as the target distribution.
This may be advantageous relative to (Ding and Freedman, 2019), as well as MCMC based approaches like (Hoffman, 2017) for training the decoder, because if the flow-based encoder can learn a good approximation of the latent posterior, then this alleviates the requirement to run MCMC for long in order to obtain good samples for the decoder. In neural transport MCMC (Hoffman et al., 2019), a flow is trained via KL divergence minimisation with a target distribution, and then is used to perform MCMC with the target in the latent space of the flow.
This improves the geometry in which MCMC takes place, allowing for better mixing.
Neural transport MCMC relies on the flow capturing the geometry of the target distribution and can worsen sampling in the tails of the target distribution if the flow fails to capture the target well.
Thus by replacing the KL divergence training objective with FAB training, our approach may be useful for improving the robustness of neural transport MCMC. We proposed FAB, a novel method of combing flows with AIS in a training procedure that allows them to improve each other in a bootstrapping manner.
For future work we hope to (1) identify how important each component of FAB is, such as the number of intermediate AIS distributions used and the form of the loss (2) perform a rigorous benchmarking of FAB’s performance and computational complexity relative to alternative approaches such as MCMC, (3) explore the aforementioned connections of FAB to other recent advances in literature, and (4) scale up FAB and apply it to to more challenging real world problems, for example Boltzmann distributions of complex molecules. Acknowledgments
GNCS and JMHL acknowledge support from a Turing AI Fellowship under grant EP/V023756/1. This work has been performed using resources operated by the University of Cambridge Research Computing Service, which is funded by the EPSRC (capital grant EP/P020259/1) and DiRAC funding from the STFC ( http://www.dirac.ac.uk/ ). LIM acknowledges support from the Cambridge Trust, the Skype Foundation, and the Oppenheimer Memorial Trust.
Appendix A. FAB derivation
FAB aims to minimise
[Equation]
We obtain a gradient estimator by differentiating Equation 6 with respect to the parameters θ of the flow model
[Equation]
where we have used the fact that since θ is independent to samples from p ( x ), we can move ∇ θ inside the expectation.
If we set f ( x ) = ∇ θ p ( x ) q θ ( x ) , we see that Equation 7 is in the form
E p ( x ) [ f ( x )], so we can estimate it with AIS.
With this goal in mind, during the training loop, we generate a batch of importance weights w (1: L ) AIS , and samples x (1: L ) AIS using AIS, with p ( x ) as the target distribution and q θ ( x ) as the proposal distribution.
We can then obtain an importance weighted estimate of the above gradient operator
[Equation]
We note that in Equation 7, q θ ( x ) is only differentiated through θ ’s contribution to the probability density function, and not 5 via ∇ θ x .
Therefore, in Equation 8 we take care to block the gradient of x (1: L ) AIS with respect to θ .
We denote the blocked gradients with ¯ x ( l ) AIS .
Thus, we can train the proposal by minimising the surrogate “loss function”
[Equation]
taking care 6 to block the gradient of ¯ w ( [Equation] L ) AIS and ¯ x (1: L ) AIS with respect to θ .
To obtain a good loss function for training it is beneficial to instead seek to write the surrogate loss (Equation 9) in terms of log probabilities and log importance weights, because inside the expectation the importance weights and fractions of probabilities will have high variance.
To do this we can re-write the surrogate loss as
[Equation]
Finally, we minimise the below loss, which by Jensen’s inequality is an upper bound of log O ( θ ),
[Equation]
We can now work with log probabilities and log importance weights, and use the “logsumexp” trick to obtain a numerically stable estimate.
Equation 11 is the exact surrogate loss implemented in practice for training.