Smooth Normalizing Flows
Authors: Jonas Köhler
Andreas Krämer
Frank Noe
Venue: NeurIPS 2021 Poster (2021)
Keywords: Normalizing Flows
Boltzmann Generators
Smoothness
Backpropagation through Root-Finding

Abstract: Normalizing flows are a promising tool for modeling probability distributions in physical systems. While state-of-the-art flows accurately approximate distributions and energies, applications in physics additionally require smooth energies to compute forces and higher-order derivatives. Furthermore, such densities are often defined on non-trivial topologies. 
A recent example are Boltzmann Generators for generating 3D-structures of peptides and small proteins. These generative models leverage the space of internal coordinates (dihedrals, angles, and bonds), which is a product of hypertori and compact intervals. 
In this work, we introduce a class of smooth mixture transformations working on both compact intervals and hypertori.
Mixture transformations employ root-finding methods to invert them in practice, which has so far prevented bi-directional flow training. To this end, we show that parameter gradients and forces of such inverses can be computed from forward evaluations via the inverse function theorem.
We demonstrate two advantages of such smooth flows: they allow training by force matching to simulation data and can be used as potentials in molecular dynamics simulations. 

Smooth Normalizing Flows
Jonas Köhler*T
Andreas Krämer*
Frank Noé 838
Department of Mathematics and Computer Science, Freie Universität Berlin
Department of Physics, Freie Universität Berlin
Department of Chemistry. Rice University, Houston, TX koehler, andreas kjonas.oohler, . frank.noe
Abstract
Normalizing flows promising tool for modeling probability distributions in physical systems. While state-of-thh-an flows accurately approximate distributions and energies, applications in physies additionally require smooth energies com- pute forces and higher-order derivatives. Furthermore, such densities are often defined non-trivial topologies. A recent example Boltzmann Generators for generating 3D-structures of peptides and small proteins. These generative models leverage the space of internal coordinates (dihedrals, angles, and bonds), which is a product of hypertori and compact intervals. In this work, we introduce class of smooth mixture transformations working on both compact intervals and hypertori. Mixture transformations employ root- not-finding methods to invert them in practice, which has so far prevented bi-directional flow training. To this end, we show that parameter gradients and forces of such inverses can be computed from forward evaluations the inverse function theorem. We demonstrate two advantages such smooth flows: they allow training by force matching to simulation data and can be used as potentials in molecular dynamies simulations.
Introduction
Generative learning using normalizing flows (NF) [50 in the physical sciences which has e.g. been used for sampling lattice models [36] 37] 31 4 M approximating the equilibrium density of molecular systems 38 29
42 39 has become a widely applicable tool
56
58 or estimating free energy
differences [55 10]
Such models approximate a target density μ via diffcomorphic maps f(; 0): Ω C Rd + Q by transforming samples po(z) of base density into samples æ f[(; 0) such that they follow the push-forward density
x (x; 0 = po ( 0)) |ett 0, (=:; 0)
(1)
Flows can be trained on data by maximizing the likelihood or via minimizing of the reverse KL divergence DKL (pr(:: 1)]p]] if known up to a normalizing constant.
While NFs are usually introduced as smooth diffeomorphisms, most applications like density estima- tion sampling only require C* -mooohh 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 d are discussed.
"J.K and A.K. contributed equally to this work.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

