Flow Annealed Importance Sampling Bootstrap
Authors: Laurence Illing Midgley
Vincent Stimper
Gregor N. C. Simm
Bernhard Schölkopf
José Miguel Hernández-Lobato
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Normalizing flow
Boltzmann distribution
Boltzmann generator
Annealed Importance Sampling
Approximate Inference

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

Published as a conference paper at ICLR 2023
FLOW ANNEALED IMPORTANCE SAMPLING BOOTSTRAP
Laurence I. Midgleyy
Vincent Stimper*
University of Cambridge InstaDeep Baurencemidgla eeeeeeccemmmllllymmmmmmmm eellley@@mai...om ccom
Max Planck Institute for Intelligent Systems University of Cambridge
v B8 www.ccammmmmmmmmmmmm..... ac uk
Gregor N. C. Simm
Bernhard Schölkopf Max Planck Institute for Intelligent Systems bsê tue. mwe..epgg.e.... de
José Miguel Hernández--obato University of Cambridge jmh2336 cam, ac ...uk
University of Cambridge gncs20cam cam .ac.uk
ABSTRACT
Normalizing flows are tractable density models that can approximate complicated target distributions, e.g. Boltzmann distributions of physical systems. However, current methods for training flows either suffer from mode-seeking behavior, use samples from the target generated beforehand by expensive MCMC methods, or stochastic losses that have high variance. To avoid these problems, aug- ment flows with annealed importance sampling (AIS) and minimize the mass- covering a-divergence with 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 approxi- mate them very accurately where previous methods fail. To the best of our knowl- edge, 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 re- sults than training via maximum likelihood on MD samples while using 100 times fewer target evaluations. After reweighting the samples, we obtain unbiased his- tograms of dihedral angles that are almost identical to the ground truth.
INTRODUCTION
Approximating intractable distributions is a challenging task whose solution has relevance in many real-world applications. A prominent example involves approximating the Boltzmann distribution of given molecule. In this case, the unnormalized density can be obtained by physical modeling and is given by eu(()), where æ are the 3D atomic coordinates and u(-) returns the dimensionless energy of the system. Drawing independent samples from this distribution is difficult (Lelièvre et al., 2010). It is typically done by running expensive Molecular Dynamics (MD) simulations (Leimkuhler & Matthews, 2015), which yield highly correlated samples and require long simulation times.
An alternative is given by normalizing flows. These are tractable density models parameterized by neural networks. They can generate batch of independent samples with single forward pass and any bias in the samples can be eliminated by reweighting via importance sampling. Flows are called Boltzmann generators when they approximate Boltzmann distributions (Noé al., 2019). Recently, there has been a growing interest in these methods (Dibak et al., 2022; Köhler et al., 2021; Liu et al., 2022) as they have the potential to avoid the limitations MD simulations. Most current approaches to train Boltzmann generators rely on MD samples since these are required for the estimation of the flow parameters by maximum likelihood (ML) (Wu et al., 2020). Alternatively, flows can be trained without MD samples by minimizing the Ballback-Leibbeee (KL) divergence with respect to the target distribution. Wirnsberger et al. (2022) followed this approach to approximate the
Equual contribution

---

Published as a conference paper at ICLR 2023
Boltzmann distribution of atomic solids with up 512 atoms. However, the KL divergence suffers from mode-sceking behavior, which severely deteriorates the performance of this approach with multimodal target distributions (Stimper et al., 2022). On the other hand, mass-covering objective such as the forward KI. divergence suffer from the high variance of the samples from the flow.
To address these challenges, we present new method for training flows: Flow AIS Bootstrap' (FAB). Our main contributions are as follows:
as our training objective, which is mass-covering and minimizes importance weight variance. At test time an importance sampling distribution with low a-divergence (with a = 2) may be used to approximate expectations with respect to the target with low variance. This objective is challenging to estimate during training. To approximate this objective we use annealed importance sampling (AIS) with the flow as the initial distribution and the target set to the minimum variance distribution for the estimation of the a-divergence. AIS returns samples that provide a higher quality training signal than samples from the flow, as focuses the regions that contribute the most to the a-divergence loss.
We propose use the a-divergence with
2. We reduce the computational cost of our method by introducing a scheme to re-use samples via a prioritized replay buffer.
3. We apply FAB to a toy 2D Gaussian mixture distribution, the 32 dimensional "Many Well" prob- lem, and the Boltzmann distribution of alanine dipeptide. In these experiments, we outperform competing approaches and, to the best of our knowledge, we are the first successfully train Boltzmann generator on alanine dipeptide using only the unnormalized target density, In particu- lar, we use over 100 times fewer target evaluations than Boltzmann generator trained with MD samples while producing a better approximation to the target.
2 BACKGROUND
Normalizing flows Given random variable z with distribution q(z), a normalizing flow (Tabak & Vanden-Eijnden, 2010; Rezende & Mohamed, 2015; Papamakarios al., 2021) uses an invertible map F Rd + Rd to transform z yielding the random variable x = F(z) with distribution
q(x) = g(z)) [det(J=(2))
(1)
where (z) OF / ddz is the Jacobian of F. If we parameterize F, we can use the resulting model to approximate a target distribution To simplify our notation, we will assume the target density p(x) is normalized, i.e., integrates to 1, but the methods described here are equally applicable when this is not the case. If samples from the target distribution are available, the flow can be trained via ML. If only the target density p(x) is given, the flow can then be trained by minimizing the reverse KL divergence? between q and p, i.e., (L(gllp) f- 4(x) (xg(g(x))p(())) which is estimated via Monte Carlo using samples from q.
Alpha divergence An alternative to the KL divergence is the a-divergence (Zhu & Rohwer, 1995; Minka, 2005; Müller et al., 2019; Bauer & Mnih, 2021: Campbell et al., 2021) defined by
Da(pll Le p(x)) (()x))((x)==00)
a(1 a))
The a-divergence is mode-seeking for a S ( and mass-covering for a 2 1 (Minka, 2005), as shown in Figure 1. When a = 2, minimizing the a-divergence is equivalent to minimizing the variance of the importance sampling weights wx) = p(x)/q(x> (xx), which is desirable if importance sampling will be used to eliminate bias in the samples from q at test time.
Annealed importance sampling AIS begins by sampling from an initial distribution X1 Po q. given by the flow in our case, and then transitioning via MCMC through a sequence of interme- diate distributions, p1 to Par-1, to produce sample x closer to the target distribution g PM (Neal, 2001). Each transition generates an intermediate sample x; by running a few steps of a Markov chain initialized with the previous intermediate sample xj-1 that leaves the intermediate
FAB uses the flow in combination with AIS to estimate a loss in order to improve the flow. Thus we use bootstrap in the name of our method to mean "using one's existing resources to improve oneself". We refer reverse KI. divergence as just "KI. divergence", following standard practice in literature.
2

---

Published as a conference paper at ICLR 2023
-0
0.5
-
KL(glpp
KL(pl1g)
Figure 1: Illustration of unnormalized Gaussian approximating distributions shown in red, that minimize the a-divergence for different values of with respect a bimodal target distribution p. shown in blue. The solutions q are mode-seeking for values of a < and they are mass coass-covringg for values of a 2 1. The cases a - and - correspond to KL(g and pg)p)pp) respectively. Figure reproduced from (Minka, 2005).
distribution Py-1 invariant. Each py is defined by interpolating between the initial and target log densities: log (x(x) 8, log (x) (1 - 3,) log
PM (x), where .B B1 BN 0.
AIS provides an importance weight for the final resulting sample XM given by
(AIs(XM)= = (x1) P2 (x2)
PM-1 (XM-1) PM (XM))
(3)
Po (x1) Pl (x2) PM-2 (XM-1) PM-1 (xar
These weights exhibit variance reduction compared to their importance sampling counterparts (xs = p(x)/g(x) (Neal, 2001). The AIS samples and importance weights may then be used where (x))
to estimate expectations over the target gxx) using Eggxx) [h(x) EAIS
is some function of interest. Hamiltonian Monte Carlo (HMC) is suitable transition operator for implementing AIS in challenging problems (Neal, 1995; Sohl-Dickster & Culpepper, 2012).
3
METHOD
3.1 FLOW ANNEALED IMPORTANCE SAMPLING BOOTSTRAP
FAB trains flow approximate target p by minimizing (pllg) which is estimated with AIS using y as initial distribution and y² /q as target. The latter the minimum variance importance sampling distribution for estimating the Da=2(pllg) loss. FAB performs form of bootstrapping since it fits the flow q using the samples generated by q after these have been improved with AIS to train mass- mass-covering flow without access to samples from the target. Below
fit p² /q. Thereby, we provide a brief derivation of our loss function and refer to Appendix for the full derivation.
We consider q to be specified by some parameters and write qo make this explicit. We aim to tune go by minimizing the loss function L(0) ix 2o=2(pl19a) where C(A). We can write our loss as an expectation over some distribution g(x) by using importance sampling:
D Da=2(pllg0) x L()) = p(x)²
x(x)²
dx FFg(x)
(x)
9o (x)g((xx)
We consider setting,g x p² p²/40 which minimizes the variance in the estimation of C(O). Sampling directly from g x p² /9g intractable. Instead, train the flow using an estimate of the loss based on samples generated by AIS when targeting 9 x p² p²/qa and using These AIS samples have higher quality than those returned by the flow as they occur regions where the integrand in Equation (4) takes high values. These are regions where p and q have high and low density, respectively. Another advantage of the AIS samples is that we can use the weights returned by AIS to obtain an unbiased estimate of L()).
the initial distribution.
To obtain the gradient of Equation (4) with respect to 6, let us denote fo(x) = p(x)2//o(x)) Without loss of generality* we assume ʃ fo (x)dxx and then set g(x) = fo (x). First, we write the gradient
'The importance sampling distribution 9 that minimizes the variance in the estimation of J [f(x))ll [p()dx is given by g(x) x lx(()) (Kahn & Marshall, 1953; Owen, 2013). We this differ- ent from the distribution that minimizes variance for self-normalized importance sampling, which given hy g(x) x (f(x) plp(x) (Hesterberg, 1988; Owen, 2013).
*Se Appendix B for the full derivation where keep track of the normalizing constant.
3

---

