Paper: Fast and unified path gradient estimators for normalizing flows (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
The Jacobian matrix  $\frac{\partial f_\theta(x_l^{trans}, x_l^{cond})}{\partial x_l^{trans}}$ is square and invertible and there is thus no subtlety in defining its inverse. In more detail, we define a coupling block in Eq 6 as 

 $$x_{l+1}^{trans} = f_{\theta} (x_{l}^{trans}, x_{l}^{cond})$$

 $$x_{l+1}^{cond} = x_{l}^{cond}$$

where $f_\theta(\bullet \,, x_l^{cond})$ is an invertible function for any choice of $x_l^{cond}$.
By bijectivity, the Jacobian matrix from above is not only square but also invertible.
We have added a remark after Proposition 3.3 to emphasize this and we have rewritten the proof, so that this becomes more clear.

We have revised the manuscript implementing your suggestions. 
Specifically, we have added Proposition 3.2, where we first state the recursion for a general flow, i.e., not necessarily a coupling flow:

$\frac{\partial \log q_{\theta,  l+1}(x_{l+1})}{\partial  x_{l+1}} = \frac{\partial \log q_{\theta, l}(x_l)}{\partial  x_{l}} \left( \frac{\partial T_{l,  \theta_l}(x_l)}{\partial x_l} \right)^{-1} + \frac{\partial \log |\det \frac{\partial T_{l, \theta_l}(x_l)}{\partial x_l}|}{\partial x_l}\left( \frac{\partial T_{l, \theta_l}(x_l)}{\partial x_l} \right)^{-1}$

As we also remark in the modified manuscript, the evaluation of this expression however involves inversion of the Jacobian and is therefore prohibitively expensive for generic flow architectures (see answer immediately below for more details).

Reviewer comment (Theory):
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.

Author response:
First, note that in the revised manuscript Proposition 3.2 has become Proposition 3.3. We will refer to proposition numbers in the revised manuscript.

For general flow architectures the recursion formula is given by the new Proposition 3.2. For generic architectures, where the inverse of the transformation $T_{l, \theta_l}$ is not easy to compute, this recursion however is too expensive. Specifically it involves building and inverting of a $d\times d$ Jacobian, which naively scales cubically in the number of dimensions $d$ (although optimized algorithms with slightly better scaling exist).  Even for autoregressive flows, which have triangular Jacobians, there is still a quadratic scaling. This compares unfavorable with the linear scaling of the standard total gradient estimator. For this reason then we derived a specialized recursion for coupling flows (see Proposition 3.3). This has the same linear complexity as the standard gradient estimators.
Furthermore in Corollary 3.4 we explicate for affine coupling blocks, which has the same linear complexity, but is slightly faster, since it reuses terms that have been used in the standard forward pass and uses less terms than the recursion in Proposition 3.3.

In our experiments, we  used the recursion from Proposition 3.3 for implicitly invertible flows, while Corollary 3.4 was used for explicitly invertible flows. The walltime cost of computation is summarized in Table 2 (for Proposition 3.3 see third row 'Implicitly Invertible' and for Corollary 3.4 see first row, 'Explictly Invertible').
Computing Proposition 3.3 has roughly twice the computational cost of the standard total gradient estimator, but compares favorable to the previous state of the art in the literature, which under the same constraints (numerical error and memory footprint) yields a time factor of at least 8. 
Even if we relax the numerical constraints in favor of the baseline, our fast estimators still lead to significantly faster computations, i.e. at least a 33\% speed up with an order of magnitude higher precision.

Reviewer comment (Theory):
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)? What 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.

Author response:
We substantially extended Appendix B.3 to address your questions in detail. Briefly summarized:

- Intuition: It is well known that the forward KL in target space becomes the reverse KL in base space (see Section 2.3.3 in [1]). This duality allows us to immediately apply all the results derived for the reverse KL case to the forward case, if we express all densities in the base space variables. To the best of our knowledge, we are the first to observe and harness this simple, yet powerful trick. 
- Regularization: We discuss this in detail in the new subsection B.3.2. In essence, the path gradient estimator involves a term containing the derivative of the log of the target density, $\nabla_{x_0} \log p(T_\theta(x_0)) \frac{\partial T_\theta^{-1}(x)}{\partial \theta}$, see Eq. 57. This term is absent in the standard total gradient, see Eq. 56. Note that in the case of Boltzmann generators, this derivative of the log target density is known in closed form by assumption. Thus, the path gradient estimator is able to harness information of the gradient of the energy at the given sample, which can be thought of as a (gradient) regularizer. 
- Term vanishing in expectation: The expected score of the pullback target density $E_{p_{0, \theta}}[\frac{\partial \log p_{0, \theta}(x_0)}{\partial \theta}]$ vanishes in expectation. This term has non-vanishing variance in the standard gradient estimator, even if the pullback $p_{0, \theta}(x_0)$ perfectly approximates the base density $q_0$
(see the revised Appendix B.3).
- Variance of the removed term: The removed variance corresponds to the Fisher Information Matrix (divided by the batch size) of the pullback density $p_{0, \theta}$ (see Appendix B3.1).
- Improvement in metric: We observe exactly the same beneficial sticking-the-landing behavior of our forward estimators as observed in the reverse case, i.e. the variance of the estimator vanishes if the model perfectly approximates the target and this metric is overall reduced particularly at the end of training (see Figure 3 in the revised manuscript).
- Useful visualization: We added the exact same plot as in Vaitl et al. 2022b (see Figure 3 and analogous explanations in the revised appendix B3).