C -mssmothhness 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 physical problem are Boltzmann Genera- tors (BG) [8) 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 ju(æz) X exp(-u(2))
Here u is the dimensionless potential energy defined by the molecular system and the thermodynamic state (e.g. Canonical ensemble at 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 density (z)). This bi-directional training scheme can achieve sampling of low- 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]
the target
In the context of BGs, the negative gradient of the log push-forward density with respect corre- sponds 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 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-squaree 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 transition rates or diffusivities.
Respecting Topological Constraints Many physical models operate on nontrivial topologies, i.e. the d-dimensional hyper-torus T4, which can become an obstacle constructing smooth flows.
An important example are BGs for peptides and small proteins which require an internal coordinate (IC) transformation achieve low-energy sampling of structures [38]. This non-trainable layer transforms Euclidean coordinates æ E RaX3 into a representation given by distances d E [a1, 61] x ... [ap-1, bn angles a [0, T| m-)- and dihedral torsion angles T E T=-3, As molccular 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 p(d, a. T) which supported on the nontrivial topological space Xic := [2n-3 3 x T denotes the closed unit interval.
33 where I [0. 1]
Recent work 38 566 suggested to model the density within open set Ω Mc., leverage _ smooth normalizing flows defined on IRd and then prevent significant mass around singular points using regularizing losses. Such an approach however can lead bias and ill-behaved training and requires a-priori knowledge of the support of the densities. Later work [9] approached the problem using Cmmstooth splines flows which leads accurate samples, however results in broken forces.
To overcome these limitations while still benefiting from the merits of prior work, such as bi- directional 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 Isd and T (B) They should be -mooth. (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 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 new -moshh transformation with compact support which can simulta- neously be used for expressive transformations on Td as well as the hypertorus Td. This satisfies (A) and (B).

---

We present novel algorithm which allows optimizing non con-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.
Related Work
While related work exist for each of the requirements (A-D) above, none of them provides frame- work that addresses all of them. Specifically, multi-modalltyy in element-wise flow transformations has been approached using splines [33 35 monotonoic polynomials [26 41), monotonic neural
networks which directly approximate an inverse CDF transformation
22 mixtures of unimodal
transformations
2 4 and stochastic transition steps [56 7
Flows on non-trivial manifolds have been discussed for hypertori and -spheres 44]. Lie groups
[16 4] hyperbolic spaces 3 as well on general manifolds [5] 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 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 -oint 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 (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 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 based models on data e.g. given in Song and Kingma 48
We discuss related work on flows for 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), (O) (Re [log py (æ; 0)] DKIlHllpp(C + const,
(2)
or by minimizing the reverse KI. divergence (KLD)
LKID := DKI. (p::: (:; 0)||m] + const.
(3)
BGs as discussed in Noć et al. 388 use a convex combination of the two in order to avoid mode- collapse while still being able to achieve low-energy samples. reference forces f(x) -0 (x)) corresponding to samples æ are available and f(=:0) 0) is at least Cm=-smothh, the optimization can naturally be augmented by the force mean- -quared error:
LFM == En=p(=) [lf(z) - 0x loggy (z; 0)12
(4)
As shown in Wang al. 553 such force-matching can lead unbiased potential surfaces samples are not sampled from equilibrium. Similarly to Noé et al. (38 we can thus define a loss function for smooth flows the convex combination
(()) = wn, CNLI + WK.LLI + wiLFm(
(5)
4 Smooth flows on the closed interval and the unit circle
We now discuss how to achieve smooth flows on 1d and T4. To unify the discussion we consider the unit circle S to be the quotient space I/ using the relation ~ x' + (x = x') v (x OAr = 1) (x' 0 A = 1). The d-dimensional hypertorus is given by the direct product T x 1t Following the usual definition we say that / -mooohh on a compact interval [a, b) iff there exists ** -m-sthh continuation of ff on an open set Ω 5 Ja, 6).

---

Smooth flows on I Any smooth diffeomorphism on R or an open interval (a. 5 I can be restricted to I and re-scaled to satisfy this requirement. Let f: Ω + Ω be C* -smooth for some open set C Ω C R. Then f(x) (f(z) - f(0)) - f())) defines a -mmooo diffeomorphism on I. Simple C transformations on I 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 modli transformations
can be achieved
autional-quadratc using rational- quadratic splines [13]. Those are -smooth and thus will result in discontinuous forces which can be disadvantageous for physical applications. Possible other C -sooohh candidate transformations with analytic forward evaluation are given by mixtures of logistic transformations 21 or deep sigmoid/deep dense sigmoid flows [2)]. 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 IK can be achieved using non-analytic methods 66 (4). 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 St Flows on the hypertorus were discussed in Rezende et al. [3 who introduced necessary conditions for the transformations to define C* continniuu 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 C'l ccnttinuityy and thus render forces discontinuous, only MoMT's define smooth densities. While MoM'Ts trivially also define -lowslows 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 third way to construct smooth transformations which works for both and S Here we follow a general construction principle for smooth bump functions e.g. explained in Tu [51] All proofs can be found in the supplementary material.
First, define a C* smooth and strictly increasing ramp function p: T II with p(0) 0 and p(1) 1. Second, define the generalized sigmoid
p(x)
<pp((x) p(z) + p(1 x)
(6)
Then [lp] will be C'-diffeomorphsiim on and furthermore all derivatives up order k vanish at the boundary. The first derivative defines a non-negative smooth bump function with support [o, 1] and maximum 0.5. We can introduce a concentration parameter a E R,o and location parameter b E [0, 1J, which makes gzz) == dpp] (a (r b) + 2)) smooth bijection from [- 20 + b, 1 + b] to [o, 1] with all values and higher-order derivatives vanishing outside the domain. Furthermore, by introducing e (0, 1] and setting
20
f(z) = (1 (1- [(()) -()) + C.- 2,,
(1)) ((1))-(0))
can define flexible unimodal C* di-difffommopsssss [0, 1]. When used within coupling layers, we let a, b, c be the output of some not further constrained neural network.
C² and C ramp functions There are many possible choices for ramp functions satisfying above mentioned properties. A simple C* ramp function is given by the k-th order monomial p(x) = z²
More interestingly, C smoothness on [0, 1] can be achieved using the ramp p(x) = exp (- a.2
for x > 0 and p(x) = 0 for x S Here we set a > 0 and /3 > We make a a trainable parameter. Optimizing 3 is possible in principle, yet fixing 3 E {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 C*-smooth density p(x) with support on R into C* C*-smooth density on S using the marginalization p(x) Ekez p(x + k). This construction is generally problematic due to two

---

reasons: (i) non von-shnig tails (e.g. if p(z) is Gaussian) require evaluating 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 the interval boundaries.
Smooth and compactly supported transformations as introduced above do not suffer from this: (i) due compact support the series will always have a finite number of non-vanishing contributions, (ii) as the transformations are C* on all of I and all derivatives are compactly supported wrapping will always result in a C density on Sta
Multi-modality via mixtures C de-difffeomorhiisss defines a O* diffeoommmmmhis unimodal transformations defined in Eq. 7 to obtain arbitrarily complex multi-modal transformations on ll and St, see Fig.
Similarly, as discussed in prior work [21 43 any convex sum of
Thus we can combine multiple of those
Figure 1: Construction of mixture transformations on compact intervals (left) and hypertori (right). The upper and lower row depict probability densities and cumulative distribution functions, respec- tively. Multiple unimodal bump functions (a) and d)] are added small but finite density (b) and e)] and combined to yield bijective multimodal transformations (c) and f)].
5 Optimizing non-analytic inverse flows
drawback of mixture-based flows is their lack of an analytic inverse. Prior work such as 21 43 suggested rely black-box inversion methods like bisection search sample from trained models. While this leads to accurate samples, the discrete nature of such black-box oracles does not allow minimize C(0) as defined in Eq. (5).
possible remedy to this can be derived using the inverse function function theorem. Let f(; 0): Ω R Ω be scalar diffeomorphism and let furthermore x: f 0) be obtained via black-box inversion algorithm. To minimize losses depending on æyy: 0) or log (y: 0) we need to compute gradients with respect to y and 9. Here we derive the following relations for the derivatives (all details in supplementary material):
0yz(v; 0) = (0.f(x;0))-
(8)
0gz(y; 0) (0.f(x;)) 0e fx;
(9)
0, log (, 0) = (0.](z;;)) log (x; 0);
(10)
Oo log 0, 0)
(). (log 10. 0o 0o O,;,
(11)
We remark that 8 corresponds to the implicit reparameterization gradient as introduced in Figurnov al. [17]. Using these rules, we can first compute via black-box oracle for any input y and then compute all necessary gradients using forward evaluations and derivatives of fl::0). Derivatives of [+; 0) are usually accessible for computing the log Jacobian of the transformation. We can obtain higher-order derivatives using automatic differentiation. It is easy to see that this construction extends to multivariate diffeomorphisms with diagonal Jacobian as used in coupling layers. can be generalized arbitrary higher order derivatives using the Faà di Bruno formula, e.g. when aiming
5

---