Published as a conference paper at ICLR 2023
expectation over g:
V (x) (x)
VgL(0) Fox())
Fy(x) [V log [o(()] -Bg(x) [e log (x)]
(5)
We can then write this as an expectation over the AIS forward pass:
VeL(0) EATS [WAIS Va log (xAis))
where XAIS and WAIS are the samples and respective importance weights generated by AIS when targeting g. The bar superscript denotes stopped gradients in the AIS samples, XAIS. with re- spect to 0. If we stop the gradients of WAIS as well, we can then use the surrogate loss function S(O) EAI [WAIS log go (A which can be estimated by Monte Carlo and differentiated to obtain unbiased estimates of the gradient. In practice, we found that using the self- norfaizrd importance weights greatly improved training stability. We refer to the surrogate loss with self- normalized importance weights as S'(0). Its estimate used for training is given by:
N
(1)
8'
Σ
W
AIS
(1) log go AIS)
where wiis and X are N samples and weights generated by AIS using g =² /go as target distri- bution. When evaluating the gradient of Equation (7) with respect to 6, gradients must be stopped during the computation of the AIS samples and weights. In practice, we obtain good performance with a relatively low number of intermediate AIS distributions, e.g., I for the Gaussian mixture model problem and 8 for the dipeptide problem, see the following section. Moreover, we can use AIS after training with target p further reduce variance when approximating expectations over p.
Here we have focused on the minimization of (p=l(pllg) using an AIS bootstrapping approach with y² /q as target distribution. However, our approach general and could be used minimize other objectives (Midgley et al., 2021) and to train other models, such as those that combine flows with stochastic sampling steps (Wu al., 2020; Arbel et al., 2021; Matthews et al., 2022; Jing et al., 2022). We provide further discussion and examples related to this in Appendix C. This includes a derivation of version of FAB that works for a divergence minimization with arbitrary values of a.
In Appendix D we provide an analysis of the quality of the estimates of the gradient of Da=2 (pllg) produced by FAB and by importance sampling with samples from q or p. We focus on the FAB gradient in the form from Equation (6). as it is easy to analyze. First, we show that in a simple scenario where both q and p are ID Gaussians, the signal-to-noise ratio of FAB with a small number of AIS distributions is far superior to that of estimating (a=2(pllg) using samples from & p.
We also study in Appendix the performance of FAB as the dimensionality of the problem grows. Similar to the analysis of AIS with increasing dimensionality by Neal (2001), we consider simple scenario where p and are factorized and the AIS MCMC transitions are perfect (output independent samples that follow the corresponding intermediate distributions). We then show the following: 1) Estimating Da=2 with importance sampling using samples from q or p results in a variance of the gradient estimate that grows exponentially with respect to the dimensionality of the problem. 2) This variance remains constant in FAB when the number of AIS distributions increases by the same factor as the dimensionality. We provide an empirical analysis of the gradient variance in FAB under the aforementioned assumptions. If we increase the number of AIS distributions by the same factor by which the dimensionality increases, the SNR of the gradient estimate remains roughly constant. This suggests that FAB should scale well to higher dimensional problems, relative to training via estimation of the Da=2 loss by importance sampling with samples from g or p. We acknowledge that our simplifying assumptions are strong and we leave more general analysis to future work.
3.2 RE-USING SAMPLES THROUGH A REPLAY BUFFER
Although AIS is relatively cheap, it is still significantly more expensive than directly sampling from the 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 prioritized replay buffer analogous to the one in (Mnih et al., 2015; Schaul et al., 2016).
Consider a replay buffer with a set of samples and corresponding AIS weights generated during a single run of AIS with target gxx) p(x)² (lx))2/g(( (x), where go(x) is the flow point in training
4

---

Published as a conference paper at ICLR 2023
Algorithm 1: FAB for the minimization of Da=2(plgg with prioritized replay buffer
Initialize flow 9 parameterized by 0
Initialize replay buffer to fixed maximum size
for iteration J K do Generate AIS samples and add to buffer Sample x (1:30) from qo and evaluate log (x()
Obtain AAS and log WAIS using AIS with target p² /qo and seed x (1:M)) and log qo (x) (2:40)
(1:M) and log qo (xXIs (1s to replay buffer Add X (1:M)) log AIS
for iteration 1 to I. do // Sample from buffer and update go
Sample (1.N) and log 90wu (xx (1.N)) from buffer with probability proportional to w
(2:N)
Calculate log (unnn)
log qaau (x()s)) stop-grad (x(s
Update log w AiS (1:N) and log as (XAIS
is in buffer log w (I:N) + log w conrection and log qo (x)
AIS
Calculate loss (0) = -1 I//N > (i())
Wommection log (x AIS)
Perform gradient descent S'(f) update
specified by 0. We can approximate the gradient of Equation (7) using
(oS
log go (xx),
where X1,..., X 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 0, denoted fald. we have to multiply their AIS weights with correction factor Weoorection = qe (x)/qs (x) before sampling, which requires to additionally store qe (x,a()) for each data point in the buffer. Note that qe in the denominator of this correction factor because gxx), the distribution we are sampling from, inversely proportional to go.
XN
The resulting procedure extracts data from the buffer in a prioritized manner: We sample according g, which favors points with low go and high p. As go is updated fit samples from the buffer, will take higher values on those samples and their weights will gradually be decreased to encourage drawing alternative samples. The buffer allows to re-use old AIS samples and does not require re-evaluating p(x), which could be expensive in some cases.
A limitation of the above approach that requires updating the AIS weights for all data points in the buffer before sampling, which is expensive. To significantly speed up computations, we instead draw minibatch from the buffer with probability proportional the old AIS weights and then reweight each sample with the corresponding Wecerection Before updating I, we update the AIS weights for the sampled points in the buffer and replace the respective go., (x) values with qo (x). The pseudocode for the final procedure is shown in Algorithm 1. In practice, we found that sampling from the buffer without replacement worked better, at the cost of introducing bias into our gradient estimates. Lastly, we set a maximum length for the buffer, and once this is reached we discard the oldest samples each time new samples are added.
EXPERIMENTS
This section contains an experimental evaluation of our proposed method. The code publicly available at ht :tttps://gin //tt:////hhtbb/ the normf 1 lows package to implement the flows (Stimper et al., 2023). The Appendix contains a detailed description of each experiment to guarantee reproducibility. In Appendix F we include an additional set of experiments on the 32-dimensional "Many Well" distribution given by the product of 16 copies of the 2-dimensional Double Well distribution from Noć et al. (2019); Wu et al. (2020).
Is written in PyTorch and uses
4.1 MIXTURE OF GAUSSIANS IN 2D
First, we consider a synthetic problem where p is a mixture of bivariate Gaussians with 40 mixture components. The too-dinensional nature of this problem allows us easily visualize the results
5

---

Published as a conference paper at ICLR 2023
Initialisation
Flow w/ ML
Flow w/ KLD
RBD w/ KLD
SNF w/ KLD
CRAFT
FAB w/o buffer (ours)
FAB w/ buffer (ours)
Figure 2: Contour lines for the target distribution and samples (blue dises) drawn from the approx- imation go obtained by different methods the mixture of Gaussians problem. The plot for the flow trained by Da which had the worst performance, is shown in Appendix E.3.
of different methods while the multimodality of p makes the problem relatively challenging. To increase the problem difficulty, we give the flow pathological initialization where samples from qo concentrate in a small region of the sampling space, as illustrated in the top left plot in Figure 2.
We compare the following methods: 1) FAB with a replay buffer as shown in Algorithm 1; 2) FAB without a replay buffer, where we directly optimize Equation (7); 3) flow model that minimizes <L(gollp) 4) a flow with a Resampled Base Distribution (RBD) (Stimper et al., 2022) that mini- mizes KL [p:: 5) Stochastic Normalizing Flow (SNF) model (Wu et al., 2020) that also min- imizes (L( lpp); 6) a Continual Repeated Flow Annealed Transport (CRAFT) model (Matthews al., 2022) that minimizes a CRAFT specific version of KL divergence specified in Appendix C.2; 7) flow model that minimizes D 2 (qu) estimated using samples from qo (Müller et al., 2019), and 8) flow trained by maximum likelihood (ML) using samples from p. In this toy problem. we have access to ground truth samples from the target, allowing us to train the flow by ML. However, we are interested in the case where samples from the target are not cheaply available. To denote this, we single out the results from this latter method in our tables with horizontal dashed line.
All methods besides CRAFTS use the same parametric form for go given by a Real NVP flow model with 15 layers (Dinh et al., 2017). For the FAB-based approaches, we run AIS with single in- termediate distribution (3 = 0.5) and MCMC transitions with Metropolis- Hastings step, being a Gaussian perturbation and then accept-reject step. For FAB with prioritized buffer, we perform
== 4 gradient updates to qa per AIS sampling step. For the SNF and CRAFT, we do Metropolis- Hastings step every? 3 flow layers. For training the flow by ML, we draw new samples from the target for each loss estimation. All models are trained for 2 10? flow evaluations. Further details on the hyper-parameters and architectures used in each algorithm
provided in Appendix E.
Figure shows that our two FAB-based methods and the flow trained by ML fit all modes in p. By contrast, the other alternative methods cover only small subset of the modes. The reason for this is that such methods are trained only on samples from go and the poor initialization of qo makes it unlikely that they will ever generate samples from undiscovered modes.
To evaluate the trained models we compute the effective sample size (ESS) obtained when doing importance sampling with qe: the average log- loglliiel of qo on samples from pt the forward KL divergence with respect to the target; and the mean absolute error (MAE) in the estimation of Ep(x) [f (x)] by importance sampling with qo where fx) is a toy quadratic function specified in
$CRAFT also uses affine transform in the flow, but with autoregressive dependency across dimensions (Kingma al., 2016) instead of coupling.

---

Published as a conference paper at ICLR 2023
Results for the mixture of Gaussians problem. Our methods are marked in italic. Best results are emphazised in bold. log-likelihood values for the first two methods are NaN because they assign zero density to samples from missing modes. Log-likelihood values for SNF and CRAFT are N/A because this method does not provide density values for the generated samples. The CRAFT implementation used does not allow for forward-KL to be estimated. Furthermore, the resampling step from CRAFT prevents the ESS from being estimated. Hence, these fields have N/A.
Table I:
ESS (%)
Ep(x) [log [x)]
KL(pl14)
MAE (%)
MAE w/o RW (%)
Flow w/ ML
54.3 10.4
-7.18 + 0.05
0.31 + 0.05
9.0
6.1 + 2.0
Flow w/ Da== 0..7+..3
N
NaN + NaN
99.5
9..
Flow w/ KLD
55.0+20.8
NaN + NNa
NaN + NaN
26.0+3.1
26.3+1.4
RBD w/ KLD
16.3
NaN = NNN
NaN = NaN
18.1
92.1
SNF w/ KLD
23.3
N/A
NaN = NaN
64.6
69.3
CRAFT
N/A
N/A
N/A = N/A
98.8
99.1
FAB w/o buffer
3..1
1.00 0.19 0.19
9.4
4.4
FAB w/ buffer
6.0
-..16
0.30 + 0.07
8.9
0.5
Appendix E. We express the MAE as 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 the importance weights. We provide further details on the evaluation setup Appendix E.2.
Table shows for each method average results and corresponding standard errors over 3 random seeds. FAB with the buffer performs similarly to the benehmark of training the flow by ML. Both of these methods are the best performing ones with the highest ESS, the highest log-likelihood on samples from p and the lowest forward KL divergence and lowest MAE FAB without a replay buffer is the next best method while the other methods perform poorly. This is especially the case regarding the forward KL divergence and log- log-likelihood values on samples from pr the non FAB/ML methods assign zero density to points sampled from undiscovered modes of which is represented by writing NaN in the table. Note that the ESS for the SNF, RBD and the flow trained by minimizing KL( (g)llp) spurious these methods
missing modes. Finally, note that the MAE in the estimation of Extx) [x) via importance sampling with 1000 samples from p is 8.6%. Also, the log- lidlhiodl of p on samples from this same distribution is -6 -6.85. These are close to the values obtained by FAB with a replay buffer, meaning that the corresponding qa is close to p.
4.2 ALANINE DIPEPTIDE
We now consider the 22 atom molecule alanine dipeptide, shown in Figure 3a, in an implicit solvent at temperature of T = 300 K and aim to approximate its Boltzmann distribution given the 3D atomic coordinates. This is popular benchmark when considering Boltzmann generators (Wu et al., 2020: Campbell et al., 2021: Dibak et al., 2022: Stimper et al., 2022). Previous works have used coordinate transformation which are normalized using their mean and standard deviation computed on samples generated by MD (Noé et al., 2019). Since we aim train models without using any data, we replace the mean by the minimum energy configuration, which can be cheaply estimated through gradient descent within less than 100 steps. Similarly, we replace the standard deviations with values reflecting the typical order of magnitude of each variable. Furthermore, we represent the molecule with internal coordinates on only, thereby implicitly satisfying the system's rotational and translational invariance.
map some but not all Cartesian coordinates to internal coordinates,
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 (Rezende et al., 2020), while the others are considered as unbound. The models trained with FAB use 8 intermediate distributions. For FAB with the replay buffer, we do L = 8 gradient updates per AIS forward pass. Alanine dipeptide is 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 16. In nature, 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 leam to only generate this form. More details are given in Appendix G.I.
7

---

Published as a conference paper at ICLR 2023
10%
Probability density
10-1
Ground truth
10
Flow w/ ML
Flow ML,
10"3
FAB w/ butter (ours)
FAB buffer,
10-5
-7
(b)
Figure 3: (a) Visualization of alanine dipeptide and the dihedral angles ø and & for the Ramachan- dran plot. (b) Marginal distribution of & in log scale as given by the ground truth, the flow trained with MI on MD samples, and FAB with replay buffer. RW indicates whether samples have been reweighted with importance sampling before generating the plot. Figure 17 is the same plot in nor- mal scale, revealing how small the mode at 1 is.
Ground truth
Flow w/ ML
FAB w/ buffer (ours)
FAB w/ buffer, RW (ours)
10
10
o
10
-
10
17
2
0
7
17
17
17
12
17
-17
π
Φ
Figure 4: From left to right, Ramachandran plots of the ground truth generated by MD, flow model trained by ML on MD samples, and by FAB using replay buffer before and after reweighting samples to eliminate bias.
To evaluate our models, we generated samples using parallel tempering MD simulations, which serve as ground truth. They are split into training and validation sets with 10t samples each and test set with 10" samples. We compare FAB several baseline methods already mentioned in the previous section, see Table 2 for the full list. The SNF method performs 10 Metropolis- Hastings steps every two layers, meaning a total of 60 additional sampling steps. All methods are trained for 2.5 10 flow evaluations except for the SNE, which uses 6 x 107 as it is more expensive in terms of target evaluations. Table provides an overview of the number of flow and target evaluations by each method. We compare methods via the ESS of importance sampling weights and the average log-likelihooe on the test set. Moreover, we generate Ramachandran plots, which are histograms for the marginal distribution of the dihedral angles á and ve illustrated in Figure 3a. We compute their KI. divergence to the ground truth with and without reweighting using the importance weights. Our experiments are repeated over random seeds and average values and standard errors are given.
Table 2 shows our results. The flow trained by minimizing Do (pll(p))g)) with samples from had convergence problems due the high gradient variance. As result, it performs very poorly in practice, especially in terms of ESS. The models trained by minimizing KL(gllp) have no conver- gence problems, but they only approximate a subset of the target modes. This results in poor test log-likelihood and KLD values, and spurious values for the ESS. The flow trained by ML on MD samples obtains very good results terms of test likellhodo and KLD values. However, strug- gles to model the dim mode at ~ 1 correctly, shown in Figures 3b and 4. Its ESS is also fairly low and, hence, reweighting worsens performance. The models trained with FAB have higher ESS and test log-likelihood than the other methods. FAB with buffer obtains lower KLD values than the flow trained by ML. When reweighting is applied to the samples generated by this version of FAB, the resulting distribution is nearly the same the ground truth, 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, shown in Table 8.
8

---

