Verlet Flows: Exact-Likelihood Integrators for Flow-Based Generative Models

Authors: ezra  erives
bowen  jing
tommi  jaakkola
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 02:28:53

Abstract: Approximations in computing model likelihoods with continuous normalizing flows (CNFs) hinder the use of these models for importance sampling of Boltzmann distributions, where exact likelihoods are required. In this work, we present \textit{Verlet flows}, a class of CNFs on an augmented state-space inspired by symplectic integrators from Hamiltonian dynamics. When used with carefully constructed \emph{Taylor-Verlet integrators}, Verlet flows provide exact-likelihood generative models which generalize coupled flow architectures from a non-continuous setting while imposing minimal expressivity constraints. On experiments over toy densities, we demonstrate that the variance of the commonly used Hutchinson trace estimator is unsuitable for importance sampling, whereas Verlet flows perform comparably to full autograd trace computations while being significantly faster.

Main paper: 
 introduction: 
 Flow-based generative models-also called normalizing flows-parameterize maps from prior to data distributions via invertible transformations. An exciting application of normalizing flows is in learning the Boltzmann distributions of physical systems (Noé et al., 2019; Midgley et al., 2023; Kim et al., 2024) . At inference time, these Boltzmann generators provide model likelihoods which can be used to reweigh samples towards the target energy with importance sampling. While nearly all existing Boltzmann generators are built from composing invertible layers such as coupling layers or splines, experiments on image domains suggest that continuous normalizing flows (CNFs)which can parameterize arbitrary vector fields mapping noise to data-are far more expressive than their discrete counterparts (Chen et al., 2018; Grathwohl et al., 2018) . Unfortunately, the exact model likelihood of CNFs can only be accessed through expensive trace computations and numerical integration, preventing their adoption in Boltzmann generators.In this work, we propose Verlet flows, a flexible class of CNFs on an augmented state-space inspired by symplectic integrators from Hamiltonian dynamics. Instead of parameterizing the flow γ with a single neural network, Verlet flows instead parameterize the coefficients of the multivariate Taylor expansions of γ in both the state-space and the augmenting space. We then introduce Taylor-Verlet integrators, which exploit the splitting approximation from which many symplectic integrators are derived to approximate the intractable time evolution of γ as the composition of the tractable time evolutions of the Taylor expansion terms. At training time, Verlet flows are a subclass of CNFs, and can be trained accordingly. At inference time, Taylor-Verlet integration enables theoretically-sound importance sampling with exact likelihoods. 
background: 
 Discrete Normalizing Flows Given a source distribution π 0 and target distribution π 1 , we wish to learn an invertible, bijective transformation f θ which maps π 0 to π 1 . Discrete normalizing flows parameterize f θ as the compositionf θ = f N θ • • • • • f i θ , from which log π 1 (f θ (x)) can be computed using the change of variables formula and the log-determinants of the Jacobians of the individual transformations f i θ . Thus, significant effort has been dedicated to developing expressive, invertible building blocks f i θ whose Jacobians have tractable log-determinant. Successful approaches include coupling-based flows, in which the dimensions of the state variable x are partitioned in two, and the each half is used in turn to update the other half (Dinh et al., 2016; 2014; Müller et al., 2019; Durkan et al., 2019) , and autoregressive flows (Kingma et al., 2017; Papamakarios et al., 2018) . Despite these efforts, discrete normalizing flows have been shown to suffer from a lack of expressivity in practice.Continuous Normalizing Flows Continuous normalizing flows (CNFs) dispense with the discrete layers of normalizing flows and instead learn a time-dependent vector field γ(x, t; θ), parameterized by a neural network, which maps the source π 0 to a target distribution π 1 (Chen et al., 2018; Grathwohl et al., 2018) . Model densities can be accessed by the continuous-time change of variables formula given byEQUATIONwhere x t = x 0 + t 0 γ(x t , t; θ) dt, Tr denotes trace, and J γ (x t , t; θ) = ∂γ(x,t;θ) ∂x | xt,t denotes the Jacobian. Compared to discrete normalizing flows, CNFs are not constrained by invertibility or the need for a tractable Jacobian, and therefore enjoy significantly greater expressivity.While the trace Tr J γ (x t , t; θ) appearing in the integrand of Equation 1 can be evaluated exactly with automatic differentiation, this grows prohibitively expensive as the dimensionality of the data grows large, as a linear number of backward-passes are required. In practice, the Hutchinson trace estimator (Grathwohl et al., 2018) is used to provide a linear-time, unbiased estimator of the trace. While cheaper, the variance of the Hutchinson estimator makes it unsuitable for importance sampling.Symplectic Integrators and the Splitting Approximation Leap-frog integration is a numeric method for integrating Newton's equations of motion which involves alternatively updating q (position) and p (velocity) in an invertible manner not unlike augmented, coupled normalizing flows.1 Leap-frog integration is a special case of the more general family of symplectic integrators, designed for the Hamiltonian flow γ H (of which the equations of motion are a special case). Oftentimes the Hamiltonian flow decomposes as γ H = γ q + γ p , enabling the splitting approximationEQUATIONwhere φ(γ, τ ) denotes the time evolution operator along the flow γ for a duration τ , and where the terms on the right-hand side of Equation 2 are possibly tractable in a way that the left-hand side is not. For example, the leap-frog integrator corresponds to analytic, invertible, and volumepreserving φ(γ {q,p} , t), whereas the original evolution may satisfy none of these properties. While Verlet flows, to be introduced in the next section, are not in general Hamiltonian, they similarly exploit the splitting approximation. A more detailed exposition of symplectic integrators and the splitting approximation can be found in Appendix A. 