to do force matching through a black-box inversion algorithm. As our experiments only require losses involving second order derivatives for the inverse direction we leave this for future work. Note that Equations 8 11 can be applied to any other element-wise flow transformation that lacks an analytic inverse and is not tied mixture models.
Another problem of bisection is its sequential execution which can become prohibitively slow during training. E.g. achieving an eтror of 10-6 requires around 20 iterations. To obtain speedup on GPUs we suggest to generalize the classic binary search to searching in a grid of K bins simultaneoulsy (see details in Supp. Mat.). For K 2 results in the usual bisection method. However, by leveraging vectorized execution we can reduce the number of iterations by factor O (1/ log K) at the expense of increasing memory by a factor O (K). Practical speedup depends on the actual number of parallel workers and thus depending on dimensionality, batch size and number of mixture components the optimal choice of K varies.
6
Experiments
The numerical experiments in this work are tailored to show the benefits of smooth flows over non-smooth flows. Therefore, we mainly compare mixtures of bump functions to neural spline flows which exhibit ofte--of-the-rt generative performance but whose densities are only first-order continuously differentiable.
6.1 Illustrative Toy Example
C)
Figure 2: a) Reference density and corresponding force field as approximated by b) a -s-moooth NSF and c) -moooth flow using the mixture of bump functions introduced in Sec. 4
To highlight the difference in terms of smoothness, the two flow architectures were applied to two-dimensional toy example. Fig. _ a) depicts the reference energy and forces on samples from the ground-truth distribution. Both flows were trained through density estimations on those samples (see SI for details). The smooth flows and spline flows matched the density well, which demonstrates their expressivity. However, the foree field of the spline flow contained dramatic outliers, while the forces the smooth flow matched the regular behavior of the reference forces.
6.2 Runtime Comparison
While the rational-quadratic splines are analytically invertible, mixture transformations obviously increase computational cost the inversion due to the iterative root-finding procedure. To quantify this gap in performance, the inverse evaluation of spline flow was compared with a modified version, where the analytic inverse was replaced by the multi-bin bisection from Section The performance of the bisection and the most efficient K was picked for each input size ("dim"), Both transformations were coupled to conditioner with the same input size (dim) and 64 hidden neurons.
evaluated for different numbers of bins K 2m L....,8, .8
For small tensor dimensions (2-32), the optimal multi-bin bisection employed up 256 bins on the GPU, which resulted in only a factor of 2-3 slowdown compared to analytic inversion. For larger dimensions (2048), the parallelization over multiple bins became less effective, leading to one order of magnitude difference in computational cost. The compute times and optimal bin sizes were comparable for the inverse network pass and its backward-pass. More details on these runtime comparisons can be found in the supplementary material.

---

