Smooth Normalizing Flows

1 Introduction . Generative learning using normalizing flows ( NF ) [ 50 , 42 , 39 ] has become a widely applicable tool in the physical sciences which has e.g . been used for sampling lattice models [ 36 , 37 , 31 , 4 , 1 ] , approximating the equilibrium density of molecular systems [ 38 , 29 , 56 , 58 ] or estimating free energy differences [ 55 , 10 ] . Such models approximate a target density µ via diffeomorphic maps f ( · ; θ ) : Ω ⊂ Rd → Ω by transforming samples z ∼ p0 ( z ) of a base density into samples x = f ( z ; θ ) such that they follow the push-forward density x ∼ pf ( x ; θ ) : = p0 ( f−1 ( x ; θ ) ) ∣∣det ∂xf−1 ( x ; θ ) ∣∣ . ( 1 ) Flows can be trained on data by maximizing the likelihood or via minimizing of the reverse KL divergence DKL [ pf ( · ; θ ) ‖µ ] if µ is known up to a normalizing constant . While NFs are usually introduced as smooth diffeomorphisms , most applications like density estimation or sampling only require C1-smooth transformations . Higher-order smoothness of flows has not been discussed so far and can become a challenge as soon as multi-modal transformations on other topologies than Rd are discussed . ∗J.K and A.K . contributed equally to this work . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Ck-smoothness of NFs with k > 1 is especially important for applications in physics . Physical models usually come in the form of differential equations , where the derivatives bear a physical meaning that is often crucial to fit , evaluate , or interpret the model . Thus , the construction of expressive , smooth flow architectures will likely open up new avenues of research in this domain . Boltzmann Generators A recent application of NFs to a physical problem are Boltzmann Generators ( BG ) [ 38 ] , which we see as the main application area for the methods introduced in this paper . They are generative models trained to sample conformations of molecules in equilibrium , which follow a Boltzmann-type distribution µ ( x ) ∝ exp ( −u ( x ) ) . Here u is the dimensionless potential energy defined by the molecular system and the thermodynamic state ( e.g . Canonical ensemble at a certain temperature ) it is simulated in . BGs can be trained by a combination of MLE on possibly biased trajectory data and simultaneous minimization of the reverse KL divergence to the target density µ ( x ) . This bi-directional training scheme can achieve sampling of low-energy equilibrium structures . After training , BGs can be used e.g . for importance sampling or for providing efficient proposals when being used in MCMC applications [ 45 ] . In the context of BGs , the negative gradient of the log push-forward density with respect to x corresponds to the atomistic forces . Access to well-behaved forces is pivotal in most classical molecular computations : they simplify optimization of molecular models , drive molecular dynamics ( MD ) simulations , enable computations of macroscopic observables , and facilitate multiscale modeling . In this work , we will primarily focus on two important implications of smooth flow forces . First , we show that they enable the training of NFs via force matching ( FM ) , i.e . minimizing the force mean-squared error with respect to reference data . In combining FM with density estimation , NFs are pushed to match the distributions and their derivatives , which implicitly leads to favorable regularization . Second , we apply flow forces to drive dynamics simulations . Such simulations are required in many applications to not only sample the target distribution but also compute dynamical observables such as transition rates or diffusivities . Respecting Topological Constraints Many physical models operate on nontrivial topologies , i.e . the d-dimensional hyper-torus Td , which can become an obstacle in constructing smooth flows . An important example are BGs for peptides and small proteins which require an internal coordinate ( IC ) transformation to achieve low-energy sampling of structures [ 38 ] . This non-trainable layer transforms Euclidean coordinates x ∈ Rn×3 into a representation given by distances d ∈ [ a1 , b1 ] × . . . × [ an−1 , bn−1 ] , angles α ∈ [ 0 , π ] n−2 and dihedral torsion angles τ ∈ Tn−3 . As molecular energies are often invariant to translation and rotation , IC transformations are useful as they are bijective with the equivalence class of all-atom coordinates that are identical up to global rotations and translations . The learning problem can then be reduced to modeling the joint distribution µ ( d , α , τ ) which is supported on the nontrivial topological space XIC : = I2n−3 × Tn−3 , where I = [ 0 , 1 ] denotes the closed unit interval . Recent work [ 38 , 56 ] suggested to model the density within an open set Ω ⊂ XIC , leverage C∞smooth normalizing flows defined on Rd and then prevent significant mass around singular points using regularizing losses . Such an approach however can lead to bias and ill-behaved training and requires a-priori knowledge of the support of the densities . Later work [ 9 ] approached the problem using C1-smooth splines flows which leads to accurate samples , however results in broken forces . To overcome these limitations while still benefiting from the merits of prior work , such as bidirectional training and fast forward/inverse transformations we formulate the following desiderata for flow transformations on-top of IC layers : ( A ) They must have support on Id and Td . ( B ) They should beC∞-smooth . ( C ) They must allow bi-directional training . ( D ) Forward and inverse direction must be efficient to evaluate . Satisfying ( D ) can be achieved using coupling layers [ 11 , 12 ] . This reduces the problem to finding element-wise conditional transformations satisfying ( A-C ) . Contributions In this work we propose the following novelties : • We present a new C∞-smooth transformation with compact support which can simultaneously be used for expressive transformations on Id as well as the hypertorus Td . This satisfies ( A ) and ( B ) . • We present novel algorithm which allows optimizing non-analytic inverse directions of flows that can only be evaluated via black-box root finding methods . This satisfies ( C ) . • We show that training of smooth flows through combinations of force matching , density estimation , and energy-based training can achieve nearly perfect equilibrium densities . • We show that forces of such smooth flows can be used in dynamical simulations . 2 Related Work . While related work exist for each of the requirements ( A-D ) above , none of them provides a framework that addresses all of them . Specifically , multi-modality in element-wise flow transformations has been approached using splines [ 13 , 35 ] , monotonoic polynomials [ 26 , 41 ] , monotonic neural networks which directly approximate an inverse CDF transformation [ 22 , 8 ] , mixtures of unimodal transformations [ 21 , 43 ] and stochastic transition steps [ 56 , 7 ] . Flows on non-trivial manifolds have been discussed for hypertori and -spheres [ 43 ] , Lie groups [ 16 , 4 ] , hyperbolic spaces [ 3 ] as well as on general manifolds [ 15 , 18 , 33 , 32 , 5 , 27 ] . Applications of flows for molecular systems include approximation of the equilibrium density [ 38 , 56 , 29 ] , generation of molecular conformations [ 58 , 44 ] and free energy differences [ 55 , 10 ] . Using forces to train neural network potentials of molecules was done e.g . in [ 53 , 24 ] . Backpropagation through black-box functions was generally discussed in [ 19 ] . More related to our work is Bai et al . [ 2 ] who discuss training models whose outputs are obtained through fix-point equations . Finally Shirobokov et al . [ 46 ] discuss backpropagation through black-box simulators using surrogate models . Using force-matching for training generative models is rarely used in machine learning due to absence of an ( unnormalized ) target density . However , a common method to train unnormalized energy-based models purely on data is given by score-matching ( SM ) [ 25 , 47 ] . SM assumes an unknown density for the data distribution and attempts to match forces implicitly , e.g . via sliced [ 49 ] or denoising [ 52 ] score-matching . A survey on score-matching and its relation to other approaches of training energy-based models on data is e.g . given in Song and Kingma [ 48 ] . We discuss related work on flows for I and the unit circle in detail in Section 4 . 3 Incorporating Forces into Flow Training . NFs are most commonly trained by minimizing the negative log likelihood ( NLL ) , LNLL ( θ ) : = −Ex∼µ ( x ) [ log pf ( x ; θ ) ] = DKL [ µ||pf ( · ; θ ) ] + const , ( 2 ) or by minimizing the reverse KL divergence ( KLD ) LKLD ( θ ) : = DKL [ pf ( · ; θ ) ||µ ] + const . ( 3 ) BGs as discussed in Noé et al . [ 38 ] use a convex combination of the two in order to avoid modecollapse while still being able to achieve low-energy samples . If reference forces f ( x ) = −∂xu ( x ) corresponding to samples x are available and f ( · ; θ ) is at least C2-smooth , the optimization can naturally be augmented by the force mean-squared error : LFM ( θ ) : = Ex∼µ ( x ) [ ‖f ( x ) − ∂x log pf ( x ; θ ) ‖22 ] . ( 4 ) As was shown in Wang et al . [ 53 ] such force-matching can lead to unbiased potential surfaces even if samples are not sampled from equilibrium . Similarly to Noé et al . [ 38 ] we can thus define a loss function for smooth flows as the convex combination L ( θ ) = ωnLNLL ( θ ) + ωkLKLD ( θ ) + ωfLFM ( θ ) . ( 5 ) 4 Smooth flows on the closed interval and the unit circle . We now discuss how to achieve smooth flows on Id and Td . To unify the discussion we consider the unit circle S1 to be the quotient space I/ ∼ using the relation x ∼ x′ ⇔ ( x = x′ ) ∨ ( x = 0 ∧ x′ = 1 ) ∨ ( x′ = 0 ∧ x = 1 ) . The d-dimensional hypertorus is given by the direct product Td = S1× . . .×S1 . Following the usual definition we say that f is Ck-smooth on a compact interval [ a , b ] iff there exists a Ck-smooth continuation of f on an open set Ω ⊃ [ a , b ] . Smooth flows on I Any smooth diffeomorphism on R or an open interval ( a , b ) ⊃ I can be restricted to I and re-scaled to satisfy this requirement . Let f : Ω→ Ω be Ck-smooth for some open set I ⊂ Ω ⊂ R. Then f̃ ( x ) = ( f ( x ) − f ( 0 ) ) / ( f ( 1 ) − f ( 0 ) ) defines a Ck-smooth diffeomorphism on I . Simple C∞ transformations on R are given by affine transformations [ 12 ] . After re-scaling any affine map will result in the same identity mapping and thus will not be able to model complex densities . Powerful and computationally efficient multi-modal transformations on I can be achieved using rational-quadratic splines [ 13 ] . Those are C1-smooth and thus will result in discontinuous forces which can be disadvantageous for physical applications . Possible other C∞-smooth candidate transformations with analytic forward evaluation are given by mixtures of logistic transformations [ 21 ] or deep sigmoid/deep dense sigmoid flows [ 22 ] . Those are only continuous on ( 0 , 1 ) and thus special care has to be taken to avoid problematic behavior on the tails . Finally , multi-modal C∞ transformations on R can be achieved using non-analytic methods [ 6 , 54 ] . Yet , computational costs ( solving and backpropagating over ODEs , inverting quadrature integrations via bisections ) and numerical accuracy ( non-anlytic forward passes ) limit their applicability e.g . when used in deep coupling layers which are required to capture multivariate correlations or when trying to match densities up to high accuracy e.g . as necessary for molecular modeling . Smooth flows on S1 Flows on the hypertorus were discussed in Rezende et al . [ 43 ] who introduced necessary conditions for the transformations to define C0-continuous densities , such that they can be used in density estimation tasks . In their work Rezende et al . [ 43 ] introduced three candidates satisfying this condition : mixtures of non-compact projections ( NCP ) , rational-quadratic circular spline flows ( CSF ) and mixtures of Moebius transformations ( MoMT ) . While NCPs and CSFs satisfy C1-continuity and thus render forces discontinuous , only MoMTs define smooth densities . While MoMTs trivially also define C∞-flows on I their periodicity would be limiting when applied to non-periodic densities . Smooth compact bump functions In addition to this previous work we leverage a third way to construct smooth transformations which works for both I and S1 . Here we follow a general construction principle for smooth bump functions e.g . as explained in Tu [ 51 ] . All proofs can be found in the supplementary material . First , define a Ck-smooth and strictly increasing ramp function ρ : I→ I with ρ ( 0 ) = 0 and ρ ( 1 ) = 1 . Second , define the generalized sigmoid σ [ ρ ] ( x ) : = ρ ( x ) ρ ( x ) + ρ ( 1− x ) . ( 6 ) Then σ [ ρ ] will be a Ck-diffeomorphsim on I and furthermore all derivatives up to order k vanish at the boundary . The first derivative defines a non-negative smooth bump function with support [ 0 , 1 ] and maximum at 0.5 . We can introduce a concentration parameter a ∈ R > 0 and a location parameter b ∈ [ 0 , 1 ] , which makes g ( x ) : = σ [ ρ ] ( a · ( x− b ) + 12 ) a smooth bijection from [ − 12a + b , 1 2a + b ] to [ 0 , 1 ] with all values and higher-order derivatives vanishing outside the domain . Furthermore , by introducing c ∈ ( 0 , 1 ] and setting f ( x ) : = ( 1− c ) · ( g ( x ) − g ( 0 ) g ( 1 ) − g ( 0 ) ) + c · x , ( 7 ) we can define flexible unimodal Ck-diffeomorphisms on [ 0 , 1 ] . When used within coupling layers , we let a , b , c be the output of some not further constrained neural network . Ck and C∞ ramp functions There are many possible choices for ramp functions satisfying above mentioned properties . A simple Ck ramp function is given by the k-th order monomial ρ ( x ) = xk . More interestingly , C∞ smoothness on [ 0 , 1 ] can be achieved using the ramp ρ ( x ) = exp ( − 1 α·xβ ) for x > 0 and ρ ( x ) = 0 for x ≤ 0 . Here we set α > 0 and β > 1 . We make α a trainable parameter . Optimizing β is possible in principle , yet fixing β ∈ { 1 , 2 } and treating it as a hyper-parameter stabilized training and led to better results . Smooth and efficient circular wrapping As discussed in Rezende et al . [ 43 ] and Falorsi et al . [ 16 ] we can turn any Ck-smooth density p ( x ) with support on R into a Ck-smooth density on S1 using the marginalization p̃ ( x ) = ∑ k∈Z p ( x+ k ) . This construction is generally problematic due to two reasons : ( i ) non-vanishing tails ( e.g . if p ( x ) is Gaussian ) require evaluating an infinite series which in most cases has no analytic expression and can be numerically challenging ( ii ) smooth densities with compact support in some interval [ a , b ] , for example sigmoidal transforms , can introduce non-smooth behavior at the interval boundaries . Smooth and compactly supported transformations as introduced above do not suffer from this : ( i ) due to compact support the series will always have a finite number of non-vanishing contributions , ( ii ) as the transformations are Ck on all of R and all derivatives are compactly supported wrapping will always result in a Ck-smooth density on S1 . Multi-modality via mixtures Similarly , as discussed in prior work [ 21 , 43 ] any convex sum of Ck-diffeomorphisms on I defines a Ck-diffeomorphism on I . Thus we can combine multiple of those unimodal transformations defined in Eq . ( 7 ) to obtain arbitrarily complex multi-modal transformations on I and S1 , see Fig . 1 .

Summary: The authors introduce a flow architecture, which they call smooth normalizing flows, that can operate on compact intervals and hypertori and has a C-infinity-smooth transformation. Furthermore, they propose to train normalizing flows via force matching, i.e. matching the gradients of the flow transformation which those of the energy function of the target distribution. They also show that a transformation with non-analytic inverse can be used to define a flow model which is trained by evaluating the inverse via black-box root finding and computing the gradients through the forward pass of each layer. On the experiment side, the authors apply their method to a toy example as well as the Boltzmann distribution of a complex molecule. They demonstrate that their flow architecture is expressive and the presented training procedure can be successfully used. When force matching is used, the gradients of the flow transformation resemble the respective gradients of the energy function of the target indeed more accurately and they can even be used as forces in molecular dynamics simulations.