Fast and unified path gradient estimators for normalizing flows
Authors: Lorenz Vaitl
Ludwig Winkler
Lorenz Richter
Pan Kessel
Venue: ICLR 2024 poster (2024)
Keywords: Normalizing Flows
Gradient Estimators
Lattice Field Theory
Variational Infernce
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: New low variance gradient estimator for normalizing flows

Abstract: Recent work shows that path gradient estimators for normalizing flows have lower variance compared to standard estimators, resulting in improved training. However, they are often prohibitively more expensive from a computational point of view and cannot be applied to maximum likelihood training in a scalable manner, which severely hinders their widespread adoption. In this work, we overcome these crucial limitations. Specifically, we propose a fast path gradient estimator which works for all normalizing flow architectures of practical relevance for sampling from an unnormalized target distribution. We then show that this estimator can also be applied to maximum likelihood training and empirically establish its superior performance for several natural sciences applications.

Published as a conference paper at ICLR 2024
FAST AND UNIFIED PATH GRADIENT ESTIMATORS FOR NORMALIZING FLOWS
Lorenz Vait!" Ludwig Winkler Lorenz Richter Pan Kessel Machine I Learning Group, TU Berlin, Zuse Institute Berlin, dida Datenschmiede GmbH, *Prescient Design, Genentech, Roche
ABSTRACT
Recent work shows that path gradient estimators for normalizing flows have lower variance compared to standard estimators for variational inference, resulting in improved training. However, they are often prohibitively more expensive from a computational point of view and cannot be applied to maximum likelihood train- ing in a scalable manner, which severely hinders their widespread adoption. In this work, we overcome these crucial limitations. Specifically, we propose fast path gradient estimator which improves computational efficiency significantly and works for all normalizing flow architectures of practical relevance. We then show that this estimator can also be applied to maximum likelihood training for which it has a regularizing effect as it can take the form of given target energy func- tion into account. We empirically establish its superior performance and reduced variance for several natural sciences applications.
INTRODUCTION
Normalizing flows (NFs) have become a crucial tool applications of machine learning in the natural sciences. This is mainly because they can be used for variational inference, i.c., for the approximation of distributions corresponding to given physical energy functions. Furthermore, they can be synergistically combined with more classical sampling methods such as Markov chain Monte Carlo (MCMC) and Molecular Dynamics, as their density is tractable.
The paradigm of using normalizing flows as neural samplers has lately been widely adopted for example in quantum chemistry (Boltzmann generators (Noé et al., 2019)), statistical physics (gen- eralized neural samplers (Nicoli et al., 2020)), as well high-energy physies (neural trivializing maps (Albergo et al., 2019)). In these applications, the normalizing flow is typically trained using a combination of two training objectives: Reverse Kullback-Leibler (KL) training is used to train the model by self-sampling (see Section 2). Crucially, this training method on its own often fails in sigh-dimensional sampling settings because self- self-sampling is unlikely to probe exceedingly concentrated high probability regions of the ground-truth distribution and can potentially lead to mode collapse. As such, reverse KL training is often combined with maximum likelihood training (also known as forward KL. training). For this, samples from the ground-truth distribution are obtained by standard sampling methods such as, e.g., MCMC. As these methods are typically costly, the samples are often of low number and possibly biased. The model is then trained to maximize its likelihood with respect to these samples. This step is essential for guiding the self-sampling towards high probability regions and, by extension, for successful training.
Since training normalizing flows for realistic physical examples is typically computationally challenging, methods to speed up the convergence have been focus of recent research. To this end, path estimators for the gradient of the reverse KL loss have been proposed (Roeder et al., 2017; Vaitl al.. 2022a;b). These estimators focus on the parameter dependence of the flow's sampling process, also known as the sampling path, while discarding the direct parameter dependency, which vanishes expectation. Path gradients have the appealing property that they are unbiased and tend to have lower variance compared to standard estimators, thereby promising accelerated convergence (Roeder et al., 2017; Agrawal et al., 2020; Vaitl et al., 2022a;b). At the same time, however, current path gradient estimation schemes have often a runtime that is several multiples of the standard
Shared first authorship.

---

Published as a conference paper at ICLR 2024
gradient estimator, thus somehow counteracting the original intention. As remedy, recently, Vaitl et al. (2022b) proposed a faster algorithm. Unfortunately, however, this algorithm is limited to continuous normalizing flows.
Our work resolves this unsatisfying situation by proposing unified and fast path gradient estimators for all relevant normalizing flow architectures. Notably, our estimators are between 1.5 and 8 times faster than the previous state-of-the-art. Specifically, we a) derive a recursive equation to calcu- late the path gradient during the sampling procedure. Further, for flows that are not analytically invertible, we b) demonstrate that implicit differentiation can be used to calculate the path gradient without costly numerical inversion, resulting in significantly improved system size scaling.
Finally, we e) prove by a change of perspective (noting that the forward KI. divergence in data space is a reverse KI. divergence in base space) that our estimators can straightforwardly be used for max- imum likelihood training. Crucially, the resulting estimators allow us to work directly on samples from the target distribution. As result of our manuscript, path gradients can now be used for all widely used training objectives - as opposed to only objectives using self-sampling - in a unified and scalable manner.
We demonstrate the benefits of our proposed estimators for several normalizing flow architectures (RealNVP and gauge-equivariant NCP flow) and target densities with applications both in machine learning (Gaussian Mixture Model) as well as physics (U(1) gauge theory, and ø² lattice model).
1.1 RELATED WORKS
Pathwise gradients take the sampling path into account and are well established in doubly stochas- tie optimization, see e.g. L''euyer (1991); Jankowiak Obermeyer (2018); Parmas & Sugiyama (2021). The present work uses path gradient estimators, subset of pathwise gradients, originally proposed by Roeder et al. (2017) in the context of reverse KL training of Variational Autoencoders (VAE), which is motivated by only using the sampling path for computing gradient estimators and disregarding the direct parameter dependency. These were subsequently generalized by Tucker et al. (2019); Finke & Thiery (2019); Geffner & Domke (2021a;b) to generie VAE self-san losses. There has been substantial work on reducing gradient variance not by path gradients but with control variates, for example in Miller et al. (2017); Kool et al. (2019); Richter al. (2020); Wang et al. (2023). For an extensive review on the subject, we refer to Mohamed et al. (2020).
Bauer& Mnih (2021) generalized path gradients to score functions of distributions which do not coincide with the sampling distribution in the context of hierarchical VAEs. As we will show, our fast path gradient for the forward KL training can be brought into the same form. However, only our formulation allows the application of a fast estimation scheme for NFs and establishes that forward and reverse path gradients are closely linked.
Path gradients for normalizing flows have recently been studied: Agrawal et al. (2020) were the first to apply path gradients to normalizing flows as part of broader ablation study. However, their algorithm has double the runtime and memory constraints it requires a full copy of the neural net- work. Vaitl et al. (2022a) proposed a method that allows path gradient estimation for any explicitly invertible flow at the same runtime cost Agrawal et al. (2020) but half the memory footprint. They also proposed an estimator for forward KL training which is however based on reweighting and thus suffers from poor system size scaling, while our method works on samples from the target density, For the rather restricted case of continuous normalizing flows, Vaitl et al. (2022b) proposed a fast path gradient estimator. Our proposal unifies their method in a framework which applies across a broad range of normalizing flow types.
2 NORMALIZING FLOWS
normalizing flow a composition of diffeomorphisms
x = To (To) = Tr.0, ....o T,,0
where we have collectively denoted all parameters of the flow by 8= (01...., 9z). Since diffeo- morphisms form group under composition, the map To is diffeomorphism as well. Samples from a normalizing flow can be drawn by applying To to samples from a simple base den- sity xo go such go (N,0,1) 1). The density of x = To (I5)), denoted by 4o is then given by the
2

---