verlet flows: 
 We consider the problem of mapping a source distribution π0 (q) on R dq at time t = 0 to a target distribution π1 (q) on (R dq ) at time t = 1 by means of a time-dependent flow γ(x, t). We will now augment this problem on the configuration-space R dq by extending the distribution π0 (q) to π 0 (q, p) = π 0 (p|q)π 0 (q) and π1 (q) to π 1 (q, p) = π 1 (p|q)π 1 (q) where both π i (p|q) are given by N (p; 0, I dp ). In analogy with Hamiltonian dynamics, we will refer to the space M = R dq+dp as phase space.2 Observe that any analytic flow γ is given (at least locally) by a multivariate Taylor expansion of the formEQUATIONfor appropriate choices of functions s q i and s p i , which we have identified in the last equality as (i, 1)tensors: multilinear maps which take in i copies of q ∈ T q R n and return a tangent vector. While s {q,p} 0 and s {q,p} 1 can be thought of as vectors and matrices respectively, higher order terms do not admit particularly intuitive interpretations. Whereas traditional CNFs commonly parameterize γ θ directly via a neural network, Verlet flows instead parameterize the coefficients s {q,p};θ k with neural networks, allowing for Verlet integration via the splitting approximation. By parameterizing all the terms in the Taylor expansion, Verlet flows are in theory as expressive as CNFs parameterized as γ(q, p, t; θ). However, in practice,we must truncate the series after some finite number of terms, yielding the order N Verlet flowEQUATIONIn the next section, we examine how to obtain exact likelihoods from these truncated Verlet flows.Denote by γ q k the flow given byγ q k (x, t; θ) = s q k (p, t; θ)(q ⊗k ) 0 ∈ T x M,and define γ p k similarly. 3 For any such flow γ ′ on M , denote by φ ‡ (γ ′ , τ ) the time evolution operator, transporting a point x ∈ M along the flow γ ′ for time τ . We denote by just φ the pseudo time evolution operator given by φ(γ ′ , τ ) : 4 Note that t is kept constant throughout integration, an intentional choice which we shall see allows for a tractable closed form. Although our Verlet flows are not Hamiltonian, the splitting approximation from Equation 11 can be applied to Verlet flows to decompose the desired time evolution into simpler, analytic terms, yieldingx t → x t + t+τ t γ ′ (x s , t) ds.φ ‡ (γ, τ ) ≈ φ(γ t , τ ) • φ(γ p N , τ ) • φ(γ q N , τ ) • φ(γ p N -1 , τ ) • φ(γ q N -1 , τ ) • • • φ(γ p 0 , τ ) • φ(γ q 0 , τ ). (5)Note here that the leftmost term of the right hand side is the time-update term φ(γ t , τ ). The key idea is that Equation 5 approximates the generally intractable φ ‡ (γ, τ ) as a composition of simpler, tractable updates allowing for a closed-form, exact-likelihood integrator for Verlet flows.The splitting approximation from Equation 5, together with closed-form expressions for the time evolution operators and their log density updates (see Figure 1 ), yields an integration scheme specifically tailored for Verlet flows, and which we shall refer to as a Taylor-Verlet integrator. Explicit integrators for first order and higher order Verlet flows are presented in Appendix D. One important element of the design space of Taylor-Verlet integration is the order of the terms within the splitting approximation of Equation 5, and consequently, the order of updates performed during Verlet integration. We will refer to Taylor-Verlet integrators which follow the order of Equation 5 as standard Taylor-Verlet integrators, and others as non-standard. While the remainder of this work focuses on standard Taylor-Verlet integrators, the space of non-standard Taylor-Verlet integrators is rich and requires further exploration. Certain coupling-based normalizing flow architectures, such as RealNVP (Dinh et al., 2016) can be realized as the update steps of non-standard Taylor-Verlet integrators, as is discussed in Appendix E.For each pseudo time evolution operator φ(γ k {q,p} , τ ), we compute its closed-form and the logdeterminant of its Jacobian. Together, these allow us to implement the integrator given by Equation 5. Results are summarized in the Table 1 for γ q k only, but analogous results hold for for γ p k as well. Note that for terms of order k ≥ 2, and for the sake of tractability, we restrict our attention to sparse tensors, denoted s k {q,p} , for which only "on-diagonal" terms are non-zero so that s k {q,p} (q ⊗k ) collapses to a simple dot product. We similarly use γ {q,p} k to denote the corresponding flows for sparse, higher order terms. Full details and derivations can be found in Appendix C. Table 1 : A summary of closed-forms for the time evolution operators φ(γ q k ; τ ), and their corresponding log density updates. Analogous results hold for for φ(γ p k ; τ ) as well.Flow γ Operator φ(γ, τ )Density Update log det |Jφ(γ, τ )| γ q 0 q p → q + τ s q 0 (p, t) p 0 γ q 1 q p → exp(τ s q 1 (p, t))q p Tr(τ s q 1 (p, t)) γ q k , k > 1 q p → (q •(1-k) + τ (s q k )i(1 -k)) •( 1 1-k ) p i k 1-k log q 1-k i + τ (1 -k)(s q k )i -k log |qi| 
