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: This paper is proposing smooth normalizing flows for modeling probability distribution in physical systems.  This is because the existing flow-based model approximates distributions and energies by computing forces and higher-order derivatives based on smoothed energies. This work aims at addressing this challenge by introducing a class of smooth mixture transformations working on both compact intervals and hypertori.  The proposed methods show advantages on two points: 1) enable training by force matching to simulation data and 2) used as potentials in MD simulations.