[1] Papamakarios, George, et al. "Normalizing flows for probabilistic modeling and inference." The Journal of Machine Learning Research 22.1 (2021): 2617-2680. https://arxiv.org/abs/1912.02762

Reviewer comment (Experiments):
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.

Author response:
For our $\phi^4$ experiments, we use the exact same hyperparameter choices as in the recent publication [1] as well as the same codebase. 
To the best of our knowledge, there are currently only two publications considering maximum likelihood training for $\phi^4$ theory: [1] and [2].
We chose [1] as it is more recent. 

For the target density, we choose the same bare quartic coupling $\lambda = 0.022$ as in [1] and hopping parameter $\kappa = 0.25$. 
The authors of [1] considered various hopping parameters and we restrict to this particular choice as it is the most physically relevant one: $\phi^4$ theory has a second order phase transition at this critical point, see Figure 1 in [2]. 
It is well-known that second order phase transitions lead to divergent correlation length and thus to long-range correlation which are challenging to learn. 
Unfortunately, [1] and [2] do not report any effective sampling sizes to compare to. However, we confirmed in private communication with the authors of [1] that our effective sampling sizes of 85.1% for the baseline are compatible with their results. We also experimentally calculated the wbar estimator which is 1.0001+-8e-5 and thus compatible with Figure 21 in [1] for the considered target parameters. 