Published as a conference paper at ICLR 2024
pushforward density of go under Tø, i.e.,
log (x)) log go (f + log det 0Ig 1 (x)
Ox
see also Appendix A for general remarks on the notation. We focus on applications for which normalizing flows trained to closely approximate ground-truth target density p(z) 2 exp( -b(x)) where the energy : [R + is known closed-fron but the partition function
1
2 JR= dz is intractable. To this end, there are two widely established training methods:
Reverse KL training relies on self-sampling from the flow and minimizes the reverse KL diver- gence
(go. p) Es [Ex) + log qo (x)] + const.
(3)
Since reverse KL training is based on self-sampling, the flow needs to be evaluated in the base-to- target direction Tg.
Forward KL training requires samples from the target density and is equivalent to maximum likelihood training
DKL (p, go) = Eesp |- log [o) + const.
(4)
Since forward KL training requires the calculation of the density (z)), the flow needs to be evalu- ated in the target-to-base direction T see (2).
As mentioned before, one typically uses a combined forward and reverse training to guide the self- sampling to high probability regions of the target density. When choosing normalizing flow archi- tecture for this task, it is therefore essential that both directions Tg and T can be evaluated with reasonable efficiency. As result, the following types of architectures are of practical relevance:
Coupling Flows are arguable the most widely used (see, e.g., Noé et al. (2019); Albergo et al. (2019); Nicoli et al. (2020); Matthews al. (2022); Midgley al. (2023); Huang al. (2020)). They split the vector 2ι E Id in two components
Ft = (x) trans 27
cond)
(5)
with a.rane E && and uooad E RRd- for {1,... ... 1}. The map 1, is then given by
trans trans xi coud) 1+1,
= 7 T
trans ha. (
ViE (1....,k),
(6a)
goond x rcond
(6b)
R are invertible maps with respect to their first argument for any choice of the second argument and ha Rd- RT is the i-th output of a neural network. Note that the function fo acts on the components of arrass element-wise.
where fo R* x Rd-k R
There are broadly two types of coupling flows with different choices for the transformation T:
1. Explicitly invertible flows have the appealing property that the inverse map Title can be calculated in closed-form and as efficiently as the forward map 1.+1,0.41 A particular example of this type of flows are affine coupling flows (Dinh et al., 2014; 2017) that use an affine transformation T, i.e.,
Iraus fo( trans 411
21 condd 00(2; @ x
Arraas + (o(
(7a)
cond
cond
(7b)
21+1
with hg = (og.Po) Another example are neural spline flows (Durkan et al., 2019) which use splines instead of an affine transformation.
2. Implicitly invertible flows use a map т whose inverse can only be obtained numerically, such as mixture of non-compact projectors (Kanwar et al., 2020; Rezende et al., 2020) or smooth bump functions (Köhler et al., 2021). This often results in more expressive flows in particular in the context of normalizing flows manifolds (Rezende et al., 2020).
Recently, it has been shown in Köhler et al. (2021) that implicit diffcrentation can be used train these types of flows using the forward KI objective.
3

---

Published as a conference paper at ICLR 2024
Continuous Normalizing Flows ordinary differential equation (ODE) which relates to the bijection Tg : R - E2 allowing for straightforward implementation cquivariances, but typically coming with high computational costs (Chen et al., 2018). Notably, autoregressive flows (Huang et al., 2018; Jaini et al., 2019) are less relevant in the context of learning target density p because they only permit fast evaluation in one direction and there is no training method based on implicit differentiation. As a result, they are not considered in this work.
3 PATH GRADIENTS FOR REVERSE KL
In this section, we introduce path gradients and show how they are related to the gradient of the reverse KL objective. The basic definition of path gradients is as follows:
Definition 3.1. The path gradient of a function (o (zo)) is given by
(0, To(x)) 0=((0xz)
(To(Fo)
(8)
dx
T==To(zo)
00
Note that the total derivative of the function yp can be decomposed in the following way:
d (0, (zo)) (0, To (zo)) a0 (0, [==))
(9)
d0
The path gradient therefore only takes the parameter dependence of the sampling path Tp into count, but does not capture any explicit parameter dependence denoted by the second term. This decomposition was applied by Roeder et al. (2017) to the gradient of the reverse KL divergence to obtain the notable result
d d0
DKL. DKI.(go.P) = Ezo-no (E(To(zo)) + log ge (To ((oo)))
(10)
where we have used the fact that Exprogo 16 the gradient of the KL divergence is given by
loggeeTT (To (2o)) 0. Thus, an unbiased estimator for
(((TT(=("))) log go (To((()))))
9path
(11)
N
where (1()) ~ qu are i.i.d. samples. This path gradient estimator has been observed to have lower variance compared to the standard gradient estimator (Roeder et al., 2017; Tucker et al., 2019; Agrawal et al., 2020).
As the total derivative of the energy agrees with the path gradient of the energy function, i.e., ae (zo)) = a (T(T((to) the first term in the estimator can be straightforwardlly calculated using automatic differentiation. The second term, involving the path score v log qe (zo)), is however non-trivial as the path gradient through the sampling path Ta has to be disentangled from the explicit parameter dependence in go. Recently, Vaiil et (2022a) proposed a method caleu- late this term using the following steps:
d
1. Sample from the flow without building the computational graph:
stop-seopp stto--gradieee gradientsddtttr (zo()
for Fo ~ qu
(12)
2. Calculate the gradient of the density with respect to the sample x' using automatic differ- entiation:
a dr'
(log go ( + log det
0 (3!)
))
G
log go('') Dx
(13)
dx'
3. Calculate the path gradient using a vector Jacobian product which can be efficiently calcu- lated by standard reverse-mode automatic differentiation'.
ÖT (xro
log qo (xo)) = G
(14)
00
Following standard convention in the autograd community, we adopt the convention that G is row vector. Øf. of function is one-form and thus an element of the co-tangent
This is because the differential ff dr, de, space.
4

---

Published as a conference paper at ICLR 2024
This method therefore requires the evaluation of both directions Tg and T For implicitly invert- ible flows, backpropagation through a numerical inversion per training iteration is thus required, which is often prohibitively expensive.
Even in the best case scenario, i.e., for flows that can be evaluated in both directions with the same computational costs, such as RealNVP (Dinh al., 2017), this algorithm has significant computa- tional overhead. Specifically, it has roughly the costs of five forward passes; one for the sampling (12) and two each for the two gradient calculations (13) and (14) (which each require a forward as well as a backward pass). This is to be contrasted with the costs of the standard gradient estimator which only requires a single forward as well as a backward pass, i.e., has the cost of roughly two forward passes. In practical implementations, typically runtime overhead of factor of two instead of 5 is observed for the path gradient estimator compared to the standard gradient estimator.
2
3.1 FAST PATH GRADIENT ESTIMATOR
In the following, we outline a fast method to estimate the path gradient. An important downside of the algorithm outlined the last section that one has to evaluate the flow both directions To and T The basic idea of the method outlined in the following is to calculate the derivative Og log (2) of the flow model recursively during sampling process. As a result, the flow only needs to be calculated in the forward direction T3 as the second step in the path gradient algorithm discussed in the previous section can be avoided. In more detail, the calculation of the path gradient proceeds in steps:
1. The sample (xo) and the gradient G Zur log qo can be calculated alongside the sampling process using the recursive relation derived below.
2. The path gradient is then calculated with automatic differentiation using a vector Jacobian product, where, however, the forward pass To (xro) does not have to be recomputed:
(To((o)
log (og
(15)
00
The recursion to calculate the derivative 02 log qo (a)) is as follows:
Proposition 3.2 (Gradient recursion). Using the diffeomorphism Tú, the derivative of the induced probability can be computed recursively as follows
Olog det o
Ô log go. ( 0x1+1
a log qo
(1 -1
(oTi,o
(8T
Om
Oxi
Fm
Oz
(16)
For general Tr, computing the inverse Jacobian ( entails a time and space complexity higher than Od), which is the complexity of the standard gradient estimator. For autoregressive flows, the total complexity
(dd), since its Jacobian is triangular. For coupling-type flows, we can simplify and speed up the recursion to have linear complexity in the number of dimensions, Le. Odd.. We state the recursive gradient computations for these kind of flows in the following proposition.
Proposition 3.3 (Recursive gradient computations for coupling flows). For a coupling flow, 1141
cond trans fo(atrans trans .2ood) and 2141
xt aondnd
(17)
the derivative of the logarithmic density can be calculated recursively as follows
i log 99.14 (91441 Oog (0)() Outtans
(ffo ..cond) Oxtraus
(18)
a f trans gooud)
b fo
trans cand x
-1
log det
Jutrrans
Drrtrans
Axtrans
i log (0.441
0 blg ((r)
2 log 1
0S (5((x) (attaass groodd
dx
Oxpond
Oxttans
daçond
(19)
0 axpond
S (a(aarau traus o coud
log det
arans
5

---

Published as a conference paper at ICLR 2024
starting with
i log (0,) 0 trans
i log go axb Oxtrans
log go,0(20) Oxgond
i log go( axgond
(20)
For proof, see Appendix B.I. We stress that the Jacobian (fo(z) (anaaa /2/ 1/0 is a h: x k square and invertible matrix, since fo 2ffnd) is bijective for any 2gand E Rdd see (6).
Implicitly Invertible Flows. An interesting property of the recursions in Proposition 3.3 that they only involve (derivatives of) fa 22 ..fond) and thus be evaluated during the sampling from the flow. As such, they are directly applicable to implicitly invertible flows. Further note that the diagonal matrix; the function f acts element- wise on 2) ztrane, see (6). Therefore, the recursion has the decisive advantage that no numerical inversions need to be performed. In particular, there is no need for prohibitive backpropagation through such an inversion.
Jacobian Hfo(z] //a can be inverted in linear time (dd),
Explicitly Invertible Flows. For explicitly invertible normalizing flows the most favorable setup for the baseline method from Vaitl et al. (2022a) the runtime reduction appears to be more mild at first sight. The algorithm has roughly the cost of three forward passes: one each for the calculation of both x and G and one more for the backward pass when calculating the path gradient in (15). This to be compared to the cost of five forward passes for the baseline method by Vaitl et al. (2022a) to calculate path gradients and two forward passes for the standard total gradient. However, this rough counting neglects the synergy between the sampling process x = Txx)) and the calculation of the score G. As will show experimentally in Section the actual runtime increase is only about forty percent compared to the standard total gradient.
Finally, let note that for the aforementioned popular of affine coupling flows recursion from Proposition 3.3 takes a particular form. Since fewer terms need to be calculated, the following recursion gives an additional improvement in computational speed.
Corollary 3.4 (Recursive gradient computations for affine coupling flows). For an affine coupling flow (7), the recursion for the derivative of the logarithmic density can be simplified to
a log (0.41((1441) i log (2,1 trans
2 (2(2 (oood))
(21)
Oxtrans
i log 40,1+1 (I141) 0ztond
log (x)) Oxcond
b Olog qa (1411)
(Don(xrond) (zgond)
Spal Qualatone (afond)
4 x
trans
Ox/41
Oryoon
dngond
0 Dxgond
k
log
0 o (rnond)))
where ztrans is a matrix with entries
trans
trans for i E {I,..., E {l.....d k). Li.
For proof, see Appendix B.2. Additionally, we show in Appendix that the fast path gradient derived by Vaitl et al. (2022b) for continuous normalizing flows can be rederived using analogous steps as Proposition 3.3. Our results therefore unify path gradient calculations of coupling flows with the analogous ones for continuous normalizing flows.
Finally, we note that a further distinctive strength of the proposed fast path gradient algorithm is that it can be performed at constant memory costs. Specifically, the calculation of G' can be done without saving any activations. Similarly, the activations needed for the vector Jacobian product (15) can be calculated alongside the backward pass To To(x)) is known using the techniques of Gomez et al. (2017).
PATH GRADIENTS FOR THE FORWARD KL DIVERGENCE
For training normalizing flows with the forward KI. divergence, previous works have mainly relied on reweighting path gradients (Vaitl et al., 2022a). Specifically, their basic underlying trick is to rewrite the expectation value with respect to the ground-truth p as an expectation value with respect to the model go
DRI(p. go) Exap log p(x:) 9o (x))))
p(x)
log p(x)
Exoage
(22)
(x(x))
qo (z)

---

Published as a conference paper at ICLR 2024
For this reweighted loss, suitable path gradient estimators were then derived in Tucker et al. (2019). Reweighting, however, has the significant downside that leads to estimators with prohibitive vari- ance especially for high-dimensional problems and in the early stages of training (Hartmann & Richter, 2021). As a result, the proposed estimators cannot be applied in a scalable fashion (Dhaka al., 2021; Geffner & Domke, 2021a).
In the following, we will propose general method to apply path gradients to forward KL training without the need for reweighting, To this end, first notice that the forward KL of densities in data space can be equivalently rewritten as reverse KL in base space, namely
DKL (p. qo) DKL (pe.o: qo).
(23)
where we have defined the pullback of the target density to base space as follows
(Tp(
Po.o(zo) p(T)( det
(24)
Oro
We refer to Papamakarios et al. (2021) for a proof. As result, all results derived for the reverse KI. case in the last sections also apply verbatim to the forward KL. case if one exchanges:
go
pa.0
p
40
I,
To( + (x).
(25)
In particular, the fast path gradient estimators can be straightforwardly applied. More precisely, the following statement holds:
Proposition 4.1 (Path gradient for forward KL). For the derivative of the forward KL divergence (plgo) w.r.t. the parameter it holds
d DKL (plgo) = Egnp
po.0 (To (x))
log
(26)
de
qu
where po,0(20) = p(To del
(To())) Ozu
is the pullback of the target density p to base space.
For proof, Appendix B.3. Note that if p is only known in unnormalized form, is its pullback P0.0- However, this has no impact on the derived result as it only involves derivatives of the log density for which the normalization is irrelevant. The following comments are in order:
The proposed path gradient for maximum likelihood training provides an attractive mecha- nism to incorporate the known closed-form target energy function into the training process. In partieular, this can help alleviate overfitting, cf. Figures 1 and 8 particularly relevant concern as the forward training often uses a low amount of samples which entails the risk of density collapse on the individual samples for standard maximum likelihood training. The information about the energy function helps the path-gradient-base. training avoid this undesired behaviour. On the other hand, forward KL path gradient training cannot be used if the target energy function is not known such as in image generation tasks.
As for path gradients of the reverse KL,, we expect lower variance of the Monte Carlo esti- mator of (26) compared standard maximum likelihood gradient estimators. In particular, we note that at the optimum go p the variance of the gradient estimator vanishes.
The proof in Appendix shows that the so-called generalized doubly reparameterized gradient proposed in Bauer & Mnih (2021) in the context of hierarchical VAEs can be brought in the the same form as the path gradient for the forward KL objective derived in this section. However, only formulation elucidates the symmetry between the forward and reverse objective and therefore allows the application for fast path gradient estimators.
5 NUMERICAL EXPERIMENTS
In this section, we compare our fast path gradients with the conventional approaches for several normalizing flow architectures, both using forward and reverse KL optimization. We consider target densities with applications in machine learning (Gaussian mixture model) as well as physics (U(1)
gauge theory, and øt lattice model). We refer to Appendix E for further details."
"Code for reproducing the experiments for GMM and ((1) at gwtthh.com//eem21/000//nneeeddee gradients.

---

Published as a conference paper at ICLR 2024
1.00
0.75
0.50
Reverse KL Geadient
Forward KL Puth Gradient Feeward Kl. Gradien
0.25
Reverse Kl.
Path Grodient
0.00
20
10k 20k 30k 40k 50k 60k Samples Usedd Forward KI.
Walltime Minutes
Walltime Minntes
Figure 1: Effective sample size (ESS) over the training iterations for Gaussian mixture model using the forward and the reverse KL. divergence. The intervals denote the standard error over 5 runs. The best performance is indicated by a dot with subsequent faded average performance the left and center figure. For forward KL, we compare multiple hyperparameter settings (see Appendix E) and plot the respective best runs in the central plot. The right plot displays stereotypical dependency on the data set size for fixed hyperparameters. see Tables 3, and for more details. We see that, typically, path gradients perform better than standard maximum likelihood gradients.
+.344 10
0.450
0.40
0.425
ELBO
0.400
Reverse KL Poth Gradient
0.35
0.375
Reverse KL.
Gradient
0.30
0.350
Walltime Days
Walltime in Days
Figure 2: Training the (1) flow for Lattice Gauge Theory. Shaded area shows standard error over runs. The Reverse KL Path Gradients reach higher performance and exhibit less erratic behavior.
Gaussian Mixture Model. As a tractable multimodal example, we consider a Gaussian mixture model in Rdd with /² 0.5, i.e. we choose the energy function
E(x) = log Σ
Nz Ia)
(27)
(-11)
Note that the number of modes of the corresponding target density increases exponentially in the dimensions, i.e. have 2d modes in total. We choose
resulting in 64 modes. As shown in Figure 1, for most choices of hyperparameters, path gradient training outperforms the standard training objectives. In Figures to 7 in the appendix we present further experiments, showing that path gradient estimators are indeed often better and never significantly worse than standard estimators. The slight overhead runtime is therefore more than compensated by better training convergence. The additional information about the ground-truth energy function included in the forward path gradient training alleviates overfitting in forward KL training, see the discussion in Section 4.
t² Field Theory can be described by a random vector E Rd, whose entries &n. represent values of the corresponding field across a 16 8 lattice. The lattice positions are encoded in the set C N We assume periodic boundary conditions of the lattice. The random vector & admits the density" p(Ø) = ((p)))())) with action
S(Ø)
Σ D.Dnnd Σ (m2oz + \
(28)
M,PAA
where Age is the lattice Laplacian. The parameters m and À are the bare mass and coupling, respectively. We choose the value these parameters such that they lie in the so-called critical region, as this is the most challenging regime. We refer to Gattringer & Lang (2009) for more details
Note that, by slightly abusing notation, plays the role of what before.
8

---

Published as a conference paper at ICLR 2024
Table 1: Results of the experiments from Section 5. We measure the approximation quality of the variational density Yo by the effective sampling size (ESS) plus standard deviations, where higher is better, i.e., 100% indicates perfect approximation, see Appendix E for details.
Reverse KL
Forward KL
Gradient
Path Gradient
Gradient
Path Gradient
E P
4.0
97.4 + 0.0
79.1 + 0.0
91.8 + 0.0
GMM
ESS,
93.0 0.0
97.4 0.0
84.1 0.0
91.8 + 0.0
FISSp
85.6 0.1
96.0 + 0.1
85.1 + 0.1
95.6 + 0.0
ESS,
85.6 + 0.1
96.0 0.1
85.1 = 0.1
95.6 + 0.0
ESS,
40.1000
41.1 = 0.0
U(1)
ELBO
1346.42 + .01
1346.43 + .00
Table 2: Factor of runtime increase (mean and standard deviation) in comparison to the standard gradient estimator, i.e., nanrime path gradien//uuniine standard gradient on an A100-80GB GPU. The upper set of experiments cover the explicitly invertible flows, applied to Ø² as treated in the experiments. The lower set covers implicitly invertible flows applied (1) theory.
Algorithm
Runtime factor with batch size
64
1024
8192
Impl Expl
Alg. I (ours)
1.6 - 0.1
1.4 - 0.1
1,4 = 0.0
Alg. 2 (Vaitl al., 2022a)
2.1 + 0.1
2.2 - 0.1
2.1 0.0
Alg. I (ours)
2.2 + 0.0
2.0 = 0.1
2.3 = 0.0
Alg. 2 + Köhler et al. (2021)
17.5 0.2
11.0 0.1
8.2 = 0.0
on the underlying physics. Training is performed using both the forward and reverse KI. objective with and without path gradients. For the flow, the same affin-coupling-basee Nicoli et al. (2020) is used. Samples for forward KL and ESS are generated using Hybrid Monte Carlo. We refer to Appendix E for more details. The path gradient training again outperforms the standard objective for both forward and reverse training, see Table 1.
architecture as in
Gauge Theory was recently widely studied in the context of normalizing flows (Kanwar et al., 2020; Albergo et al., 2021; Finkenrath, 2022; Bacchio et al., 2023; Cranmer et al., 2023) as it provides an ideal setting for illustrating the power of inductive biases. This is because the theory's action has a gauge symmetry, i.c., symmetry which acts with independent group elements for each lattice site, see Gattringer& & Lang (2009) for more details. Crucially, the field takes values in the circle group U(1). Thus, flows on manifolds need to be considered. We the flow architecture proposed by Kanwar et al. (2020) which is only implicitly invertible. Sampling from the ground-truth distribution with Hybrid Monte Carlo is very challenging due critical slowing down and we therefore refrain from forward KL training and forward ESS evaluation. Table and Figure 2 demonstrate that path gradients lead overall better approximation quality.
Runtime Comparison. In Table 2, we compare the runtime of our method to relevant baselines both implicit for the ex- and implicitly invertible flows. To obtain strong baseline for the latter, numerical differentiation as in Köhler et al. (2021) to avoid costly backpropagation through the inversion. Our method is significantly faster than the baselines. We refer to Appendix E for a detailed analysis of how this runtime comparison scales with the chosen accuracy of the numerical inversion. Briefly summarized, we find that our method compares favorable to the baseline irrespective of the chosen accuracy.
CONCLUSION
We have introduced a fast and unified method to estimate path gradients for normalizing flows which can be applied to both forward and reverse training. We find that the path gradient training consistently improves training for both the reverse and forward case. An appealing property of path- gradient maximum likelihood is that it can take information about the ground truth energy function into account and thereby acts as a particularly natural form of regularization. Our fast path gradient estimators several multiples faster than the previous saate-of-the-ase they are applicable accross a broad range of NF architectures, and considerably narrow the runtime gap to the standard gradient while preserving the desirable variance reduction.

---

Published as a conference paper at ICLR 2024
Acknowledgements. L.V. thanks Matteo Gätzner for his preliminary work GDReGs. L.W. thanks Jason Rinnert for visualization help and acknowledges support by the Federal Ministry of Education and Research (BMBF) for BIFOLD (011S18037A). The research of L.R. has been par- tially funded by Deutsche Forschungsgemeinscaatt (DFG) through the grant CRC 1114 "Scaling Cascades in Complex Systems" (project A05, project number 235221301). P.K. wants to thank Andreas Loukas for useful discussions.
REFERENCES
Abhinav Agrawal, Daniel R. Sheldon, and Justin Domke. Advances in black-box VI: normaliz- ing flows, importance weighting, and optimization. Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Michael S Albergo, Gurtej Kanwar, and Phiala Shanahan. Flow-based generative models for Markov chain Monte Carlo in lattice field theory. Physical Review D, 100(3)003451,, 2019.
Michael S Albergo, Denis Boyda, Daniel C Hackett, Gurtej Kanwar, Kyle Cranmer, Sébastien Racaniere, Danilo Jimenez Rezende, and Phiala E Shanahan. Introduction to normalizing flows for lattice field theory. arXiv preprint arXiv:2101.08176 2021.
Simone Bacchio, Pan Kessel, Stefan Schaefer, and Lorenz Vaitl. Learning trivializing gradient flows for lattice gauge theories. Physical Review D, 107(5), 2023.
Matthias Bauer and Andriy Mnih. Generalized doubly reparameterized gradient estimators. In Marina Meila and Tong Zhang (eds.), Proc. of ICMI., Proceedings of Machine Learning Research, 2021.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicold Cesa-Bianchi and Roman Garnett (eds.). in Neural Information Processing Systems 31; Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, 2018
Kyle Cranmer, Gurtej Kanwar, Sébastien Racanière, Danilo Rezende, and Phiala E Shanahan. Advances in machine-learning-barr sampling motivated by lattice quantum chromodynamics. Nature Reviews Physics, pp. 1-10, 2023.
Akash Kumar Dhaka, Alejandro Catalina, Manushi Welandawe, Michael Riss Andersen, Jonathan H. Huggins, and Aki Vehtari. Challenges and opportunities high dimensional varia- tional inference. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, 2021.
Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear independent components
estimation. arXiv preprint 17Xiv:1410.8516, 2014.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In Proc. of ICLR, 2017.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' A'chh--BBir Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Con- ference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, 2019.
Axel Finke and Alexandre H. Thiery. On importance-weighted autoencoders, 2019.
Jacob Finkenrath. Tackling critical slowing down using global correction steps with equivariant flows: the case of the Schwinger model. arXiv preprint arXiv:2201.022116 2022.
10

---

Published as a conference paper at ICLR 2024
Christof Gattringer and Christian Lang. Quantum chromodynamics the lattice: an introductory presentation, volume 788. Springer Science & Business Media, 2009.
Tomas Geffner and Justin Domke. On the difficulty of unbiased alpha divergence minimization. In Marina Meila and Tong Zhang (eds.), Proc. of ICML, Proceedings Machine Learning Research, 2021a.
Tomas Geffner and Justin Domke. Empirical evaluation of biased methods for alpha divergence minimization. ArXiv preprint, 2021b.
Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger Grosse. The reversible residual net- work: Backpropagation without storing activations. Advances in neural information processing systems, 30, 2017.
Carsten Hartmann and Lorenz Richter. Nonasymptotic bounds for suboptimal importance sampling. ArXiv preprint, 2021.
Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive
flows. In International Conference on Machine Learning, pp. 2078-2087. PMLR, 2018.
Chin-Wei Huang, Laurent Dinh, and Aaron Courville. Augmented normalizing flows: Bridging the gap between generative flows and latent variable models. ArXiv preprint, 2020.
Priyank Jaini, Kira A Selby, and Yaoliang Yu. Sum-of-squares polynomial flow. In International Conference on Machine Learning, pp. 3009-3018. PMLR, 2019.
Martin Jankowiak and Fritz Obermeyer. Pathwise derivatives beyond the reparametcrization trick. In International conference on machine learning, pp. 2235-2244. PMLR, 2018.
Gurtej Kanwar, Michael S Albergo, Denis Boyda, Kyle Cranmer, Daniel Hackett, Sébastien Racaniere, Danilo Jimenez Rezende, and Phiala Shanahan. Equivariant flow-based sampling for lattice gauge theory. Physical Review Letters, 125(12), 2020.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), Proc. of ICLR, 2015.
Jonas Köhler, Andreas Krämer, and Frank Noé. Smooth normalizing flows. In Marc' Aurelio Ran- zato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Informa- fion Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, 2021.
Wouter Kool, Herke van Hoof, and Max Welling, Buy 4 reinforce samples, get a baseline for free! In DeepRLStructPreee ICLR, 2019.
Pierre L L'Ecuyer. An overview of derivative estimation. Proceedings of the 1991 Winter Simulation Conference, 1991.
Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Arnaud Doucet. Continual repeated annealed flow transport Monte Carlo. In International Conference Machine Learning, pp- 15196-15219. PMLR, 2022.
Laurence I Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernaneez-Lobato. SE(3) equivariant augmented coupling flows. ArXiv preprint, 2023.
Andrew C. Miller, Nick Foti, Alexander D' Amour, and Ryan Adams. Reducing reparameteriza- tion gradient variance. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Anmual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, 2017.
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte Carlo gradient estimation in machine learning. J. Mach. Learn. Res., 2020.

---

Published as a conference paper at ICLR 2024
Kim A Nicoli, Shinichi Nakajima, Nils Strodthoff, Wojeiech Samek, Klaus-Robert Müller, and Pan Kessel. Asymptotically unbiased estimation of physical observables with neural samplers. Physical Review E, (2) 2020.
Kim A Nicoli, Christopher Anders, Tobias Hartung, Karl Jansen, Pan Kessel, and Shinichi Naka- jima. Detecting and mitigating mode-collapse for flow-based sampling of lattice field theories. ArXiv preprint, 2023.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium
states of many-body systems with deep learning. Science, 365(6457), 2019.
George Papamakarios, Eric T. Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. J. Mach. Learn.
Res., 2021.
Paavo Parmas and Masashi Sugiyama. A unified view of likelihood ratio and reparameterization gradients. In International Conference on Artificial Intelligence and Statistics, pp. 4078-4086. PMLR, 2021.
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racanière, Michael Albergo, Gurtej Kanwar, Phiala E. Shanahan, and Kyle Cranmer. Normalizing flows tori and spheres. In Proc. of ICML, Proceedings of Machine Learning Research, 2020.
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco J. R. Ruiz, and Ömer Deniz Aky- ildiz. VarGrad: low-variance gradient estimator for variational inference. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin (eds.), Advances in Neural Information Processing Systems 33; Annual Conference on Neural Information Pro- cessing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Geoffrey Roeder, Yuhuai Wu, and David Duvenaud. Sticking the landing: Simple, lower-variance gradient estimators for variational inference. In Isabelle Guyon, Ulike von Luxburg, Samy Ben- gio. Hanna M. Wallach, Rob Fergus, S. V. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Pro- cessing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, 2017.
George Tucker, Dieterich Lawson, Shixiang Gu, and Chris Maddison. Doubly reparameterized gradient estimators for Monte Carlo objectives. In Proc. of ICLR, 2019.
Lorenz Vaitl, Kim A Nicoli, Shinichi Nakajima, and Pan Kessel. Gradients should stay path: bet- ter estimators of the reverse-and forward KI. divergence for normalizing flows. Machine Learn- ing: Science and Technology, 3(4), 2022a.
Lorenz Vaitl, Kim Andrea Nicoli, Shinichi Nakajima, and Pan Kessel. Path math-gradient estimators for continuous normalizing flows. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference Machine Learning, ICML. 2022, 17-23 July 2022, Baltimore, Maryland, USA, Proceedings of Machine Learning Re- search, 2022b.
Xi Wang, Tomas Geffner, and Justin Domke. A dual control variate for accelerated black-box variational inference. In Fifth Symposium on Advances in Approximate Bayesian Inference, 2023.
12

---

Published as a conference paper at ICLR 2024
APPENDIX
A NOTATION
For function v Rd - R we denote with di its total derivative, for function y R x R' - R we denote with @p(x.) th its partial derivative w.r.t. the first argument x E R and for a function 8T its Jacobian. For ease of notation, we sometimes use the shorthand T: R + I2d we denote by Br notation dz, for 8_ T 1+1,0 (I441), keeping in mind that I1+1 1,+1,0 (21). We write Fzi+1 3
de
v for the path gradient of jo as defined in Definition 3.1. The symbol @ denotes elementwise multiplication and @ denotes elementwise division.
B PROOFS
In this section we collect the proof' of our propositions and corollaries, which we will recall here for convenience.
B. RECURSIVE GRADIENT COMPUTATIONS FOR COUPLING FLOWS
First, let us recall the following proposition from Section 3.1.
Proposition 3.2 (Gradient recursion). Using the diffeomorphism Tr, the derivative of the induced probability can be computed recursively follows
o (21) Bar
i log q0 (1141)
b log go,(
[0Ti,,
ô log det
(TT (
0z141
dr
dr
Om
(16)
Proof. The basic definition of flows consisting of compositions
log det (1141)
log qo, (0,141 (1141) log go. (477 (z1+1))
(29)
011+1
implies the following recursion
① log (0.141
aog qo (oo (1141))
log det
(30)
0m
0r1+1
0r
8x
Since for general normalizing flows, the inverse T, (x141) is not efficiently computable, we apply the chain rule and the inverse function theorem
i log q0,14 (
log q0,((x)
0T,o) (F141)
(ogg det
ATi, (zi) Om
0T (#1+1)
05
Ox
Oe1t1
Om
0m1t1
(3la)
Ooo (xx)
(0 (x1)
[og del
0T, 0 (TT)(10))
(
(oTio, (mr)
Oxq
Oay
Om
Om
(31b)
We can break down the recursive operation further the next proposition. Proposition 3.3 (Recursive gradient computations for coupling flows). For a coupling flow,
21+1
trans
fo(
21
condy
and
F141
cond
x
cond
(17)
13

---

Published as a conference paper at ICLR 2024
the derivative of the logarithmic density can be calculated recursively follows
i log (0.44
① Olog (a.(
[af xroud)
0x141
Oxtrans
Oxt Arraans
(18)
0 axtraus
ff (z) (atran 27 220nd)
@fe
(ff zi
goand)
log det
Ortrans
Oztrans
i log g0.141(11+1)
Ölog qo.
a log 90,140 (11+11 ofo (af
Oxttud
Oxyond
Oxttis log det
Oxyond
(19)
0 Oxpond
Oxtrans
starting with
i log (o,0(<r))
2 log (2 Ozb trans
log go,n(
2 log (2ro
(20)
Ouo trans
Oxyond
Oxgond
Proof. For coupling flows nant of the Jacobian is
(atraus gound) and goond
27 goond This implies that the determi-
Su trans teans
Ou trane
Outrans
det
= det
= det
(32)
0z1+1
Outrans
Using this result and splitting into the transformed and conditional components, the recursion can be rewritten as follows
i log 11.14111100011 (F141)
log (3,)(
a log det Ox Öxtrans
+
(33)
Oxtrans
Oxtraan
Oztrans
@ log (1141) Oxfind
a log (3:)))
θ det
log
(34)
Ouftn
Onnd
0x
trans
We want to rewrite the right-hand side of this recursion terms of quantities that involve derivatives with respect to zy. To this end, let us study derivatives w.r.t. ztrans and goond respectively,
For generic function up TR Rd-k + TR. it holds via the chain rule that
aeplutrans . ,roond) aeplutrans .roond axtrans Oxtrans
Oeplutrans, Oxgond
x gondy axcond
(35)
Oxtrans
Oxtrans
transs
Noting that 2) trrans tor (affar ( and noting that xit1 2oond
we can compute
aiplatrans x goondd
Oepatrans . (ooodd 05o 2 cond
(36a)
Ontrans
Oxtrans
ap(xttans
x cond)
f (ot trans zrond
-1
(36b)
Ontrans
Oxtrams
where we used the fact that the Jacobian of the inverse is the inverse of the Jacobian due to the inverse function theorem.
Similarly, we can write
aylatraass 2gond) Ox
aplatrrns zcond)
Oztrans On
aeplttrans acondy Ongond ap(atrans agonnd Ongond
Ozcond
(37a)
Drtrans
Oz
ap(atrans trans ggond) Oztrans
(37b)
traus
dnoond
where we have used the decomposition (6) of coupling flow.
14

---

Published as a conference paper at ICLR 2024
Ozj ax]
Note that the Jacobian have assumed that urreas fo (1/+1 - is invertible only with respect to its first argument ff 2f 2fond) Duy
θecīr not necessarily invertible and may not even be square matrix. We
S
(for any choice of its last). So while the Jacobians a
and
square and
invertible matrices, the same cannot be said for the Jacobian following trick
However, we can use the aur
0 = [+1
( :zfond)
(38a)
0x1+1
Oufond
0 Sso Wofffrane 11 .ond)
dxprane
+ 05o , soond)
Oxyond
(38b)
aztrans
Expond
dxyond
Outtn
( 2n sooond autrans
O S sggond)
(38c)
Oztrans
Oxfond
Oryoud
0 (ala))nn toae on gound
Since the Jacobian
is invertible, the above statement is equivalent to
Oz Leuue
Oxtrans
Salaatrans aroond -1 fo(attrans at acond)
(39)
Ozcond
Oztrans
Ozgond
Substituting this result into (37b) yields
Op(xtraass arund)
ap(a ccond) Intrans
ffolztrans sfand) Ontrans
2 ff(arrans xrond az
0xft1
(40)
(p( trans 17
cond
goond
Next, we note that the determinant of the Jacobian can be rewritten as
0fg (zxtrans condd
0.f (x) trans x gond
log det
log det
(41)
Oxtraas
using again the inverse function theorem.
Now, plugging (36), (40) and (41) into (33) for suitable choices of yo, we can rewrite the recursion in the desired form:
i log go (x1+1) 0 log (x1)
Offe (ztraas, , xt xf Oxtraus
Oxtrans Oxtraaus
(42)
a artrans This is precisely the form stated in the proposition. Analogously, plugging (36), (40) and (41) into (34), the conditional component can be rewritten as
fz trans Ox.
20 gondd
Of (o x: -1
log det
Ox
i log (21-1)
2 log (2))) azcond
a log (0()) Oxtirans
(Ofolattans rondd Oxtraus
a ff 1
cond)
Oarond
Dagond
Oz log det a trans poond)
@fo(o)rams .z .roodd Oxtrans
-1 Ofo (ood Inyoud
+
Oxtrans
① log det 0 grond) Ozcond
Oz
(43)
15

---

Published as a conference paper at ICLR 2024
This be brought in more compact form by noticing that the term expression for the conditional component. Using this, we can rewrite the conditional component
8 log de (2141) 0
appears in the
log go (11111) a log go a log (o(())) ofo(ztrane axfond
Ozgond
0zt+1
Ozyoud
(44)
log det 0So(x)rame 2f aoondy
Orgond
Oxtrans
which is precisely the form stated in the proposition.
B.2 RECURSIVE GRADIENT COMPUTATIONS FOR AFFINE COUPLING FLOWS
Let us recall the following corollary from Section 3.1.
Corollary 3.4 (Recursive gradient computations for affine coupling flows). For an affine coupling flow (7), the recursion for the derivative of the logarithmic density can be simplified
a log g0,41((444)) θ log (0,( Oritns
2 (00
(21)
Ozirans
i log (0,+1(1141)) Olog g0 (x11)
b blog (1
(Z0g(z)on))
ppo (p (2fand)
Oxfond
Orgond
trars
Orgond
Iz
Oxgond
a Oxcond
log
Il 0 o0 (zpond))
where Ztrans is matrix with entries
(Itrans)
Lis gtrans for i [l,....k),,7 E {l,.... d k).
Proof. Affine coupling flows are defined by
#1 So (xx
(arfrans ln (fondd oo (2 (zfand) 2 xi trans
uo (x(ond),
F141 27
cond cond
(45)
This implics for the Jacobian
ff faann . ggond) Oxtrans
diag (2f(nd)))
(46)
and for its determinant it holds
Of fallztran p roond det
k II iml
Ta.s (2) (2cond).
(47)
Oz drata
which notably does not depend on attans, Therefore, the recursion for the transformed component simplifies to
i log (9.441 (0.44((444) ① log (ai)
(85 (Dfo(at (atrane, (ygond) -1 Oxtraus (
Oxtrans
ixtraus
Ofe(atrans at aronnd Oxtrans
log del Ofe(ztrans grnnd) Inttran
(48a)
axttraus
i log (0.1(1z)
2 (2(2z (2gond)
(48b)
0z trans
16

---

Published as a conference paper at ICLR 2024
Forward KL Path Gradient
Forward KL Gradient
Gradient Norm
2-
1-
10
20 Wall time (h)
40
50
Figure 3: Gradient norm during training of the o' ex-eerriments. The norm of the path gradient estimator closer to zero than the norm of the standard gradient estimator when the target density is well approximated, indicating lower variance.
which is precisely the form stated in the proposition. The recursion for the conditional component simplifies to
i log 90 (1141) d log go (x1)
θ Ölog qo (1142) 0fo grndd)
(49a)
@rgond
Oztrrns
anyond
a augond log (1))
Sf ( gondd)
log det
Oat trans
alog q0.14 (1(11)
@o (Don(ztondl Orcoud
Ops Qu (2r
Itrans +
Orcond
Oulrans
dzcond
i Dacond
log
o ( cond
(49b)
where zfrans is a matrix with entries (z)rans) = I1,
trans for i {l,...,,),, k), E (l,...,d ddd--) k}.
This shows the claim.
B.3 PATH GRADIENTS FOR THE FORWARD KL DIVERGENCE
In this section, prove Proposition 4.1 and lay out its implications for path gradient estimators, both for the forward and reverse KL divergence. In particular, we will discuss the favorable vari- ance properties of path gradients in the case of the forward KI. divergence, which also verify experimentally in Figure 3.
Proposition 4.1 (Path gradient for forward KL). For the derivative of the forward KL divergence DKL (plge) w.nt. the parameter it holds
DKL (plge) = Eony
v log
Pe.0
(T
(x))
(26)
d0
where pa.o(ao) = (zo)) det
(To))) Uzo
is the pullback of the target density p to base space.
17

---

Published as a conference paper at ICLR 2024
Proof. Let first note that
log logp(ar) log det (x)
DKL (pla0)
Exap
log qo (To (x))
(50a)
dx
Enorpooo log logg(To(<D)) + log det
Ta((v))
- log 20 (20)
(50b)
Ozo
(Pe.049)))))
(50c)
Note that in (50) we have essentially transformed a forward into reverse KL divergence. By looking at the problem of minimizing the forward KL divergence (xl(p) - from target density to the variational density Qo
reverse KI. divergence DKL (po.olo12 (o0 from variational density po.o to target density Oo - we can employ the tools that exist for optimizing the reverse KI. divergence. In particular, we can apply the standard path gradients derived by Roeder et al. (2017) and Tucker et al. (2019). We compute
DKL (po.ololgg) d0
d
Bunp
log
P8.0
(Ta
(x))
(51a)
00
40
0
P8.0 (2o)
0T, 'Tx'(x)
2 log pe (2o)
En-p
log
(51b)
Ozo
40
00
To (a))
Exp
Vo log Q0
P8.0
(To (x))
+ PETOOrPOOO
a log (a.)
(51c)
00
Enep
log
P6,0
(To (x))
(51d)
q0
where we used the fact that the score
log pe.o (zo) vanishes in expectation over Pe,0. since
00
0
E Eizopoo.
log (xo)
log Pog pog.
(52a)
00
Ra
00
a
Pe,o (o,0
(52b)
Rad 00
/ po,o(ro)dog
(52c)
00
Ô O.
(52d)
00
Now, the statement follows by combining (50) and (51).
B.3.1 VARIANCE: STICKING THE LANDING PROPERTY
The covariance matrix of the score term is known to the Fisher Information (cf., e.g., Vaitl et al. (2022a)). In this case, is the Fisher Information (o) of the pullback density p8,o, defined
à
(o =
log
T 00
log pa.pa..o (000)
(53)
00
the model perfectly approximates the target density, i.e. po,o(zu) = (u(zo) for all xo E Rd, then the path gradient term
log
P9.0 (T) (x))
(log (P..((o)) - log qo (zo))
aT. (x) 00
Ixo
Ioo=T
(x)
is zero almost surely. This implies that the path gradient estimator has zero variance in the limit of perfect approximation (which sometimes called sticking the landing), while the standard gradient estimator has non-vanishing covariance Io(O)/N, where N is the sample size of the Monte Carlo estimator. Note that these results are exactly analogous to the reverse KL path gradients, noting the relations in (25). Experimentally, we find for the forward KL divergence that the gradient norm indeed behaves as in previous works (Roeder et al., 2017; Vaitl et al., 2022a), ie.. it exhibits the vanishing gradient properties. This is illustrated in Figure 3.
18

---

Published as a conference paper at ICLR 2024
B.3.2 MOTIVATION FOR REGULARIZATION
The favorable behavior of path gradients for the forward KL divergence, which we have defined in Proposition 4.1, can potentially be understood by identifying the additional terms appearing in the gradient as having regularizing effect. To this end, let us compare the standard maximum likelihood gradients with the path gradients. The former is given by
d
de DKL -Beup
log go()))
(54a)
d0
d
log (Tp (x) + log det 0T' (x)
)
= - BBep
(54b)
d0
Ox
whereas the latter can be computed as
DKL( = Exnp d0
log
(Tg (x))
(55a)
P8.0
MTa (xo
Expp
log (x())
log det
Uzo
dxp
01 (x)
log (To((())))))
(55b)
00
Tx=
Note that only in the path gradient version (55), the target density p appears and we conjecture that incorporating this information helps to not overfit to the given data sample. Crucially, in many ap- plications, is (up to normalization) given explicitly such that (55) can indeed be readily computed. Because of the duality of the KL divergence, the regularization property also appears for the reverse KL, where the term involving do does not appear in the standard gradient estimator, whereas for the path gradients it does.
B.3.3 RELATION TO GDREG
The above gradient formula (26) can be seen as special case of the Generalized Doubly- Reparameterized Gradient Estimator (GDReG) derived in Bauer Mnih (2021), by noting that
P9.0 (To (x)) = - Enpp
To log P (To(a)))) qo
Expup
log
(56)
Qo
xo=T,
This can be seen as follows:
e log (To(2))
(57a)
qo
20 To
a (logp(z) log det
0r, ((x) Ox
log 2o (To (x))
(x))
(57b)
Or
30
To
a [ogp(2) log det
aT, (x)
- log qo (TØ (z)))
dx
(2o)
OT, (x)
(57c)
Ōr
Oro
00
a (og (To(())) + log det
(To) Ozo
log (0(
0T (x)
(57d)
Ozo
00
P8.0 (Tg (x)),
log
(57e)
Go
where in (57c) we have used the identity
(T, (x))
(x)()) OT, (x)
[Tg(z)
(58)
00
Ø=
30
00
(tr)
However, only our derivation allows interpret the estimator an instance of the standard sticking- lan-ianding trick of eliminating score function, which then motivates variance reduced estimators and allows to the proposed Algorithm 1. Further, the time for computing (26) is expected to be significantly lower than for the GDReG estimator since in this form we can employ our proposed algorithm for efficiently computing the path gradient.
19

---

Published as a conference paper at ICLR 2024
C RELATION TO FAST PATH GRADIENT FOR CONTINUOUS NORMALIZING FLOWS
In this section, we demonstrate that the recently proposed fast path gradient for continuous normal- izing flows (CNFs) proposed by Vaitl et al. (2022b) can be obtained using completely analogous reasoning as our derivation for the fast path gradient of the coupling flows.
We first note that the algorithm described in Section 3.1 can be applied verbatim to the CNF case with
X, (X()) = Xo /' vo ds,
(59)
where IR x R - Rd is the generating vector field of the CNF. Note that To defined in (1) can be interpreted as a discretization of the flow To and x in (1) can be seen as discrete approximation of X,. The only difference between the CNF and coupling case arises in the recursive relation to obtain the derivative
log qo 0x,
(60)
Vaitl et al. (2022b) proposed an ODE that can be evolved along with the sampling process. As we will demonstrate subsequently, this ODE can easily be recovered by using the same reasoning as we applied for the coupling flows.
We start from the observation that the ODE can be discretized as
11+1 = xt + 2o
(31, (31,
At.
At.
(61)
for I E (O,..., 1 11 the sense that Xia: = 2) + (( Here, we assume for convenience that the time increment At > 0 is chosen such that f is an integer.
Az
As for the coupling flows, start from equation (31b), which for a discretized CNF is given by i log g0(
a lloggo(a))
à log det
dz 00141
(62)
00141
00141
Ja1+1
Using the chain rule, this can be rewritten as log go(x++))
a log (o(a))
Oxi
Oxi
+
log
det
(63a)
0ri+1
0ri+1
Om
0zi+1
0ri+1
i log 90(xz)
(0x++1
Oxi
))
0xl+1
+
log
det
(63b)
dx
Dzl+1
Ox
i log (m)))
021+1
0x1+1
Oxi )
00141
log
del
(63c)
dz
am
Ox
From the discretized ODE (61), it follows that
(x1+1
Oes
Duo
1 + At
(x),
/tt)
=1 1 -
At
(x), 1At).
(64)
dzi
Oxi
Omi
Using this expression, we obtain
i log go (1141))
( log (m()))
Oue dmi41 )) (1-2/ A
log
det
(51, 121)
(65)
0m141
Oa
Oai
Oay
Using the standard series expansion of the matrix-valued logarithm, we can check that
Iva 11 At
(-1)++1
(Ovg) k At*
[ovo
log det
(ox, 1bt)
Σ
Tr
ArTr
(I,, lAt)
+O(AP),
dm
k
Oxi
(66)
20

---

Published as a conference paper at ICLR 2024
where in the second equality, have used that log det A Tr A and log(1 + B) Ex=1 (-1)++1 B for matrices A and B. Substituting this expression, in the expression above,
we obtain
(olg
0 log go(mz)
i log (o())) dvo (sn, (A,,At))
a dz
Tr
(xi, (si, O(A)).
At
0m++1
Dzi
dz
dx
Ör
Taking the limit At ++ O, we obtain precisely the ODE derived by Vaitl et (2022b):
log qo Xx) log go( (x) Oug (X,,t) ax, ax,
Oug
Tr
(Xt,t).
(67)
dt 0x,
0x,
ax,
As result, the fast path gradient derived in this manuscript unifies path gradient calculations of coupling flows with the analogous ones for CNFs.
D ALGORITHMS
In this section we state the different algorithms used our work. As discussed in Section 4, due to the duality of the KL divergence. employ the algorithms for both the forward and the reverse KI.
The different algorithms treated in this paper are:
1. The novel fast path gradient algorithm - shown Algorithm 1.
2. As a baseline, the method proposed in Vaitl et al. (2022a) shown in Algorithm 2.
3. As further baseline, the same algorithm, amended the GDReG estimator (Bauer & Mnih, 2021) - shown in Algorithm 3. We stress that this algorithm is an original proposal of this paper, however, we did not introduce in detail in the main text it slower than our fast path gradient and therefore more of a side-produe serving as strong baseline.
The algorithms do not contain the path gradients for the target density, since those can be done readily. A graphical visualisation of the respective terms of the algorithms is provided in Figure
Algorithm 1: Fast Path Gradient: computation of log qo (TT((T)) 2o qo
Input:
base sample To ~ 90
for in (0...., L 1]: Apply 1:41.0 to compute #141
D joint forward pass and recursive equations
Compute 49, (01+))))))))) according to Proposition 3.3
(
return
DE
FL.
compute vector-Jacobian products
Algorithm 2: Path gradient: computation of V log go (To (Z))
~ qu
Input: base sample To 40
stoo-graddeet((a (xo)) Qa (xr) (x')) det or, (c))
D forward pass of mp through the flow without gradients
reverse pass to density
(
I. compute gradient with respect to at'
m < (2)))
standard forward pass
return G G
compute vector-Jacobian products
E COMPUTATIONAL DETAILS
In this section elaborate computational details.
21

---

Published as a conference paper at ICLR 2024
Algorithm 3: Path gradient: computation of Vo log go (To
zo= '(x)) x
p
Input: target sample x' ~
4o (2') + Lo(TT (x)) (det (w')
(a')
D reverse pass to calculate dousity
0x
G (2222
compute gradient with respect x'' P. copy To (ar)
de
m6 +- stop coo..grdien((T (x))
To(x))
standard forward pass
return o
compute vector Jacobian products
AD
p(aso Bore [80 log go(())) (2)]2== doL [[[O,, logp(ze) Oz, log go (xa)) dsze]
doc(0)= E Or, log (ata)))
Puth Gradient
Score
Path Gradient
Ze
log gu((ss)
log yo (20)
8 log ge (219)
Normalizing Flow
Normaliaing Flow
log g(z)
log ((2)
log (())
Path Gradients and Score
Fast Path Gradients without Score
Figure 4: The gradient of the loss function de (() consists of the path gradient (blue) and a score term (red) - the latter vanishes in expectation, but has vanishing variance. The path gradient framework computes the necessary quantities to perform stochastic gradient descent with only the path gradients, eliminating the impact of the score term which tends to increase the variance of the gradient leading suboptimal gradient estimation. We propose fast algorithm for computing the Path Gradient estimator which evaluates 0zg log qo (xx)) (green) alongside the forward pass. The term On. log p(x₂) is assumed to be available for given energy function within the problem formulation. We can then compute the derivative of the log ratio, Øz, log (p(ze)/go(zo)) which can be interpreted a scaling of the gradient Ögzg, the gradient with respect to the parameters. Thus the gradient dgL (0) only consist of the path gradient, eliminating the negative influence of the score term.
Evaluation Metries. As common, we assess the approximation quality of the variational density qo by the effective sampling size (ESS), defined as
ESS
(68)
Eggago [=)(x)] E (w()
where w(x) == p(z) [c) are the importance weights.
As can be seen in (68), the ESS can be computed with samples from either qo or p (if available), which leads to two different estimators, namely
N
(n) go
ESS
((()) ~ ~p (69)
(2(())))
EN
(z((())))
m==
where the normalization constant Z appearing in the importance weights w can as well be approxi- mated either with samples from go or p. respectively, namely
S((r
())
wx(x) = (z)))
2,
(n) qo.
(70)
N n=
go (ar(m))
22

---

Published as a conference paper at ICLR 2024
E -S(x)
Nv
yo (x:(())
-1
(p(x)=
qo (z) (a(x)2z Zp
7p
N
S(afro))
x!n)
(71)
Note that E might be biased due to potential mode collapse of qo, which is not the case for ESS,. The ESS is measure of how efficiently one can sample from the target distribution. E.g., an ESS of 0.5 means that we have N samples from the sampling distribution 9o, the variance of the reweighted estimator is large an estimator using 0.5NV samples from the target distribution. maximum ESS of occurs when the importance weights w are exactly for every sample, the minimum ESS is O.
Gradient Estimators. As baselines for the path gradients we used the Maximum Likelihood gra- dient estimator as the Forward KL Standard Gradient estimator
.V. Σ
DKL (ploo) = d0
log go (z)(),),
(1)) P.
(72)
N de ==1
and for the reverse KL we used the reparamettiization trick gradients
d a DKL
d
log
98 (Ia(=b"))) p
(n) T0
qo
(73)
Note that both gradients are independent of the normalization constants of both qo and p.
Gaussian Mixture Model. We use RealNVP Dinh et al. (2017) flow with weight normalization.
The ReaINVP flow contains 1,000 hidden neurons per layer and six coupling layers, each of which consists of 6 hidden neural networks layers with Tanh activation to generate the affine coupling parameters. We draw 10,000 samples from the Gaussian mixture model (GMM) for the forward KL training. thus mimicking finite yet large sample set.
The superior performance of path gradients for Gaussian Mixture Models can be observed in Figure 1. In of the reverse KL in the left plot, both path gradients and standard gradients achieve the same forward ESS, yet the path gradients converge faster in wall time, despite their increased computational cost of 40% compared the standard gradients.
The middle figure shows the improved performance of the forward KI, path gradients compared to the standard non-path gradient. The forward KI. path gradients converge faster high forward ESS and are able to maintain the achieved sampling efficiency compared to the standard gradients which decrease a forward ESS of zero.
Finally, the right plot examines the different performance of forward KL path gradients in more detail by increasing the number of training samples on which the forward KL optimized. The advantage of path gradients becomes evident even more as they generally surpass standard gradients earlier and achieve higher ultimate forward effective sampling sizes even for very large forward KL training sets.
Besides the forward ESS, the corresponding negative Loglikelihood (NLL) for both gradients and path gradients is plotted in Figure 9 for an increasing number of training samples. One observes how the NLL for the training and test data set differ only slightly compared to the performance gap of standard gradients measured in terms of NLL. Path gradients are robust even in low data regimes where standard gradients diverge strongly in terms their training and test performance, indicating tendeney to overfit.
The results in the experiments Figures and 9 are obtained after 10, 000 optimization steps, with learning rate of 0.00001 with the Adam optimizer (Kingma & Ba, 2015) with batch size of 4, 000. The target distribution is identical to the setup in Section and the forward ESS and NLL evaluated with 10, 000 test samples.
Finally, Figure 8 illustrates the performance gap between standard gradients and path gradients on simple two dimensional multivariate Normal distribution. The visualization hints stronger
23

---

Published as a conference paper at ICLR 2024
regularization for path gradients which are able to incorporate gradient information of the underlying ground truth energy function.
the course of the optimization for varying numbers of linear layers per coupling block, number of hidden neurons per linear layer and the batch size used for optimization trained with Forward KL Path Gradients and non-path Forward KL Gradients. While standard Forward KL gradients can relatively surpass Forward KL Path Gradients mildly in few cases for smaller models, larger model trained with Path Gradients doubles the mean of the ESS, in direct comparison and improves upon the best Forward KL Gradients in absolute terms. Ultimately, the best performance as measured with the E achieved with larger models trained with path gradients. Importantly, path gradients provide increased robustness against overfitting as can seen in the performance during training. Standard gradients for larger models tend to deteriorate their ESSpp
Figures 5, 6 and 7 show the ESSp
The Tables and collect the best ESSp achieved during optimization for the same combination of number of linear layers, number of hidden neurons per linear layer and batch size. This corre- sponds to saving checkpoints during training and choosing the respective model with the best ESSp. The Forward KL Path Gradients outperform their non-path counterpart except for few instances, which perform worse than larger models as measured in ESS, and trained with path gradients.
Table 3: We summarize the highest achieved average ESSp over the entire optimization which corresponds to the best attainable performance terms of ESS, possible for that model capacity. The flow consists of 6 coupling layers, each with 2 linear layers, each of which has an increasing width (number of neurons).
Layer Width
10
50
100
250
500
1000
Batch Size
Std
Path
Std
Path
Std
Path
Std
Path
Std
Path
Std
Path
10
55.7
56.9
58.5
58.8
58.5
58.8
60.0
62.4
71.3
72.7
72.1
71.7
50
58.0
58.3
58.9
59.2
59.3
60.0
69.8
75.2
76.0
77.2
66.6
80.5
100
58.3
58.6
59.1
59.4
60.2
62.1
73.4
79.7
75.1
78.3
66.3
84.1
200
58.5
58.7
59.4
59.7
62.7
66.2
75.1
81.7
73.7
78.4
68.3
86.1
500
58.6
58.8
60.0
60.8
68.9
66.8
75.7
81.8
73.0
80.5
67.0
86.9
1000
58.7
58.9
61.1
63.0
70.8
66.8
75.0
81.5
72.5
86.2
54.7
86.9
2000
58.7
58.9
63.5
66.7
69.9
72.0
75.2
85.7
69.8
88.0
44.5
86.6
4000
58.8
59.0
65.3
69.5
73.4
77.0
75.4
88.9
72.1
88.9
41.8
88.4
Table 4: summarize the highest achieved average ESS, over the entire optimization which corresponds to the best attainable performance in terms of ESSp possible for that model capacity. The flow consists of 6 coupling layers, each with linear layers, each of which has an increasing width (number of neurons).
Layer Width
10
50
100
250
500
1000
Batch Size
Std
Path
Std
Path
Std
Path
Std
Path
Std
Path
Std
Path
10
56.2
57.4
59.0
59.2
59.0
59.1
59.1
59,2
59.5
64.7
66.5
76.8
50
58.5
58.8
59.0
59.2
59.0
59.1
61.3
68.1
73.1
78.9
74.7
84.7
100
58.7
58.9
59.0
59.2
59.0
59.1
68.1
76.6
74.0
83.4
72.4
86.6
200
58.9
59.0
59.0
59.2
59.1
59.7
73.7
81.1
75.9
86.0
69.8
88.4
500
58.9
59.1
59.0
59.2
59.9
65.5
77.8
84.7
74.7
88.3
68.3
90.5
1000
59.0
59.1
59.1
59.3
66.8
73.2
79.1
85.2
75.7
90.3
60.4
88.3
2000
59.0
59.1
59.2
60.3
71.3
73.4
78.4
86.9
71.8
90.0
52.9
87.6
4000
59.0
59.1
59.6
62.9
73.4
78.0
77.7
89.0
68.8
89.9
50.7
88.0
24

---

Published as a conference paper at ICLR 2024
Number of Linear Lagere per Couplng Blocic
Batch Size 10
Dotch Size 50
Bareh Size T0o
Datch Sioe 2001
Donch Sice
Batci Size
Dotch Size
Boteh Stac
500
1000
2000
4000
Hiodm: 30
Hilde 200
Eiiioo 300
Eilbeer 1000
Optinsookis Siew
Opennoniinn Supp
Optintotion Step
Ope inbooaa
Ostiization
Optintoskin
Openncsian
Optinication
Supp
Saye
Saye
Suqu
Stape
The ESS, RealNVP flow with two linear layers in each its six couplings blocks trained with Forward KL Gradients shown in red and Forward KL Path Gradients shown in blue. For higher model capacity and larger batch sizes, the Forward KL Path Gradients achieve higher absolute ESSp while the Forward KL Gradients collapse with increasing model capacity. The rows increase width of the linear layers (hidden neurons) and the columns increase the batch size used during optimization.
Figure 5:
6 Field Theory. For flow architecture we slightly modified NICE (Dinh al., 2014) architecture, called Z2Nice (Nicoli et al., 2020), which is equivariant with respect to the Ze symme- try of the &² action (28). We use lattice of extent 16
learning rate of 0.0005, batch size 8, 000, AltFC coupling, 8 coupling blocks with 4 hidden layers each. A learning rate decay with pa- tience of 3, 000 epochs is applied. We used global scaling, Tanh activation and hidden width 1, 000. As base-density we chose a Normal distribution applied. Just like Nicoli et al. (2023), training is done on 50 million samples. Optimization performed for 48h on single 1000 each, which leeds to up 1.5 million steps for the standard gradient estimators and 1.1 million epochs with the fast path algorithm.
(00 1). Gradient clipping with norm=1 is
U(1) Gauge Theory, The U1) experiments are based on the experimental design Kanwar et al. (2020). We consider lattice with 162 sites with batch size of 12, 288, learning rate 0.0001, 24 coupling blocks with a NCP (Rezende et al., 2020) with 6 mixtures, hidden size x and kernel size 3 and uniform base-density do (0, 2T). We train models A A100 for one
25

---

Published as a conference paper at ICLR 2024
Number of Linear Lagere per Couplng Blocic
Batch Size 10
Dotch Size 50
Bareh Size T0o
Datch Sioe 3001
Donch Sice
Batci Size
Datch Size
Botch Stac
500
1000
2000
4000
1akier 36
Hiodm: 30
Hileec 200
Miiiom 300
Eibbeer 1000
Optintookis Siew
Opennnnian
Optintotion Stepe
Ostiizatio Saye
Optistoskin Saye
Opennostan
Optinication
Supu
Suppe
Suqu
Stape
Graiánk
Figure 6: For many combinations increasing parameterization and increasing batch size, Forward KL Path Gradients able to maintain their high ESSp, while performance of Forward KL Gradi- ents increasingly deteriorates. The rows increase width of the linear layers (hidden neurons) and the columns increase the batch size used during optimization.
week, the batch size was chosen, so to maximize GPU-RAM for single GPU. Because training from random initialization led to very high variance in performance, we pre-train the model for
3.0 for 200,000 epochs. The shown benchmarks show the training after initializing from the pretrained model for target beta 3 = 3.0. For the standard reverse KL gradient estimator (using reparametrization trick) this yields 700,000 and for the fast path gradients 300,000 epochs.
The results can be seen in Figure 2; they are averaged over running average of window size and 4 repetitions, the mean and standard error shown. The ESS is estimated 10 batch size samples.
Runtimes. In order to give a fair comparison for the runtimes, we use the hyperparameters of experiments in this work testing ground. Namely for the affine coupling flow, we use the aa experiments and for the implicitly invertible flows, we use the (11 experiments. Walltime runtimes measured For the affine flows, we use the setup from the o experiments and look at the existing Algorithm for computing path gradient, the proposed fast path Algorithm 1, well as the equally new Algo-
A100 GPU with 1,000 repetitions.
26

---

Published as a conference paper at ICLR 2024
Number of Linear Lagere per Couplng Blocic
Batch Size 16
Dotch Size 50
Bareh Size T0o
Datch Sioe 300
Dotch Sice
Batci Size
Datch Size
Botch Stac
500
1000
2000
4000
Hiodm: 36
Hildec 20C
Miiiom 300
Eibbeer 1000
Optinsookks Siew
Openvnnian
Optintotion Stepe
Opeinbot
Ostiizatio
Optistoskan Saye
Openncsian
Optinization
Supu
Supp
Saae
Suqu
Slape
Graizai
Figure 7: The higher model capacity compared to models with fewer linear layers per coupling block increases the performance of Forward KI. Path Gradients in terms of ESS, while Forward KL Gradients deteriorate, resulting in a final ESSp 0.0% for the model with the largest capacity trained with the largest batch size. The rows increase width of the linear layers (hidden neurons) and the columns increase the batch size used during optimization.
rithm which uses the insights of Section 4 speed up Algorithm 2. Here, Algorithm I for the fast path gradient uses the recursive equation (49b). The results can be seen in the upper rows of Table
For the implicitly invertible flows, we use the setup of the U(1) experiments. Since the flows are easily invertible, significant percentage of the time spent on the root finding algorithm, which is implemented as the bisection method. The root finding algorithm employs an absolute enror tolerance which determines when the recursive search stops. Each iteration of the bisection method requires evaluation of the function, in case normalizing flow.
Using backpropagation through the bisection is not only error-prone, but also costly in compute and memory. Recently, Köhler et al. (2021) proposed circumventing the backpropagation through the root finding via the implicit function theorem. Both of these methods can be combined with the existing path gradient Algorithm 2. Due to the increase of computational cost and numerical error of the root finding algorithm, these methods are outperformed in runtime and precision by our proposed Algorithm 1. As the tolerance in root-finding, Köhler al. (2021) chose value of 1e
27

---

Published as a conference paper at ICLR 2024
Forward KL Gradient
Forward KL Path Gradient
-1
-2
0
2
1
Figure 8: We visualize the contours of the learned probability distribution ge over the true contour plot of p which is a two dimensional multivariate Gaussian distribution centered at with 0.5 on the diagonal entries of the covariance matrix and 0.25 on the off-diagonal entries. The ReaINVP is trained with 750 samples from with the forward KL divergence. While the forward KL Path Gradient able to recover the true distribution well, the standard gradient exhibits numerous irreg- ularities for the same training samples.
Forward KL Path Gradient Test NLL
Forward KL Path Gradient Train NLL
13
Forward KL Gradient Test NLL
Forward KL Gradient Train NLL
NLL
12
9-
5k
10k
15k
20k
25k
30k
Samples Used in Training
Figure 9: We plot the train and test negative Loglikelihood (NLL) intervals denote the minimum and maximum performance over runs. We see that the path gradient training consistently outperforms standard training in terms of NLL. The Forward KL Path Gradients is less prone to overfitting and able to maintain a steady test NLL. Already with comparitively few data samples, path gradients converge to the test set NLL. In data regimes with little data, forward KI, gradients are prone to overfitting on the training data compared to forward KI. path gradients.
28

---

Published as a conference paper at ICLR 2024
Table 5: We summarize the highest achieved average ESS, over the entire optimization which corresponds to the best attainable performance in terms of ESSp possible for the corresponding model capacity. The flow consists coupling layers, each with linear layers each of which has an increasing width (number of neurons).
Layer Width
10
50
100
250
500
1000
Batch Size
Std
Path
Std
Path
Std
Path
Std
Path
Std
Path
Std
Path
10
57.9
58.5
59.1
59.3
59.1
59.2
59.0
59.1
59.0
59.1
59.0
60.6
58.9
59.1
59.1
59.3
59.1
59.2
59.0
60.8
59.7
75.9
64.8
80.8
100
59.0
59.1
59.1
59.3
59.1
59.2
59.0
65.2
71.1
84.0
75.9
84.8
200
59.1
59.2
59.1
59.3
59.1
59.3
63.6
75.1
76.7
87.8
75.9
87.5
500
59.1
59.2
59.1
59.3
59.1
61.9
72.2
80.1
77.6
89.0
75.6
90.1
1000
59.1
59.2
59.1
59.2
61.3
68.9
75.7
82.3
78.6
91.1
75.6
86.5
2000
59.1
59.2
59.1
59.2
65.8
74.4
74.7
80.7
76.4
91.8
76.4
87.0
4000
59.1
59.2
59.1
60.2
69.4
77.8
73.1
85.8
75.7
91.6
74.6
88.9
Table 6:
Factor of runtime increase in comparison to standard gradient estimator, i.e.,
rantime path gradient/ runtime standard gradient on A100-80GB. (For error sterr and for runtime std is shown.)
Algorithm
Error
Runtime increase (batch size)
X 1e-7
64
1024
8,192
Explicilly
Alg. (ours)
1.6 + 0.1
1.4
+
0.1
1.4 + 0.0
Alg. (Vaitl al., 2022a)
2.1
0.1
2.2
- 0.1
2.1
0.0
Alg. (ours)
1.8 - 0.0
1.9
- 0.1
1.8 + 0.0
Alg. (ours)
2.0 + 0.4
2.2 0.0
2.0 + 0.1
2.3 0.0
Implicilly
Black box root tinding
abs tol
Alg. Autodiff
2c-6
> ~10,0000
19.2
11.8 0.1
Mem
Alg. 2 + Köhler et al. (2021)
2c-6
29.5 + 13.5
4.8 0.0
3.4
- 0.0
3.4 - 0.0
Alg. + Köhler et al. (2021)
le-6
4.1 1.1
17.5 0.2
11.0
+ 0.1
8.2 0.0
for testing the runtime and during training on other benchmarks, they chose tolerance of 2e-6 Results for both of these tolerances are shown in the lower rows of Table 6. Numerical errors are computed 50 batches, each with 64 samples.
We show the behavior of the baseline Algorithm 2 with bisection and gradient computation as pro- posed by Köhler et al. (2021) in Figure 10. We can see that our proposed method outperforms the baseline irrespective of the chosen hyperparameters. For single-precisio floating point, the numer- ical error can only be reduced far, after double-precision has be employed which leads drastic increase in runtime and memory.
F CONTRIBUTIONS
LV
LW
Conceptualization
X
Methodology
X
Formal Analysis
X
Software
X
X
Investigation
Visualization
X
X
Writing Original Draft
Writing Review & Editing
X
29

---

Published as a conference paper at ICLR 2024
Black Box (Köhler et al.)
10
FastPath (ours)
Erro
10
Runtime increase
Figure 10: Runtime vs Precision trade-off Köhler al. (2021) U(1) experiments with lattice 162
30


Reviews:
[{"review_id": "S0GjQTr3qO", "reviewer": "Reviewer_yXNA", "summary": "This work deals with improving the pathwise gradient\nestimator in the context of variational inference\nusing normalizing flow based models (i.e., they\nwant a fast method for computing the \"sticking the landing\"\nestimator by Roeder).\nIn particular, they are looking at deriving pathwise gradients\nfor the log probability term of the normalizing flow.\nComputing this efficiently is non-trivial due to the\nmodification to the probability caused by a change of coordinates\nthat requires computing the determinant of the Jacobian.\n\nThey derive a faster method for computing the pathwise gradient in\nthis setting for coupling flows (the most widely used normalizing\nflow). The improvement in computational speed ranges between 1.3 times\nto 8 times (takes 1.4 - 2.3 times the standard estimator that has a\nhigher variance, so doesn't work as well). The improvement is\nespecially large for implicitly invertible coupling flows, but more\nmodest for explicitly invertible coupling flows.\n\nTheir formulation allows computing the pathwise gradient for\nboth the forward and reverse KL, allowing to also perform\nmaximum likelihood training.\n\nExperiments were performed on a multimodal Gaussian distribution as\nwell as physics settings: U(1) gauge theory and $\\phi^4$ lattice\nmodel.", "questions": "I have a naive question about computing the pathwise gradient of the\nreverse KL. In equation (2), it seems to me that we could rewrite the\nequation by using the Jacobian of the forward transform based on the\ninverse function theorem, so that the $+\\log |\\textup{det} ~ dT^{-1}/dx|$ term\nbecomes $- \\log |\\textup{det}~dT/dx_0|$. Then we could compute the quantity and\nuse backprop to get the pathwise gradient. Am I misunderstanding, or\nwhy would this not work? Is the computation of the Jacobian too\ncostly?\n\n\"Path gradients have the appealing property that they are unbiased and\nhave lower variance compared to standard estimators, thereby promising\naccelerated convergence (Roeder et al., 2017; Agrawal et al., 2020;\nVaitl et al., 2022a;b).\"  -> Other estimators are also unbiased. The\nsentence makes it seem like they aren't. Also, the \"have lower\nvariance\" is not always true. I suggest revising to make the sentence\ncorrect, e.g., making it \"tend to have lower variance\".", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+Fast pathwise gradients are certainly necessary for normalizing flows,\nand the current work provides this with a large improvement over the\nprior work in terms of computational speed.\n\n+The method improves in both walltime and efficiency.\n\n+The method allows both forward and reverse KL training.", "weaknesses": "-The literature review is a bit misleading, as pathwise\ngradients have been around for a long time, e.g., see [L'Ecuyer,\nP. (1991). An overview of derivative estimation] where it is\nreferred to as \"infinitesimal perturbation analysis\". Moreover,\nreparameterization gradients are a type of pathwise gradient, and\nthere are other works discussing it, e.g., [Jankowiak & Obermeyer, 2018]\nor [Parmas & Sugiyama, 2021]. The current work is mainly referring\nto pathwise gradients in the context of normalizing flows and\nvariational modeling, but the broader picture of pathwise gradients\nshould be briefly mentioned, and probably the terminology should\nbe clarified because the current paper refers to \"pathwise\" gradients\nas the narrow application of it to normalizing flows, whereas there\nare many other estimators that have been around for decades that are\nalso referred to as pathwise estimators.\n\n-The experiments are a bit toy, or at least their significance\nwas not explained. \n\nJankowiak, M., & Obermeyer, F. (2018, July). Pathwise derivatives\nbeyond the reparameterization trick. In International conference on\nmachine learning (pp. 2235-2244). PMLR.\n\nParmas, P., & Sugiyama, M. (2021, March). A unified view of likelihood\nratio and reparameterization gradients. In International Conference on\nArtificial Intelligence and Statistics (pp. 4078-4086). PMLR.", "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": "", "extra_scores": {}}, {"review_id": "3pNuEvd8Ey", "reviewer": "Reviewer_vn2u", "summary": "This paper considers the problem of learning a distribution $p$ given an oracle\nfor log probabilities plus a constant (i.e., $\\log p(x) + c$ at sample $x$). It\nproposes a method for estimating the gradients of forward and reverse KL\ndivergence that dispenses with a term known to have zero expectation value, thus\nallowing lower variance estimators of the gradient, with less computational\ncomplexity than prior work. In particular, this method deployed beyond previous\nresults for continuous flows to include coupling flows.", "questions": "No questions.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper technically precise and, to my knowledge, presents valuable original\nwork with immediate applications. The experiments were generally informative.\nIts major contribution is reducing the computational complexity for calculating\npath gradients of both forward and reverse KL when $\\log p(x) + c$ is queriable.\n\nThe theoretical results appear sound after some inspection.\n\nI believe the overall contribution is valuable enough to share with the broader\nICLR community, though I was surprised that the proposed \"fast\" gradient\nestimator was not already established. Perhaps like many key results, it seems\nobvious in hindsight. The suggestion that removal of the $\\frac{\\partial}{\\partial \\theta} \\log q$\nterm from the gradient estimate makes learning empirically robust to overfitting\nis quite interesting and provocative, but unexplored in detail.", "weaknesses": "I had some difficulty reading this work, despite some prior exposure to the\nsubject matter. It took me several passes to make sense of what the key\ncontribution was, and I wished for additional clarity.  The key idea behind\n\"path gradients\" (dropping a term that has zero expectation value) from the\nempirical estimation of the gradient is easy enough to understand, but took some\ntime to distill from the intro [1].\n\nRegarding the experiments, at least one sentence introducing effective sample\nsize would also have been appreciated.\n\n[1] It took me far too long to realize that the expectation value in Equation\n(10) was for $x_0 \\sim q_0$, not $x \\sim q_{\\theta}$. This might have been\nmore clear if different symbols were used for inputs $x_0 \\to x$ and outputs\n$x \\to y$ of the transformation, since layer indexing was only used in the\ncontext of coupling flows.", "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": "", "extra_scores": {}}, {"review_id": "R6DWZzxh4A", "reviewer": "Reviewer_UvVP", "summary": "The authors propose a technique for improving the efficiency of the calculation of path-gradients for both the forwards and reverse KL loss. \nTypically, the path gradient is lower variance but has a significantly higher computational cost, preventing scalability to large problems. Their method avoids having to evaluate the flow in both the forwards and reverse directions by recursively calculating the gradient during the forward pass using JVPs. \nThe speedup is especially significant for flows that require implicit differentiation for inversion. The main contributions are (1) efficient calculation of the path gradient based losses and (2) path gradient version of the forwards KL loss.", "questions": "- How come the flow trained via the standard maximum likelihood objective achieves such poor performance on the MGM problem (Table 1)?. It seems possible that poor hyper-parameters have been used as training by maximum likelihood should be able to obtain reasonable results.\n\n- In the case of forwards KL with flows that require implicit differentiation for inversion, is it not more efficient to set the forwards direction of the flow to map from the target to the flow’s base (rather than base to target), such that implicit differentiation is required for sampling, but not density evaluation)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The method obtains significant improvement in speed in practice, especially for the case of flows that require implicit differentiation for inversion. \n- The method obtains improved generalization for the forward KL training relative to \n- Incoporating the energy function of the target in the forward KL training is novel. And having a loss with the “sticking the landing” property for the forward KL is useful.", "weaknesses": "- The speedup for explicitly invertible flows (which are more common) is relatively minor. \n- The authors emphasise that an advantage of their method relative to those from Vaitl et al. for the estimation of the forward KL is that their method does not require reweighting. However, their method uses samples from the target, while the method from Vaitl et al. uses samples from the flow - hence the two methods are not directly comparable as they are for different situations. I think this is somewhat misleadingly presented in the text (it is presented as an improvement relative to the forward KL objective from Vaitl).", "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": "", "extra_scores": {}}, {"review_id": "iIIUU7nYOB", "reviewer": "Reviewer_ZWTi", "summary": "The paper concerns path gradient estimators for normalizing flows, a reduced variance estimator for the KL divergence gradient in normalizing flows. They come at the cost of additional forward and backward passes through the normalizing flow at hand. The present paper reduces this computational overhead to compute path gradient of reverse KL, while being analytically equal to previous work. It then provides a new path gradient estimator for the forward KL. Experiments demonstrate that the resulting path gradient estimators work both in the forward and reverse KL setting on physical sciences data sets (where the unnormalized $p(x)$ is known).", "questions": "see Weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "*Originality*\n\n- The iterative procedure for computing the path gradient has no memory overhead over non-path gradients and is potentially faster (see Weakness 3).\n- Path gradients are applied to the forward KL with reduced variance by applying the same algorithm to .\n- The approach has the potential to be generically applied to abitrary coupling blocks, if clarified.\n\n\n*Quality*\n\nThe theoretical results might be correct, but I cannot judge at this point (see below). I have some doubts on the baseline experiments (see below).\n\n\n*Clarity*\n\nThe motivation and main chain of reasoning are clear, but several parts of the manuscript lack clarity and detailed explanations (see below).\n\n\n*Significance*\n\nMaking use of path gradients in order to regularize for the known unnormalized density of training data has the potential to greatly reduce compute over classical methods, so this chain of work is relevant to the machine learning + natural sciences community. Allowing the forward KL to make use of the unnormalized density is attractive, as the forward KL may have better properties than reverse KL (mode covering instead of mode seeking).", "weaknesses": "Generally, the presentation interpretation of the results can be greatly improved. I also have concerns on some of the results.\n\nIn detail:\n\n1. The notation of Proposition 3.2 and its proof in the appendix are sloppy and I cannot determine the correctness: what is the inverse of the rectangular matrix $\\frac{\\partial f_\\theta(x_l^t, x_l^c)}{\\partial x_l^t}$? Is it a pseudo-inverse, or is it a part of the network Jacobian? I suggest to greatly rewrite this proposition as a Theorem that outlines the general idea of the recursion (that the path gradient can be constructed iteratively by vector-Jacobian products with the inverse of each block, if I am right). Then proceed to derive concrete realizations for coupling blocks and affine couplings in particular if they allow for unique results.\n2. What is the cost of computing Proposition 3.2? As I mentioned in the first point, by rewriting the recursion more generally, this could easily be showcased.\n3. What is the intuition behind Proposition 4.1? What is the regularization obtained from including the unnormalized density (probably something like the corrected relative weight of each sample according to the ground truth density)?\nWhat derivative vanishes in expectation? How large is the variance of the removed gradient term? Is your result an improvement in this metric? What is the regularizing effect? Vaitl et al. 2022b have useful visualizations and explanations in this regard.\n4. The baseline Algorithm 2 should not be used to measure baseline times. The second forward pass through the network is unneccessary, as one can simply store the result from the first forward pass, once with stop_gradient and once without. Please report Table 2 again with this change.\n5. I have strong doubts on the validity of the experiment on the multimodal gaussian model. It is hard to believe that a standard RealNVP network cannot be trained effectively on this data, with an ESS_p of 0.0(!). I see several warning signs that a bad performing network has been selected in order to have a badly performing baseline:\n\t- the network is huge, with a number of parameters bounded from below by six coupling blocks $\\times$ five hidden subnetworks $\\times$ (1000 $\\times$ 1000 entries in each weigh matrix) amounting to more than 30 million parameters;\n\t- the batch size of 4,000 given 10,000 samples makes the network see almost the entire data set in every update.\n  This indicates that the training is set up in a way that training from samples only must fail. Given that training yields useful models in only five minutes, it is reasonable to expect hyperparameter tuning of the baseline model from the authors.\n6. In this light, how much parameter tuning was involved in the other experiments $\\phi^4$ and $U(1)$? Please compare your numbers to the state of the art results on these benchmarks.\n\n\nGiven that the theoretical results need improved presentation and explanation, and given the doubts on the numerical experiments, the manuscript does not reach the quality ICLR in the current form. Many of the proposed changes can be achieved with additional explanations and better notation. I am looking forward to the author's rebuttal, happy to be corrected on my understanding.\n\n\n\n## Minor comments:\n\n- Eq. (13) is missing a logarithm.\n- The caption for Figure 1 is on page 21 in the appendix, took me some time.\n- The statement that building up the computation graph takes measurable time is false, as this simply means storing already  computed activations in a dictionary (right before section 3.1).\n- Eq. (25) is missing that $p_{\\theta, 0}$ can be computed from the unnormalized density.\n- If a reader is not familiar with the terms forward and reverse KL, it is hard to understand the introduction. Point the reader to Section 2 or drop it here, leaving space for more explanations on theoretical results.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)