Flow Annealed Importance Sampling Bootstrap
Anonymous Author(s)
AffiliationAddress email
Abstract
Normalizing flows are tractable density models that can approximate complicatedtarget 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 by expensive MCMC simulations, or usestochastic losses that have high variance.
To avoid these problems, we augmentflows with annealed importance sampling (AIS) and minimize the mass-coveringα -divergence with α = 2 , which minimizes importance weight variance.
Ourmethod, Flow AIS Bootstrap (FAB), uses AIS to generate samples in regionswhere the flow is a poor approximation of the target, facilitating the discoveryof new modes.
We apply FAB to complex multimodal targets and show that wecan approximate them accurately where previous methods fail.
To the best ofour knowledge, we are the first to learn the Boltzmann distribution of the alaninedipeptide molecule using only the unnormalized target density, without access tosamples generated via Molecular Dynamics (MD) simulations: FAB produces better results than training via maximum likelihood on MD samples while using 100times fewer target evaluations.
After reweighting samples, we obtain unbiasedhistograms of dihedral angles that are almost identical to the ground truth.
Introduction
Approximating intractable distributions is a challenging task whose solution has relevance in manyreal-world applications.
A prominent example involves approximating the Boltzmann distribution ofa given molecule.
In this case, the unnormalized density can be obtained by physical modeling and isgiven by e − u ( x ) , where x are the 3D atomic coordinates and u ( · ) returns the dimensionless energy ofthe system.
Drawing independent samples from this distribution is difficult (Leli ` evre et al., 2010).
It is typically done by running expensive Molecular Dynamics (MD) simulations (Leimkuhler &
Matthews, 2015), which yield highly correlated samples and, hence, require long simulation times.
An alternative is given by normalizing flows.
These are tractable density models parameterized byneural networks.
They can generate a batch of independent samples with a single forward pass andany bias in the samples can be eliminated by reweighting via importance sampling.
Flows are called
Boltzmann generators when they approximate Boltzmann distributions (No´e et al., 2019).
Recently,there has been a growing interest in these methods (Dibak et al., 2021; K¨ohler et al., 2021; Liu et al.
,2022) as they have the potential to avoid the limitations of MD simulations.
Most current approachesto train Boltzmann generators rely on MD samples since these are required for the estimation ofthe flow parameters by maximum likelihood (ML)
(Wu et al., 2020).
Alternatively, flows can betrained without MD samples by minimizing the Kullback–Leibler (KL) divergence with respect tothe target distribution (Wirnsberger et al., 2022).
However, the KL divergence suffers from modeseeking behavior, which severely deteriorates the performance of this approach with multimodaltarget distributions (Stimper et al., 2022).
On the other side, mass-covering objective such as theforward KL divergence suffer from the high variance of the samples from the flow.
To address these challenges, we present a new method for training flows: F low A IS B ootstrap( FAB ).
Our main contributions are as follows:
We propose to use the α -divergence with α = 2 as our training objective, which is mass-coveringand minimizes importance weight variance.
We approximate this objective using annealed importance sampling (AIS) with the flow as the initial distribution and the target set to the minimumvariance distribution for the estimation of the α -divergence via importance sampling.
AIS returnssamples that provide a higher quality training signal than samples from the flow, as it focuses onthe regions that contribute the most to the α -divergence loss.
We reduce the computational cost of our method by introducing a scheme to re-use samples viaa prioritized replay buffer.3.
We apply FAB to a toy 2D Gaussian mixture distribution, the 32 dimensional “Many Well” problem 1 , and the Boltzmann distribution of alanine dipeptide.
In these experiments, we outperformcompeting approaches and, to the best of our knowledge, we are the first to successfully train aBoltzmann generator on alanine dipeptide using only the unnormalized target density.
In particular, we use over 100 times fewer target evaluations than a Boltzmann generator trained with MDsamples while producing a better approximation to the target.
BackgroundNormalizing flows
Given a random variable z with distribution q ( z ) , a normalizing flow (Tabak& Vanden-Eijnden, 2010; Rezende & Mohamed, 2015; Papamakarios et al., 2021) uses an invertiblemap F : R d → R d to transform z yielding the random variable x = F ( z ) with distribution
[Equation]
(1) where J F ( z ) = ∂F/∂ z is the Jacobian of F .
If we parameterize F , we can use the resulting model toapproximate a target distribution p .
To simplify our notation, we will assume that the target densityp ( x ) is normalized, i.e., it integrates to 1, but the methods described here are equally applicablewhen this is not the case.
If samples from the target distribution are available, the flow can betrained via ML.
If only the target density p ( x ) is given, the flow can then be trained by minimizingthe reverse KL divergence 2 between q and p , i.e., KL ( q ∥ p ) =
(cid:82) x q ( x ) log { p ( x ) /q
( x ) } d x , which isestimated via Monte Carlo using samples from q .
Alpha divergence
An alternative to the KL divergence is the α -divergence (Zhu & Rohwer,1995; Minka, 2005; M ¨ uller et al., 2019; Bauer & Mnih, 2021;
Campbell et al., 2021) defined by
[Equation]
The α -divergence is mode-seeking for α smaller than 0 .
5 and mass-covering for α larger than 0 .
5(Minka, 2005), as shown in Figure 1.
When α = 2 , minimizing the α -divergence is equivalentto minimizing the variance of the importance sampling weights w IS ( x ) =
p ( x ) /q
( x ) , which isdesirable if importance sampling will be used to eliminate bias in the samples from q .
Annealed importance sampling
AIS begins by sampling from an initial distribution x 1 ∼ p
=q , given by the flow in our case, and then transitioning via MCMC through a sequence of intermediate distributions, p 1 to p N − 1 , to produce a sample x N closer to the target distribution p N = p
(Neal,2001).
Each transition generates an intermediate sample x j by running a few steps of a Markovchain that is initialized with the previous intermediate sample x j − 1 and that leaves the intermediatedistribution p j − 1 invariant.
Each p j is defined by interpolating between the initial and target logdensities: log p
i ( x ) =
β i log p
[Equation]
x ) +
[Equation] [Equation]
β i ) log p N ( x ) , where 1 = β 0 > β 1 > ... > β N = 0 .
AISprovides an importance weight for the final resulting sample x N given by
[Equation]
These weights exhibit variance reduction compared to their importance sampling counterpartsw IS ( x )
= p ( x ) /q
( x ) (Neal, 2001).
Hamiltonian Monte Carlo (HMC) is a suitable transition operator for implementing AIS in challenging problems (Neal, 1995; Sohl-Dickstein & Culpepper,2012).
Method
[Figure]
3.1 Flow annealed importance sampling bootstrap
[Figure]
FAB trains a flow q to approximate a target p by minimizing D α =2 ( p ∥ q ) , which is estimated with
AIS using q as initial distribution and p 2 /q as target.
The latter is the minimum variance importancesampling distribution for D α =2 ( p ∥ q ) .
FAB performs bootstrapping since it fits the flow q using itsown samples after they have been improved with AIS to fit p 2 /q .
FAB allows us to overcome thechallenges of training a mass-covering flow without access to samples from the target.
Below weprovide a brief derivation of our training loss function and refer to Appendix A for the full derivation.
We consider q to be specified by some parameters θ and write q θ to make this explicit.
We aim totune q θ by minimizing L ( θ ) ∝ D α =2 ( p ∥ q θ ) where L ( θ ) is our loss function.
We can write our lossas an expectation over some distribution g ( x ) by using importance sampling:
[Equation]
We consider setting g ∝ p 2 /q θ which minimizes 3 the variance in the estimation of L ( θ ) .
Samplingdirectly from g ∝ p 2 /q θ is intractable.
Instead, we train the flow using an estimate of the loss basedon samples generated by AIS when targeting g ∝ p 2 /q θ and using q θ as the initial distribution.
These AIS samples have higher quality than those returned by the flow as they occur in regionswhere the integrand in Equation (4) takes high values.
These are regions where p and q have highand low density, respectively.
Another advantage of the AIS samples is that we can use the weightsreturned by AIS to obtain an unbiased estimate of L ( θ ) .
To obtain the gradient of Equation (4) with respect to θ , let us denote f θ ( x )
= p ( x ) 2 /q θ ( x ) .
Withoutloss of generality 4 we assume that (cid:82) f
θ ( x )d x = 1 and then set g ( x ) =
f θ ( x ) .
First, we write thegradient as an expectation over g :
[Equation]
We can then write this as an expectation over the AIS forward pass:
[Equation]
where ¯ x AIS and w AIS are the samples and corresponding importance weights generated by AIS whentargeting g .
We use the bar superscript to denote stopped gradients in the samples generated by
AIS, ¯ x AIS , with respect to θ .
If we also stop the gradients of w AIS , we can then use the surrogateloss function S ( θ ) =
− E AIS [ ¯ w AIS log q
θ
(¯ x AIS )] , which may then be estimated by Monte Carlo anddifferentiated to obtain unbiased estimates of the gradient.
In practice, we found that using the selfnormalised importance weights greatly improved training stability.
We refer to the surrogate loss x )
( x ) | p ( x ) (Kahn & Marshall, 1953; Owen, 2013).
We that this is different from the distribution that minimizes variance for self-normalized importance sampling, which is given by g ( x ) ∝ | f ( x ) − µ | p ( x ) (Hesterberg, 1988; Owen, 2013)
Initialize base q parameterized by θ
Initialize replay buffer to a fixed maximum size for iteration = 1 to K do //
Generate AIS samples and add to buffer
Sample x (1: M ) q from q θ and evaluate log q θ ( x (1: M ) q )
Obtain x (1: M ) AIS and log w (1: M ) AIS using AIS with target p 2 /q θ and seed x (1: M ) q and log q θ ( x (1: M ) q )
Add x (1: M ) AIS , log w (1: M ) AIS and log q θ ( x (1: M ) AIS ) to replay buffer for iteration = 1 to L do //
Sample from buffer and update q θ
Sample x (1: N ) AIS and log q θ old ( x (1: N ) AIS ) from buffer with probability proportional to w (1: N ) AIS
Calculate log w (1: N ) correction = log q
θ old ( x (1: N ) AIS ) − stop-grad(log q θ ( x (1: N ) AIS ))
Update log w (1: N ) AIS and log q θ old ( x 1: N AIS ) in buffer to log w (1: N ) AIS + log w (1: N ) correction and log q θ ( x 1: N AIS )
Calculate loss S ′ ( θ ) = − 1 /N (cid:80) Ni w ( i ) correction log q θ ( x ( i ) AIS )
Perform gradient descent on S ′ ( θ ) to update θ with self-normalised importance weights as S ′ ( θ ) .
Thus, we can use the following estimate of thesurrogate loss function for training:
[Equation]
where ¯ w ( i ) AIS and ¯ x ( i ) AIS are N samples and weights generated by AIS using g
= p 2 /q θ as target distribution.
When evaluating the gradient of Equation (7) with respect to θ , gradients must be stoppedduring the computation of the AIS samples and weights.
In practice, we obtain good performancewith a relatively low number of intermediate AIS distributions, e.g., 1 for the Gaussian mixturemodel problem and 8 for the dipeptide problem, see the following section.
Moreover, we can use
AIS after training with target p to further reduce variance when approximating expectations over p .
Here we have focused on the minimization of D α =2 ( p ∥ q ) using an AIS bootstrapping approach withp 2 /q as target distribution.
However, our approach is general and could be used to minimize otherobjectives and to train other models, such as those that combine flows with stochastic sampling steps(Wu et al., 2020; Arbel et al., 2021; Matthews et al., 2022;
Jing et al., 2022).
We provide furtherdiscussion and examples related to this in Appendix B.
In Appendix C we provide an analysis of the quality of the estimates of the gradient of D α =2 ( p ∥ q )produced by FAB and by importance sampling with samples from q or p .
We focus on the FABgradient in the form from Equation (6), as it is easy to analyze.
First, we show that in a simplescenario where both q and p are 1D Gaussians, the signal-to-noise ratio of FAB with a small numberof AIS distributions is far superior to that of estimating D α =2 ( p ∥ q ) using samples from q or p .
Wealso study in Appendix C the performance of FAB as the dimensionality of the problem grows.
Ina simple scenario, we show that the variance of our objective remains constant if we increase thenumber of AIS distributions by the same factor by which the dimensionality increases.
3.2 Re-using samples through a replay buffer
Although AIS is relatively cheap, it is still significantly more expensive than directly sampling fromthe flow as it requires additional flow and target evaluations.
To speed up computations, we re-use
AIS samples during the flow updates by making use of a prioritized replay buffer analogous to theone in (Mnih et al., 2015; Schaul et al., 2016).
Consider a replay buffer with a set of samples and corresponding AIS weights generated during asingle run of AIS with target g ( x )
=
p ( x ) 2 /q θ ( x ) , where q θ ( x ) is the flow at a point in trainingspecified by θ .
We can approximate the gradient of Equation (7) using
[Equation]
where x 1 , . .
, x N are sampled from the buffer with probability proportional to their AIS weights.
However, if the buffer data points have been generated with a previous value of θ , denoted θ old ,we have to multiply their AIS weights with a correction factor w correction =
q θ old ( x ) /q θ ( x ) beforesampling, which requires to additionally store q θ old ( x ) for each data point in the buffer.
When drawing a minibatch from the buffer, we do this with probability proportional to the old AIS weightsand reweight each sample afterwards so that we do not have to correct every sample in the buffer.
Further, we set a maximum length for the buffer to limit its memory footprint.
The resulting procedure, shown in Algorithm 1, extracts data from the buffer in a prioritized manner:
We sample according to g , which favors points with low q θ and high p .
As q θ is updated to fitsamples from the buffer, it will take higher values on those samples and their weights will graduallybe decreased to encourage drawing alternative samples.
The buffer allows us to re-use old AISsamples and does not require re-evaluating p ( x ) , which could be expensive in some cases.
Experiments
[Figure]
4.1 Mixture of Gaussians in 2D
[Figure]
First, we consider a synthetic problem where p is a mixture of bivariate Gaussians with 40 mixturecomponents.
We give the flow a pathological initialization where samples from q θ concentrate in asmall region of the sampling space, as illustrated in the top left plot in Figure 2.
We compare 1) FAB with a replay buffer as shown in Algorithm 1; 2) FAB without a replay buffer,where we directly optimize Equation (7); 3) a flow model that minimizes KL ( q θ ∥ p ) ; 4) a flowwith a Resampled Base Distribution (RBD)
(Stimper et al., 2022)
that minimizes KL ( q θ ∥ p ) ; 5) a
Stochastic Normalizing Flow (SNF) model (Wu et al., 2020) that also minimizes KL ( q θ ∥ p ) ; 6)
a
Continual Repeated Flow Annealed Transport (CRAFT) model (Matthews et al., 2022) that minimizes a CRAFT specific version of KL divergence specified in Appendix B.1; 7) a flow model thatminimizes D α
=2 ( p ∥ q θ ) estimated using samples from q θ (M¨uller et al., 2019), and 8) a flow trainedby ML using samples from p .
We single out the results from this latter method in our tables with ahorizontal dashed line because these samples might not be available for complicated targets.
All methods use the same Real NVP flow architecture (Dinh et al., 2017).
For the FAB-basedapproaches, we run AIS with a single intermediate distribution ( β = 0 . 5 ) and MCMC transitionsgiven by a single Metropolis-Hastings step, being a Gaussian perturbation and then an accept-rejectstep.
For FAB with prioritized buffer, we perform L = 4 gradient updates to q θ per AIS samplingstep.
For the SNF and CRAFT, we do 1 Metropolis-Hastings step every 3 flow layers.
All models
[Table]
and architecturesused in each algorithm are provided in Appendix D.1.
Figure 2 shows that our two FAB-based methods and the flow trained by ML fit all modes in p .
Bycontrast, the other alternative methods cover only a small subset of the modes.
The reason for thisis that such methods are trained only on samples from q θ and the poor initialization of q θ makes itunlikely that they will ever generate samples from undiscovered modes.
For evaluation of trained models we compute the effective sample size (ESS) obtained when doingimportance sampling with q θ ; the average log-likelihood of q θ on samples from p ; the forward
KL divergence with respect to the target; and the mean absolute error (MAE) in the estimation of 180 E p ( x ) [ f ( x )] by importance sampling with q θ , where f ( x ) is a toy quadratic function specified in
Appendix D. We express the MAE as a percentage of the true expectation to ease interpretability.
Finally, we also report the MAE that is obtained when we do not reweight samples according to theimportance weights.
We provide further details on the evaluation setup in Appendix D.2.
Table 1 shows for each method average results and corresponding standard errors over 3 randomseeds.
FAB with the buffer performs similarly to the benchmark of training the flow by ML.
Bothof these methods are the best performing ones with the highest ESS, the highest log-likelihood onsamples from p and the lowest forward KL divergence and lowest MAE.
FAB without a replay bufferis the next best method while the others perform poorly.
The non FAB/ML methods assign zerodensity to points sampled from undiscovered modes of p , which is represented by writing NaN asforward KL divergence and log-likelihood in the table.
Note that the ESS for the SNF, CRAFT, RBDand the flow trained by minimizing KL ( q θ ∥ p ) are spurious as these methods are missing modes.
4.2 Alanine dipeptide
We now consider the 22 atom molecule alanine dipeptide, shown in Figure 3a, in an implicit solventat a temperature of T = 300 K and aim to approximate its Boltzmann distribution given the 3D
[Figure]
atomic coordinates.
This is a popular benchmark when considering Boltzmann generators (Wu et al.
,2020; Campbell et al., 2021; Dibak et al., 2021; Stimper et al., 2022).
Previous works have useda coordinate transformation to map some but not all Cartesian coordinates to internal coordinates,which are normalized using their mean and standard deviation computed on samples generated by
MD (No´e et al., 2019).
Since we aim to train models without using any data, we replace the mean bythe minimizer of the energy, which can be cheaply estimated through gradient descent with less than100 gradient evaluations.
Similarly, we replace the standard deviations with values reflecting thetypical order of magnitude of each variable.
Furthermore, we represent the molecule with internalcoordinates only, thereby implicitly satisfying the system’s rotational and translational invariance.
We use Neural Spline Flows with 12 rational quadratic spline coupling layers (Durkan et al., 2019).
Dihedral angles of those bonds that can move freely are treated as circular coordinates (Rezendeet al., 2020), while the others are considered as unbound.
The models trained with FAB use 8intermediate distributions.
For FAB with the replay buffer, we do L = 8 gradient updates per AIS forward pass.
Alanine dipeptide is a chiral molecule, meaning that it can exist in two distinct forms(L-form and D-form) that are mirror images of each other, as illustrated in Figure 13.
In nature,we find almost exclusively the L-form which is why only this form is considered in the literature.
During training, we filter the samples generated by our flows and keep only those for the L-form,whereby the flow models learn to only generate this form.
More details are given in Appendix F.1.
[Figure]
To evaluate our models, we generated samples using parallel tempering MD simulations, whichserve as ground truth.
They are split into training and validation sets with 10 6 samples each and atest set with 10 7 samples.
We compare FAB to several baseline methods already mentioned in theprevious section, see Table 2 for the full list.
The SNF method performs 10 Metropolis-Hastingssteps every two layers, meaning a total of 60 additional sampling steps.
All methods are trained for2 .
5 × 10 8 flow evaluations except for the SNF, which uses 6 × 10 7 as it is more expensive in termsof target evaluations.
We compare methods using the ESS of importance sampling weights andthe average log-likelihood on the test set.
Moreover, we generate Ramachandran plots, which arehistograms for the marginal distribution of the dihedral angles ϕ and ψ illustrated in Figure 3a.
Wecomputed their KL divergence to the ground truth with and without reweighting.
Our experimentsare repeated over 3 random seeds and average values and corresponding standard errors are given.
[Table]
Table 2 shows our results.
The flow trained by minimizing D α =2 ( p ∥ q θ ) with samples from q θ didnot converge due to the high gradient variance.
Hence, it performs very poorly in practice, especiallyin terms of ESS.
The models trained by minimizing KL ( q ∥ p ) only approximate a subset of the targetmodes, leading to poor test log-likelihood and KLD values
[Table]
and spurious values for the ESS.
Theflow trained by ML on MD samples obtains very good results in terms of test log-likelihood and
KLD values but struggles to model the dim mode at ϕ ≈ 1 correctly, as shown in Figures 3b and4.
Its ESS is low and, therefore, reweighting worsens performance.
The models trained with FABhave a higher ESS and test log-likelihood than all other methods.
FAB with a buffer obtains lower
KLD values than the flow trained by ML.
When the samples generated by this FAB version arereweighted, the resulting distribution is nearly the same as the ground truth, as illustrated in Figures
3b and 4.
These results show that FAB with a replay buffer outperforms the flow trained by ML on
MD samples while using 100 times fewer evaluations of the target density, as shown in Table 6.
Discussion and Related Work
Stochastic Normalizing Flows (SNFs) combine flows with MCMC methods by introducing samplinglayers between flow layers to improve model expressiveness (Wu et al., 2020; Nielsen et al., 2020).
SNFs have been extended in Continual Repeated Annealed Flow Transport Monte Carlo (CRAFT)(Matthews et al., 2022; Arbel et al., 2021), where flows are combined with Sequential Monte Carlo(SMC).
In CRAFT, flows are used to transport SMC samples between consecutive intermediateannealing distributions, with each flow being trained by minimizing a KL divergence with respect tothe next intermediate distribution.
As the CRAFT model is more general and has more expressivepower than the flow models used in this paper, combining FAB with CRAFT would be a promisingavenue for future work.
We outline a new FAB-flavored version of CRAFT in Appendix B.1.
Within the MCMC/AIS literature, significant work has focused on improving transition kernels(Levy et al., 2018; Gabri ´ e et al., 2022), intermediate distributions (Brekelmans et al., 2020) andthe extended target distribution (Doucet et al., 2022a,b) of AIS.
These techniques are applicable to
FAB for improving the simple AIS procedure used within this paper.
FAB does not differentiatethrough AIS to obtain the gradient with respect to the flow, but doing so as in (Geffner & Domke,2021; Zhang et al., 2021;
Metz et al., 2021; Doucet et al., 2022a) may allow for a lower variancegradient estimate.
FAB does not use gradients of the target distribution when optimizing its lossfunction, although such gradients are used in the sampling process by HMC.
This is in contrast withthe alternative approach of training the flow by minimizing KL ( q θ ∥ p ) , which does use these gradients.
Such gradient information could be included in FAB through force matching (Wang et al.
,2019; K¨ohler et al., 2021; K¨ohler et al., 2022) or the addition of a KL divergence loss term.
Conclusions
We have proposed FAB, a method for training flows to approximate complicated multimodal targetdistributions.
FAB combines α -divergence minimization with α = 2 with an AIS bootstrappingmechanism for improving the samples used for the loss estimate.
By focusing on this divergence, wefavor mass-covering of multimodal distributions and minimize importance weight variance.
Using
AIS, FAB targets the ratio between the squared target density and the flow density, which providesa high quality training signal by focusing on the regions where the flow is a poor approximation ofthe target.
We have also proposed to use a prioritized replay buffer, which reduces the cost of FABand improves performance.
Our experiments show that FAB can produce accurate approximationsof complex multimodal targets without using samples from such distributions, while alternativeapproaches fail.
Remarkably, for the alanine dipeptide, FAB produces better results than trainingthe flow by ML on samples generated via MD simulations while still using 100 fewer evaluationsof the target than the MD simulations.
In future work, we hope to scale up our approach to morechallenging problems, such as the modelling of the Boltzmann distribution of large proteins.
A Derivation of the loss
We consider the general case of training a parameterized probability distribution q θ to minimize aloss function L ( θ )
= (cid:82) f
( x , θ )d x , where f ( x , θ ) ≥ 0 .
Later, we will focus on the specific case of
D α =2 ( p ∥ q θ ) minimization, where f ( x , θ )
= p ( x ) 2 /q θ ( x ) , however, the general case is interestingas
well and simplifies notation.
Let us consider the gradient of L ( θ ) written as an expectation oversome distribution g ( x ) :
[Equation]
We select g ( x ) to be the minimum variance importance sampling distribution given by g ( x ) =
f
( x , θ ) /Z
f where Z f = (cid:82) f ( x , θ )d x is the normalizing constant.
Then, plugging in the identity∇ θ log f ( x , θ ) = ∇ θ f ( x , θ ) /f
( x , θ ) , i.e., applying the log-derivative trick similar to REINFORCE(Williams, 1992), we obtain
[Equation]
We may not generally be able to sample directly from g ( x ) , but we can use AIS to estimate theright part of Equation (10).
To do this, we consider running AIS with p 2 /q θ as the target and q θ asthe initial distribution.
We note that, when the target density is unnormalized, the AIS weights arescaled by the target normalizing constant:
[Equation]
where ˜ p i denotes the unnormalised density for the i -th intermediate AIS distribution and Z i is thecorresponding normalizing constant such that p i
= ˜
p
i /Z
i .
We can then write Equation (10) as an expectation over the AIS forward pass
[Equation]
(12) where ¯ x AIS and w AIS are the samples and corresponding importance weights generated by AIS whentargeting g .
We use the bar superscript to denote stopped gradients of the samples generated by
AIS, ¯ x AIS , with respect to the parameters θ .
Now, returning to the case where we minimize L ( θ ) =
D α =2 ( p ∥ q θ ) , we set f ( x, θ ) = p ( x ) 2 /q θ ( x ) to obtain
[Equation]
where w ( i ) AIS and x ( i ) AIS are samples and weights generated by AIS with p 2 /q θ as target.
Equation (13)provides an unbiased estimate of ∇ θ L ( θ ) .
If we also stop the gradients of w AIS , we can then use thesurrogate loss function
[Equation]
where ∇ θ L ( θ ) = ∇ θ S ( θ ) .
The surrogate loss function may then be estimated using Monte Carlo.
In practice, we found that replacing the unnormalized weights in Equation (14) with the normalizedweights greatly improved training stability.
We refer to the surrogate loss with normalized importance weights as S ′ ( θ ) .
To normalize the weights we divide them by ¯ Z f = ¯ L ( θ ) =
E AIS [ ¯ w AIS ] such that (cid:82) ¯ w AIS d x = 1 , where the bar superscripts denotes stopped gradients.
The relationship between
S ( θ ) and S ′ ( θ ) is therefore S
′ ( θ ) =
S ( θ ) / ¯ L ( θ ) .
The gradient of S ′ ( θ ) has the same direction asthe gradient of the original surrogate loss, S ( θ ) , but a has different magnitude.
Using S ′ ( θ ) insteadof S ( θ ) improves training stability by removing the effect of large fluctuations in the magnitude of
L ( θ ) =
D α =2 ( p ∥ q θ ) without changing the direction of the gradient as training proceeds.
Thus, we use the following estimate of the surrogate loss function for training:
[Equation]
(15) where ¯ w ( i ) AIS and ¯ x ( i ) AIS are the samples and importance weights generated by AIS but evaluated inpractice using stopped gradients when computing the gradient of Equation (15).
The use of selfnormalization in the loss function introduces bias for finite N for the estimation of S ′ ( θ ) .
We useq θ as the initial distribution for AIS, and a relatively small number of intermediate distributions toprevent the AIS forward pass from becoming too computationally expensive.
B Variations of FAB
As mentioned above, in this paper we focus on the minimization of D α
=2 ( p ∥ q θ ) as estimated withour AIS bootstrap approach targeting p 2 /q θ .
However, the general approach of improving gradientestimation through the addition of the AIS bootstrap process may be applied in other settings.
Firstly,in Appendix A above, the loss function in Equation (12) is written in a general manner and couldtherefore be used for any f ( x , θ ) ≥ 0 and not only D α =2 ( p ∥ q θ ) .
We could simply plug in into
Equation (12) other divergence measures in the form L ( θ ) = (cid:82) f ( x , θ )d x and satisfying f ( x , θ )
≥0 .
We can also apply the proposed approach to other types of models.
For example, in the sectionbelow we show how we can obtain a FAB flavored version of Continual Repeated Annealed Flow
Transport Monte Carlo (CRAFT) (Matthews et al., 2022).
B.1 FAB applied to CRAFT
Continual Repeated Annealed Flow Transport Monte Carlo (CRAFT) is an extension of SNFs (Wuet al., 2020; Nielsen et al., 2020) proposed by Matthews et al.
(2022) which combines normalizingflows with sequential Monte Carlo (SMC).
Specifically, flows are used to transport samples betweenconsecutive annealing distributions in combination with SMC.
CRAFT trains each of the flows byminimizing the reverse KL divergence with respect to the next annealing distribution and gradientsare estimated using the samples generated by the flow transport step.
FAB could be combined with the CRAFT model.
For example, the reverse KL divergence could bereplaced with a mass-covering divergence such as the α -divergence with α = 2 .
Furthermore, wecould improve the estimation of gradients by targeting with AIS the minimum variance distributionfor importance sampling instead of just the next annealing distribution.
A FAB style version of the CRAFT algorithm is described in algorithms 2 and 3.
Each flow transportstep is now trained by minimizing the α -divergence with α = 2 and using samples generated by thenext immediate MCMC step ahead of the flow in the SMC process.
Below, we briefly introduce
CRAFT with an emphasis on the loss function used.
Next, we describe how FAB may be used toimprove training of the flow transport steps.
We follow the notation from the CRAFT paper exactlyand use their pseudo code as a basis for our proposed algorithm indicating changes clearly.
We referto (Matthews et al., 2022) for further details on the CRAFT algorithm.
B.1.1 Continual Repeated Annealed Flow Transport Monte Carlo
As in FAB, the aim of CRAFT is to approximate an intractable target distribution that we cannotsample from and whose density can only be evaluated up to a normalizing constant.
This targetdistribution is denoted π K ( x ) (equivalent to p ( x ) in our notation).
In CRAFT, SMC is run withinterleaved flow transport steps through a sequence of annealed distributions ( π k ( x ))
Kk =0 , each withnormalization constant Z k .
The base distribution π 0 is a tractable distribution (e.g., a Gaussian)from which we can sample.
Similarly to AIS, ( π k ( x ))
K − 1 k =1 are defined by interpolating between base and target log-densities, where the target density may be unnormalized.
The SMC process in
CRAFT begins by sampling from the base distribution X
i 0 ∼ π 0 .
Then, for each distribution fromk = 1 to k = K , a flow T k is trained to transport samples from π k − 1 ( x ) to π k ( x ) .
Additionally,at each step from k = 1 to k = K , CRAFT utilises resampling and MCMC to bring the samplescloser to π k ( x ) .
Similarly to AIS, the CRAFT algorithm returns a set of points and normalized importance weights( X iK , W iK ) Ni =1 , which may be used for approximating expectations with respect to the target.
Eachpoint X
ik has an associated normalized importance weight W iK for importance sampling with respectto the intermediate target distribution π k ( x ) .
We refer back to the CRAFT paper, and to the pseudocode in Algorithm 2 for how these importance weights are calculated.
CRAFT minimizes the following training objective:
[Equation]
where # denotes the push forward between distributions.
The above objective trains each flowtransport step T k to minimize the KL divergence between T # k π k − 1 , i.e., the distribution of outputsof the flow when given as input samples from π k − 1 , and the next intermediate distribution π k .
Thegradient estimate used to train each flow transport step T k is given by
[Equation]
(17) where γ k ( x ) ∝ π k ( x ) .
The flow is trained by passing it samples
X k − 1 from the previous SMC step,computing the corresponding output samples from the flow T k ( X k − 1 ) and using these to estimatethe gradient of KL [ T # k π k −
|| π k ] .
The normalized importance weight in the loss W ik − 1 account forthe fact that the samples X ik − 1 passed from the previous step in the SMC forward pass come froman approximation to π k − 1 .
B.1.2 FAB-CRAFT
We now propose a FAB flavored version of CRAFT.
First, we re-introduce some notation from ourpaper: We use q to denote the initial distribution used in AIS and p to denote the target distributionthat we wish to approximate.
Recall that q is trained to fit p .
All other notation in this section followsthe CRAFT paper’s notation.
In our FAB-CRAFT method, we use the MCMC samples following each flow transport step in
CRAFT to update the flow to minimize D α =2 ( p ∥ q ) , where q = T # k π k − 1 and p = π k .
To do thiswith minimal changes to the original CRAFT algorithm, we make the observation that samplingfrom the initial distribution q
and then running MCMC targeting p is equivalent to running AIStargeting p 2 /q with 1 intermediate distribution at β = 0 .
Thus, the samples generated by the
MCMC steps following each flow transport step in CRAFT can be repurposed for an AIS bootstrapestimate of the flow training loss.
For training, the only adjustment to the SMC forward pass of
CRAFT is then to move the resampling step to occur after each MCMC step, where previously itoccurred after each flow transport step.
At inference time the original CRAFT algorithm can be runwith the flows trained with our method in its exact original form.
We describe this in more detailbelow and provide pseudo code in Algorithm 2 and 3.
We begin by deriving the AIS importance weights when targeting p 2 /q with 1 intermediate distribution and setting β = 0 .
For only 1 intermediate distribution, the AIS weights are given by
[Equation]
As before, we set the intermediate distributions as interpolations between the base and the target:log p i ( x ) =
β i log p
( x ) + (1 − β i ) log p N ( x ) .
Now, if we set β 1 = 0 .
5 , then plugging in p
= q ,p 1 = q 0 .
5 ( p 2 /q ) 0 .
5 = p and p 2 =
p 2 /q , we obtain the following AIS weights:
[Equation]
Recall that in CRAFT we set q = T # k π k − 1 and p = π k .
AIS is then run by first sampling x 1 fromq , which is done in practice by setting Y ik ← T k
( X ik − 1 ) where X ik − 1 ∼ π k − 1 and then generatingx 2 from x 1 by MCMC, which is done in practice by setting X ik ∼ K k
( Y ik ) , where K k is an MCMCtransition kernel that leaves π k invariant.
The importance weights of X ik with respect to the AIStarget p 2 /q are then given by
[Equation]
i
[Equation] [Equation]
Y
[Equation]
Using the normalized importance weights W i AIS ,k = w i AIS ,k / (cid:80) Nj =1 w j AIS ,k , we can calculate the
FAB gradient estimate given by
[Equation]
This assumes that the samples passed to the flow are from the distribution π k − 1 .
However, inpractice, these samples are passed from the previous SMC step which is an approximation to π k − 1 .
Similarly as in the original CRAFT loss, see Equation 17, we can correct for this by instead using
[Equation]
where W ik − 1 accounts for X ik − 1 coming from an approximation to π k − 1 .
Calculating the normalized AIS weights requires all the samples from the flow to be passed to the
MCMC step.
Because of this, we move the SMC resampling step to take place after the MCMC stepinstead of just after the flow transport step, see Algorithm 2.
Note that the weights W k for resampling
X k are equal to the weights for resampling the corresponding flow outputs Y k that generate suchsamples.
This result is due to the MCMC kernel K k leaving π k invariant.
The resulting FAB flavorof CRAFT is shown in algorithms 2 and 3.
Some final remarks : our goal has been to create a FAB flavored version of CRAFT while keepingthe algorithm as similar to the original version as possible.
However, in practice, it would be betterto make further changes.
For example, using a prioritized replay buffer would significantly decreasethe computational requirements of the algorithm.
Furthermore, for updating each flow, it may alsobe beneficial to consider samples across the whole chain of intermediate distributions, instead ofusing only samples from the local MCMC step immediately following the flow.
C Analysis of FAB
C.1 Gradient estimation performance
We first analyze the quality of the noisy gradients provided by the proposed AIS bootstrap method.
For this, we consider a toy problem in which q θ and p are unit variance 1D Gaussians with means 0 .
5and − 0 . 5 , respectively, as shown in Figure 5a.
We estimate the gradient of D α =2 ( p ∥ q θ ) with respectto the mean of q θ and compare different methods: first, importance sampling (IS) with samples fromq θ ; second, IS with samples from p ; third, AIS with p as target and q θ as initial distribution; andfourth, our proposed method using AIS with p 2 /q θ as target and q θ as initial distribution.
For AISwe use 3 intermediate distributions and, as transition operator, HMC with 5 leapfrog steps andresampling of momentum variables once per intermediate AIS distribution.
Figure 5b shows the Signal-to-noise ratio (SNR) for the different gradient estimators as a functionof the number of samples used.
AIS bootstrap is clearly the best method.
IS with q θ performs verypoorly and it is outperformed by both IS with p and AIS targeting p , with these two latter techniquesperforming similarly but way worse than AIS bootstrap.
Figure 5c shows that the quality of theproposed method increases fast as the number of intermediate AIS distributions grows, with IS withsamples from p being outperformed quite early in the plot while still using a rather small number ofdistributions.
It is important to note, however, that in more challenging problems, it is unlikely thatour AIS bootstrap method will outperform IS with samples from p , especially early in training whenq
θ is a poor approximation to p .
where p = π
=
# k
− setting
=
5 we simply run the original CRAFT
MCMC transition kernel with p as a target.
[Equation]
objective ˆ
[Equation]
T − 1 θ k
( X ik ) (cid:3) 12: if ESS Nk ≤ NA then 13:
Resample N particles denoted abusively also X ik according to the weights W ik , then set
W ik = 1 N .
14: end if 15: Return (cid:16) π Nk , Z Nk , ˆ h k (cid:17) .
C.2 Scaling FAB to higher dimensions
Now we consider how the performance of FAB is affected by an increasing problem dimensionality.
To investigate this, we analyse a simple scenario where we assume p and q to be factorised, witheach marginal of q having its own separate parameters (no parameter sharing between dimensions).
We acknowledge that these are strong simplifying assumptions and leave a more general analysis tofuture work.
Given these assumptions, we show the following: 1) The variance in the estimates ofthe gradient of D α =2 by importance sampling with samples from q and p increases exponentiallywith respect to the dimensionality of the problem.
The variance in the corresponding estimatesobtained with FAB can remain constant if the number of AIS intermediate distributions increaseslinearly with the dimensionality of the problem.
C.2.1 Theoretical analysis on factorized p and q
We consider the problem of estimating the gradient of D α =2 where both p and q are factorizeddistributions: they are equal to the product of their marginals.
To further simplify the analysis, weconsider the gradient with respect to the parameters of the j -th dimension of q and assume that thereis no parameter sharing across dimensions of q .
In this case, D α =2 is given by
[Equation]
[Figure]
where p d and q d are the marginal distributions for dimension d .
The gradient of this quantity withrespect to the parameters θ j of the j -th marginal of q is given by,
[Equation]
[Equation]
We are interested in studying how the variance in our estimate of ∇ θ j µ scales with D .
Equation (24)shows that each additional dimension d ̸ = j adds an extra factor µ d
= (cid:82) p d ( x d ) 2 /q d ( x d ) d x d in thegradient expression.
To eliminate the effect of this change in the gradient as D increases, we focusour analysis on Var [ ∇ θ j ¯ µ/ (cid:81)
Dd ̸ =
j µ d ] , where ¯ µ is an estimate of µ and (cid:81)
Dd ̸ =
j µ d is a normalizationfactor that cancels the effect of the additional dimensions.
Importance Sampling with q : We consider first how increasing D affects the variance in theestimation of the gradient of D α =2 by importance sampling with samples from q .
The importancesampling estimate of D α =2 with N samples from q is given by
[Equation]
[Equation]
where w ( x n ) = p ( x n ) /q
( x n ) and x n ∼ q .
Now, let us define w d ( x n,d ) =
p d ( x n,d ) /q
d ( x n,d ) ,where x n,d is the d -th entry in x n .
We then obtain
[Equation]
The variance of this estimate, after dividing by the aforementioned normalization factor, is given by
N
[Equation] [Equation]
Since Var q d (cid:2) w d ( x d,n ) 2 (cid:3) =
E q d (cid:2) w d ( x d,n ) 4 (cid:3) − µ 2 d > 0 , we have that E q d (cid:2) w d ( x d,n ) 4 (cid:3) /µ
2 d ≥ 1 .
Thus, the first factor in Equation (27) is multiplied in this equation by D − 1 factors all larger than1, which implies that the variance of this estimator increases exponentially as a function of D .
Thisis a well-known problem of importance sampling.
Importance Sampling with p : Interestingly, we get a similar result when estimating D α =2 byimportance sampling with samples from p .
The estimate for the gradient of D α =2 with respect tothe j -th dimension of q is now
[Equation]
where x n,d ∼ p d .
The variance of this gradient estimate, after dividing by the normalization factor,is given by
[Equation]
Since Var p d [ w d ( x d,n )] =
E p d (cid:2) w d ( x d,n ) 2 (cid:3) − µ 2 d > 0 , we have that E p d (cid:2) w d ( x d,n ) 2 (cid:3) /µ
2 d ≥ 1 .
Thus,the first factor in Equation (29) is again multiplied by D − 1 factors all larger than 1, which impliesthat the variance of this estimator increases exponentially as a function of D , albeit at a lower ratethan in the case of importance sampling with q .
This implies that, even with access to ground truthsamples from p , the number N of samples required to keep the variance of the gradient estimates of
D α =2 constant grows exponentially as a function of the problem dimensionality D .
FAB :
We now apply the same type of analysis to the estimates of the gradient given by FAB.
Weconsider the FAB gradient estimate from Equation (13), which is equal in expectation to ∇ θ µ .
Thisestimate relies on the raw importance weights from AIS rather than the self-normalized importanceweights, which makes it easier to analyze.
The FAB estimate of the gradient of D α =2 with respectto the parameters of the j -th marginal of q is given by
[Equation]
where ¯ w are the importance weights from AIS with stopped gradients and we have decomposedthe AIS weights into the contributions from each dimension: ¯ w ( x n )
=
(cid:81)
d ¯ w d ( x n,d ) .
Note that,since p and q are factorized, we have that all the intermediate AIS distributions are factorized aswell.
If we assume that the MCMC transition kernels in AIS produce independent samples from theground truth intermediate target distributions, we have that x n, 1 , . . .
, x n,D are independent randomvariables.
The variance of ∇ θ j ¯ µ FAB after dividing by the normalization factor is then given by
N Var
[Equation]
where E AIS denotes the expectation with respect to the AIS forward pass.
The first expectationin the equation above is constant as D increases.
Therefore, we focus on the contributions of the other expectations for the importance weights of dimensions d ̸ = j .
As in the previous cases wherewe used importance sampling with samples from q and p , the variance in Equation (31) will againincrease exponentially with D .
However, under the assumption that the MCMC transitions produceindependent samples from the intermediate AIS distributions, Neal (2001) shows that the variancein the log importance weights of AIS is proportional to D/K where K is the number of intermediate
AIS distributions.
This implies that (cid:81)
Dd ̸ =
j E AIS d [ ¯ w d ( x n,d ) 2 ] /µ
d will remain roughly constant if weincrease the number K of AIS distributions by the same factor as the dimensionality of the problem.
In this case, the variance of ∇ θ j ¯ µ FAB will remain constant as we increase D and the cost of the FABgradient estimator will only increase linearly as D increases.
C.2.2 Empirical analysis on toy problem
We now run an empirical analysis to assess the performance of the FAB gradient estimator as Dincreases.
We consider the case where q and p are both factorized Gaussians with unit marginalvariances and with mean vectors equal to 0 .
5 × 1 D and − 0 .
5 × 1 D , where 1 D is a vector of dimension
D with all of its entries equal to one.
For the FAB gradient estimate, we increase the number of AISdistributions by the same factor as the dimensionality of the problem.
We analyze the SNR in theestimates of the gradient of the mean for the first marginal of q as the dimensionality of the problemincreases.
The AIS transition operators are performed by running a single iteration of HMC, with5 leapfrog steps, with a step size of 0.5.
We found our results to be sensitive to the choice of stepsize.
In practice, we selected this value by trial and error, assessing the quality of the AIS samplesby looking at their empirical histogram, as shown in Figure 7.
Figure 6a shows that, when using importance sampling with samples from p , the log-weight varianceincreases linearly as the number of dimensions increases.
By contrast, this variance remains constantwith FAB.
This is achieved by fixing the number of AIS intermediate distributions to be equal to thenumber of dimensions.
This result is consistent with the analysis from the previous section and withthe results of Neal (2001).
Furthermore, in Figure 6b we see that the SNR remains roughly flat for
FAB (stays within a single order of magnitude), while it quickly decreases for importance samplingwith samples from p .
If the same results were to hold in more complex problems, this would implythat we could safely apply FAB in those settings by linearly increasing the number of intermediate
AIS distributions as the dimensionality of the problem grows.
[Figure]
D Mixture of Gaussians experiments
D.1 Training Setup
[Figure]
All flow models have 15 RealNVP layers (Dinh et al., 2017), with a 2 layer (80 unit layer width)
MLP for the conditioner.
The flow is initialized to the identity transformation, so q θ is initially astandard Gaussian distribution.
Training is performed with a batch size of 128, using the Adamoptimizer (Kingma & Ba, 2015) with a learning rate of 1 × 10 − 4 and we clip the gradient norm to amaximum value of 100.
For training the flow by ML, we draw new samples from the target for eachloss estimation.
The model that uses a RBD uses an acceptance function composed of a residualnetwork with three blocks containing 512 hidden units per layer.
The truncation parameter is setto the common value T = 100 .
For the SNF and CRAFT methods, we do 1 Metropolis-Hastingsstep every three flow layers.
We used a fixed step size of σ = 5 . 0 for the Gaussian perturbation ofthe Metropolis-Hastings step, which is the same as what is used within AIS for FAB.
This meansthat the SNF and CRAFT models has 5 stochastic Metropolis-Hastings steps in total.
The CRAFTmodel uses 6 annealing temperatures with geometric spacing, and a resampling threshold of 0.3.
Weuse the code provided by (Matthews et al., 2022) at https://github.com/deepmind/annealed_flow_transport for training the CRAFT model.
We train all models for 2 × 10 7 flow evaluations.
For each method, we train 3 models, each with a different random seed, and results are reported asaverages over these seeds.
FAB specific details : The batch size for both the AIS forward pass ( M ) and sampling from thebuffer ( N ) is equal to 128.
We run AIS with a single intermediate distribution ( β = 0 . 5 ) and
MCMC transitions are given by a single Metropolis-Hastings step: a Gaussian perturbation andthen an accept-reject step.
We used a fixed step size of σ = 5 . 0 for the Gaussian perturbation.
We initialize the buffer with 1280 samples from the initialized flow-AIS combination and use amaximum buffer length of 12800.
We do not use any clipping when computing w correction .
D.2 Evaluation Setup
For each method, we compute after training the effective sample size (ESS) obtained when doingimportance sampling with q θ ; the average log-likelihood of q θ on samples from p ; the forward
KL divergence with respect to the target; and the mean absolute error (MAE) in the estimation of 650 E p ( x ) [ f ( x )] by importance sampling with 1000 samples from q θ , where f ( x ) =
a T ( x − 2 b ) +2 ( x − 2 b ) T C ( x − 2 b ) , with the entries in vectors a and b and matrix C randomly initialized bysampling from a standard Gaussian then kept fixed to such values during all the experiments.
Weexpress the MAE as a percentage of the true expectation to make it easier to interpret.
We also reportthe MAE that is obtained when we do not reweight samples according to the importance weights.
[Table]
The ESS is calculated using 5 × 10 4 samples from q θ .
The MAE is calculated by averaging over100 repetitions.
D.3 Further Results
Figure 8 shows a plot of samples from each trained model on the mixture of Gaussians problem, withthe target contours in the background.
We see that the FAB based methods and the flow trained with
ML cover all the modes in the target distribution.
All the other methods fit a subset of the modes.
The flow trained with D α =2 minimization exhibited highly unstable behavior during training and,thus, is the worst performing model.
E Many Well experiments
E.1 Description and Results
We consider another synthetic problem that is significantly more difficult than the GMM problem:approximating the 32-dimensional “Many Well” distribution given by the product of 16 copies ofthe 2-dimensional Double Well distribution 5 from Wu et al.
No´e et al.
[Equation]
22 + constant , (32) where each copy of the Double Well is evaluated on a different pair of the 32 inputs to the Many
Well.
The original Double Well has two modes as shown in the top-right contour plot in Figure 9.
Therefore, our 32-dimensional Many Well has 2 16 = 65536 modes, one for each possible choice ofmode in each of the 16 copies of the Double Well.
We obtain exact samples from the Many Wellby sampling from each independent copy of the Double Well.
Exact samples from the Double Wellare obtained by sampling from each independent marginal distribution.
The first marginal p ( x 1 ) canbe sampled from exactly using rejection sampling (see Appendix E.2), while the second marginaldistribution p ( x 2 ) can be sampled from directly as it is a (unnormalized) standard Gaussian.
Thesesamples are cheap to produce.
We use them for training a flow by ML as well as for the evaluationof the different methods.
Additionally, we created an artificial test set for evaluation purposes bymanually placing a point on each of the 2 16 modes.
By computing log-likelihoods on this test set, wecan then check if a method is covering the entire target distribution, as any missing mode will resultin very low log-likelihood values.
We can calculate the normalizing constant for each marginalof the Double Well problem via numerical integration (for p ( x 1 ) ) and analytical integration (forp ( x 2 ) ), and use this to obtain the normalizing constant of the Many Well distribution (see Appendix
E.2).
This may then be used to obtain the normalized probability density function of the Many Welldistribution, which is useful for model evaluation.
We can also compare models on how accuratelythey estimate the normalizing constant as the average unnormalized importance weights.
For eachmodel, we report the MAE in the estimation of the Many
Well’s normalizing constant using 1000samples, averaged over 100 runs.
We express this as a percentage of the true value of the normalizingconstant.
[Equation]
[Table]
We compare FAB to the same alternative approaches as in the mixture of Gaussians problem excluding CRAFT 6 and use also the same Real NVP flow architecture as before.
For the the modelthat uses a RBD, we use the same architecture as before, with an acceptance function composed of aresidual network with three blocks containing 512 hidden units per layer.
The truncation parameteris set to the common value T = 100 .
For FAB based methods, we use AIS with 4 intermediatedistributions (linearly spaced) and with a HMC transition operator containing a single iteration with5 leapfrog steps.
For FAB with prioritized buffer we use L = 8 .
The SNF model uses 1 step of HMCwith 5 inner leapfrog steps every 2 layers.
As before, when training the flow by ML we draw newsamples from the target for each loss estimation.
All models are trained for 10 10 flow evaluationswith the number of target evaluations by each method being reported in Table 5.
Further details onthe hyper-parameters and architectures used by each algorithm are provided in Appendix E.3.
Figure 9 shows contour plots for several two-dimensional marginals of the Many Well target.
Eachplot is obtained by scanning two variables that are inputs to different
Double
Well factors in the
Many Well distribution while the other variables are kept fixed to zero.
We also show in this figuresamples generated by FAB with a replay buffer (left) and by the method that tunes q θ by minimizing
KL ( q ∥ p ) (right).
We see that FAB generates samples on each of the contour modes while this isnot the case for the alternative baseline, which misses several modes.
Additional plots for all othermethods can be found in Appendix E.4.
Table 4 shows for each method 1) the ESS when doing importance sampling with q θ ; the averagelog-likelihoods for q θ 2) on samples from the target and 3) on test points placed on the modes of the Many Well distribution; 4) the forward KL divergence with respect to the target; and 5)
the
MAE in the estimation of the May Well normalizing constant.
Average log-likelihoods and ESS arecalculated with 5 × 10 4 samples.
All the results in the table are averages across 3 random seeds.
We see similar results as in the previous experiment: FAB with a buffer performs similarly to thebenchmark of training the flow by ML.
These are the two best performing methods, obtaining thehighest ESS and average log-likelihoods and the lowest forward KL divergence and MAE values.
The next best method is FAB without buffer, while the other methods perform very poorly sincethey miss many modes from the Many Well distribution.
The method that minimizes D α =2 ( p ∥ q θ )as estimated by sampling from q θ diverged early in training and always returned NaN values.
The
ESS for the flow trained by minimizing KL ( q θ ∥ p ) , RBD and the SNF are spurious, as they aremissing modes (see Figure 9 and 11).
After training, we may combine the trained flows with AISto further improve the ESS.
If we run AIS as during training but targeting p instead of p 2 /q , the
ESS is 89 .
9 % and 13 .
6% for the FAB flows trained with and without a buffer, respectively.
Thelog-likelihood of p on samples from p and on the test set with points at the modes are -27.4 and-20.9, respectively.
These values are very close to the ones obtained by FAB with buffer, showingthat this method produces highly accurate approximations to the target distribution.
captures the far better
E.2 Obtaining the normalizing constant and exact samples
In this section, we describe how to obtain the exact normalizing constant and exact samples fromthe Double Well distribution.
These allow us to obtain the normalizing constant and exact samplesfrom the Many Well distribution.
The Double Well log-density is given by
[Equation]
[Equation]
By noting that x 1 and x 2 are independent, we see that their distribution factorises as p ( x 1 , x 2 ) =
p ( x 1 ) p ( x 2 ) .
Thus, the normalizing constant of p ( x 1 , x 2 ) is given by the product of the normalizingconstants of each marginal distribution.
Furthermore, samples from p ( x 1 , x 2 ) may be obtainedby sampling independently from each marginal.
By inspection, we see that the marginal p ( x 2 ) isstandard Gaussian.
Thus, its normalizing constant is given by √ 2 π , and samples from this marginal may be obtained trivially.
The normalizing constant of the second marginal distribution may becalculated via numerical integration Z 1 = 11784 .
We obtain exact samples from p ( x 1 ) byusing rejection sampling, a visual summary of this is provided in Figure 10.
For the rejectionsampling proposal distribution, denoted q , we use a two-component Gaussian mixture distributionwith mixture weights (0.2, 0.8), means ( − 1 . 7 , 1 . 7) and standard deviations equal to 0.5.
For thecomparison function kq ( x 1 ) , we set k =
Z 1 to ensure that kq ( x 1 ) > p ( x 1 ) .
E.3 Setup
All flow models have 10 RealNVP layers (Dinh et al., 2017), with a 2 layer (320 unit layer width)
MLP for the conditioner.
The flow is initialized to the identity transformation and, consequently,q θ is initially a standard Gaussian distribution.
Training is performed with a batch size of 2048and using the Adam optimizer (Kingma & Ba, 2015) with a learning rate of 3 × 10 − 4 .
We clipthe gradient norm to a maximum value of 100.
For the SNF method, we do 10 Metropolis-Hastingssteps every two flow layers.
We train all models for 10 10 flow evaluations.
For each method, we train3 models using different random seeds.
Results are reported as averages over these three models.
FAB specific details : In FAB, the batch sizes for the AIS forward pass ( M ) and sampling fromthe buffer ( N ) are both equal to 2048.
We run AIS using four intermediate distributions with linearspacing.
Each MCMC transition is given by a single Hamiltonian Monte Carlo step consisting of leapfrog steps.
The momentum variable for HMC is sampled from a standard Gaussian and it isnot tuned throughout training.
However, an important parameter to tune in HMC is the step sizeparameter for the leapfrog integrator.
We do tune this parameter for each intermediate distribution.
This is done by using a parametrization of step sizes that includes coefficients that are specific andshared across intermediate distributions.
In more detail, we define the HMC step size for the nth intermediate AIS distribution as ϵ n = ϵ shared + ˆ ϵ n , where we define ϵ shared and ˆ ϵ n as follows:
ϵ shared is a shared parameter across all AIS transitions,
and
is updated at the transition for every
intermediate distribution
This
allows
for faster
adaption of
the
step size
if
the step sizes
for
all
transition kernels are “too big” or “too small”, which is common at the start of training
ˆ ϵ n is a
parameter specific to each n -th intermediate distribution
and is only updated during its specific
transition
This
allows for ϵ n to be
tailored
to the specific n -th
intermediate
distribution
transition.
We found this parameter sharing to improve performance in practice.
The HMC transition kernelfor each intermediate distribution is initialized with a step size of 1.0, where we set ϵ shared = 0 .
1and ˆ ϵ 1: N − 1 = 0 .
The step size is then tuned to target a Metropolis acceptance probability of0.65.
For the transition corresponding to each intermediate distribution, if the average acceptanceprobability across the batch is greater than 0 .
65 , we set ˆ ϵ n = 1 . 05ˆ
ϵ n , and ϵ shared = 1 .
02 ϵ shared .
If the average acceptance probability across a batch is lower than 0 .
65 , we set ˆ ϵ n = ˆ ϵ n / 1 . 05 , andϵ shared =
ϵ shared / 1 . 02 .
In FAB with prioritized buffer, we use a total of L = 8 gradient update steps per AIS sampling step.
We initialise the buffer with 65 , 536 samples from the initialized flow-AIS combination and use amaximum buffer length of 512 , 000 .
We do not use any clipping for w correction .
[Table]
E.4 Further Results
Figure 11 and Figure 12 show contour plots of 2D marginals from the 32 dimensional Many Welltarget.
The contours are for pairs of variables in the first four elements of x belonging to differentcopies of the Double Well distribution.
Each plot is obtained by scanning two variables while theother ones are kept fixed to zero.
These figures also shows samples from each analyzed method.
FAB based methods and the flow trained by ML place samples at each of the modes in the contourplots, while the other methods fail to do so.
F
Alanine dipeptide experiments
F.1 Setup
Coordinate transformation
Boltzmann generators usually do not operate on Cartesian coordinates.
In particular, No´e et al.
(2019) introduced a coordinate transformation whereby a subset ofthe coordinates are mapped to internal coordinates, i.e., bond lengths, bond angles, and dihedral angles.
The internal coordinates are normalized and the respective means and standard deviations forthese coordinates are computed on samples from the target distribution generated with MD.
For theremaining Cartesian coordinates, principal component analysis is applied to the samples and the sixcoordinates with the lowest variance are eliminated.
The rationale behind this is that the Boltzmanndistribution is invariant in six degrees of freedom, i.e., three of translation and three of rotation, andconsequently, the corresponding unnecessary coordinates should be removed.
However, the mapping of vectors onto a fixed set of principal components is generally neither invariant to translationsnor to rotations, and, therefore, the transformed coordinates do not satisfy these invariances.
Whentraining Boltzmann generators with samples, this is not a problem since the flow will learn to generate molecular configurations for a specific rotation or translation, but when only using the targetdistribution to train the flow, the model will spend some of its capacity to sample different translational and rotational states, which is unnecessary since they can easily be sampled independently.
This will harm performance.
Instead, we transform all Cartesian coordinates to internal coordinates, which is a representationinvariant to translations and rotations.
Since we do not want to use MD samples for our model,we use the position with the minimum energy instead as shift and fix values for the scale whennormalizing the coordinates.
The former can be easily estimated with gradient descent using lessthan 100 steps.
As scale parameters, we used 0 .
005 nm for the bond lengths, 0 .
15 rad for the bondangles, and 0 .
2 rad for the dihedral angles.
Coordinates which are treated as circular are not scaled.
Model architecture
We use Neural Spline Flows with rational quadratic splines having 8 binseach.
The parameter mapping is done through coupling (Durkan et al., 2019).
Dihedral angles whichcan freely rotate, e.g., because it is not a double bond, are treated as periodic coordinates (Rezendeet al., 2020).
For these coordinates, we use a uniform base distribution, while we pick a Gaussianfor the other ones.
The flow has 12 layers and the parameter maps are residual networks with oneresidual block, while the two linear layers in each block have 256 hidden units.
The flow layers wereinitialized in a way that they correspond to the identity map.
One model uses a RBD, which has a residual network with two blocks having 512 hidden units perlayer as acceptance function.
The truncation parameter is set to the common value T = 100 .
The models trained with FAB do AIS with 8 intermediate distributions, which are linear interpolations between the flow and target log-densities (Neal, 2001), where the latter one is unnormalized.
We use HMC with 4 Leapfrog steps as the MCMC operator in AIS.
The same procedure is usedwhen we use AIS in the other baseline models, see Table 7, Figure 17, Figure 18, and Figure 21.
The HMC parameters are initialized and tuned using the same procedure as the Many Well problem,see Appendix E.3.
The SNF model does additionally 10 Metropolis-Hastings steps every two flowlayers.
Since this renders sampling from this model already expensive, we do not do AIS with thismodel.
Dataset
Since the energy surface of alanine dipeptide in an implicit solvent has several modes ofdifferent sizes with large energy barriers between them, see Figure 16, a very long MD simulationwould be required to obtain samples that represent the target distribution well.
To get around thisproblem, which is well known in computational physics and chemistry, we do a replica exchange
MD simulation (Mori & Okamoto, 2010), which is a parallel tempering technique (Earl & Deem,2005).
We use 21 replicas starting at a temperature of 300 K and increasing the temperature byan increment of 50 K .
The replicas are exchanged every 200 iterations and use the state at eachmultiple of 1000 time steps as samples.
To reduce the time it takes to generate the data, we runmany of these simulations in parallel with different seeds.
Since the initial condition is always thesame, i.e., the position with minimum energy as it is usually done, we let the system equilibrate for2 × 10 5 iterations and run the simulation subsequently for 2 × 10 6 iterations.
We split the data into 1) a training set, which consists of 10 6 samples and is only used to train thebaseline flow model with ML; 2) a validation set consisting of 10 6 samples as well, which is used to nature, exclusively and, likewise, we to only generate a suitable set
To generate the training data alone we had to evaluate the target distribution and its gradients2 .
3 × 10 10 times, which is what we report as cost in terms of target evaluations in Table 6.
Filtering chiral forms
As mentioned in Section 4.2, alanine dipeptide is a chiral molecule, i.e.,it can occur in two different forms that are mirror images of each other, see Figure 13.
They cannotbe easily converted into each other as this would involve breaking existing and forming new bonds.
In nature, we find almost exclusively, while the D-form typically only exists in synthetically createdcompounds.
Hence, whenever alanine dipeptide is considered in the literature, it is almost alwaysas the L-form (Wu et al., 2020;
Campbell et al., 2021; Stimper et al., 2022; Dibak et al., 2021;
K¨ohler et al., 2022).
Therefore, we aim to train our model on only this form as well.
However, sincethe energy of the molecule does not change when creating a mirror image of it, models trained toapproximate its Boltzmann distribution will a priori learn to generate both forms.
The two forms can be separated by using the following procedure.
The two chiral forms differ bythe positioning of the neighboring atoms at a chiral center, i.e., the center carbon atom.
Hence, thedifference of the dihedral angles of those atoms with respect to a fixed reference will differ relative toeach other, i.e., their difference will change.
Hence, we compute this difference and check whetherit is close to a reference configuration for which we know that it corresponds to the L-form.
Asreference configuration, we use the position with minimum energy, which we already determinedfor the coordinate transformation.
We use this procedure to filter the configurations generated by the flow model during training andincluded only the samples that correspond to the L-form when computing the loss.
Thereby, themodel learns to only generate this chiral form.
In Appendix F.2 we will investigate a model trained on both chiral forms and compare it to one thatwas only trained on the L-form.
Training
All models were trained using the Adam optimizer (Kingma & Ba, 2015) with a batchsize of 1024.
A learning rate of 5 × 10 − 4 is initially linearly warmed up over 1000 iterations anddecayed with a cosine annealing schedule over the course of training.
We use a weight decay of10 − 5 and clip gradients at a value of 10 3 .
When training the models with the prioritized replaybuffer, we ensured a minimum buffer length of 64 batches and started replacing the oldest samplesonce its length reached 512 batches.
Evaluation
To evaluate the models, we draw 10 7 samples from the models with and withoutthe use of AIS.
Since there were some outliers of the importance weights due to flow numerics, wetook the 10 3 highest weights and clipped them to the lowest value in this set to compute the ESS(Koblents & M´ıguez, 2015; Dibak et al., 2021).
This corresponds to a fraction of 10 − 4 , or 0 .
01% ,of the weights.
For the flow trained with the α -divergence with α = 2 , the resulting ESS is close to
10 − 4 , indicating that the true ESS is even lower.
We estimated the Ramachandran plots, i.e., madea histogram of the dihedral angles ϕ and ψ , see Figure 3a, with 100 × 100 bins, and used them tocompute the KL divergence between the test samples and the samples from the model.
We repeatedthis with the reweighted samples, whereby we also used the clipped weights.
We computed thelog-likelihood on the test set with the models, but we did not do so for the SNF, as it only computesimportance weights and does not directly estimate the density.
[Table]
Computational cost
The two main contributors to the computational expenses necessary to trainflows approximating Boltzmann distributions are the number of evaluations of the flow and thetarget.
Typically, we need both the value and the gradient and, hence, we regard obtaining them asone operation.
The number of flow and target evaluations for each model and training procedure arelisted in Table 6.
In general, we trained all models using an equal number of flow evaluations, withthe exception of SNF.
SNF requires a large number of target evaluations due to the sampling layers.
Because of this, we reduced the number of flow evaluations done in total by this method.
Computational resources
To generate the MD dataset, we ran the replica exchange MD simulations on servers with an Intel Xeon IceLake-SP 8360Y processors having 72 cores and 256 GB
RAM.
We used a total of 100 nodes of the course of 16 h adding up to about 115k CPU h.
The flowmodels were trained on servers with an NVIDIA V100 GPU and an Intel Xeon Gold 6148 processorwith 20 cores and 96 GB RAM.
Training each model took around 2 days except the SNFs, whichtook around 8 days, totaling to around 1.3k GPU h.
F.2 Further results
Model trained on both chiral forms
To demonstrate the importance of filtering for the L-formduring training, we trained a model on both chiral forms using FAB with a replay buffer with thesame setting as in the other experiments.
We drew 10 7 samples from the model and found exactly50% of them correspond to the L- and 50% to the D-form.
As can be seen in Figure 15, the marginal distributions of the dihedral angles ψ and ϕ for the two forms are mirror images of each other, whilethe flow model, generating both forms, is a mixture of the two.
The log-likelihood on the test set is 210 .
80 , which is 0 .
74 less than the corresponding model onlytrained on the L-form.
This is close to log(2)
≈ 0 .
69 , i.e. the density is roughly by a factor of twolower, confirming once more that the flow density is a mixture of the density of the two forms.
Model performance with AIS
As mentioned in the previous section, we do AIS with all ourtrained models except the SNF, which already involves sampling layers.
We adopt the same AISsetting used for training the flow models with FAB, i.e., we use 8 intermediate distributions given bylinear interpolations between the flow and the target distributions and sample from them with HMCperforming 4 Leapfrog steps.
The results are shown in Table 7.
When comparing tables 2 and 7, weobserve that AIS improves performance for those models that approximate the target distribution atleast fairly well.
Again, the flow trained with FAB with a replay buffer outperforms the baselines.
[Table] [Table]
[Table]
Flow w/ D α =
[Table] [Table]
Flow w/ KLD
[Table]
RBD w/ KLD
[Table] [Table]
w/
[Table]
buffer
[Table]
FAB w/ buffer
[Table]
Ramachandran plots
Figure 16 shows the Ramachandran plot of the test set and Figure 17,
Figure 18, Figure 19, Figure 20, Figure 21, Figure 22, and Figure 23 show the Ramachandran plotsof all the models we trained for the first run, including the samples drawn from them via AIS.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]