6.3 Smooth Flows Enable Boltzmann Generator Training through Force Matching
To demonstrate the advantages of smooth flows, we train Boltzmann generator (BG) for small molecule, alanine dipeptide, which is described in the supplementary material. This is common test case for molecular computations and was previously considered for sampling with normalizing flows in [56, 9 30] The molecular system has 60 degrees of freedom (global rotation and translation excluded). Its potential energy surface is highly sensitive to small displacements (30 and contains singularities.
a)
b)
c)
Count [#Samples /00000] Count | 100000
2000
Flow Forces [kg7/nml
Target
1000
TEST SET (MD)
Flow
800
400
200
-200
Flow Forces [kgTinm]
Target Flow
1000
BG SAMPLES
1000
800
600
1100
400
2000
200
22000
20
Target (kg7/mm)
Energy
[Ke7]
Figure 3: a) Generated sample structures. b) Torsion distribution after training through root-finding. c) Smooth normalizing flow trained on alanine dipeptide through combination of force matching and density estimation. The top and bottom row show the performance on 10,000 samples each from the holdout test set and from the flow, respectively. Left: joint distribution of backbone torsion angles. Center: scatter plot of flow vs. target force components. Right: energy histograms for the flow energy and the target energy. (Flow energies were shifted by constant so that the minimum energy matched with the minimum energy from the molecular potential).
In previous work 56 30,, BGs for alanine dipeptide used stochastic augmentation of the base space. The introduction of noise variables facilitates creating expressive normalizing flows but prevents computation of deterministic forces. Some have also used spline flows to sample molecular
configurations [55] [0]
As proof-of-concept for the force matching loss, we trained smooth flows for alanine dipeptide using 1000:1 weighting between the force matching residual and the negative log likelihood, see supplementary material for details on the flow and training setup. No stochastic augmentation or energy-based training was used and reweighting was conducted for postprocessing.
Figure 3 compares the flow distribution with the target Boltzmann distribution on the test data from MD and on flow samples. The left column shows that the flow has almost perfectly learned nontrivial joint distribution berween the two backbone torsion angles ø and v. including sparsely populated metastable region around & 1. In contrast to previous work that used affine coupling layers [56), the modes are cleanly separated, see also Fig. [S-3 for direct comparison. The center column compares flows forees with target forces. Those forces are highly sensitive to small

---

perturbations of the molecular configurations. Nevertheless, they matched up to root-mean square deviation of 25 and 46 3/7mmmm in the target ensemble and the generated ensemble, respectively. Note that the training was stopped after 10 epochs, where the validation force matching residual had not yet fully converged, so that even further improvements may be possible with longer training or extensive hyperparameter optimization. A more detailed analysis of the forces and sampling efficiency of the BGs is shown in supplementary material, Figure [SI44] Smooth flows trained by a combination of density estimation and FM attain a favorable sampling efficiency of 42%, compared to 25% and < 1% for spline and RealNVP flows, respectively.
Finally, the right column of Figure c) depicts the distribution of flow and target energies evaluated on the same samples. The flow energies were shifted so that the minimum energy matched with the molecular mechanics energy. It has not escaped our notice that this constant offset (189 kn on the test set and 190 7 on the flow samples) corresponds to the log partition function, a quantity whose intractability has complicated research in statistical mechanics for decades.
The energy distribution of the flow tracked the ground truth exponential distribution almost perfectly even though no target energies or forces were evaluated on the flow samples during training. This glose-to-perfect match demonstrates that including force information into the training process presents an efficient means for regularization. Figure 3]a) a) shows representative conformations generated by the BG, which confirm the high quality of the molecular structures.
6.4 Boltzmann Generator Training through Black-Box Root-Finding
Table 1:
Test set metrics of alanine dipeptide training with different losses: negative log likelihood (NLL), force matching error (FME), and reverse Kullback- Leibler divergence (KLD). Weighted loss functions were used as defined in Eq. (5) with weight factors Wk, Wr and W,, 1 w/ wy. Metries were recorded after 10 training epochs. Statistics are means and 2x standard errors over 10 replicas for each experiment. The lowest value with respect each metric is highlighted in bold type.
Wk 0
Wk =0 0
Wk 0.1
Wi 0.1
Metric
Method
Wp = 0
W/ = 0.001
W/ 0
wr 0.001
210.32
-196.40
48.13
spline
(=0.16)
(= 1.53)
(# 95.37)
NLL
-211.04
-211.40
-206.33
-208.70
smooth
(#: 0.09)
(# 0.05)
(+2.04)
(# 1.93)
12.13
313.64
1498.98
spline
(+: 4.94)
(+: 107.92)
(+ 1935.72)
FME x10
1.04
0.32
5.80
0.48
smooth
(#: 0.08)
(+: 0.02)
(= 2.30)
(# 0.05)
263.77
230.08
1205.69
spline
(d: 3.21)
(+: 12.38)
(+: 36.25)
KLD
193.91
195.17
219.03
207.81 (+: 11.64)
smooth
(+: 0.39)
(+: 0.43)
(+: 22.16)
Two experiments were conducted to demonstrate that flows can be trained through a black-box root-finding algorithm.
First, the BG from section 6.3 was trained with all smooth transformations operating in the inverse direction. Consequently, inverse problems had to be solved when computing the negative log likelihood with respect to data, while the sampling occurred analytically. Figure 3) shows the joint marginal distribution of the backbone torsions the BG samples. The distribution matches the data and the smooth flows that were trained in forward mode (see Fig. 3 This result shows that training with gradients computed from the inverse function theorem is feasible.
Second, flows were trained by different combinations of density estimation, force matching, and energy-based training, where computing Op CKLDD requires the gradients of the inverse flow. The training of spline flows with inclusion of the force matching error (wi > 0) was notoriously unstable. This led to 10/10 and 7/10 failed runs for We 0 and W3 0.1, respectively. This is understandable
8

---

from the discontinuity of the spline forces. Even the combination of NLL and KLD led to instabilities for 2/10 runs. (Gradients were not clipped during training to enable a mostly unbiased comparison). In contrast, all training runs with our smooth flows concluded successfully.
Table shows the metrics the test set after 10 training epochs using spline flows and smooth flows (see SI for specifics about the flow architecture and training). With pure density estimation (we wp 0), both architectures achieved similar NLL. However, the smooth flow achieved much lower FME and KLD. The NLL and FME were further improved when force data was presented to the flow during training. Including reverse KLD in the loss yielded reasonable metrics for the smooth flows, indicating that the backpropagation through root-finding was numerically stable. However, the metries were consistently worse than with pure NLL training for both types of flows (spline and smooth) and were subject to large fluctuations. In contrast, including the FME proved to be a more stable approach to include information about the target energy into the training.
6.5
Smooth Flows as Potentials in Molecular Dynamics
Total Energy [Kj/mol]
Amber Force Field
Smooth Flow
Spline Flow
Spline (Rescaled)
5.0
3x10
2x10
2.5
1x10"
0.0
0
-1x10
-2.5
a)
b)
d)
-2x10
-5.0
3x10
0
1
2
4 5 t [ps]
2
3 [pss]
[ps]
[ps]
Figure 4: Total energy fluctuation in simulations with a classical force field and with two flows that employed smooth transforms and neural spline transforms, respectively. For each potential, 10 simulations were run starting from 10 different initial configurations (grey color). One run each is highlighted in black. Subfigures c) and d) show identical data on different scales.
challenging test for the smoothness of a flow its use as a potential energy in a dynamics simulation. Especially simulations in the microcanonical ensemble are extremely sensitive to numerical errors as the symplectic integration does not impose any additional stabilizing mechanism on the total energy. This means that any inconsistencies between energies and forces on the scale of one integration step easily cause drifting, strongly fluctuating, or exploding energies.
Therefore, we ran MD simulations using the flow energy = - log (.5 0), from Section 6.3 as the potential. For comparison, simulations were also run with spline flows that were trained purely by density estimation. Simulations started from stable structures of an MD simulation and were equilibrated in each potential for 1 ps using a Langevin thermostat with 10/ps friction coefficient.
depicts the evolution of the total energy over ps simulations in the microcanonical ensemble. As expected, the classical simulation kept the total energy roughly constant within a standard deviation of 6 x 10-4 kJ/mol per degree of freedom. Astoundingly, the smooth flow potentials also maintained the energies within 8 10-4 kJ/mol per degree of freedom using the same 1 fs time step. In contrast, the simulations with spline flow potentials quickly fell apart with potential and kinetic energies growing out of bounds. Those instabilities persisted even with step (not shown).
order-of order-of-magninne maggitude smaller time
While the infeasibility of spline flows for this task was expected, the competitive behavior of smooth flows with molecular mechanics force fields that were tailored for dynamics highlights their regularity and the consistency of their forces.
7 Discussion
Limitations and outlook Despite the promising results we mention the following possible limita- tions of the method in the present state and discuss possible improvements for future work:
First, smooth flows as implemented in this work impose numerical overhead in comparison to non-smooth alternatives such as spline flows. While this can be considered a question of concrete engineering it also opens search for more efficient smooth alternatives.

---

In addition, the differentiation rules for black-box inverses provide correct algebraic gradients. Yet for considerably multi-modal distributions maintaining numerical stability can become challenge. Future work should thus focus on studying and improving numerical stability of the bidirectional training.
While potentials can be trained up to an accuracy where they allow for energy- -orrgeicnne simulations, is yet to show that they also provide accurate equilibrium distribution for long-run simulations. At this point major bottleneck is the inferior evaluation performance per step of a flow-based force-field compared to common force-field implementations.
Furthermore, internal coordinates are a very efficient representation space for smaller polyy-eptides. However, they are difficult to scale to large proteins, protein complexes or systems with con-trivial topologies, e.g. proteins with disulfide bonds. Integrating recent advances in graph-based molecular representations with force-matched flows will likely be required for succesful modeling of such systems.
Finally, training with forces improves the resulting potential compared to just training with samples. However, for many MD data forces have not been stored and recomputing them requires a significant computational overhead.
Conclusion We introduced a range of contributions which can improve upcoming work on applying flows to physical problems.
The proposed approach for backpropagation through black-box root-finders can help to obviate the search for analytically invertible transformations if bi-directional training is necessary. As we show the method works especially well for relatively fow-dimensional problems. Scaling the approach, generalizing it to non-diagonal Jacobians and improving its numerics are interesting questions for future work.
The MD simulation example shows that C -mooth flows on non-trivial topologies can open new applications for flows. By carefully respecting the topological
avenues of research well domain as well as the smoothness of the target potential we can approximate the equilibrium density of small peptide nearly perfectly. Finally, we showed that incorporating force information together with MLE outperforms the state-of-the-An approach of combining MLE with minimizing of the reverse KL divergence. It does not suffer from mode-collapse, yet includes information of the target potential which is required to achieve low-energy samples.
Combined, these results could improve methods used for learning the mean potential surface of coarse-grained molecules, pave the way to multi-scale flows for modeling large protein systems and eventually help to accelerate simulations and sampling.
10

---

Acknowledgements
We thank the anonymous reviewers at NeurIPS 2021 for their insightful comments. Special thanks to Yaoyi Chen for setting up the alanine dipeptide system and Manuel Dibak, Leon Klein, Oana-Iuliana Popescu, Leon Sixt, and Michael Figurnov for helpful discussions and feedback on the manuscript and pointers to related work. We acknowledge funding from the European Commission (ERC CoG 772230 ScaleCell), Deutsche Yorschungssemmensccaatt (GRK DAEDALUS, SFB1 14/A04), the German Ministry for Education and Research (Berlin Institute for the Foundations of Learning and Data BIFOLD), and the Berlin Mathematics center MATH+ (Project AA 1- and EF1-2).
References
[1] M. Albergo, G. Kanwar, and Shanahan. Flow-based generative models for markov chain monte carlo in lattice field theory. Physical Review D, 100(3):034515, 2019.
[2] S. Bai, J. Z. Kolter, and V. Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, 32:690-701, 2019.
[3] J. Bose, A. Smofsky, R. Liao, P. Panangaden, and W. Hamilton. Latent variable modelling with hyperbolic normalizing flows. In International Conference on Machine Learning, pages 1045-1055. PMLR, 2020.
[4] D. Boyda, G. Kanwar, S. Racanière, D. Rezende, M. S. Albergo, K. Cranmer, D. C. Hackett, and P. E. Shanahan. Sampling using (n) gauge equivariant flows. Physical Review D, 103(7): 074504, 2021.
[5] J. Brehmer and K. Cranmer, Flows for simultancous manifold learning and density estimation. Advances in Neural Information Processing Systems, 33, 2020.
[6] Q. Chen, Y. Rubanova, Bettencourt, and K. Duvenaud. Neural ordinary differential equations. In Advances in neural information processing systems, pages 6571-6583, 2018.
[7] R. Cornish, A. Caterini, G. Deligiannidis, and A. Doucet. Relaxing bijectivity constraints with continuously indexed normalising flows. In International Conference on Machine Learning, pages 2133 2143. PMLR, 2020.
[8] N. De Cao, W. Aziz, and I. Titov. Block neural autoregressive flow. In Uncertainty in Artificial Intelligence, pages 1263-1273. PMLR, 2020.
[9] M. Dibak, I.. Klein, and F. Noé. Temperaature-sseerralee flows. arXiv preprint arXiv:2012.00429, 2020.
[10] X. Ding and B. Zhang, Deepbar: fast and exact method for binding free energy computation. The Journal of Physical Chemistry Letters, 12(10):2509-255 2021.
[11] L. Dinh, D. Krueger, and Y. Bengio. Nice: Non-linear independent components estimation. arXiv preprint 1rXivv1410.855 2014.
[12] L. Dinh, J. Sohl-Dickstein and S. Bengio. Density estimation using real NVP. In 5th Inter- national Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. URL net/forum? -oorumm?ddH
[13] C. Durkan, A. Bekasov, I. Murray, and G. Papamakarios. Neural spline flows. In Advances in Neural Information Processing Systems, pages 7509-7520, 2019.
[14] P. Eastman, J. Swails, J. D. Chodera, R. T. McGibbon, Y. Zhao, K. A. Beauchamp, L.-P. Wang, A. C. Simmonett, M. P. Harrigan, C. D. Stern, al. Openmm Rapid development high performance algorithms for molecular dynamics. PLoS computational biology, 13(7)ee0056559 2017.
[15] I.. Falorsi. Continuous normalizing flows on manifolds. arXiv preprint arXiv:2104.1499 2021.

---

[16] L. Falorsi, P. de Haan, T. R. Davidson, and P. Forré. Reparameterizing distributions on lie groups. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3244 3253. PMLR, 2019.
[17] M. Figurnov, S. Mohamed, and A. Mnih. Implicit reparameterization gradients. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 439-450, 2018.
[18] M. C. Gemici, D. Rezende, and S. Mohamed. Normalizing flows on riemannian manifolds. arXiv preprint arXiv: 161 (0..1...034444 2016.
[19] W. Grathwohl, D. Choi, Y. Wu, G. Roeder, and D. Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018, URL htttp:// net/ group?id=ICLR. group7/dd-IUaa.cc/2111//ooooeee
[20]
C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. Smith, R. Kern, M. Picus, Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and E. Oliphant. Array programming with Sept. 2020. doi: 10.10388<41186-000-26666 020 2649-2
NumPy. Nature, (885(7825)3357-255
URL
hatpp::/doii wttp:://.....r.......o
[21] Ho, X. Chen, A. Srinivas, Y. Duan, and P. Abbeel. Flow++: Improving flow-based generative models with variational dequantization and architecture design. In International Conference on Machine Learning, pages 2722-2730. PMLR, 2019.
[22] C.-W. Huang, D. Krueger, A. Lacoste, and A. Courville. Neural autoregressive flows, In International Conference on Machine Learning, pages 2078-2087. PMLR, 2018.
[23] J. D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering, (3):90-95, 2007. doi: 10.1 101109//CSE.200..55.
[24] B. E. Husic, N. E. Charron, D. Lemm, J. Wang, A. Pérez, M. Majewski, A. Krämer, Y. Chen, Olsson, G. de Fabritiis, et Coarse graining molecular dynamics with graph neural networks. The Journal of Chemical Physics, 153(19):194101 2020.
[25] A. Hyvärinen and P. Dayan. Estimation of normaooaal statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
(26) P. Jaini, K. A. Selby, and Y. Yu. Sum-of-squares polynomial flow. In International Conference on Machine Learning, pages 3009-3018. PMLR, 2019.
[27] D. Kalatzis, J. Z. Ye, J. Wohlert, and S. Hauberg. Multi-chart flows. arXiv preprint arXiv:2106.03500 2021.
[28] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 14X.v::4/2.698/// 29800 2014.
[29] Köhler, L. Klein, and F. Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning, pages 5361-5370. PMLR, 2020.
[30] A. Krämer, J. Köhler, and F. Noé. Training invertible linear layers through rank-one perturba- tions. arXiv preprini arXiv:2010.07033 2020.
[31] S.-H. Li and L. Wang. Neural network renormalization group. Physical review letters, 121(26): 260601, 2018.
[32] A. Lou, D. Lim, I. Katsman, L. Huang, Q. Jiang, S. N. Lim, and C. M. De Sa. Neural manifold ordinary differential equations. H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 17548- 17558. Curran Associates, Inc., 2020. URL / 1020/file/Coe 22t00//iiLaaaiet1oaaatasssssass3
neuriis.. cttpsscccccpr/////// Paperr .dff
12

---

[33] E. Mathieu and M. Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33, 2020.
[34] R. T. McGibbon, K. A. Beauchamp, M. P. Harrigan, C. Klein, M. Swails, C. X. Hernández, C. R. Schwantes, L.-P. Wang, T. Lane, and V. S. Pande. Mdtraj: a modern open library for the analysis of molecular dynamics trajectories. Biophysical journal, 109(8):1528- 1532, 2015.
[35] T. Müller, B. McWilliams, F. Rousselle, M. Gross, and J. Novák. Neural importance sampling. ACM Transactions on Graphics (TOG), 38(5):1-19, 2019.
[36] K. A. Nicoli, S. Nakajima, N. Strodthoff, W. Samek, K.-R. Müller, and Kessel. Asymptotically unbiased estimation of physical observables with neural samplers. Physical Review E, 101(2): 023304, 2020.
[37] K. A. Nicoli, C. Anders, I.. Funcke, T. Hartung, K. Jansen, P. Kessel, Nakajima, and P. Stor- nati. Estimation of thermodynamic observables in lattice field theories with deep generative models. Physical review letters, 126(3):03201,, 2021.
[38] F. Noé, S. Olsson, J. Köhler, and H. Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457))eaaw/ 147, 2019.
[39] G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan. Nor- malizing flows for probabilistic modeling and inference. Jourmal of Machine Learning Research, 22(57):1-64, 2021. URL http::/jmlr. .rg. pptpppppppp/111--epepp 1028 html
[40] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury. G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang. Z. DeVito, M. Raison, A. Te- jani, S. Chilamkurthy, B. Steiner, L. Fang. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, H'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: pttptp/::///.oooeenggg.... neurips cc http:////////eeaaaaaaaeeeeeeeee0
Paper. .dff
[41] S. Ramasinghe, K. Fernando, S. Khan, and N. Barnes. Robust normalizing flows using bernstein- type polynomials. arXiv preprint arXiv:2 102.03509, 2021.
[42] D. Rezende and S. Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning, pages 5530-153.. PMLR, 2015.
[43] D. J. Rezende, G. Papamakarios, S. Racaniere, M. Albergo, G. Kanwar, P. Shanahan, and K. Cranmer. Normalizing flows on tori and spheres. In International Conference on Machine Learning, pages 8083-8092. PMLR, 2020.
[44] V. G. Satorras, E. Hoogeboom, F. B. Fuchs, I. Posner, and M. Welling. E (n) equivariant normalizing flows for molecule generation in 3d. arXiv preprint arXiv:210S.0001 2021.
[45] L. Sbailò, M. Dibak, and F. Noé. Neural mode jump monte carlo. The Journal of Chemical Physics, 154(7):074101, 2021.
[46] Shirobokov. V. Belavin, M. Kagan, A. Ustyuzhanin, and A. G. Baydin. Black-box optimiza- tion with local generative surrogates. Advances in Neural Information Processing Systems, 33, 2020.
[47] Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. In Proceedings of the 33rd Annual Conference on Neural Information Processing Systems, 2019.
[48] Y. Song and D. P. Kingma. How to train your energy- based models. CoRR, 0iss/101.033288 2021. URL http::ttttt//att /ttp::///////v.//aa http:/vv..rr/baa/////11// abt///bbb....../////////////111/ https////0.0028.88331
[49] Y. Song, S. Garg, J. Shi, and S. Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence, pages 574-584, PMLR, 2020.
[50] F. G. Tabak, E. Vanden-Fijnden et al. Density estimation by dual ascent of the log-likelihood. Communications in Mathematical Sciences, 8(1):217-233, 2010.
13

---

[51] L. W. Tu. Bump Functions and Partitions of Unity, pages 127-134. Springer New York, New York, NY, 2008. ISBN 978-0-387-48101-2 doi: 10.1007/973.00387-8810----- 13. URL
/ /ttp://owo.00.0000//////00000050
[52] P. Vincent. A connection between score matching and denoising autoencoders. Neural Compu- tation, 23(7):1661-167) 2011. doi: 10.1162/NECO a 00142.
[53] Wang, S. Olsson, C. Wehmeyer, A. Pérez, N. E. Charron, G. De Fabritiis, F. Noé, and C. Clementi. Machine learning of coarse-gras soarse-raiied molecular dynamics force fields. ACS central science, 5(5):755-777, 2019.
[54] A. Wehenkel and G. Louppe. Unconstrained monotonic neural networks. Advances in Neural Information Processing Systems 32 (2019), 2019.
[55] Wirsherger, A. Ballard, Papamakarios, Ahercrombie, S. Racanière, A. Pritzel, Blun- dell, et al. Targeted free energy estimation via learned mappings. The Journal of Chemical Physics, 153(14):144 14111224441111444 14, 2020.
[56] H. Wu, J. Köhler, and F. Noe. Stochastic normalizing flows. H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neu- ral Information Processing Systems, volume 33, pages 5933-5944. Curran Asso-
ciates, Inc., 2020. URL https://pproceedinggs www.eeeddinnnneuis.ss.... hew..ips.cc/ppppr/000000 (1000 Paper. pdff
[57] C. Xenophontos. A formula for the nth derivative of the quotient of two functions. http:
//ww.ww... mas. www.maas.cc...o ac ct/////eooooo/uuubs....eo.e httpppub//ppplppppppppppp 2007. Accessed: 2021-05-28.
[58] M. Xu, S. Luo, Bengio, Peng, and Tang. Learning neural generative dynamics for molecular conformation generation. arXiv preprint arXiv:2 10240, 2021.
14