Published as a conference paper at ICLR 2023
Table ESS, log- logeliiel on the test set, and KL divergence (KLD) of Ramachandran plots with and without reweighting (RW) for each method. Our methods are marked in italic and best results are emphasized in bold.
ESS (%)
Ep(x) (log (x))
KLD
KLD w/ RW
Flow w/ ML
2.8
209.22
(7.5733.80)
10 10-3
(2.58 0.80) X 10-2
Flow w/ D 2
0.000
73.5 1.3
0.13
1..
Flow w/ KLD
12
100
0.27
0.19
RBD w/ KLD
44+18
143 +
3.00
3.00
SNF w/ KLD
0.16+
N/A + N/A
8.36
9.58
FAB w/o buffer
5..
211.13 + 0.03 0.03
(6.28
10
(2.66=0.90)
10
FAB w/ buffer
92.8+
211.54 + 0.00
(3.42 +.455
0.45)
x 10-3
(2.51 (2351
x 10-3
DISCUSSION AND RELATED WORK
SNFs combine flows with MCMC methods by introducing sampling layers between flow layers to improve model expressiveness (Wu et al., 2020; Nielsen et al., 2020). SNFs have been extended to 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 in- termediate distributions, with each flow being trained by minimizing a KL divergence with respect to the next intermediate distribution. CRAFT improves the issue of mode seeking relative SNFs, which can be seen in the Many Well problem in Appendix F where it performs well. However, CRAFT fails catastrophically
the GMM problem, as the CRAFT loss can still favour mode seek- ing and uses samples directly from the flow for its estimation which can provide poor training signal. In Appendix C.2 we describe a new FAB-flavored version of CRAFT using the 2 as objective, estimating it with the MCMC samples ahead of the flow in the SMC process.
Within the MCMC/AIS literature, significant work has focused on improving transition kernels (Levy et al., 2018; Gabrié et al., 2022), intermediate distributions (Brekelmans et al., 2020), and the extended target distribution (Doucet al., 2022a;b) of AIS. These techniques are applicable the AIS procedure used in FAB. FAB does not differentiate through AIS to obtain the gradient with re- spect to the flow. Combining FAB with methods that allow for differentiation through AIS (Geffner & Domke, 2021; Zhang al., 2021: Doucet et al., 2022a) may allow for lower variance gradient estimate. Works on differentiation through iterated systems would be relevant for this, notably (Metz et al., 2021). FAB does not use gradients of the target distribution when optimizing its loss although such gradients are used in the sampling process by HMC. This is in contrast with the alter- native approach of training the flow by minimizing KL(gollp), which does use these gradients. Such gradient information could be included in FAB through force matching (Wang et al., 2019; Köhler al., 2021; Köhler al., 2022) the addition of a KL divergence loss term.
function,
CONCLUSION
We have proposed FAB, method for training flows to approximate complicated multimodal target distributions. FAB combines a-divergence minimization with a = with an AIS bootstrapping mechanism for improving the samples used for the loss estimate. By focusing this divergence, we favor 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 provides a high-quality training signal by focusing on the regions where the flow is a poor approximation of the target. We have also proposed use prioritized replay buffer, which reduces the cost of FAB and improves performance. Our experiments show that FAB can produce accurate approximations of complex multimodal targets without using samples from such distributions. By contrast, other alternative approaches fail in this challenging setting. Remarkably, for the alanine dipeptide, FAB produces better results than training the flow by ML on samples generated via MD simulations while still using 100 fewer evaluations of the target than the MD simulations. In future work, we hope to scale up our approach to more challenging problems, such as the modelling of the Boltzmann distribution of large proteins.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
We thank Emile Mathieu, Kristopher Miltiadou, Alexandre Laterre, Clément Bonnet, and Alexander Matthews for the helpful discussions. José Miguel Hernández-Lobace acknowledges support from Turing AI Fellowship under grant EP/V023756/1. This work was supported by the German Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 01IS! 8039BB and by the Project number 390727645.
Machine Learning Cluster Excellence, EXC number 2064/1
REFERENCES
Michael Arbel, Alex Matthews, and Arnaud Doucet. Annealed flow transport Monte Carlo. In Proceedings of the 38th International Conference Machine Learning, pp. 318-330. PMLR, 2021.
Matthias Bauer and Andriy Mnih. Generalized doubly reparameterized gradient estimators. In Ma- rina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139, pp. 738- 747. PMLR, 2021.
Rob Brekelmans, Vaden Masrani, Thang Bui, Frank Wood, Aram Galstyan, Greg Ver Steeg, and Frank Nielsen. Annealed importance sampling with q-paths. arXiv preprint arXiv:2012.07823, 2020.
Andrew Campbell, Wenlong Chen, Vincent Stimper, Jose Miguel Mernandez-Lobatt and Yichuan Zhang. A gradient based strategy for Hamiltonian Monte Carlo hyperparameter optimization. In Proceedings of the 38th International Conference on Machine Learning, pp. 1238-1248, PMLR, 2021.
Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and Boltzmann generators. Phys. Rev. Research, 4, 2022.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. In- ternational Conference on Learning Representations, 2017.
Arnaud Doucet, Will Grathwohl, Alexander G. D. G. Matthews, and Heiko Strathmann. Score-based diffusion meets annealed importance sampling. arXiv preprint 2XXv:2208.00698 2022a.
Arnaud Doucet, Will Sussman Grathwohl, Alexander G de G Matthews, and Heiko Strathmann. Annealed importance sampling meets score matching. In ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022b.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios, Neural spline flows. Ad- vances in Neural Information Processing Systems, 32:7511-7522, 2019.
David J. Earl and Michael W. Deem. Parallel tempering: Theory, applications, and new perspectives.
Phys. Chem. Chem. Phys., 7:3910-3916, 2005.
Marylou Gabrié, Grant M Rotskoff, and Eric Vanden-Eijnden. Adaptive Monte Carlo augmented with normalizing flows. Proceedings of the National Academy Sciences, 119(10):c2109220119 2022.
Tomas Geffner and Justin Domke. MCMC variational inference via uncorrected Hamiltonian an- nealing. Advances in Neural Information Processing Systems, 34:639-655 651, 2021.
Timothy Classen Hesterberg. Advances in importance sampling. Stanford University, 1988.
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffu- sion for molecular conformer generation. arXiv preprint arXiv:2206.01729 2022.
Herman Kahn and Andy W. Marshall. Methods of reducing sample size in Monte Carlo computa- tions. Journal of the Operations Research Society of America, 1(5):263-278 1953.
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. 3rd Interna- lional Conference Learning Representations, 2015.
10

---

Published as a conference paper at ICLR 2023
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Im- proved variational inference with inverse autoregressive flow. Advances in neural information processing systems, 29:4743-4477 47511, 2016.
Eugenia Koblents and Joaquín Míguez. A population Monte Carlo scheme with transformed weights and its application stochastic kinetic models. Statistics and Computing. 25:407-425, 2015.
Jonas Köhler, Andreas Krämer, and Frank Noé. Smooth normalizing flows. In Advances in Neural Information Processing Systems 34, 2021.
Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé. Force-matching coarse-graining without forces. arXiv preprint arXiv:2203. 1167,, 2022.
Ben Leimkuhler and Carles Matthews. Molecular Dynamics With Deterministic and Stochastic Numerical Methods. Number 39 in Interdisciplinary Applied Methematics. Springer, 2015.
Tony Lelièvre, Mathias Rousset, and Gabriel Stoltz. Free Energy Computations: A Mathematical Perspective. Imperial College Press, 2010.
Daniel Levy, Matthew D) Hoffman, and Jascha Sohl-Dickstein. Generalizing Hamiltonian Monte Carlo with neural networks, In 6th International Conference on Learning Representations, 2018.
Tianyi Liu, Weihao Gao, Zhirui Wang, and Chong Wang. PathFlow: A normalizing flow generator that finds transition paths. In The 38th Conference on Uncertainty in Artificial Intelligence, 2022.
Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Arnaud Doucet. Continual repeated annealed flow transport Monte Carlo. In Proceedings of the 39th International Conference on Machine Learning, pp. 15196-15219. PMLR, 2022.
Luke Metz, Daniel Freeman, Samuel S Schoenholz, and Tal Kachman. Gradients are not all you need. arXiv preprint 1tt.x::21111008000 2021.
L. I. Midgley, V. Stimper, G. N. C. Simm, and J. M. Hemández-LLobto.. Bootstrap your flow. In 1st ELLIS Machine Learning for Molecule Discovery Workshop, December 2021. URL https: / //aaxiivvor org. abaa 2211 11510.
Tom Minka. Divergence measures and message passing. Technical report, Microsoft Research, 2005.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 ((7754)):529-533 2015.
Yoshiharu Mori and Yuko Okamoto. Replica-exchange molecular dynamics simulations for various constant temperature algorithms. Journal of the Physical Society of Japan, 79(7):074001, 2010.
Thomas Müller, Brian McWilliams, Fabrice Rousselle, Markus Gross, and Jan Novák. Neural im- portance sampling. ACM Transactions on Graphics (TOG), 38(5):1-19, 2019.
Radford M. Neal. Bayesian learning for neural networks. PhD thesis, Department of Computer Science, University of Toronto, 1995.
Radford M. Neal. Annealed importance sampling. Statistics and computing, 11(2):125-13) 2001.
Didrik Nielsen, Priyank Jaini, Emiel Hoogeboom, Ole Winther, and Max Welling. SurVAE flows: Surjections to bridge the gap between VAEs and flows. Advances in Neural Information Process- ing Systems 33, 2020.
Frank Noé. Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of mody systems with deep learning. Science, 365(6457), 2019.
Art B. Owen. Monte Carlo theory, methods and examples. 2013.
11

---

Published as a conference paper at ICLR 2023
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lak- shminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1-64 2021.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on Machine Learning, pp. 1530-1538. PMLR, 2015.
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racanière, Michael S. Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer. Normalizing flows on and spheres. In Proceedings of the 37th International Conference Machine Learning, volume 119, pp. 8083- 8092. PMLR, 2020.
Tom Schaul, John Quan, loannis Antonoglou, and David Silver. Prioritized experience replay. In 4th International Conference Learning Representations, 2016.
Jascha Sohl-Dickstein and Benjamin J Culpepper. Hamiltonian annealed importance sampling for partition function estimation. arXiv preprint arXiv: 2012.
Vincent Stimper, Bernhard Schölkopf, and José Miguel Hernández-Lobate Resampling Base Distri- butions of Normalizing Flows. Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pp. 4915-4936, 2022.
Vincent Stimper, David Liu, Andrew Campbell, Vincent Berenz, Lukas Ryll, Bernhard Schölkopf, and José Miguel Hernández-Loobto. normflows: A PyTorch Package for Normalizing Flows. arXiv preprint arXiv:2302. /2014, 2023.
Esteban G. Tabak and Eric Vanden Vanden-Eijnden. Density estimation by dual ascent of the log- likeliihoodd. Communications in Mathematical Sciences, 8(1):217 233, 2010.
Jiang Wang, Simon Olsson, Christoph Wehmeyer, Adrià Pérez, Nicholas E Charron, Gianni De Fab- ritii,, Frank Noé, and Cecilia Clementi. Machine learning coarse-grained molecular dynamics force fields. ACS central science, 5(5):755-767, 2019.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3):229-256, 1992.
Peter Wimnsbrrgr, George Papamakarios, Borja Ibarz, Sébastien Racanière, Andrew J Ballard, Alexander Pritzel, and Charles Blundell. Normalizing flows for atomic solids. Machine Learning: Science and Technology, 3(2)0225009, 2022.
Hao Wu, Jonas Köhler, and Frank Noe. Stochastic normalizing flows. In Advances in Neural Information Processing Systems, volume 33, pp. 5933-5944, 2020.
Guodong Zhang, Kyle Hsu, Jianing Li, Chelsea Finn, and Roger Baker Grosse. Differentiable annealed importance sampling and the perils of gradient noise. In Advances in Neural Information Processing Systems, pp. 19398-194101 2021.
Huaiyu Zhu and Richard Rohwer. Information geometric measurements of generalisation. Technical report, Aston University, 1995.
12

---

Published as a conference paper at ICLR 2023
A EXTENDED DISCUSSION
Gabrié et al. (2022) use normalizing flows to learn the transition kernels for MCMC. These transition kernels are used to perform large MCMC steps between meta-stable states, improving the notorious issue of mixing in MCMC. In FAB, the flow has similar function, although is used as the base distribution for AIS rather than for a transition kernel. Namely, in FAB, the flow learns to balance mass between meta-stable states. This done by using AIS and its importance weights to reweight inaccuracies in the mass allocated across different meta-stable states. In (Gabrié et al., 2022), priori knowledge of the meta-stable states is required to obtain good performance. For example, they show that their approach fails on a 2D bimodal mixture of Gaussians problem if mode state is missing in the model's s initialization. FAB contrasts this, as the flow is able to incorporate modes discovered by AIS into the flow that were not present during initialization (see Figure 2). Notably, a single sample from the AIS bootstrap process comes from a new mode, the flow will be updated strongly towards it immediately.
In our experiments, we have used a relatively low number of intermediate AIS distributions and minimal hyperparameter search. However, the performance of FAB could be improved by further tuning. For example. by trading off the reduction in variance from more intermediate distributions with the corresponding increase in compute cost. Furthermore, since the loss variance decreases throughout training, it may be beneficial to reduce the number of AIS distributions in an online fashion.
Wimnserrgrr et al. (2022) train flows to accurately approximate the Boltzmann distribution of same- atom atomic solids with up to 512 atoms just using the target distribution's density, Similarly to our approach on the alanine dipeptide molecule, they incorporated physical knowledge about the system into their base distribution and flow architecture. In their case, this corresponds to permutation and translation invariance, periodic boundary conditions and meta-stable states. FAB could help to scale their approach to larger systems with different atom species and more complex potentials. Moreover, incorporating the chiral structure of a molecule into the model architecture might simplify training and aid in applying FAB to larger proteins. In general, including prior knowledge of the system into the model is important for FAB, it increases the effective sample size initially during training, which decreases the computational burden of AIS in reducing loss variance. Incorporating symmetries into the model has two key benefits. First, it often lets us operate in lower dimensional space, alleviating the curse of dimensionality. Second, it often greatly reduces the number of modes in the distribution, as not incorporating symmetries causes multiple repeats of "the same" mode.
This work has focused on the application of normalizing flows. However, diffusion models have also shown great promise for learning Boltzmann generators. Jing et al. (2022) are able train a single diffusion model to learn the Boltzmann distribution over the torsional angles of multiple molecules, while using cheminformatics methods for the bond lengths and angles. They perform energy- training via estimation of a score matching loss using samples from the model. As with flows, this will exhibit high variance for complex target distributions, especially during initialization when the model is a poor match for the target. Thus, incorporating an AIS bootstrap process similar to FAB may improve training in these methods.
B DERIVATION OF THE LOSS
We consider the general case of training a parameterized probability distribution qo to minimize a loss function C(O) = f fxx 0)dx, where fxx 0) > 0. Later, we will focus on the specific case of D Da=2(pllq0) minimization, where f(x, 0) = p(x)² 2 (x). however, the general case is interesting as well and simplifies notation. Let us consider the gradient of L()) written as an expectation over some distribution g(x):
VoL(O) / V Vef(( 0)dx -Eg(x)
Vof(x, 0) g(x)
We select g(x) to be the minimum variance importance sampling distribution given by g(x) fx, 0)/Z, where Z, = f(x, f)xx is the normalizing constant. Then, plugging in the identity No log (x vof(x,0)/ 0) 0), i.e., applying the log-derivative trick similar REINFORCE
13

---