experiments: 
 Across all experiments in this section, and unless stated otherwise, we train an order-one Verlet flow γ θ , with coefficients s {q,p};θ 0,1 parameterized as a three-layer architecture with 64 hidden units each, as a continuous normalizing flow using likelihood-based loss. Non-Verlet integration is performed numerically using a fourth-order Runge-Kutta solver for 100 steps. EQUATIONWe train an order-one Verlet flow γ θ targeting a trimodal Gaussian mixture in two-dimensional qspace, and an isotropic Gaussian N (p 1 ; 0, I 2 ) in a two-dimensional p-space. We then perform and time importance sampling using Equation 6to estimate the natural logarithm log Z in two ways: first numerically integrating γ θ with a fourth-order Runge-Kutta solver and using automatic differentiation to exactly compute the trace, and secondly using Taylor-Verlet integration. We find that integrating γ θ using a Taylor-Verlet integrator performs comparably to integrating numerically while being significantly faster. Results are summarized in Figure 1 .The poor performance of the Hutchinson trace estimator can be seen in Figure 2 , where we plot a histogram of the logarithm log π(x) π θ (x) of the importance weights for x ∼ π θ (x). The presence of just a few positive outliers (to be expected given the variance of the trace estimator) skews the resulting estimate of Z to be on the order of 10 20 or larger.  
