normflows: A PyTorch Package for Normalizing Flows

TL;DR

normflows is a comprehensive PyTorch package that facilitates building and integrating various normalizing flow models for density estimation and generative modeling tasks.

Contribution

It provides an easy-to-use, extensible library supporting multiple flow architectures, streamlining their application in machine learning workflows.

Findings

Supports diverse flow architectures like Real NVP, Glow, and Neural Spline Flows.

Simplifies integration of normalizing flows into larger models.

Open-source package available on GitHub.

Abstract

Normalizing flows model probability distributions through an expressive tractable density. They transform a simple base distribution, such as a Gaussian, through a sequence of invertible functions, which are referred to as layers. These layers typically use neural networks to become very expressive. Flows are ubiquitous in machine learning and have been applied to image generation, text modeling, variational inference, approximating Boltzmann distributions, and many other problems. Here, we present normflows, a Python package for normalizing flows. It allows to build normalizing flow models from a suite of base distributions, flow layers, and neural networks. The package is implemented in the popular deep learning framework PyTorch, which simplifies the integration of flows in larger machine learning models or pipelines. It supports most of the common normalizing flow architectures, such as Real NVP, Glow, Masked Autoregressive Flows, Neural Spline Flows, Residual Flows, and many more. The package can be easily installed via pip and the code is publicly available on GitHub.

Full text

normflows: A PyTorch Package for Normalizing Flows

Vincent Stimper

University of Cambridge, Cambridge, United Kingdom

Max Planck Institute for Intelligent Systems, Tübingen, Germany

Corresponding author: vstimper@tue.mpg.de

David Liu

University of Cambridge, Cambridge, United Kingdom

Andrew Campbell

University of Cambridge, Cambridge, United Kingdom

Vincent Berenz

Max Planck Institute for Intelligent Systems, Tübingen, Germany

Lukas Ryll

University of Cambridge, Cambridge, United Kingdom

Bernhard Schölkopf

Max Planck Institute for Intelligent Systems, Tübingen, Germany

José Miguel Hernández-Lobato

University of Cambridge, Cambridge, United Kingdom

Summary

Normalizing flows model probability distributions through an expressive tractable density (Tabak and Vanden-Eijnden, 2010; Tabak and Turner, 2013; Rezende and Mohamed, 2015). They transform a simple base distribution, such as a Gaussian, through a sequence of invertible functions, which are referred to as layers. These layers typically use neural networks to become very expressive. Flows are ubiquitous in machine learning and have been applied to image generation (Kingma and Dhariwal, 2018; Grcić et al., 2021), text modeling (Wang and Wang, 2019), variational inference (Rezende and Mohamed, 2015), approximating Boltzmann distributions (Noé et al., 2019), and many other problems (Kobyzev et al., 2021; Papamakarios et al., 2021).

Here, we present normflows, a Python package for normalizing flows. It allows to build normalizing flow models from a suite of base distributions, flow layers, and neural networks. The package is implemented in the popular deep learning framework PyTorch (Paszke et al., 2019), which simplifies the integration of flows in larger machine learning models or pipelines. It supports most of the common normalizing flow architectures, such as Real NVP (Dinh et al., 2017), Glow (Kingma and Dhariwal, 2018), Masked Autoregressive Flows (Papamakarios et al., 2017), Neural Spline Flows (Müller et al., 2019; Durkan et al., 2019), Residual Flows (Chen et al., 2019), and many more. The package can be easily installed via pip and the code is publicly available on GitHub111https://github.com/VincentStimper/normalizing-flows.

Statement of need

normflows focuses on flows that are composed of discrete transformations, as opposed to continuous normalizing flows (Chen et al., 2018; Papamakarios et al., 2021). There are several other packages implementing discrete normalizing flows, such as TensorFlow Probability (Dillon et al., 2017) for TensorFlow, distrax (Babuschkin et al., 2020) for JAX, and nflows (Durkan et al., 2020) and FrEIA (Ardizzone et al., 2018-2022) for PyTorch. However, none of them support the two popular flow architectures, residual and autoregressive flows, within a single package, while we do so.

Moreover, normflows stands out by providing tools that are often used when approximating Boltzmann distributions. First, sampling layers needed for Stochastic Normalizing Flows (Wu et al., 2020; Nielsen et al., 2020) are included. Second, Neural Spline Flows on circular coordinates are supported (Rezende et al., 2020), which can be combined with standard coordinates on bounded or unbounded intervals. They are needed when modeling the internal coordinates of molecules consisting of angles and lengths (Midgley et al., 2023). Furthermore, there is an extension for normflows that adds Boltzmann distributions as targets as well as flow layers converting between Cartesian and internal coordinates (Stimper et al., 2023).

Our package has already been used in several scientific projects and publications (Campbell et al., 2021; Stimper et al., 2022; Midgley et al., 2023). Due to its modular nature, normflows can be easily extended to house new flow layers, base distributions, or other tools. For instance, Stimper et al. (2022) extends the package by adding resampled base distributions, which overcome an architectural weakness of normalizing flows and make them more expressive.

Examples

In the GitHub repository of our package@footnotemark, we provide various examples illustrating how to use it. We show how to build a flow model from a base distribution, a list of flow layers, and, optionally, a target distribution. They can be trained by computing a loss through the respective methods provided and minimizing it with the standard PyTorch optimizers. We show how to approximate simple 2D distributions, but, moreover, apply flows to images through the multiscale architecture (Dinh et al., 2017), which normflows provides as well. Furthermore, there is an example of how to build a variational autoencoder with normalizing flows as well.

Here, we want to illustrate a strength of normflows, i.e. that it can deal with combinations of standard and circular coordinates. Therefore, we consider a distribution of two random variables, $x$ and $\phi$. $x$ follows a Gaussian distribution with density $p(x)=\mathcal{N}(x|0,1)$ and $\phi$ has a circular von Mises distribution such that $p(\phi|x)=\mathcal{M}(\phi|\mu(x),1)$ with $\mu(x)=3x$. We train a Neural Spline Flow with an unbound and a circular coordinate to approximate the target distribution $p(x,\phi)=p(x)p(\phi|x)$ by minimizing the reverse Kullback-Leibler divergence. As shown in Figure 1, the density of the flow is almost indistinguishable from the target.

Acknowledgements

We thank Laurence Midgley and Timothy Gebhard for their valuable contributions to the package. Moreover, we thank everyone who contacted us via mail or on GitHub for the valuable feedback and spotting bugs.

José Miguel Hernández-Lobato acknowledges support from a Turing AI Fellowship under grant EP/V023756/1. This work was supported by the German Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 01IS18039B; and by the Machine Learning Cluster of Excellence, EXC number 2064/1 - Project number 390727645.