Published as a conference paper at ICLR 2023
(Williams, 1992),
obtain
Vof(x, 01 Z, gix()) g(x)
Vef(x. 0) f(x,0)
g(x)
= 2, By(x) VV log fxx 0)]
(10)
We may not generally be able to sample directly from g(x), ) but we can use AIS to estimate the right part of Equation (10). To do this, we consider running AIS with p² /qq as the target and qe as the initial distribution. We note that, when the target density is unnormalized, the AIS weights are scaled by the target normalizing constant:
(AIS(XM) Pi (x1) ) P2 (x2)
PM (XM-1) PM (XM)
Po (x1) P1 (X2)
PM-2 (XM-1) PM-1 (XM)
Z1Pı (x1) Zopa (x2)
ZM- (XM-)
ZMPM (XM)
(11)
Po (x1) Zipi (x2)
2PM-2M -2 (XM-1) -MS-IPM- (XM)
-M P1 (x1) P2 (X2)
PM-1 (XM-1) PM (XM)
Po (x1) P1 (x2)
PM-2 (XM-1) PM-1 (xN)
where p; denotes the unnormalized density for the i-th intermediate AIS distribution and Z; is the corresponding normalizing constant such that p; p:/Zz.
Given that we have set PM - g. PM = f and Z = Z, expectations over the AIS forward pass
hold the following relationship to expectations over g: Eyix) [b(x)] = EAIS h(x) function of interest.
wai((x) ((x)
where
Zt
Using this we can then write Equation (10) as expectation over the AIS forward pass
Z Exxx) Vo log f(x, 0)] = z EAIS
WAIS
Ve log f(AAIS; 0)
Zj
(12)
EAIS [WAIS V log f(AAIS 0)]
where AIS and WAIS the samples and corresponding importance weights generated by AIS when targeting f. We use the bar superscript to denote stopped gradients of the samples generated by AIS, RAIS, with respect the parameters Now, returning to the case where we minimize C(0)
D Da=2(pllgo), we set f(x 0) = p(x)² 2 //o to obtain
Vo log p(x)2
V VoDa=e(olloo) = 2 Fx(x)
Zg Ex(x)
VV log [x))
40 (x))
(13)
1
Σ
AIS Vo log go (x (IS))
N
where (6) WAs and A are samples and weights generated by AIS with p² p²/go target. Equation (13) provides an unbiased estimate of (o)( we also stop the gradients of WAIS, we can then use the surrogate loss function
(i)
S(0) = - EAIS [WAis log go (XAIs))
(14)
where VoL(0) = (0)( 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 normalized weights greatly improved training stability. We refer to the surrogate loss with normalized impor- tance weights as (²))). To normalize the weights we divide them by Z, [(() EAIS [WAIS] such that f WAsdxx = 1, where the bar superscripts denotes stopped gradicnts. The relationship between S(0) and S'(0) therefore (0) = ((0)/C(0) The gradient of S' (0) has the same direction as the gradient of the original surrogate loss, S(Ø), but has different magnitude. Using S'(I) instead of S()) improves training stability by removing the effect of large fluctuations in the magnitude of [(() D ( without changing the direction of the gradient training proceeds. Thus, we use the following estimate of the surrogate loss function for training:
N
(1)
S'(0) = Σ
log qo (x
(15)
W
AIS
14

---

Published as a conference paper at ICLR 2023
where xix and A are the samples and importance weights generated by AIS but evaluated in practice using stopped gradients when computing the gradient of Equation (15). The use of self- normalization the loss function introduces bias for finite N for the estimation of (0)).. We use qe as the initial distribution for AIS, and a relatively small number of intermediate distributions to prevent the AIS forward pass from becoming too computationally expensive.
It is possible for (Pl go) to have very large value, and it can even be infinite. For example if q assigns zero density to regions in p. then (p=l(pllqg) is infinite. In practice we drop any points that have infinite/NaN density under during training. Sometimes early in training when q is a poor approximation to the target, importance weights with NaN values arose, and thus were dropped. However, towards the end of training, when & relatively accurate sampler, the importance weights were stable and infinit//NaN did not typically occur. This issue therefore did not result in any practical problems during training. Another example of when =a=2 (p]lg) can be very large, or even infinite is if target distribution is heavy tailed, and the tails of q are light. Using the normalized importance weights (for FAB without the buffer), or sampling points from the buffer in proportion to their importance weights means that FAB is only effected by the direction of the gradient of (p=l(p)) with respect to the parameters of the flow. This helps improve the robustness of FAB in situations where the magnitude of (pl qo) is very large or has very high variance in its estimates. Using an architecture for q that includes a defensive mixture component distribution with heavy tails would help address this issue, and this could be way of improving the stability of FAB further (Owen, 2013).
C VARIATIONS OF FAB
As mentioned above, in this paper we focus on the minimization of D==2 as estimated with our AIS bootstrap approach targeting p² /qg. However, the general approach of improving gradient estimation through the addition of the AIS bootstrap process may be applied in other settings. For example, in previous version of this work, we used bound on (==2(pl(qo) objective, which we estimated with AIS targeting p (Midgley et al., 2021).
In Appendix B above, the loss function in Equation (12) is written in a general manner and could therefore be used for any [x, 0) 2 and not only Do=2(pl (1=2(pll))) [go)). We could simply plug other divergence measures of the form (()) = f f(x, d) and satisfying f 0) > 0 into Equation (12). In the below section we show how we can apply the FAB approach for a-divergence minimization with other values of a.
We can also apply the proposed approach to other types of models. For example in Appendix C.2 we show how Monte Carlo (CRAFT) (Matthews et al., 2022).
obtain FAB flavored version of Continual Repeated Annealed Flow Transport
C.I FAB FOR GENERIC DIIVERRGENCE MINIMISATION
Below we provide a derivation of how the FAB may be generalized to a-divergence minimization with arbitrary values of a. We restate the definition of a-divergence:
Da(pll)) (
(16)
a(1 - ax)
We can write the gradient of the above expression as an expectation over an importance sampling distribution following the same approach as in Appendix B.
VoDa()llo)) = Vo Vo p(x) 90 9x()))))))))) "x
1 a a) gg(x)
Vo (p(x)) (x)) 1-0)
a(1 aa)
g(x)
(17)
Setting f(x, 0) = p(x)* (x)
2-) 1-a, we note that the integral in the above equation is in the form
f(x f) satisfying f(x, 8) 2 0. Thus, we consider setting gxx) to the minimum variance importance sampling distribution for estimating (p(ll)g) given by g(x) = (x()) where 7 = -_ (x))))((x))==x))) Using the result from Equation 12 can then estimate the
15

---