For U(1), we base our experiments on [4] using their code base (which was released in this separate publication [5]). 
To the best of our knowledge, this is the only publicly available code for a normalizing flow with U(1) gauge symmetry. 
The chosen bare coupling parameter is $\beta=3$ as in this publication (in addition, the authors also fine-tuned from this target value to further bare coupling values). 
We followed the same hyperparameter choices as in the publication. 
Unfortunately, important hyperparameter choices such as details on the architecture and training setup were not reported in the publication and we therefore had to rely on the PhD thesis [5] of one of the authors to extract more details. 
We choose exactly the same architecture and hyperparameteres as listed in the thesis except for the batch size and learning rate. 
This is because the thesis is not clear on the precise batch size stating that “batches of size ranging from 16,384 to 131,072” were used, see Section 4.6.3 on page 118. 
We used a batch size of 12,288 since it was the biggest one that we could reach without gradient accumulation. 
We also chose a learning rate of 1e-4 as opposed to 1e-3 reported in the thesis since training proved unstable with a higher learning-rate (most likely due to the difference in batch size). 
Furthermore, [4] does not seem report any effective sampling sizes to compare to but our results are consistent with the integrated autocorrelation time shown in Figure 4 of [4]. 
Note that in neural MCMC, the integrated autocorrelation time is closely related to the effective sampling size, see [https://arxiv.org/abs/2111.10189](https://arxiv.org/abs/2111.10189) for a rigorous discussion.

[1] Detecting and Mitigating Mode-Collapse for Flow-based Sampling of Lattice Field Theories, Nicoli et al. 2023, [https://arxiv.org/abs/2302.14082](https://arxiv.org/abs/2302.14082)

[2] Flow-based sampling for multimodal distributions in lattice field theory, Hacket et al, 2021, [https://arxiv.org/abs/2107.00734](https://arxiv.org/abs/2107.00734)

[3] Estimation of Thermodynamic Observables in Lattice Field Theories with Deep Generative Models, Nicoli et al, [https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.126.032001](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.126.032001)

[4] Equivariant flow-based sampling for lattice gauge theory, Kanwar et al,  Physics Review Letters 2020, [https://arxiv.org/abs/2207.08945](https://arxiv.org/abs/2003.06413)

[5] Introduction to Normalizing Flows for Lattice Field Theory, Albergo et al, [https://arxiv.org/abs/2101.08176](https://arxiv.org/abs/2101.08176)

[6] Machine Learning and Variational Algorithms for Lattice Field Theory, Gurtej Kanwar [https://arxiv.org/abs/2106.01975](https://arxiv.org/abs/2106.01975)

[7]Analysis of autocorrelation times in Neural Markov Chain Monte Carlo simulation, Bialas et al, 2023, Physical Reviews E, [https://arxiv.org/abs/2111.10189](https://arxiv.org/abs/2111.10189)

Reviewer comment (Theory):
Eq. (13) is missing a logarithm.

Author response:
We fixed the typo.

Reviewer comment (Presentation):
The caption for Figure 1 is on page 21 in the appendix, took me some time.

Author response:
We rephrased the caption.

Reviewer comment (Writing):
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.

Author response:
We have pointed the reader to Section 2 in the introduction as suggested.

Reviewer comment (Evaluation):
The speedup for explicitly invertible flows (which are more common) is relatively minor.

Author response:
Our estimators for explicitly invertible flows have about 60 percent the runtime of the previous state-of-the-art. Thus the speed-up is significant. We however agree that proposing path gradient estimators for implicitly invertible normalizing flows is probably the more important contribution of our manuscript since many state-of-the-art architectures in the Boltzmann generator context are of this form. Particular examples are Smooth Flows [1] for Quantum Chemistry as well as the gauge invariant coupling flows [2-4]. To the best of our knowledge, we are the first to propose fast path gradient estimators for this model class of high practical relevance. A further important contribution of our manuscript is that we provide the first path-gradient estimators for sample-based training (see discussion below).

[1] Smooth Normalizing Flows, Jonas Koehler et al, NeurIPS 2021, [https://arxiv.org/abs/2110.0035](https://arxiv.org/abs/2110.00351)

[2] Equivariant Flow-Based Sampling for Lattice Gauge Theory, Kanwar et al, Physics Review Letters 2020, [https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.125.121601](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.125.121601)

[3] Sampling using SU(N) gauge equivariant flows, Boyda et al, Physical Review D, [https://journals.aps.org/prd/abstract/10.1103/PhysRevD.103.074504](https://journals.aps.org/prd/abstract/10.1103/PhysRevD.103.074504)

[4] Flow-based sampling in the lattice Schwinger model at criticality, Albergo et al, 2022, Physical Review D [https://journals.aps.org/prd/abstract/10.1103/PhysRevD.106.014514](https://journals.aps.org/prd/abstract/10.1103/PhysRevD.106.014514)

Reviewer comment (Evaluation):
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).

Author response:
We agree that this is an important difference between our method and the one proposed by Vaitl et al. We have revised the manuscript to make this clearer.
We stress however that their reweighting method comes with an important downside: it fails as the system size grows because the probability mass of the target density becomes increasingly concentrated, see for example Figure 4 right in Vaitl et al. As a result, the importance weights suffers from large variance. 
Futhermore, their reweighting-based approach does not allow to incorporate samples from MD or MCMC into the training. This is unfortunate as Boltzmann generators are often trained using such samples. We discussed extensively with the authors of Vaitl et al. and they agree with this assesment.
Even if the bulk of the training is energy-based, samples are essential as a pretraining setp because the self-sampling nature of the energy-based training will fail to find any modes when the flow is randomly initialized and the target distribution is high-dimensional. 
Only our approach allows for path-gradients of sample-based maximum likelihood training. It is therefore a substantial improvement upon the method of Vaitl et al. and of great interest to the Boltzmann generator community.

Reviewer comment (Experiments):
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.

Author response:
Note that our experiments merely establish that path gradients facilitate *more sample efficient training* and help avoid overfitting. 
This can be seen from the rhs of Fig 1, namely that standard maximum likelihood training can easily fit the MGM if enough samples are provided. For smaller number of samples, maximum likelihood training leads to overfitting. A non-vanishing ESS_p could be achieved by early stopping. Early stopping is however challenging in the context of Boltzmann generators, as their training often relies on a limited number of biased samples which are only used for pre-training. In the revised manuscript, we nevertheless report results with early stopping to adopt the most charitable setting for the baseline. We stress that sample efficiency is a crucial requirement in learning unnormalized distributions for which training set generation involves costly MD or MCMC simulations. 

To avoid any impression of hyperparameter tuning and to provide a more detailed analysis, we include an extensive analysis of results for various batch and flow sizes in the revised Appendix E. In more detail, Figures 5, 6 and 7 show the corresponding Forward ESS during the course of both path and non-path forward training for varying numbers of 
- linear layers per coupling block,
- number of hidden neurons per linear layer,
- batch sizes.
  
While standard Forward KL gradients can *relatively* surpass Forward KL Path Gradients for smallest models trained with the two smallest batch sizes by a small margin, a larger model trained with path gradients doubles the mean of the $ESS_p$ in a direct comparison and improves upon the best forward KL gradients by almost 10 percentage points in absolute terms.

Tables 3, 4 and 5 summarize the best possible Forward ESS over the course of training (corresponding to the results obtained by early stopping) for direct comparison.

As the additional experiments demonstrate, the superior performance of our method does not depend on the particular hyperparameter choices.

Reviewer comment (Theory):
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)?

Author response:
You are correct in that one can choose the “directionality” of the flow such that density estimation is fast. 
In such as situation, implicit differentiation is not necessary. 
However, such a choice is strongly disfavored in the context of Boltzmann generators: one wants to use these flows to facilitate fast sampling from a given target distribution. 
As such, the sampling direction has to be sufficiently fast. Furthermore, one often combines sample-based pretraining with energy-based training which uses the base-to-target direction of the flow. As a result, sample-based forward training has to rely on implicit differentation. Therefore, our implicit path gradients are of great significance to improve training of Boltzmann generators.

Reviewer comment (Experiments):
The experiments are a bit toy, or at least their significance was not explained.

Author response:
We politely disagree with this statement. Lattice field theory provides the mathematical framework underlying many parts of modern theoretical physics, in particular, high-energy physics, gravitational physics, condensed matter and statistical physics. Indeed all known fundamental forces of nature can be described by quantum field theory (gravity only in an effective field theory sense). Applications of normalizing flows to lattice field theory is an emerging and highly promising field, see, for example, this recent review ([https://www.nature.com/articles/s42254-023-00616-w](https://www.nature.com/articles/s42254-023-00616-w)). In fact, a significant part of the global supercomputing resources are spent on lattice field theory simulations.

From a machine-learning point of view, this problem setting is extremely challenging because of the extremely large symmetry of the target distribution. Specifically, the energy (or, more precisely, action) of many lattice field theories is invariant under local symmetry groups, i.e., the field at each lattice site can be rotated with a separate symmetry operation. As such the size of the symmetry group scales with the number of lattice sites. Without making the model architecture manifestly invariant under this symmetry, the generative model would not be able to learn. Furthermore, lattice field theories have to be extrapolated to the continuum limit. In other words, they have to be considered at a second order phase transition for which it is well known that the correlation length of the system diverges. Therefore, the generative model has to learn long-ranged interactions.

We also emphasize that our experiments considered state-of-the-art architectures for both $\phi^4$ and U(1) gauge theory. Thus, our manuscript demonstrates improvements of models that are at the forefront of this exciting and highly dynamic field of research.

Reviewer comment (Theory):
I have a naive question about computing the pathwise gradient of the reverse KL. In equation (2), it seems to me that we could rewrite the equation by using the Jacobian of the forward transform based on the inverse function theorem, so that the $+\log |\textup{det} ~ dT^{-1}/dx|$ term becomes $- \log |\textup{det}~dT/dx_0|$. Then we could compute the quantity and use backprop to get the pathwise gradient. Am I misunderstanding, or why would this not work? Is the computation of the Jacobian too costly?

Author response:
The stated identity is, of course, correct but unfortunately unhelpful as we are interested in the derivative with respect to $x,$ i.e.,   $\frac{\partial \log \det |\frac{d T_{\theta}^{-1}(x) }{ d x }|}{\partial x}$. The term $\log |\det dT_{\theta}/d x_{0}|$ manifestly only depends $x_0$ and only implictly on $x$ through $x_0 = T^{-1}_{\theta}(x)$. Of course, the gradient can straightforwardly be calculated using the chain rule

$\frac{\partial \log \det |\frac{d T_{\theta}^{-1}(x)}{d x}|}{\partial x} = - \frac{\partial \log \det |\frac{d T_{\theta}(x_0)}{d x_0}|}{\partial x_0} \frac{\partial T_{\theta}^{-1}(x)}{\partial x}.$

However this involves a Jacobian in the target-to-base direction which we wanted to avoid in the first place.

Reviewer comment (Writing):
"Path gradients have the appealing property that they are unbiased and have lower variance compared to standard estimators, thereby promising accelerated convergence (Roeder et al., 2017; Agrawal et al., 2020; Vaitl et al., 2022a;b)." → Other estimators are also unbiased, but the sentence makes it seem like they aren't. Also, the "have lower variance" is not always true. I suggest revising to make the sentence correct, e.g., making it "tend to have lower variance".

Author response:
We fully agree and have rephrased the relevant sentences.