Fast and unified path gradient estimators for normalizing flows

Strengths

Originality

• The iterative procedure for computing the path gradient has no memory overhead over non-path gradients and is potentially faster (see Weakness 3).
• Path gradients are applied to the forward KL with reduced variance by applying the same algorithm to .
• The approach has the potential to be generically applied to abitrary coupling blocks, if clarified.

Quality

The theoretical results might be correct, but I cannot judge at this point (see below). I have some doubts on the baseline experiments (see below).

Clarity

The motivation and main chain of reasoning are clear, but several parts of the manuscript lack clarity and detailed explanations (see below).

Significance

Making 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).

We thank the reviewer for the positive feedback. We agree that our method is generically applicable to coupling-based flows and is of relevance to natural sciences applications.

Weaknesses

1. 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
$\partial f_{\theta}(x_{l}^{trans}, x_{l}^{cond})/\partial x_{l}^{trans}$  ? 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.

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).

Strengths

• The method obtains significant improvement in speed in practice, especially for the case of flows that require implicit differentiation for inversion.
• The method obtains improved generalization for the forward KL training relative to
• 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.

We thank the reviewer for recognizing these strengths. We fully agree that regularization of maximum likelihood training with the ground-truth energy function is an important contribution of our manuscript. We also appreciate that the reviewer rightly pointed out the significance of path gradients for implicitly invertible normalizing flows in the context of Boltzmann generators.

Weaknesses

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

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

[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

[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

[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

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).

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.

Strengths

The paper technically precise and, to my knowledge, presents valuable original work with immediate applications. The experiments were generally informative. Its major contribution is reducing the computational complexity for calculating path gradients of both forward and reverse KL when log⁡p(x)+c is queriable.

The theoretical results appear sound after some inspection.

I believe the overall contribution is valuable enough to share with the broader ICLR community, though I was surprised that the proposed "fast" gradient estimator was not already established. Perhaps like many key results, it seems obvious in hindsight. The suggestion that removal of the $\frac{\partial}{\partial \theta} log(⁡ q )$ term from the gradient estimate makes learning empirically robust to overfitting is quite interesting and provocative, but unexplored in detail.

We thank the reviewer for the positive feedback. We agree that path gradients lead to more robust training for normalizing flows in the context of tractable ground-truth energy function $\log ⁡p(x)+c$.

Weaknesses

I had some difficulty reading this work, despite some prior exposure to the subject matter. It took me several passes to make sense of what the key contribution was, and I wished for additional clarity. The key idea behind "path gradients" (dropping a term that has zero expectation value) from the empirical estimation of the gradient is easy enough to understand, but took some time to distill from the intro [1].

We have rephrased the introduction to make this point and our key contributions clearer.

Regarding the experiments, at least one sentence introducing effective sample size would also have been appreciated.

We have expanded the discussion of the Effective Sampling Size in appendix E.

It took me far too long to realize that the expectation value in Equation (10) was for $x_0 \sim q_0$, not $x \sim q_\theta$ This might have been more clear if different symbols were used for inputs

$x_0 \to x$ and outputs $x \to y$ of the transformation, since layer indexing was only used in the context of coupling flows.

We will update the draft following your suggestions for the camera-ready version. We however prefer to use $z=x_0$ as it is more standard notation.

Strengths

• Fast pathwise gradients are certainly necessary for normalizing flows, and the current work provides this with a large improvement over the prior work in terms of computational speed.
• The method improves in both walltime and efficiency.
• The method allows both forward and reverse KL training.

We thank the reviewer for pointing out these strenghts. We agree that that path gradients are an important tool in successfully training normalizing flows.

Weaknesses

The literature review is a bit misleading, as pathwise gradients have been around for a long time, e.g., see [L'Ecuyer, P. (1991)]. An overview of derivative estimation] where it is referred to as "infinitesimal perturbation analysis". Moreover, reparameterization gradients are a type of pathwise gradient, and there are other works discussing it, e.g., [Jankowiak & Obermeyer, 2018] or [Parmas & Sugiyama, 2021]. The current work is mainly referring to pathwise gradients in the context of normalizing flows and variational modeling, but the broader picture of pathwise gradients should be briefly mentioned, and probably the terminology should be clarified because the current paper refers to "pathwise" gradients as the narrow application of it to normalizing flows, whereas there are many other estimators that have been around for decades that are also referred to as pathwise estimators.

We thank the reviewer for pointing this out. We indeed focused narrowly on normalizing flow applications in the literature review. In light of the reviewers comments, we have extended the discussion with the suggested references providing a broader perspective on the field.

The experiments are a bit toy, or at least their significance was not explained.

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). 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.