Published as a conference paper at ICLR 2023
gradient of VgDa(Pll9e) with AIS targeting f.
Vo VoDn(pllov) a(1 a) EAIS [WAIS Vo log [x(()[-a)] (xxx)
(18)
EAIS ACA V log gx())
We that plugging in FAB gradient from Equation 13. Furthermore, if we plug in a = which is equivalent to minimiz- ing forward KL divergence, generated by AIS with = as the target, multiplied by the AIS importance weight correction factor. This gradient exactly equal the gradient of the forward KL estimated with AIS,
the above equation gives gradient proportional to the original
that this results in maximizing the log probability of samples
Ve KL (pllgo) = Ve Ep [logp(x) log do (x)] = - Ep [Vo log (x))
(19)
- EAIS [LAIS Vo log [x(x)]
We note that this method does not hold Da-o(ollue) = KL (ullp).
We can combine the generalised FAB loss from Equation 18 with the prioritised buffer train- ing procedure from Algorithm by instead simply setting the AIS target to (x))""((x)* and (1 a) go log )0m)). We provide the pseudo code for this in Algorithm
log Weurrection
below.
Using this algorithm, in Appendix E.3 and G.2 analyse the performance of FAB with varying values of a. We find that a does best, which provides empirical support for this choice.
Algorithm 2: FAB for the minimization of D. (plg) with prioritized replay buffer
Set target p
Initialize flow 9 parameterized by 0
Initialize replay buffer to fixed maximum size for iteration
K do // Generate AIS samples and add to buffer Sample X; (1:M) from go evaluate log (x() (1:M)
Obtain A
(1.M) and log WAIs (1:Af) using AIS with target p 03 a- and seed x
(1:M) and log (x(g (1:N))
Add (1:M) log WAIS
(1:M) and log (x
(1:M) replay buffer
for iteration 1 do Sample from buffer and update go
Sample x AIS
(1:N) and log 20,0 (xc
AIS (1:N)) from huffer with probability proportional i
(1:N)
Calculate log Womeesion
(1 (stop--grad qu (x (1s)))) log qa (x (I:A))
Update log AAA
(1:N) and log u (xxss) in buffer to log was (1.N) + log w coneeioo and log go (x
(1-N)
Calculate loss S' (0) - 1/N E. Womeection log qu (x((x)
(()
Perform gradient descent on S' (0) to update
C.2 FAB APPLIED TO CRAFT
Continual Repeated Annealed Flow Transport Monte Carlo (CRAFT) is an extension of SNFs (Wu et al., 2020; Nielsen et al., 2020) proposed by Matthews et al. (2022) which combines normalizing flows with sequential Monte Carlo (SMC). Specifically, flows are used to transport samples between consecutive annealing distributions in combination with SMC. CRAFT trains each of the flows by minimizing the reverse KL divergence with respect to the next annealing distribution and gradients 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 be replaced with mass- mass-covering divergence such as the a-divergence with 2. Furthermore, could improve the estimation of gradients by targeting with AIS the minimum variance distribution for importance sampling instead of just the next annealing distribution.
A FAB style version of the CRAFT algorithm is described in algorithms and 4. Each flow transport step now trained by minimizing the a-divergence with 2 and using samples generated by the
16

---

Published as a conference paper at ICLR 2023
next immediate MCMC step ahead of the flow in the SMC process. Below, we briefty introduce CRAFT with an emphasis on the loss function used. Next, we describe how FAB may be used to improve training of the flow transport steps. We follow the notation from the CRAFT paper exactly and use their pseudo code as a basis for our proposed algorithm indicating changes clearly. We refer to (Matthews al., 2022) for further details the CRAFT algorithm.
C.2.1 CONTINUAL REPEATED ANNEALED FLOW TRANSPORT MONTE CARLO
As in FAB, the aim of CRAFT is to approximate intractable target distribution that we cannot sample from and whose density can only be evaluated up to a normalizing constant. This target distribution is denoted T (x) (equivalent
p(x) in notation). In CRAFT, SMC is run with interleaved flow transport steps through a sequence of annealed distributions (NH (x)) each with normalization constant z The base distribution
is tractable distribution (e.g., Gaussian) AIS, (TTx 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 XO-TO- Then, for each distribution from k = 1 to k K, flow Tk is trained to transport samples from 1 to T& (x). Additionally, each step from &
from which we can sample. Similarly
k K, CRAFT utilises resampling and MCMC to bring the samples
closer to T& (x).
Similarly AIS, the CRAFT algorithm returns set of points and normalized importance weights (Xi. W which may be used for approximating expectations with respect to the target. Each point X has an associated normalized importance weight W for importance sampling with respect the intermediate target distribution T& (x). We refer back to the CRAFT paper, and the pseudo code in Algorithm 3 for how these importance weights are calculated.
CRAFT minimizes the following training objective:
H
Σ
|L. TR T [T])]
(20)
where # denotes the push forward between distributions. The above objective trains each flow transport step T& to minimize the KL divergence between F5 πk-1. i.c. the distribution of outputs of the flow when given input samples from mT-11 and the next intermediate distribution πy. The gradient estimate used to train each flow transport step Ta is given by
Vo H
Σ W Wh-iVos Vo, [- log 7h (Tx(XE- log V (xt((x)-))),
(21)
where (x) x (x). The flow is trained by passing samples X from the previous SMC step, computing the corresponding output samples from the flow Tx (XR-1) and using these to estimate the gradient of K - The normalized importance weight in the loss W W account for the fact that the samples X passed from the previous step in the SMC forward pass come from an approximation to T
C.2.2 FAB-CRAFT
We now propose a FAB flavored version of CRAFT. First, we re-introduce some notation from our paper: We use g denote the initial distribution used in AIS and p to denote the target distribution that we wish to approximate. Recall that is trained to fit p. All other notation in this section follows the CRAFT paper's '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 Do=2(p)))) where _ T and p To do this with minimal changes to the original CRAFT algorithm, we make the observation that sampling from the initial distribution and then running MCMC targeting is equivalent to running AIS targeting p² //q with 1 intermediate distribution at 3 0.5. Thus, the samples generated by the MCMC steps following each flow transport step in CRAFT can be repurposed for an AIS bootstrap estimate of the flow training loss. For training. the only adjustment to the SMC forward pass of CRAFT is then to move the resampling step occur after each MCMC step, where previously it
17

---

Published as a conference paper at ICLR 2023
occurred after each flow transport step. At inference time the original CRAFT algorithm can be run with the flows trained with our method in its exact original form. We describe this in more detail below and provide pseudo code in Algorithm 3 and 4.
We begin by deriving the AIS importance weights when targeting p²/q with intermediate distri- bution and setting /3 0.5. For only 1 intermediate distribution, the AIS weights given by
(x1) P2 (x2)
WAIS (x2)
(22)
Po (xi) P1 (x2)
As before, we set the intermediate distributions as interpolations between the base and the target: log (x) 3, log (x( (1 - {.) log (x)) Now, if we set 61 0.5, then plugging in po 9. P1 q².5 (1-(pp//g) (p²/q) p and P2 p² /q, we obtain the following AIS weights:
(x²) = p (x1) P (x2)
(23)
q (x1) q (x2)
Recall that CRAFT we set q = _ Th-1 and = Tke AIS then run by first sampling x from 4. which done in practice by setting Y + ( where X ~ Th-1 and then generating X2 from X₁ by MCMC, which is done in practice by setting X ~ K&x)), where Kg an MCMC
transition kernel that leaves Ty invariant. The importance weights of X with respect to the AIS target p² q then given by
" AIS p(22)/((()) p(X))/((Xx).
(24)
Using the normalized importance weights W Wis.kK FAB gradient estimate given by
/ Ei
w x we can calculate the
h
Σ
W & Vol log qo, (Xi)
(25)
N
Σ
&A Vo,
+ log V
This assumes that the samples passed to the flow are from the distribution T4-1- However, in practice, these samples are passed from the previous SMC step which is approximation to TT- Similarly as in the original CRAFT loss, see Equation 21, we can correct for this by instead using
N
hus = Σ W WAIS, Vo, [lognk-1
7k- 7k- (T (xx) + log / V (Xi)
(26)
where k-1 accounts for X coming from 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 step instead of just after the flow transport step, see Algorithm 3. Note that the weights W3 for resampling X₁ are equal to the weights for resampling the corresponding flow outputs Y; that generate such samples. This result is due to the MCMC kernel IC leaving ग&. invariant. The resulting FAB flavor of CRAFT is shown in algorithms 3 and 4.
Some final remarks:
our goal has been to create a FAB flavored version of CRAFT while keeping the algorithm as similar to the original version possible. However, in practice, would be better to make further changes. For example, using a prioritized replay buffer would significantly decrease the computational requirements of the algorithm. Furthermore, for updating each flow, it may also be beneficial to consider samples across the whole chain of intermediate distributions, instead of using only samples from the local MCMC step immediately following the flow.
D ANALYSIS OF FAB
D.1 GRADIENT ESTIMATION PERFORMANCE
We first analyze the quality of the noisy gradients provided by the proposed AIS bootstrap method. For this, consider toy problem in which go and p are unit variance ID Gaussians with means 0.5
18

---

Published as a conference paper at ICLR 2023
Algorithm 3: SMC-NF-step for FAB-CRAFT: Additions are in green and removals in red
1: Input: Approximations ( 1: Z to (x--1, normalizing flows Tk. unnormalized annealed targets 7 and "Yk and resampling threshold A E [1/N, 1).
2: Output: Gradient hgg of FAB loss w.r.t 8g, particles at iteration k: Tk (X) W
approximation _ to Zz.
3: Transport particles: Y = Ta (X
4: Compute IS weights:
w/.- < G (Xi k-1) unnormalized
W < 14///m1 w // normalized
5: Estimate normalizing constant Zμ:
N _ 1 (Ei= -
6: Compute effective sample size ESS.
7:
if ESS & .N then
8: Resample N particles denoted abusively also Y according the weights W then set
W N
9:
end if
Generate samples and IS weights via AIS targetting p² /qe, with intermediate distribution,
where = Te and q = T Th-1- By setting / = 0.5 we simply run the original CRAFT MCMC transition kernel with p target.
Sample X Ky k1 -). MCMC
AASS - /(x)))/((X()) x ((x))/((X))
W AIS.A Ae Σj=n EXISA.
Estimate gradient of FAB objective
hu - E; W WÄss, Vou [logk-1((2) (T (x+)) + log V T
(X))
02
12: if YESS NA then
13:
Resample N particles denoted abusively also X according to the weights W then set
W N
14: end if
Return
Tk N 1a.).
Algorithm 4: CRAFT-training: Additions are shown in green and removals in red
Input: Initial NFs [Tk]..N number of particles N, unnormalized annealed targets {rk 1== with T0= and =k= resampling threshold A E [1/N, 1).
2: Output: Learned flows T. and length J sequence of approximations (m,, ( Zx). do 4: Sample Xi~TT and set Wi- and Z -1.
3: for = 1,
5: for k 1,..., K do
6:
he + flow-grad (Tk, ) using eqn (21).
F N + SMC-NF step (T 2 Tk)
7:
8:
Update the flow T using gradient hx. end for
9:
10: Yield (T ) and continue for loop.
end for
12:
Return learned flows K
and -0.5, respectively, as shown in Figure 5a. We estimate the gradient of Da=z(pl go) with respect to the mean of qo and compare different methods: first, importance sampling (IS) with samples from qo: second, IS with samples from pt third, AIS with p as target and qo as initial distribution; and fourth, proposed method using AIS with p² /qo as target and go initial distribution. For AIS
19

---

Published as a conference paper at ICLR 2023
use 3 intermediate distributions and, transition operator, HMC with 5 leapfrog steps and resampling of momentum variables once per intermediate AIS distribution.
Figure 5b shows the Signal-to-noise ratio (SNR) for the different gradient estimators function of the number of samples used. AIS bootstrap is clearly the best method. with do performs very poorly and it outperformed by both IS with p and AIS targeting p, with these two latter techniques performing similarly but way worse than AIS bootstrap. Figure 5c shows that the quality of the proposed method increases fast the number of intermediate AIS distributions grows, with IS with samples from being outperformed quite early in the plot while still using rather small number of distributions. is important to note, however, that in more challenging problems, is unlikely that our AIS bootstrap method will outperform IS with samples from especially early in training when qo is poor approximation to p.
0.40
100
9
IS with
17.5
0.35
IS with
80
15.0
AIS with =
ATS with g /4/4
12.5
0.25
PDR
0.20
SNR
SNNP
10.0
0.15
7.5
0.10-
5.0
20-
2.5
0.00
0
0.0
Number of samples
Number of AIS distributions
(a)
(b)
(c)
Figure 5: (a) Target and approximating densities p and q. respectively. (b) Signal-to-noise ratio (SNR) of various gradient estimators as function of the number of samples used. (c) SNR as a function of the number of intermediate AIS distributions when using 100 samples. Legend as in (b).
D.2 SCALING FAB HIGHER DIMENSIONS
Now we consider how the performance of FAB is affected by an increasing problem dimensionality. To investigate this, we analyse simple scenario where we assume p and q to factorised, with each marginal of having its separate parameters (no parameter sharing between dimensions). We acknowledge that these are strong simplifying assumptions and leave more general analysis to future work. Given these assumptions, the gradient of Da=2 by importance sampling with samples from q and p increases exponentially with respect the dimensionality of the problem. 2) The variance in the corresponding estimates obtained with FAB can remain constant the number of AIS intermediate distributions increases linearly with the dimensionality of the problem.
show the following: 1) The variance in the estimates of
D.2.1 THFORETICAL ANALYSIS ON FACTORIZED p AND 4
We consider the problem of estimating the gradient of Da=2 where both p and are factorized distributions: they are equal to the product of their marginals. To further simplify the analysis, consider the gradient with respect to the parameters the j-th dimension and assume that there no parameter sharing across dimensions of q. In this case, Da is given by
D p == ==/nbe [(x)2 dx = II ((x)
pa(sa)? dxa, qa(ara)
(27)
20

---

Published as a conference paper at ICLR 2023
where Pa and qu the marginal distributions for dimension d. The gradient of this quantity with respect to the parameters f of the j-th marginal of is given by,
p) (1()g))
pa(xa)² (a(ra)
Vo Vo,
day
II dzj"
dra
(28)
We are interested in studying how the variance in our estimate of V0,,4 scales with D. Equation (28)
shows that each additional dimension d o j adds extra factor μla = pa(ar))² (alira))/ga(zaa) /qa(xa) dza in the gradient expression. To eliminate the effect of this change in the gradient as D increases, we focus our analysis on Var Vo, /1/ ITaz; pa), where já is an estimate of μ. and ITzz Ha is a normalization factor that cancels the effect of the additional dimensions.
Importance Sampling with q: We consider first how increasing D affects the variance in the estimation of the gradient of Dx=2 by importance sampling with samples from The importance sampling estimate of Da=2 with N samples from g given by
N
Fg
w uxx))²
(29)
N
n=1
where w(xx) p(xn)/g(xx.) and x,, ~ 9. Now, let us define wa(En,d) = where andd the d-th entry in X,. We then obtain
1 Vo, N
ΣΠ
walna(n 1 Vo, D II walanna)
Vo,Pg
(30)
d
The variance of this estimate, after dividing by the aforementioned normalization factor, given by
VoyKo
Egg
(aa.) (Vo, (31)
N Var
Egg [V, wy (D 2/
II d
Ilat Md
1. Thus, the first factor in Equation (31) is multiplied in this equation by D factors all larger than 1, which implies that the variance of this estimator increases exponentially function of D. This is a well-known problem of importance sampling.
Since Varga [wa(Fra,))2] Ey [ (z 9
O, have that Egg [wa( )1 /
Importance Sampling with p: Interestingly, we get a similar result when estimating Da=2 by importance sampling with samples from p. The estimate for the gradient of Da=2 with respect to the j-th dimension of q is now
1 N
Vo
Σ V0
II wa(Fnn)) dØf
(32)
where æn,d ~ Pd- The variance of this gradient estimate, after dividing by the normalization factor, given by
N Var
Vo,Kp D I1 Md
= Ep, [[Vo, (zy(m)))))) x II dz.
D
Epa [wa(snaa)
(Vo,1152.
(33)
1
Since Varp, (wa(xan))) = Epg [waa(rd,..]]] - > 0, we have that Ep,, [wa(x...)) /r3 2 1. Thus, the first factor Equation (33) is again multiplied by that the variance of this estimator increases exponentially factors all larger than 1, which implies function of D, albeit at lower rate than in the case of importance sampling with q. This implics that, even with access to ground truth samples from p, the number N of samples required to keep the variance of the gradient estimates of Da=2 constant grows exponentially as a function of the problem dimensionality D.
We apply the same type of analysis to the estimates of the gradient given by FAB. We consider the FAB gradient estimate from Equation (13), which is equal in expectation to Vop. This estimate relies the raw importance weights from AIS rather than the self-normalized importance
FAB:
21

---

Published as a conference paper at ICLR 2023
weights, which makes it easier to analyze. The FAB estimate of the gradient of Da=2 with respect to the parameters of the j-th marginal of q is given by
1 Σ log (xx))
Vo, /FAB
Voz N
(34)
N Σ
Vo, log oj (Zn.g) )uy (Zn.)) x II Wa(=ndd)
where w are the importance weights from AIS with stopped gradients and we have decomposed the AIS weights into the contributions from each dimension: t (xx,) = IIa @z(Zn.a). Note that, since p and y are factorized, we have that the intermediate AIS distributions are factorized as well. If we assume that the MCMC transition kernels in AIS produce independent samples from the ground truth intermediate target distributions, we have that 2n,1, Dn are independent random variables. The variance of Vo, (FAB after dividing by the normalization factor is then given by
Var
Vo,yens PEAB Toey Hid
EAIS,
[[V,, log z (zz,n) hy ( ()
x II
EAIS [el(zna))]] (Vo,14)), (35)
where EAIS denotes the expectation with respect to the AIS forward pass. The first expectation in 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 where used importance sampling with samples from q and p. the variance in Equation (35) will again increase exponentially with D. However, under the assumption that the MCMC transitions produce independent samples from the intermediate AIS distributions, Neal (2001) shows that the variance in the log importance weights of AIS is proportional to D/K where K is the number of intermediate AIS distributions. This implies that Iaz; EAIS, [Wa(=n.a)) )2/pp) will remain roughly constant if we increase the number K of AIS distributions by the same factor as the dimensionality of the problem. In this case, the variance of e, MFAB will remain constant as we increase D and the cost of the FAB gradient estimator will only increase linearly as D increases.
D
D.2.2 EMPIRICAL ANALYSIS ON TOY PROBLEM
We now run an empirical analysis to assess the performance of the FAB gradient estimator as D increases. We consider the case where q and p are both factorized Gaussians with unit marginal variances and with mean vectors equal to 0.5 D and -0.5x where P is a vector of dimension D with all of its entries equal to one. For the FAB gradient estimate, we increase the number of AIS distributions by the same factor as the dimensionality of the problem. We analyze the SNR in the estimates of the gradient of the mean for the first marginal of q as the dimensionality of the problem increases. The AIS transition operators are performed by running a single iteration of HMC, with 5 leapfrog steps, with a step size 0.5. We found our results to be sensitive to the choice of step size. In practice, we selected this value by trial and error, assessing the quality of the AIS samples by looking at their empirical histogram, as shown in Figure 7.
Figure 6a shows that, when using importance sampling with samples from p. the wei-ge variance increases linearly as the number of dimensions increases. By contrast, this variance remains constant with FAB. This is achieved by fixing the number of AIS intermediate distributions to be equal to the number of dimensions. This result is consistent with the analysis from the previous section and with the 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 sampling with samples from p. If the same results were hold in more complex problems, this would imply that we could safely apply FAB in those settings by linearly increasing the number of intermediate AIS distributions as the dimensionality of the problem grows.
E MIXTURE OF GAUSSIANS EXPERIMENTS
E.1 TRAINING SETUP
All flow models have 15 RealNVP layers (Dinh et al., 2017), with 2 layer (80 unit layer width) MLP for the conditioner. The flow is initialized to the identity transformation, go is initially
22

---

Published as a conference paper at ICLR 2023
Number of AIS distributions
Number of AIS distributions
0
20
40
20
40
50
AIS with g 9²/9
40
IS with p
Var (log w)
30
SNR
4-
20
10
0
0
20
40
20
40
Number of dimensions
Number of dimensions (b)
(a)
Figure 6: Analysis of efficiency of AIS bootstrap with increasing dimensionality. For AIS, we set the number of AIS distributions be equal to the number of dimensions of the problem-which results in a linearly increasing compute cost as the dimension scales. (a) Variance in the log importance weights. for the AIS bootstrap vs importance sampling with p. (b) Signal-to-noise ratio (SNR) for gradient estimation with the AIS bootstrap vs importance sampling with Legend as in (a).
0.4
0.3
q
PDF
0.2
p
g cc p²
0.1
AIS with g =-p2/ p2/9
AIS with =
0.0
-5.0
-2.5
0.0
2.5
5.0
X
Figure 7: Histogram of samples from AIS targetting and p² compared to the PDF of p. q and (normalized) p² /9/q. The AIS samples are generated with the tuned HMC step size of 0.5. This tuning was performed by trial and error using the displayed histogram. AIS is run with 4 intermediate distributions, and the transition each intermediate distribution performed by running a single iteration of HMC with 5 leapfrog steps.
standard Gaussian distribution. Training is performed with batch size of 128, using the Adam optimizer (Kingma & Ba, 2015) with learning rate of 10-4 and we clip the gradient norm to maximum value of 100. For the model that uses a RBD, we use acceptance function composed of a residual network with three blocks containing 512 hidden units per layer. The truncation parameter is set to the common value T 100. For the SNF and CRAFT methods, we do Metropolis- Hastings step every three flow layers.
used fixed step size a = 5.0 for the Gaussian perturbation of the Metropolis-Hasting step, which is the same what is used within AIS for FAB. This means that the SNF and CRAFT models has stochastic Metropolis-Hastings Hastings steps in total. The CRAFT model uses 6 annealing temperatures with geometric spacing, and resampling threshold of 0.3. We use the code provided by (Matthews et al., 2022) at ht https: /github. com/ dewpmiiepmmmmmmmmm nd annea led flow transport for training the CRAFT model. We train all fed_sfot_tttaasppp
23

---

Published as a conference paper at ICLR 2023
models for 2 10 flow evaluations. For each method, we train 3 models, each with different random seed, and results are reported as averages over these seeds.
FAB specific details: The batch size for both the AIS forward pass (M) and sampling from the buffer (N) is equal to 128. We run AIS with a single intermediate distribution (3 0.5) and MCMC transitions
given by single Metropolis-Hastingg step: Gaussian perturbation and then an accept-reject step. We used a fixed step size a = 5.0 for the Gaussian perturbation. We initialize the buffer with 1280 samples from the initialized flow-AIS combination and maximum buffer length of 12800. We do not use any clipping when computing Wccorection- The log density of the flow occasionally gave NaV values to points sampled from the buffer, resulting in NaN values for Woorretion- As this resulted in NaN loss values, the parameter update was skipped in iterations where this occurred. Furthermore, since the Weooeetion adjustment for these points in the buffer is invalid, the weights and a0 (x) values in the buffer were left as their previous values.
Table 3: Number of flow and target evaluations during training for each method on the mixture of Gaussians problem. For SNF and CRAFT we report the number of model forward passes each of which contains multiple flow transport and MCMC steps.
Number of flow/model evaluations
Number of target evaluations
Flow w/ ML
2. 107
2 10"
Flow w/ Do=2
2. 2
2- 10
Flow w/ KLD
2-107
2 107
RBD w/ KLD
2-107
2 2-10
SNF w/ KLD
2-107
10
CRAFT
2. 2.107
10
FAB w/o buffer
20
2-
FAB w/ buffer
2.107
6.6 x10
E.2 EVALUATION SETUP
For each method, we compute after training the effective sample size (ESS) obtained when doing importance sampling with qo: the average log-likelihood of on samples from p; the forward KL divergence with respect to the target; and the mean absolute error (MAE) in the estimation of Ep(x) [f(x) by importance sampling with 1000 samples from qo. where f(x) = a (x- 1xx-bb 2b) + (x - 2-)bbb (x - 2b). with the entries in vectors a and b and matrix C randomly initialized by sampling from standard Gaussian then kept fixed to such values during all the experiments. We express the MAE as percentage of the true expectation to make it easier to interpret. We also report the MAE that is obtained when do not reweight samples according to the importance weights. The ESS is calculated using 5 x 104 samples from qo. The MAE is calculated by averaging over 100 repetitions.
E.3 FURTHER RESULTS
Figure 8 shows a plot of samples from each trained model on the mixture of Gaussians problem, with the 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 Da=2 minimization exhibited highly unstable behavior during training and, thus, is the worst performing model.
FAB with varying values of a In Appendix C.I, we derived a variant of FAB that works with an arbitrary value for the parameter of the a-divergence. Here, we want to investigate how the performance of models trained with FAB changes as we vary a. Therefore, we leave the setup the same as used in Section 4.1 and only changed the parameter of FAB. The results when using FAB without the replay buffer are given in Table 4 and Figure 9. The results when using FAB with the replay buffer are reported in Table 5 and Figure 10. For FAB without the replay buffer,
2 is slightly superior in performance to FAB with other values of a. For FAB with the replay buffer, we see that all of the methods with 2 1 are able obtain good fit for the target, with 2 and
24

---

Published as a conference paper at ICLR 2023
Flow w/ ML
Flow w/ Da-2
Flow w/ KLD
RSB w/ KLD
SNF KLD
CRAFT
FAB w/o buffer (ours)
FAB w/ buffer (ours)
Figure 8: Contour lines for the target distribution and samples (blue discs) drawn from the approx- imation qo obtained by different methods the mixture of Gaussians problem.
achieving the best performance. For these runs, the limits of the expressiveness of the flow is most likely the limiting factor to improving performance even further. The same style of analysis for FAB with varying values of performed with the Alanine Dipeptide problem in Appendix G.2 which finds a is best, with a larger differences in performance between different values of a.
α 0.25
0.5
a 1=5.5
a = 2
==
Figure 9: Contour lines for the target distribution and samples (blue discs) drawn from the ap- proximation qo obtained by FAB without the replay buffer for varying values of on the mixture of Gaussians problem.
25

---

Published as a conference paper at ICLR 2023
0 = 0.25
a = 0.5
α2
a 1.5
a a=
α =
Figure 10: Contour lines for the target distribution and samples (blue discs) drawn from the ap- proximation qo obtained by FAB with the replay buffer for varying values of on the mixture of Gaussians problem.
Table 4: Results for the mixture of Gaussians problem for FAB without the replay buffer for varying values of ligeliiid values for the first two methods are NaN because they assign density to samples from missing modes. Best results are emphazised in bold.
ESS (%)
Ep(x) [log (xx)
(L((p1g)
MAE (%)
MAE w/o RW (%)
0.25
52.4
10.0
NaN NaN
NaN + NaN
44.2
48.8
0.5
14.7
NaN + NaN
NaN + NaN
14.6+3.66
15.9
1.0
30.8 2.2
-7.59 0.08
0.73 + 0.08
4.8
4.3
1.5
19.0
-..2
1.10
6.2
13.4
2.0
382
-..6
0.73
3.7
1.3
3.0
2.9
-8.09
1.23
5.8
19.3
F
MANY WELL EXPERIMENTS
F.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 of the 2-dimensional Double Well distribution from Wu et al. (2020); Noć et al. (2019):
(2) -zt 6z² + 1/2x - 1/2x2 + constant,
(36)
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 shown in the top-right contour plot in Figure I.. Therefore, our 32-dimensional Many Well has 216 65536 modes, one for each possible choice of
'We the Double Well distribution from the code provided in Wu al. (2020), which has different coefficients the Noć al. (2019).
26

---

Published as a conference paper at ICLR 2023
Table 5: Results for the mixture of Gaussians problem for FAB with the replay buffer for varying values of a. log-liielihood values for the first two methods are NaN because they assign zero density to samples from missing modes. Best results are emphazised in bold.
ESS (%)
Fx(x) [log [x))
(l.(pllg)
MAE (%)
MAE w/o RW (%)
0.25
44.4 18.2
NaN + NaN
NaN + NaN
NaN
91.7+3.4
915
0.5
4.0
NaN NaN
NaN NaN
3..
3..1
1.0
59.1
0.00
0.33 0.30
10.1
11.4
1.5
40.8
0.02
3.6+:.11
2.8
2.0
63.1
-.014=0.02
0.28
3.0
3.0
3.0
6.7
0.03
0.00
3.2
2.8
mode in each of the copies of the Double Well. We obtain exact samples from the Many Well by sampling from each independent copy of the Double Well. Exact samples from the Double Well are obtained by sampling from each independent marginal distribution. The first marginal p(x1) can be sampled from exactly using rejection sampling (see Appendix F.2). while the second marginal distribution p(z2) can be sampled from directly as is a (unnormalized) standard Gaussian. These samples cheap produce. We use them for training flow by maximum likelihood as well for the evaluation of the different methods. Additionally, we created an artificial test set for evaluation purposes by manually placing a point each of the 216 modes. By computing log-likelihoods on this test set, we can then check if a method is covering the entire target distribution, as any missing mode will result in very low log-likelihood values. We can calculate the normalizing constant for each marginal of the Double Well problem via numerical integration (for p(z1)) )) and analytical integration (for p(rr)), and use this obtain the normalizing constant of the Many Well distribution (see Appendix F.2). This may then be used to obtain the normalized probability density function of the Many Well distribution, which useful for model evaluation. We can also compare models on how accurately they estimate the normalizing constant the average unnormalized importance weights. For each model, we report the MAE in the estimation the Many Well's 's normalizing constant using 1000 samples, averaged over 50 runs. We express this percentage of the true value the normalizing constant.
Table 6: Results the 32 dimensional Many Well Problem. Our methods are marked in italic. Best results are emphasized in bold. CRAFT (config 1) refers to the CRAFT model trained with a similar flow and MCMC config FAB. CRAFT (config 2) uses the configuration provided in the CRAFT repository, that uses more expressive neural spline flows, and a larger number of intermediate distri- butions within SMC.
ESS (%)
Ep(x) [log qxx)]
Mean log (X
[L[pllq]
MAE (%)
Flow w/ ML.
2.1
27.6 + 0.01
21.3 0.0
0.1
1.2 0.0
Flow w/ D==2
0.0
NaN
NaN=ENaN
NaN=NaN
NaN
Flow w/ KLD
27.7
-176.9 18.08
-536.6+44.22
14..4+18.1
8.2
RBD w/ KLD
51.9
-183.8 21.8
-533.7
156.2+21.8
8..9
SNF w/ KLD
19.1
N/A N/A
N/A
167.9
88.9
CRAFT (config
N/A N/A
N/A + N/A
N/A
N/A
1.6.3
CRAFT (config 2)
N/A + N/A
N/A + N/A
N/A N/A
N/A
3.0 1.7
FAB w/o buffer
-2..8 0.7
43.0
2.3
12.7
FAB w/ buffer
7..
- -77.6+
-2.
0.1
14+0.1
We compare FAB to the same alternative approaches as in the mixture of Gaussians problem and use also the Real NVP flow architecture but with 10 layers. The MLP used for the conditioner is composed of 2 layers each with 320 units. For the the model that uses a RBD, we use the same ar- chitecture before, with acceptance function composed of residual network with three blocks containing 512 hidden units per layer. The truncation parameter is set to the common value T = 100. AIS with 4 intermediate distributions (linearly spaced) and with For FAB based methods,
27

---

Published as a conference paper at ICLR 2023
HMC transition operator containing a single iteration with 5 leapfrog steps. For FAB with pri- oritized buffer we use L = 8. The SNF model uses step of HMC with inner leapfrog steps every 2 layers. As before, when training the flow by maximum likelihood, from the target for each loss estimation. All models except CRAFT are trained for 1010 flow evalu- ations with the number of target evaluations by each method being reported in Table 7. For CRAFT we run our experiments with two different setups, which we simply refer to as CRAFT (config 1) and CRAFT (config 2). The first uses a configuration that is similar the that of FAB in terms of the flow architecture, and MCMC. This CRAFT model uses 5 temperatures, and uses an auto- regressive affine flow. The second setup uses the configuration provided in the CRAFT implemen- tation at sttps://gtt hig..b................ com/ deepmi nd/ annnaledd__. www.llowwwssssseeee neural spline flows, with 11 temperatures (10 flow/MCMC steps) and HMC containing 10 leap- frog steps. This model is significantly more expensive both in terms of the flow, and in terms of the MCMC performed in each forward pass. For CRAFT we train for 10 target evaluation budget- this is the same number used in FAB without the buffer, and slightly more than FAB with the buffer. Further details on the hyper-parameters and architectures used by each algorithm are provided in Appendix F.3.
draw new samples
which uses
Figure 11 shows contour plots for several two-dimensional marginals of the Many Well target. Each plot is obtained by scanning two variables that 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 figure samples generated by FAB with a replay buffer (left) and by the method that tunes qo by minimizing (L(()))) (right). We see that FAB generates samples on each of the contour modes while this is not the case for the alternative baseline, which misses several modes. Figure 12 shows the same contour plot as above for the two CRAFT models, both of which successfully sample from all the modes. Additional plots for all other methods can be found in Appendix F.4.
Table 6 shows for each method 1) the ESS when doing importance sampling with qa; the average log-likelihoods for qo 2) on samples from the target and 3) on test points placed on the modes of the Many Well distribution; 4) the forward KI. divergence with respect to the target; and 5) the MAE in the estimation of the May Well normalizing constant. Average log-likelihoods and ESS are calculated with 5 x 10at 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 the benchmark of training the flow by maximum likelihood. These are the two best performing methods, obtaining the highest ESS and average ligeiioiss and the lowest forward KL divergence and MAE values. The next best method is the CRAFT (config 2) model, and then FAB without buffer, while the other methods perform very poorly. The method that minimizes 2 as estimated by sampling from qe diverged early in training and always returned NaN values. The ESS for the flow trained by minimizing KL(gellp) RBD and the SNF are spurious, as they are missing modes (see Figure 11 and 14). After training, we may combine the trained flows with AIS to further improve the ESS. If we run AIS during training but targeting p instead of p² /q. the ESS is 89.9% and 13.6% for the FAB flows trained with and without a buffer, respectively. The log-likelihood of p samples from 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, showing that this method produces highly accurate approximations the target distribution.
Table shows that the CRAFT (config 2) model, which uses the more expressive spline flow ar- chitecture, and larger number of intermediate distributions in SMC, performs well and is able to provide accurate estimates of the normalizing constant for the Many Well. However, the flow trained with FAB provides more accurate estimates, even though the CRAFT model runs large amount of HMC at evaluation time. This comparison could be made more fair by taking the flow trained with FAB, and then at inference time combining it with AIS targetting p.
F.2 OBTAINING THE NORMALIZING CONSTANT AND EXACT SAMPLES
In this section, we describe how to obtain the exact normalizing constant and exact samples from the Double Well distribution. These allow us to obtain the normalizing constant and exact samples from the Many Well distribution.
The Double Well lon-density is given by
log (21, 11 Gaq + 1/2m1 1/2x3 + constant.
(37)
28

---

Published as a conference paper at ICLR 2023
FAB w/ huffer (ours)
Flow KL.D
2
2
X4
Figure 11: Samples from go and target contours for marginal distributions over the first four elements of in the 32 dimensional Many Well Problem. The flow trained with FAB with buffer (left) captures the target far better than the flow trained by minimizing KL(gllp) (right), which misses several modes.
By noting that and T2 are independent, we see that their distribution factorises as p( 222 p(z1)p(12) Thus, the normalizing constant of p(x1, is given by the product of the normalizing constants of each marginal distribution. Furthermore, samples from p(11, 22) may be obtained by sampling independently from each marginal. By inspection, we see that the marginal p(x2) is standard Gaussian. Thus, its normalizing constant given by V2T, and samples from this marginal may be obtained trivially. The normalizing constant of the second marginal distribution may be calculated via numerical integration Z = 11784.51. We obtain exact samples from p(z1) by using rejection sampling, a visual summary of this is provided in Figure 13. For the rejection sampling proposal distribution, denoted g. we use a two-component Gaussian mixture distribution with mixture weights (0.2, 0.8), means (-1.7, 1.7) and standard deviations equal to 0.5. For the comparison function kg(z1), we set k 3Z, to ensure that kg(x1) > p(x1).
F.3 SETUP
All flow models besides CRAFT 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, batch size of 2048 and using the Adam optimizer (Kingma & Ba, 2015) with a learning rate of 3 x 10-4 We clip the gradient norm to a maximum value of 100. For the SNF method, we do 10 Metssnolis-Hastings
is initially a standard Gaussian distribution. Training is performed with
Hastings steps every two flow layers. We train all models for 10 flow evaluations. For each method, we train 3 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 from the buffer (N) are both equal to 2048. We run AIS using four intermediate distributions with linear spacing. Each MCMC transition is given by single Hamiltonian Monte Carlo step consisting of 5 leapfrog steps. The momentum variable for HMC is sampled from a standard Gaussian and is not tuned throughout training. However, important parameter to tune in HMC is the step size parameter 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 and
29

---

Published as a conference paper at ICLR 2023
CRAFT (config I)
CRAFT (config 2)
2
-3
2
X4
Figure 12: Target contours and model samples for 2D marginals over the first four elements of the 32 dimensional Many Well Problem for the CRAFT models. CRAFT (config 1) refers the CRAFT model trained with similar flow and MCMC config to FAB. CRAFT (config 2) uses the configuration provided in the CRAFT repository, that uses more expressive neural spline flows, and larger number of intermediate distributions within SMC. Both CRAFT models sample well from the modes in the pairwise marginal distributions.
(normalised)
20000
kq
1.50
sample density
1.25
15000
PDF
1.00
10000
0.75
0.50
5000
0.25
0
0.00
-2.5
0.0
2.5
-2.5
0.0
2.5
X1
Figure 13: Using rejection sampling to obtain exact samples from the first marginal of the Double Well distribution. (LHS) We see that kg(z1) > p(x1). (RHS) Sample density (normalized histogram using 10000 samples) the normalized probability density of p. We that rejection sampling provides exact samples from p(z1).
shared across intermediate distributions. In more detail, we define the HMC step size for the th intermediate AIS distribution as 6, = 6stared + én. where we define estared and é, as follows: Estared is shared parameter across all AIS transitions, and is updated at the transition for every intermediate distribution. This allows for faster adaption the step size the step sizes for all transition kernels are "too big" "too small", which is common at the start of training. é, is
30

---

Published as a conference paper at ICLR 2023
parameter specific to each n-th intermediate distribution, and is only updated during its specific transition. This allows for En to be tailored to the specific n-th intermediate distribution transition. We found this parameter sharing for each intermediate distribution is initialized with a step size of 1.0, where we set ésared = 0.1 and 61:N-1 0.9. The step size is then tuned target Metropolis acceptance probability of 0.65. For the transition corresponding to each intermediate distribution, if the average acceptance probability across the batch greater than 0.65, we set é, I.05é,, and éshared 1.a2eshared If the average acceptance probability across a batch is lower than 0.65, we set én é, 1.05, and Stared Cstured Adapting shared parameters across the AIS forward pass violates Markov property, the transitions late in the MCMC chain will have weak dependency on the earlier transitions. However, the effect of this minor as the step size changes are relatively small for each run. For evaluation the AIS parameters are frozen, so it respects the Markov property.
improve performance in practice. The HMC transition kernel
In FAB with prioritized buffer, we use a total of We initialise the buffer with 65, 536 samples from the initialized flow-AIS combination and use maximum buffer length of 512, 000. We do not use any clipping for Weonrection
gradient update steps per AIS sampling step-
CRAFT specific details: For CRAFT we our experiments with two different setups, which we simply refer to as CRAFT (config 1) and CRAFT (config 2). The first uses a configuration that is similar the that of FAB in terms of the flow architecture, and MCMC. This CRAFT model uses 5 temperatures, and uses an auto-regressive affine flow. The second setup uses the configuration provided in the CRAFT implementation at ht htt///itttul t.ub.uub..... - com deepmi nd annea 1 ed flow_traae www_wtrassppee transport, which uses neural spline flows, with 11 temperatures (10 flow/MCMC steps) and HMC containing 10 leap-frog steps. We use the default HMC implementation and configuration for CRAFT, which uses a fixed step sizes of 0.3 for the first half of the intermediate distributions and 0.2 for the rest. Both models use flow layers per temperature.
Table 7: Number of flow and target evaluations during training for each method on the Many Well Problem. For SNF and CRAFT we report the number of model forward passes each of which contains multiple flow transport and MCMC steps.
Number of flow/model evaluations
Number of target evaluations
Flow w/ ML
o0
107
Flow w/ Da=2
10
10
Flow w/ KLD
1020
1020
Flow w/ RBD
1000
1010
SNF w/ KLD
10
2.5 10
CRAFT (config 1)
2.5 x 109
1000
CRAFT (config 2)
109
1000
FAB w/o buffer
10
1010
FAB w/ buffer
10
7.2 x 109
F.4 FURTHER RESULTS
Figure 14 and Figure 15 show contour plots of 2D marginals from the 32 dimensional Many Well target. The contours are for pairs of variables in the first four elements of x belonging to different copies of the Double Well distribution. Each plot is obtained by scanning two variables while the other ones are kept fixed to zero. These figures also shows samples from each analyzed method. FAB based methods and the flow trained by maximum likelihood place samples at each of the modes in the contour plots, while the other methods fail to do so.
31

---

Published as a conference paper at ICLR 2023
Flow w/ ML
Flow Da-2
-2-
-2
X4
Flow w/ KLD
SNF w/ KLD
0-
R
w/o buffer (ours)
w/ buffer (ours)
R
Figure 14: Target contours and model samples for 2D marginals over the first four elements of x in the 32 dimensional Many Well Problem. The plots are for pairs of variables belonging to different copies of the Double Well distribution. For D (p=2(pllge) minimization with samples from the flow, we plot results at iteration 56 of training as the final model samples were outside of the plotting regions due training instabilities.
32

---

Published as a conference paper at ICLR 2023
RBD KLLD
-33
Figure 15: Target contours and model samples for 2D marginals over the first four elements of x in the 32 dimensional Many Well Problem for the Resampling Base Distribution (RBD) flow model trained with the KL divergence.
33

---

Published as a conference paper at ICLR 2023
G ALANINE DIPEPTIDE EXPERIMENTS
G.1 SETUP
Coordinate transformation Boltzmann generators usually do not operate Cartesian coordi- nates. In particular, Noé et al. (2019) introduced a coordinate transformation whereby a subset of the coordinates mapped to internal coordinates, i.e., bond lengths. bond angles, and dihedral angles, see also Appendix G.I and Figure 11 in (Stimper et al., 2022). The internal coordinates are normalized and the respective means and standard deviations for these coordinates are computed on samples from the target distribution generated with MD. For the remaining Cartesian coordinates, principal component analysis is applied the samples and the six coordinates with the lowest vari- ance are eliminated. The rationale behind this is that the Boltzmann distribution is invariant in six degrees of freedom, i.e., three of translation and three of rotation, and consequently, the correspond- ing unnecessary coordinates should be removed. However, the mapping of vectors onto a fixed set of principal components generally neither invariant to translations nor to rotations, and, therefore, the transformed coordinates do not satisfy these invariances. When training Boltzmann generators with samples, this is not a problem since the flow will learn to generate molecular configurations for specific rotation or translation, but when only using the target distribution 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 representation invariant to translations and rotations. Since we do not want use MD samples for our model, we use the position with the minimum energy instead as shift and fix values for the scale when normalizing the coordinates. The former can be easily estimated with gradient descent using less than 100 steps. As scale parameters, we used 0.005 nm for the bond lengths, 0.15 rad for the bond angles, and 0.2 rad for the dihedral angles. Coordinates which are treated as circular are not scaled.
We use Neural Spline Flows with rational quadratic splines having 8 bins each. The parameter mapping is done through coupling (Durkan et al., 2019). Dihedral angles which can freely rotate, e.g., because it is not a double bond, are treated as periodic coordinates (Rezende al., 2020). For these coordinates, we use a uniform base distribution, while we pick a Gaussian for the other ones. The flow has 12 layers and the parameter maps are residual networks with one residual block, while the two linear layers each block have 256 hidden units. The flow layers were initialized in way that they correspond to the identity map.
Model architecture
One model uses a RBD, which has a residual network with two blocks having 512 hidden units per layer as acceptance function. The truncation parameter is set to the common value T 100.
The models trained with FAB do AIS with intermediate distributions, which are linear interpola- tions 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 used when use AIS in the other baseline models, see Table 10, Figure 20, Figure 21, and Figure 24. The HMC parameters are initialized and tuned using the same procedure as the Many Well problem, see Appendix F.3. The SNF model does additionally 10 layers. Since this renders sampling from this model already expensive, we do not do AIS with this model.
Hastings steps every two flow
Since the energy surface of alanine dipeptide in an implicit solvent has several modes of different sizes with large energy barriers between them, see Figure 19, a very long MD simulation would be required to obtain samples that represent the target distribution well. To get around this problem, which is well known in computational physies 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 by an increment of 50 K. The replicas are exchanged every 200 iterations and use the state at each multiple of 1000 time steps as samples. To reduce the time it takes to generate the data, we run many of these simulations in parallel with different seeds. Since the initial condition is always the same, i.c., the position with minimum energy as it is usually done, we let the system equilibrate for 2 x 10a iterations and run the simulation subsequently for 2 x 106 iterations.
Dataset
34

---

Published as a conference paper at ICLR 2023
L-form
D-form
Figure 16: Visualization of alanine dipeptide in its two chiral forms. In nature, we see almost exclusively the L-form and, likewise, we aim only generate samples of this form.
We split the data into 1) training set, which consists of 10t samples and is only used to train the baseline flow model with ML; 2) a validation set consisting of 10th samples well, which used to find a suitable set of hyperparameters for our experiments; and 3) test set with 10? samples, which used to evaluate all models.
To generate the training data alone we had to evaluate the target distribution and its gradients 2.3 10 times, which what we report as cost in terms of target evaluations in Table
Filtering chiral forms As mentioned in Section 4.2. alanine dipeptide is chiral molecule, i.e., can occur two different forms that are mirror images of each other, see Figure 16. They cannot be 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 created compounds. Hence, whenever alanine dipeptide is considered in the literature, it is almost always as the L-form (Wu et al., 2020; Campbell et al., 2021; Stimper et al., 2022: Dibak et al., 2022; aim to train model only this form well. However, since
Köhler et al., 2022). Therefore, the energy of the molecule does not change when creating a mirror image of it, models trained to approximate its Boltzmann distribution will priori learn generate both forms.
The two forms can be separated by using the following procedure. The two chiral forms differ by the positioning of the neighhoring atoms chiral center, i.e., the center carbon atom. Hence, the difference of the dihedral angles of those atoms with respect to a fixed reference will differ relative to each other, i.e., their difference will change. Hence, we compute this difference and check whether it is close to a reference configuration for which we know that it corresponds to the L-form. As reference configuration, we use the position with minimum energy, which we already determined for the coordinate transformation.
We use this procedure filter the configurations generated by the flow model during training and included only the samples that correspond to the L--orm when computing the loss. Thereby, the model learns to only generate this chiral form.
In Appendix G.2 we will investigate model trained on both chiral forms and compare it to one that was only trained on L-form.
Training All models were trained using the Adam optimizer (Kingma & Ba, 2015) with batch size of 1024. A learning rate of 5 x 10-1 initially linearly warmed up over 1000 iterations and decayed with cosine annealing schedule over the course of training. We 10-3 and clip gradients at a value 103, When training the models with the prioritized replay buffer, ensured a minimum buffer length of 64 batches and started replacing the oldest samples once its length reached 512 batches.
weight decay of
Evaluation To evaluate the models, draw 107 samples from the models with and without the use of AIS. Since there were some outliers of the importance weights due to flow numerics, we took the 10 highest weights and clipped them to the lowest value in this compute the ESS
35

---

Published as a conference paper at ICLR 2023
(Koblents & Míguez, 2015: Dibak et al., 2022). This corresponds to a fraction of 10-4, or 0.01% % of the weights. For the flow trained with the a-divergence with a = 2. the resulting ESS is close to 10 or 0.01%. indicating that the true ESS is even lower. We estimated the Ramachandran plots, i.e., made a histogram of the dihedral angles & and 1h see Figure 3a, with 100 x 100 bins, and used them to compute the KI. divergence between the test samples and the samples from the model. We repeated this with the reweighted samples, whereby we also used the clipped weights. We computed the log-likelihood on the test set with the models, but we did not do so for the SNF, as it only computes importance weights and does not directly estimate the density.
Computational cost The two main contributors to the computational expenses necessary to train flows approximating Boltzmann distributions are the number of evaluations of the flow and the target. Typically, we need both the value and the gradient and, hence, we regard obtaining them as one operation. Moreover, the flows that we use take the same time for sampling and for pure likelihood computation, which is why we regard them as the same operation as well. The flow with RBD is an exception, as sampling from it is more expensive due to learned rejection sampling being used in the base distribution. The number of flow and target evaluations for each model and training procedure are listed in Table 8.
In general, we trained all models using an equal number of flow evaluations, with the exception of SNF. SNF requires 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.
Table 8: Number of flow and target evaluations needed to train the models. For the flow being trained with ML on MD samples, we report the number of target evaluations that are needed to generate the training dataset with MD. Our methods are marked in italic.
Number of flow evaluations
Number of target evaluations
Flow w/ ML
2.5 x 10$
2.3 x 10
Flow w/ Da=2
2.5 x I0
2.5 x 10a
Flow w/ KLD
2.5 X 108
2.5 x 10
RBD w/ KLD
2.5 X 108
2.5 x 10$
SNF w/ KLD
6.0 10?
3.6 x 10
FAB w/o buffer
2.5 X 108
2.5 x 10a
FAB w/ buffer
2.5 x 108
2.0 x 10a
Computational resources and runtime 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 which needed roughly 15.7h each adding up to about 113 kCPUh.
The flow models were trained on servers with an NVIDIA A100 GPU and an Intel Xeon IceLake- SP 8360Y processor with 18 cores and 128 GB RAM. The training time for each model is listed in Table 9. In total, we invested around 1.02 kGPUh in the experiments. Although training the flow with ML is faster than training with FAB, note that generating the data for MI. training with MD) requires an additional 9.4 kCPUh, which would take 131h when executed on one server.
Table 9: Runtime for training the models on the same server type specified in the text. Our methods are marked in italic.
Runtime
Flow w/ ML
13.8h
Flow w/ D
20.0h
Flow w/ KLD
20.0h
RBD w/ KLD
82.5h
SNF w/ KLD
170h
FAB w/o buffer
18,8h
FAB w/ buffer
15.7h
36

---

Published as a conference paper at ICLR 2023
1.0
Ground truth
Flow w/
0.8
Flow w/ RW
Probability density
FAB w/ buffer (ours)
FAB w/ buffer, RW (ours)
0.6
0.4
0.2
0.0
0
Φ
Figure 17: Marginal distribution the dihedral angle for selected models. The visualized data is the same as in Figure 3b, but here normal scale instead of log scale is used for the density.
1.0
Ground truth, L-form Flow, and f.form
Ground truth. L-form
Flow. L. and D form
0.7
only L-form
0.8
only L-form
Probability density
Flow, only D-form
Probability density
Floo, only D-form
0.6
0.4
0.4
0.3
0.2
0.0
0.0
0
W
Figure 18: Marginal distribution of the dihedral angles tb and o of a model which has been trained on both chiral forms of the alanine dipeptide, i.e., the L- and D-form. We plot the density obtained with all samples from the model in blue and separate the samples in the forms, yielding the pink and green curves. The distributions of the two forms are mirror images of each other.
G.2 FURTHER RESULTS
Model trained both chiral forms To demonstrate the importance of filtering for the L-form during training, we trained a model on both chiral forms using FAB with replay buffer with the same setting as in the other experiments. We drew 10" samples from the model and found exactly 50% of them correspond to the L- and 50% to the D-form. As can be seen in Figure 18, the marginal distributions of the dihedral angles 1 and ø for the two forms are mirror images of each other, while the flow model, generating both forms, is a mixture of the two.
The log- likelihood on the test set 210.80, which is 0.74 less than the corresponding model only trained on the I.-frrm. This is close to log(2) ~ 0.69, i.e. the density is roughly by a factor of two lower, confirming once more that the flow density a mixture of the density of the two forms.
Model performance with AIS As mentioned in the previous section, we do AIS with all our trained models except the SNF, which already involves sampling layers. We adopt the same AIS setting used for training the flow models with FAB, i.e., use intermediate distributions given by linear interpolations between the flow and the target distributions and sample from them with HMC performing 4 Leapfrog steps. For comparison, untrained base distribution of our flow as proposal for AIS. The results are shown Table 10. When comparing tables 2 and 10, observe that AIS improves performance for those models that
provide the performance when using the
37

---

Published as a conference paper at ICLR 2023
approximate the target distribution least fairly well. Again, the flow trained with FAB with replay buffer outperforms the baselines.
Table 10: ESS, and the KL divergence of the Ramachandran plots with and without reweighting (RW) for flow models when sampling from them with AIS. The results are averages over 3 runs and the standard error is given uncertainty. Our methods marked in italic and the highest ESS or log-likelihood or lowest KI. divergence values are emphasized in bold.
ESS (%)
KLD
KLD w/ RW
Base untrained
0.013 0.000
1.96 0.00
8.5
Flow w/ ML
0..
(4.92
x 10-
(1.88+0.75)
x 10-
Flow w/ Do=
0.01210.0000
2.32
11.30..
01.31:.2
Flow w/ KLD
80
0.99
0.19
RBD KLD
61
2.84
0.01
FAB w/o buffer
0.7
(2.65 10
(2.45 x 10
FAB w/ buffer
9.2
(2.53 x 10-3
(2.17 - 0.15) x 10 -3
FAB with varying values of a In Appendix C.I, we derived a variant of FAB having an arbitrary value for the a parameter of the a-divergence. Here, we want to investigate how the performance of models trained with FAB changes as we vary a. Therefore, we leave the setup the same as used in Section 4.2 and only changed the a parameter FAB. The results when using a replay buffer given in Table 11 and without it they are reported in Table 12. We that 2 outperforms the other values in almost all performance metrics, no matter whether replay buffer is used or not. empirically.
This justifies our theoretical arguments for picking
Table 11: ESS, log-likelihood on the test set, and KL divergence (KLD) of Ramachandran plots with and without reweighting (RW) for each method. Here, all models trained with FAB using replay buffer, but we varied the value of a being used, see Appendix C.I. The results are averaged over runs and the standard error is given uncertainty. Best results emphasized in bold.
Q
ESS (%)
Ep(x) [log {(x)
KLD
KLD w/ RW
0.25
0.021
0.0000
-546
0.72
1.49
0.5
0.023
0.0001
-606
0.94
1..4
1
0.027
0.000
60.1 9.0
0.81
6.73
1.5
8.3
211.49+0.00
(1.44+0.90) x 10-2
(1.92+0.33) x 10-2
2
0.1
211.54+0.00
(3.42 x 10
(2.51 = 0.399 X 10 10-3
3
1.1
$11.49+0.00
(2.58 1.58+0.90 x 10
(2.19 X 10
Table 12: ESS, log-likelihood on the test set, and KI. divergence (KLD) of Ramachandran plots with and without reweighting (RW) for each method. Here, all models are trained with FAB without replay buffer, but we varied the value of a being used, see Appendix C.I. The results are averaged over runs and the standard error given as uncertainty. Best results are emphasized in bold.
ESS (%)
Ep(x) [log {(x)
KLD
KLD w/ RW
0.25
1.7 0.3
198.71 + 0.09
3.94 10440..00 0.40
8.45 + 0.20
0.5
4.4
-192
7.70
7.99
15.8
210.16+0.440
(1.99 x 10
(5.60 x 10-2
1.5
34.
210.87
(7.22+0.72) x 10-2
(4.04+0.04) x 10
2
5.3
211.13+0.03
(6.28+0.33) x 10
(2.66+0.90) x 10
24.8+4.6 4.6
210.48+0.14
+1.1++0005) x 10-1
(3.88 +0003) x 10-2
Ramachandran plots Figure 19 shows the Ramachandran plot of the test set and Figure 20, Figure 21, Figure 22, Figure 23, Figure 24, Figure 25, and Figure 26 show the Ramachandran plots of all the models trained for the first run, including the samples drawn from them via AIS.
38

---

Published as a conference paper at ICLR 2023
10
10
1003
10
0
Nis
Φ
Figure 19: Ramachandran plot of the test data.
10
10-1
10
10-2
10-1
10-3
10-4
(a) Flow
(b) Flow, reweighted
10
10-1
10
10
10-1
10-3
(c) AIS
(d) AIS, reweighted
Figure 20: Ramachandran plots of a flow trained with the a = 2-divergence.
39

---

Published as a conference paper at ICLR 2023
10
10-1
1002
10
10
10
10
10
17
(a) Flow
(b) Flow, reweighted
10
10-1
10-
10-2
10-1
10-3
10
10-4
(c) ATS
(d) AIS, reweighted
Figure 21: Ramachandran plots of a flow trained with the KL. divergence.
40

---

Published as a conference paper at ICLR 2023
10
10"1
10
10
10
10
10
10
77
(a) Flow
(b) Flow, reweighted
10
1001
1002
100
10
10-3
10
10-4
12
-7
(c) AIS
(d) AIS, reweighted
Figure 22: Ramachandran plots of flow with resampled base distribution trained with the KI. divergence.
10a1
1011
10
10
10
10
10
10
(a)
(b) SNE, reweighted
Figure 23: Ramachandran plots of SNF trained with the KL divergence. Since the SNF already has layers which do sampling, we did not do AIS with it.
41

---

Published as a conference paper at ICLR 2023
10-1
10-1
-107
10
10
10
10
10
-7
(a) Flow
(b) Flow, reweighted
10-1
10-1
10-2
10-2
10-1
10-3
10
10-4
(c) AIS
(d) AIS, reweighted
Figure 24: Ramachandran plots of the flows trained with ML.
42

---

Published as a conference paper at ICLR 2023
10
10-1
1002
10
10
10
10
10
-7
12
(a) Flow
(b) Flow, reweighted
10-1
10-1
10-2
10-2
10-1
10-3
10
10
(c) ATS
(d) AIS, reweighted
Figure 25: Ramachandran plots of the flows trained with FAB without the use of replay buffer.
43

---

Published as a conference paper at ICLR 2023
10
10-1
10
10
0
10
10
10
10
-7
12
(a) Flow
(b) Flow, reweighted
10-1
10-1
10-2
10-2
10
10-
10
10
(c) ATS
(d) AIS, reweighted
Figure 26: Ramachandran plots of the flows trained with FAB with the use of replay buffer.
44


Reviews:
[{"review_id": "u7eYLuRV527", "reviewer": "Reviewer_tRrj", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors use the 2-divergence (Chi^2) as the objective function for sampling target distribution. And they parameterize the distribution by a normalizing flow. They use importance sampling ideas to approximate the Chi^2 divergence. By solving the proposed optimization method numerically, they directly simulate the Gibbs distribution from Molecular Dynamics. Numerical results demonstrate the simplicity of the method.  ", "summary_of_the_review": "Overall, the paper is written well. Some analytical examples of Chi^2 divergences can be studied. This could explain or motivate why it works better in numerics. ", "strength_and_weaknesses": "Strength: The paper is written well with clear mathematics.\n\nWeakness: The authors need to demonstrate (motivate) why alpha=2 works better.  Does the author try other f-divergences? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ji7DbDqVESK", "reviewer": "Reviewer_1QZD", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The normalizing Flow-based method is used to estimate complex densities.  However, current methods for training flows have some drawbacks, including expensive computing with MCMC simulations. In this work, the authors propose low AIS Bootstrap method to generate samples with a theoretical and numerical guarantees. \n\n\n\n", "summary_of_the_review": "\nOverall, the writing is easy to follow and the proposed method seems to work well in 2-d case.\nHowever, the method does not demonstrate improvements over the baseline method in higher\ndimensional cases.", "strength_and_weaknesses": "1. In section 3.1, they claim “p^2/q is the minimum variance importance sampling distribution for D_{α=2}(p||q)”. Can the authors explain further about this point?  \n2. In Many Well Experiments(Appendix E), it seems that their method can not beat the Flow w/ ML method based on Table 4. The Ep(x)[log q(x)], Mean log q (xmodes ) and DKL[p||q] results of those two methods are exactly the same. \n3. Could the authors provide computation time for different methods? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "eUzk5BPZW0", "reviewer": "Reviewer_A5bQ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a strategy to augment flows with annealed importance sampling (AIS) with the objective to minimize importance weight variance. AIS is responsible to generate samples in regions where the flow struggles. The proposed FAB was shown to produce\nbetter results than training via maximum likelihood on MD samples with less computational demand. ", "summary_of_the_review": "In summary, I think this paper is novel both in terms of the method itself and in terms of the specific use case. And authors have shown strong empirical evidence of the superiority of the method. I would recommend accept.", "strength_and_weaknesses": "Strength: I like the idea of using AIS to cover for regions in the sample space where the current flow model struggle. The experimental results also supported the claim.\n\nWeakness: My concern is over the particular choice of α=2 objective. Why does minimizing variance indicate that this objective works better for the flow?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "Dx-y5MuDtD7", "reviewer": "Reviewer_8zPW", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a method for training normalizing flows using a mass-covering objective which does not require samples from the target distribution.   Existing methods often optimize a mode-seeking objective when target samples are not available, although the experiments section highlights applications (such a learning mixtures of Gaussians or a molecular Boltzmann distribution) where mode-seeking optimization fails to match a multi-modal target.\n\nThe current work overcomes these limitations by evaluating flow density training gradients on approximate samples from the the optimal proposal for estimating the α-divergence with α=2.    Minimizing this divergence as a function of q corresponds to minimizing the variance of the importance weights when sampling under q.   Annealed Importance Sampling (with a replay buffer) is used to transform samples generated by the flow into samples from the optimal α=2 proposal, which emphasizes regions in the sample space where the current flow model does not match the target.\n\n", "summary_of_the_review": "I am happy with this paper overall.   The experiments show success in proposing a mass-covering loss for normalizing flow training which can be used when target samples are not available.    Some concerns regarding presentation have been raised above.  ", "strength_and_weaknesses": "**Strengths:**  The paper is well-motivated in proposing a method for training mass-covering flows without access to target samples.   I am happy to see work utilizing the importance sampling variance-minimization interpretation of the α=2 divergence.   \n\nThe authors essentially solve estimation of the Boltzmann dist. for alaine dipeptide molecule.   I do not fully appreciate the significance, but this does represent a 'state-of-the-art'-style result, in which the motivating aspects of the proposed method figure prominently.\n\n\n**Weaknesses:**\nOne question is to what extent the proposed α=2 objective is preferred to MLE with $KL[p : q]$ in general / practice.   Similar derivations should apply using AIS to sample from p and minimize $KL[p : q]$(the only thing that should change is the AIS target and importance weights, if I understand correctly).\n\nResults are nearly identical for α=2 vs. MLE with exact or MD samples (both essentially solve the problems), so it doesn't seem absolutely necessary to run additional baselines (or to open up the discussion to less justified values of α).     Still, the exposition should state the generality of the AIS bootstrap approach in applying to these divergences, especially since accurate sampling for higher values of α should be more challenging in practice.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "BvhIMhItuS", "reviewer": "Reviewer_xTnc", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes to train normalizing flows with annealed importance sampling (AIS). The loss function, which is the 2-divergence, is approximated using importance weighted AIS samples targeted at the distribution with the zero-variance IS estimator. A replay buffer is introduced to reduce the computational burden of the AIS step. The algorithm has good performance on approximating a 2-d Gaussian mixture distribution and the Boltzmann distribution of alanine dipeptide.", "summary_of_the_review": "The paper provides a novel algorithm for training normalizing flows. The algorithm is competitive in the reported experiments. The writing is clear. The weakness includes the complexity of tuning parameters and the lack of comparison with the baseline of AIS alone.", "strength_and_weaknesses": "**Strength**\n- The proposed algorithm introduces AIS in the training of normalizing flows. With importance sampling, one is able to sample from the region with larger loss. The computation cost of AIS is reduced by introducing a replay buffer.\n- The algorithm is competitive in the reported numerical experiments.\n\n**Weakness**\n- There are a lot of tuning parameters in the AIS part of the algorithm: the number of intermediate distributions in AIS, $\\beta_i$ for each intermediate distribution; number of MCMC steps used to sample from each intermediate distribution; number of gradient steps per AIS step $L$; $M$. The paper does not discuss how sensitive the performance is to these parameters or how should one choose the parameters.\n- It seems that the algorithm should work for other $\\alpha$. It is not clear why the paper only considers $\\alpha=2$. Some discussion or empirical results with different $\\alpha$ would help.\n- The AIS step is trying to fit $p^2/q$, which is as challenging as the original problem of fitting $p$. One might wonder how the proposed method compares with pure AIS. For example, flow without AIS has bad performance as shown in Figure 8. So one might ask what's the performance of AIS alone.\n- The \"bootstrap\" in the paper is not the bootstrap commonly seen. There is only one line of explanation in Section 3.1. I think the paper could benefit from more elaborate explanations of why the procedure is called \"bootstrap\", since it's in the name of the algorithm.\n\n**Minor comments**\n- The notation $N$ is abused. In Equation (3), $N$ is the number of intermediate distributions. In other places, $N$ is the number of samples generated by AIS.\n- One might expect a plot with $\\alpha=2$ in Figure 1 since the whole paper is about $\\alpha=2$.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "A0Jw4cdAXJH", "reviewer": "Reviewer_itRu", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed a method, Flow AIS Bootstrap (FAB) that uses AIS to train flows to approximate complex, unnormalized distributions. Authors used. Alpha divergence with alpha=2 as the objective as it encourages mode covering and minimizes importance weight variance. Authors further reduce the computational cost of FAB by introducing a replay buffer with various tricks, and the authors show empirical evidence of the effectiveness of FAB on toy distributions and on the more sophisticated Boltzmann distribution of alanine dipeptide. \n\n", "summary_of_the_review": "In summary, I think this paper is novel both in terms of the method itself and in terms of the specific use case. And authors have shown strong empirical evidence of the superiority of the method.  I would recommend accept. ", "strength_and_weaknesses": "Strength:\n\n1. I find FAB to be well-motivated. I wasn’t familiar with Boltzmann generators but I find the adoption of AIS to train Flows for this problem to be clever. AIS samples are usually pretty good, but mostly it was only used to estimate the partition function. I have previously thought that the AIS samples could be put to better use, and in this paper the alanine dipeptide problem is a perfect use case because we have access to only the unnormalized density, but to train an explicit generative model we need samples. Yet directly training with the ML objective has various problems, AIS comes to the rescue.\n\n2. I find the use of Alpha divergence with alpha=2 to be clever and appropriate as it encourages mode covering and minimizes importance weight variance at the same time. \n\n3. I find the decomposition of HMC step sizes to be clever (In Appendix E.3). I don’t recall seeing it in prior works, is it a novel contribution by the authors? \n\n4. The FAB method is theoretically sound.\n\n5. It seems authors have sufficiently explored various ways to empirically improve FAB, tricks including the replay buffer (and various nuanced decisions went into improving that including using mini batching with probability proportional to old AIS weights and reweighting etc) are well-motivated and empirically shown to be useful.\n\n6. As a result, the FAB method is empirically strong and convincing. On the toy example, Figure2 is convincing at showing the effectiveness of FAB with buffer, without ground truth samples, it performs as well as Flow trained with ML objective with access to ground truth samples. Similarly Table 1, Table 2 and Figure 3, Figure 4 show strong evidence of the effectiveness of the method. \n\nWeakness:\n\n1. The step sizes tuning for the HMC is tricky, because to conform to mathematical correctness, it is necessary to load the step sizes from a given run and repeat the same run with different random seeds, so that the Markovian property is not violated. I don’t think the authors did this? Granted it probably doesn’t make a big difference empirically, but authors could consider at least verifying how much of a difference this would make. For more details, see Appendix C.3 of https://arxiv.org/pdf/2008.06653.pdf \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%