Smooth Normalizing Flows

Jonas Köhler, Andreas Krämer, Frank Noe

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.

Code Of Conduct: I certify that all co-authors of this work have read and commit to adhering to the NeurIPS Statement on Ethics, Fairness, Inclusivity, and Code of Conduct.

TL;DR: A novel smooth normalizing flow architecture on nontrivial topologies designed for applications in physics that benefit from training with respect to samples, energies, and reference forces