Reviews:
[{"review_id": "qNvW5ITiN_", "reviewer": "Reviewer_QQCj", "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. \n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper poses an important issue in physical systems and tries to introduce new flow-based approaches to overcome the challenges. Overall, the paper is well-written and clear to follow. However, I have some concerns:\n\n1. The computational cost and scalability would be critical concerns. The authors also mentioned this limitation but I would like to highlight it again. Since the paper aims to solve challenging problems in physical systems, the toy examples or synthesis datasets are not very meaningful. In other words, the impact of the proposed methods would be limited if the algorithm is too computationally intensive for large-scale systems.  The used example in this paper might not be a large-scale example so I can not clearly see the performance in real-world applications.  \n\n2. Due to physical motivation, I am not sure how to apply the proposed methods for general problems in the ML community, for example, image datasets. For the baseline methods, have you compared with stochastic NFs [1]? In that paper, some examples may share similar physical motivation and challenges?  Right now, only compared with NSF may not be strong enough to demonstrate the advantages of the proposed smooth NFs. \n\n[1] \"Stochastic normalizing flows.\" arXiv preprint arXiv:2002.06707 (2020).\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No", "review_text": "", "extra_scores": {"ethical_concerns": "NA", "needs_ethics_review": "No", "ethics_review_area": ["I don’t know"], "time_spent_reviewing": "8.5"}, "extra_text": {}}, {"review_id": "jlQheYILtpo", "reviewer": "Reviewer_TKYQ", "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.\nOn 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.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors are motivated by tackling the challenging problem of fitting the Boltzmann distribution of complex molecules. With their work, they made two significant contributions, which overcame problems previous articles had:\n(1) They proposed a flow architecture based on mixtures of smooth bump functions, which are C-infinity-smooth, act on a compact interval, and can easily model multimodal distributions. Furthermore, they showed that these kinds of flows can be trained although their inverse is non-analytic by evaluating the inverse through black-box root finding and computing the gradients via the push-forward density.\n(2) They trained their models via force matching which makes the gradients of the transformation closer to those of the energy function of the target. Furthermore, this stabilizes training when using the reverse Kullback-Leiber divergence as objective an allows the model to be used as potential in a molecular dynamics simulation.\nThe first contribution is novel and although force matching is a known training method, it has not been used in this context yet so the second contribution is novel as well. \nThe experimental results are impressive. However, there should be more comparisons to competing methods. For instance in section 6.3, there is no clear comparison to training without force matching. They mention a competing method, i.e. stochastic normalizing flows, but they only refer to the article and do not show the results when stochastic normalizing flows are applied to the same problem.\nOverall, the paper is well written and has clear illustrations.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Although the authors are focused on approximating Boltzmann distributions, their framework is in theory not limited to this setting so it would be interesting to see how well other densities could be approximated.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "dig_lEoLgex", "reviewer": "Reviewer_9eZb", "summary": "The paper considers a set of problems that come up when fitting normalizing flows to datasets of interest in physical sciences, and proposes a set of solutions to those problems.\n\nTo begin with, when it comes to datasets in physical sciences we might be interested in matching the *gradient* of the log density of the target distribution, as well as the log density itself. This is because this gradient can correspond to the *forces* acting on the physical objects, and as such the gradient is important to get right, in order to make the model useful in downstream applications, e.g. simulation. Furthermore, we often have access to the *true* forces at the points in the training sample, and would like to regularize the model to make its gradients match the available true gradients. Authors consider a straightforward way to implement such regularization: adding the MSE between the model's gradient and the true gradient to the training objective. However, doing this requires a normalizing flow model that is at least $C^2$-smooth, as the second-order gradients will be used during optimization.\n\nMoreover, the data of interest in this context often has non-trivial topology, containing quantities such as angles, directions or axis. There's been prior work on building normalizing flows suitable for such data, but authors point out that there is no existing model that would be\na) restricted to a manifold of interest (say, a circle or a torus) \nb) $C^{\\infty}$ (or at least $C^k$ for some $k > 1$) continuous\nc) *expressive*, i.e. able to represent complex densities on the manifold\n\nAuthors propose a way to construct a normalizing flow that *would* have the properties above. More specifically, authors parametrize a $C^k$-continuous univariate diffeomorphism on $[0,1]$, which could then be turned to a diffeomorphism on a circle, torus etc. using methods by Rezende et al. [34], and/or used as the elementwise transformation in coupling/autoregressive layers. \n\nUnfortunately, to make the transformation above expressive, authors have to consider a mixture of simpler transformations, hence the transformation has no analytic inverse. This is not an uncommon limitation when it comes to normalizing flows, where inverting some of the flows (e.g. for sampling) requires performing numerical root-finding. This presents a problem in the context of this work: in literature it has been demonstrated that it is desirable to train the flow on a convex combination of a forward and an *inverse* KL divergences. The latter is the one that presents a problem as it requires sampling from the flow *and* evaluating the gradients of its inverse during training. The cost of doing this can be prohibitive when the inverse is not analytic. Authors derive a way to compute said gradients by *only* using samples from the model and the gradients of the forward transformation, which are typically readily available.\n\nFinally, authors evaluate the proposed solutions on a battery of benchmarks, both synthetic and real, showing that the final model demonstrates desirable behaviour when compared to a strong baseline.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I should begin by saying that I am not an expert in the domain-specific aspects of the paper, so a few things went over my head. This also means I can't be the judge of the importance of the outlined problems in the context of physical sciences. Having said that, I've been able to follow the motivation behind the work --- authors have done a good job laying out the story behind why the different proposed pieces are necessary. I've enjoyed reading the paper, although it's not the easiest one to wrap your head around. \n\nOn the first look the univariate diffeomorphism proposed in section 4 very much resembles the deep sigmoid flow by Huang et al., being a mixture of sigmoid-like functions. Am I correct in understanding that the primary difference is in the proposed transformation being restricted to some interval [a, b] (e.g. [0, 1]), while also being well-behaved/smooth at the boundaries of that interval? \n\nI did not fully understand the reason authors dismiss references [6] and [41] due to their \"computational costs\". I am not *too* familiar with that literature and the trade-offs of associated methods, but I've seen them applied to large, high-dimensional datasets, so their cost must be at least somewhat manageable? Given that the method authors propose is also \"non-analytic\", at least in terms of its inverse, I wonder if these references would be worth investigating, especially given their desirable smoothness properties?\n\nThe convex combination of the forward and the inverse KL divergences in eq. (5) reminds me of the $\\alpha$-divergence: could this be a more principled way to express the desired \"distance\"?\n\nSection 5 is very cool, especially because the results in eqs. (8)-(11) are immediately applicable to all normalizing flows without an analytic inverse, in cases where it is necessary to evaluate the gradient of the inverse of the flow during training, which I have seen some examples of before.\n\nThe experiments are nice: a strong baseline (the rational-quadratic spline) is used, and the differences between it and the proposed method are clear in both the carefully-constructed toy data and real data.\n\nOverall, a well-written paper with a strong motivation, plenty of technical novelty, and robust evaluation. I strongly recommend accepting it to the conference.\n\n---\n\nI thank the authors for their detailed response. I believe the write-up, while already excellent, will benefit from the proposed clarifications. Furthermore, I agree with reviewer *q1W2* that visualisations of individual \"ramp\" functions, in addition to Figure 1, would be great to aid understanding.\n\nAfter reading the author's response and the other reviews/responses, I maintain my recommendation for acceptance, although lower my score somewhat to be better calibrated with other reviews, and in recognition of the fact that even though there are multiple useful ideas in this paper, few of them are ground-breaking *by themselves*.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations of the method are discussed by the authors throughout the paper, especially when it comes to the computational costs, though a paragraph in the discussion summarizing said limitations would be useful. The explored applications make it difficult to imagine negative societal impact coming from this work.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "MG9AwPll_5s", "reviewer": "Reviewer_q1W2", "summary": "The paper introduces a class of $C^k$-smooth ($k>1$) transformations working on compact intervals and hypertori that can be used to construct normalising flows.\nSuch type of smooth transformations are important in physical systems as (1) they respect topological constraints and (2) higher-order derivatives have physical interpretations; however, existing transformations either don't respect the constraints or are limited to have only $C^1$-smoothness.\nLearning flows consisting mixtures of proposed transformations that requires black-box root-finders is made efficient through gradients derived by the inverse function theorem.\nThe benefit of accessing to higher-order derivatives is demonstrated by training normalizing flows with force-matching regularisation and dynamics simulations of trained flows.", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: \nThe proposed transformations are novel and the application of the inverse function theorem is neat.\nThis work clearly fills the gap of $C^k$-smooth ($k>1$) flows on intervals and hypertori and provides necessary discussion on related works in the end of section 1.\n\nQuality:\nThe paper is a complete piece and the techniques proposed (both transformations and training method) are technically sound.\nThe intuition behind each transformation is given when it is introduced,\nand the experiments clearly shows that the proposed method can be trained successfully while having well defined higher-order derivatives.\nThe authors mostly focus on the advantages of smooth flows in compare to spline flows but limitations are not adequately discussed (only discussion on computation cost is given).\n\nClarity:\nThe paper is well-written: the motivation, how the transformations and training are explained and how the method is evaluated are all clear and the structure makes sense.\nWhile it might be due to page limit, it could be helpful to visualise each transformations (and their alternatives) when they are introduced and contrasted in section 4.\nIn section 5, it was mentioned in L245 that \"In contrast to previous work that used affine coupling layers [43], the modes are cleanly separated.\" but not visualisation was shown.\nAgain it would be helpful to see it in the paper.\n\nSignificance: \nThe results clearly show on $C^k$-smoothness can help model training (section 6.4) and model usage (section 6.5).\nAs said, it is not achievable by previous methods.\nAs motivated and demonstrated by the authors, smooth flows can be widely used in modelling physics systems in which topological constraints are needed, so I can see smooth flows would be beneficial to the community especially who works on scientific modelling.\n\nMisc:\n- How would other types of generative models perform on the toy example in section 6.1, or in other experiments? For example I wonder if [1] would a strong competitor as it can learn the support and the distribution on it quite well. Extra experiments would be great to have but no necessary I believe, but it's worth some discussion on why flows are the right choice of generative models in the beginning.\n\n[1] Arbel, Michael, Liang Zhou, and Arthur Gretton. \"Generalized energy based models.\" arXiv preprint arXiv:2003.05033 (2020).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There is no single paragraph or section dedicated to limitations.\nI suggest the author(s) to add one, along with future works.\n\nThere is no societal impact either but I don't think it's applied to this paper as it is a technical manuscript.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces a flow architecture referred to as smooth normalizing flows. These are C^K-smooth maps that work on compact intervals and hypertori. 
One of the main contributions is to propose a smooth transformation on the unit-interval.

The paper is overall well-written and technically sound. Conceptually, there is not much innovation in the training setup. But the experimental results are quite promising both with toy and real data. Testing learning models by using them on MD simulations (Langevin diffusions) is a neat idea as it requires that the model log-likelihood's gradient to be sufficiently accurate.

I am quite surprised to see no connections made between "force matching" and the whole literature of "score matching" [e.g. 1, 2, 3, 4, 5, 6] in section 3. I strongly advise the authors to add further discussions and references in this regard.

[1] Hyvärinen, A. and Dayan, P., 2005. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4).

[2] Song, Y., Sohl-Dickstein, J., Kingma, D.P., Kumar, A., Ermon, S. and Poole, B., 2020. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456.

[3] Song, Y. and Ermon, S., 2019. Generative modeling by estimating gradients of the data distribution. arXiv preprint arXiv:1907.05600.

[4] Bordes, F., Honari, S. and Vincent, P., 2017. Learning to generate samples from noise through infusion training. arXiv preprint arXiv:1703.06975.

[5] Song, Y. and Ermon, S., 2020. Improved techniques for training score-based generative models. arXiv preprint arXiv:2006.09011.
Decision: Accept (Poster)