conclusion: 
 In this work, we have presented Verlet flows, a class of CNFs in an augmented state space whose flow γ θ is parameterized via the coefficients of a multivariate Taylor expansion. The splitting approximation used by many symplectic integrators is adapted to construct exact-likelihood Taylor-Verlet integrators, which enable comparable but faster performance to numeric integration using expensive, autograd-based trace computation on tasks such as importance sampling.  
 
 Appendix: Given a mechanical system with configuration space R d , we may define the phase space of the system to be the cotangent bundle M = T * R d ≃ R 2d . Intuitively, phase space captures the intuitive notion that understanding the state of M at a point in time requires knowledge of both the position q ∈ R d and the velocity, or momentum (assuming unit mass), p ∈ T * R d .Hamiltonian mechanics is a formulation of classical mechanics in which the equations of motion are given by differential equations describing the flow along level curves of an energy function, or Hamiltonian, H(q, p). Denote by X (M ) the space of smooth vector fields on M . Then at the point (q, p) ∈ M , the Hamiltonian flow γ H ∈ X (M ) is defined to be the unique vector field which satisfiesγ T H Ωγ ′ = ∇H • γ ′ (7)for all γ ′ ∈ X (M ), and whereEQUATION)In other words, our state (q, p) evolves according to dq dt = ∂H ∂p and dp dt = -∂H ∂q .The time evolution φ ‡ (γ H , τ ) of γ H satisfies two important properties: it conserves the Hamiltonian H, and it conserves the symplectic form Ω.Proposition A.1. The flow γ H conserves the Hamiltonian H.Proof. This amounts to showing that d dτ φ ‡ (γ H , τ )| τ =0 = 0, which follows immediately from ∇H • γ H = 0.Proposition A.2. The flow γ H preserves the symplectic form Ω.Realizing Ω as the (equivalent) two-form i dq i ∧dp i , the desired result amounts to showing that the Lie derivative L γ H Ω = 0. With Cartan's formula, we find thatL γ H Ω = d(ι γ H Ω) + ι γ H dΩ = d(ι γ H Ω)where d denotes the exterior derivative, and ι denotes the interior product. Here, we have used that dΩ = i d(dq i ∧ dp i ) = 0. Then we compute thatd(ι γ H Ω) = d(ι γ H i dq i ∧ dp i ) = d i ∂H ∂p i dp i + ∂H ∂q i dq i = d(dH).Sinced 2 = 0, L γ H = d(dH) = 0, as desired.Flows which preserve the symplectic form Ω are known as symplectomorphisms. Proposition A.2 implies that the time evolution of γ H is a symplectomorphism.We have seen that the time-evolution of γ H is a symplectomorphism, and therefore preserves the symplectic structure on the phase space M . In constructing numeric integrators for γ H , it is therefore desirable that our integrators are, if possible, themselves symplectomorphisms. In many cases, the Hamiltonian H decomposes as the sum H(q, p) = T (q) + V (p). Then, at the point z = (q, p) ∈ M , we find thatγ T = ∂T ∂p -∂T ∂q = 0 -∂T ∂q ∈ T z (R 2 ) and γ V = ∂V ∂p -∂V ∂q = ∂V ∂p 0 ∈ T z (R 2 ).Thus, the flow decomposes as well toγ H = ∂H ∂p -∂H ∂q = ∂V ∂p -∂T ∂q = 0 -∂T ∂q + ∂H ∂p 0 = γ T + γ V .Observe now that the respective time evolution operators are tractable and are given by φ ‡ (γ T , τ ) :q p → q + τ ∂T ∂p p and φ ‡ (γ V , τ ) : q p → q p -τ ∂T ∂q .Since γ T and γ V are Hamiltonian flows their time evolutions φ ‡ (γ T , τ ) and φ ‡ (γ T , τ ) are both symplectomorphisms. As symplectomorphisms are closed under composition, it follows that that φ ‡ (γ T , τ ) • φ ‡ (γ V , τ ) is itself a symplectomorphism. We have thus arrived at the splitting approximationEQUATIONEquation 9 allows us to approximate the generally intractable, symplectic time evolution φ ‡ (γ H , τ ) as the symplectic composition of two simpler, tractable time evolution operators. The integration scheme given by Equation 9 is generally known as the symplectic Euler method.So-called splitting methods make use of more general versions of the splitting approximation to derive higher order, symplectic integrators. Using the same decomposition H(q, p) = T (q) + V (p), and instead of considering the two-term approximation given by Equation 9, we may choose coefficients {c i } N i=0 and {d i } N i=0 with c i = d i = 1 and consider the more general splitting approximationEQUATIONA more detailed exposition of higher order symplectic integrators can be found in (Yoshida, 1993) .B JUSTIFICATION FOR TREATING φ(γ, τ )'S AS TIME EVOLUTION OPERATORSIn the following discussion, we will use x t = (q t , p t ) for brevity. The splitting approximation from Equation 5, which we recall below asEQUATIONrequires some clarification. Recall that while the true time evolution operator φ ‡ (γ, τ ) is given by φ ‡ (γ, τ ) :EQUATIONthe pseudo time operator φ(γ, τ ) is given by φ(γ, τ ) :EQUATIONwhere t is kept-constant throughout the integration.To make sense of the connection between φ ‡ and φ, we will augment our phase-time space S = R dp+dq × R ≥0 (within which our points (x t , t) live), with a new s-dimension, to obtain the space S ′ = S × R ≥0 . Treating x t and t as the state variables x s and t s which evolve with s, the flow γ q k (as a representative example) on R dp+dq can be extended to a flow γ q k on S given by γ q k (x s , t s ) = ∂xs ∂s ∂ts ∂sEQUATION)where the zero t s -component encodes the fact that the pseudo-time evolution φ(γ q k , τ ) from Equation 13 does not change t. The big idea is then that this pseudo time evolution φ(γ q k , τ ) can be viewed as the projection of the (non-pseudo) s-evolution φ ‡ ( γ q k , τ ), given by φ ‡ ( γ q k , τ ) :EQUATIONonto S. The equivalency follows from the fact that for γ q k , t s+τ ′ = t s for τ ′ ∈ [0, τ ]. A similar statement can be made about the t-update γ t from Equation 11.Denoting by Proj : S ′ → S the projection onto S, we see that the splitting approximating using pseudo-time operators from Equation 11 can be rewritten as the projection onto S of an analogous splitting approximation using non-pseudo s-evolution operators, viz.,EQUATIONC DERIVATION OF TIME EVOLUTION OPERATORS AND THEIR JACOBIANS Order Zero Terms. For order k = 0, recall that γ q 0 (x) = s q 0 (p, t)(q ⊗0 ) 0 = s q 0 (p, t) 0 , so that the operator φ(γ 0 q , τ ) is given by φ(γ q 0 , τ ) :EQUATIONwith Jacobian J q 0 given byEQUATIONThe analysis for s p 0 is nearly identical, and we omit it.Order One Terms. For k = 1, we recall thatEQUATIONThen the time evolution operator φ(γ q 1 , τ ) is given by φ(γ q 1 , τ ) :EQUATIONand the Jacobian J q 1 is simply given byEQUATIONThen log det(J 1 q ) = log det(exp(τ a 1 (p, t))) = log exp(Tr(τ a 1 (p, t))) = Tr(τ a 1 (p, t)).Sparse Higher Order Terms. For k > 1, we consider only sparse tensors given by the simple dot product s q k (q ⊗k ) = i (s q k ) i q k i = s q k (q ⊗k ) T q •k where q •k denotes the element-wise k-th power of q. Then the q-component of time evolution operator γ q k is given component-wise by an ODE of the form dq dt = s q k (p, t)q k , whose solution is obtained in closed form via rearranging to the equivalent formqt+τ qt 1 s q k (p, t) q -k dq = t+τ t dt = τ.Then it follows that q t+τ is given component-wise by(q 1-k t,i + τ s q k (p, t) i (1 -k))1 1-k . Thus, the operator φ(γ q k , τ ) is given by φ(γ q k , τ ) :EQUATION)The Jacobian is then given byEQUATIONwith log det |J q k | given bylog det diag q •-k q •(1-k) + τ s q k (p, t)(1 -k) •( k 1-k ) = i k 1 -k log |q 1-k i -τ s q k (p, t) i (1-k)|-k log |q i |.Taylor-Verlet integrators are constructed using the splitting approximation given in Equation 5 of an order N Verlet flow γ θ , which we recall below asEQUATION)The standard Taylor-Verlet integrator of an order N Verlet flow γ θ is given explicitly in Algorithm 1 below.Algorithm 1 Integration of order N Verlet flow 1: procedure ORDERNVERLETINTEGRATE(q, p, t 0 , t 1 , steps, γ θ , N ) 2:τ ← t1-t0 steps , t ← t 0 3:∆ log p = 0 ▷ Change in log density.4:s q 0 , s p 0 , . . . s q N , s p N ← γ θ 5:while t < t 1 do 6:k ← 0 7:while k ≤ N do 8:q ← φ(γ q;θ k , τ ) ▷ q-update.∆ log p ← ∆ log p -log det Jφ(γ q;θ k , τ ) 10:p ← φ(γ p;θ k , τ ) ▷ p-update.∆ log p ← ∆ log p -log det Jφ(γ p;θ k , τ )12:k ← k + 1 13:t ← t + τ 14:return q, p, ∆ log p Closed-form expressions for the time evolution operators γ q;θ k , τ ) and log density updates log det Jφ(γ q;θ k , τ ) can be found in Table 1 . Algorithm 2 details explicitly standard Taylor-Verlet integration of an order one Verlet flow.Algorithm 2 Integration of order one Verlet flow 1: procedure ORDERONEVERLETINTEGRATE(q, p, t 0 , t 1 , steps, γ θ ) 2: τ ← t1-t0 steps , t ← t 0 3:∆ log p = 0 ▷ Change in log density.s q 0 , s p 0 , s q 1 , s p 1 ← γ θ 5:while t < t 1 do 6:q ← q + τ s q 0 (p, t; θ), ▷ Apply equation 17 7:p ← p + τ s p 0 (q, t; θ) ▷ Apply equation 17 8:q ← exp(τ s q 1 (p, t; θ))q ▷ Apply equation 20 9:∆ log p ← ∆ log p -Tr(τ s q 1 (p, t; θ)) ▷ Apply equation 23 10:p ← exp(τ s p 1 (q, t; θ))p ▷ Apply equation 20 11:∆ log p ← ∆ log p -Tr(τ s p 1 (q, t; θ)) ▷ Apply equation 23 12:t ← t + τ 13: return q, p, ∆ log pIn this section, we will show that two coupling-based normalizing flow architectures -NICE (Dinh et al. (2014) ) and RealNVP (Dinh et al. (2016) ) -can be realized as the Taylor-Verlet integrators for zero and first order Verlet flows respectively. Specifically, for each such coupling layer architecture f θ , we may construct a Verlet flow γ θ whose Taylor-Verlet integrator is given by successive applications of f θ .Closely related to leap-frog integration is Verlet integration, from which our method derives its name.Note that we do not require that dq = dp.When there is no risk of ambiguity, we drop the subscript and refer to γN simply by γ.Justification for use of the pseudo time evolution operator φ can be found in Appendix B.In our Euclidean context, a symplectic form is more generally any non-degenerate skew-symmetric bilinear form Ω ′ on phase space. However, it can be shown that there always exists a change of basis which satisfies ΛΩ ′ Λ -1 = Ω, where Λ denotes the change of basis matrix. Thus, we will only consider Ω. 



Reviews: review: title: title: Verlet flows: exact-likelihood integrators for flow-based generative models
main_review: main_review: ## General comments
Upon review, the manuscript on Verlet Flows should be accepted without revisions.
The paper demonstrates the use of Taylor-Verlet integrators in continuous
normalizing flows, presenting a method for exact-likelihood generative models. The
theoretical basis, methodological approach, and experimental validation are well
articulated. The overall contribution is deemed significant for the field of
generative modeling. The work is ready for publication, offering valuable insights
and advancements without the need for further modifications.
However, if the authors wish, they may increase the font size of the plot labels and
ticks of Figure 1.
title: title: Verlet Flows: Exact-Likelihood Integrators for Flow-Based Generative Models
main_review: main_review: This work is focused on Verlet flows and parameterization of the coefficients in both the state space and the augmenting space. The theory is not new however the implementation for CNFs can be interesting. It would be interesting to see, what is the optimisation and design approach if the authors claim that it's better than parametrization against neural network. It would be wise to present evidence of training and benchmark comparison.