Awesome Normalizing Flows

A list of awesome resources for understanding and applying normalizing flows (NF): a relatively simple yet powerful new tool in statistics for constructing expressive probability distributions from simple base distributions using a chain (flow) of trainable smooth bijective transformations (diffeomorphisms).

^{Figure inspired by Lilian Weng. Created in TikZ. View source.}

 Table of Contents

1. 📝 Publications
2. 🛠️ Applications
3. 📺 Videos
4. 📦 Packages
   1.  PyTorch Packages
   2.  TensorFlow Packages
   3.  JAX Packages
   4.  Julia Packages
5. 🧑‍💻 Repos
   1.  PyTorch Repos
   2.  JAX Repos
   3.  TensorFlow Repos
   4.  Other Repos
6. 🌐 Blog Posts
7. 🚧 Contributing

📝 Publications (60)

1. 2024-06-20 - Transferable Boltzmann Generators by Klein, Noé
   Boltzmann Generators, a machine learning method, generate equilibrium samples of molecular systems by learning a transformation from a simple prior distribution to the target Boltzmann distribution via normalizing flows. Recently, flow matching has been used to train Boltzmann Generators for small systems in Cartesian coordinates. This work extends this approach by proposing a framework for transferable Boltzmann Generators that can predict Boltzmann distributions for unseen molecules without retraining. This allows for approximate sampling and efficient reweighting to the target distribution. The framework is tested on dipeptides, demonstrating efficient generalization to new systems and improved efficiency compared to single-system training. [Code]

2. 2023-01-03 - FInC Flow: Fast and Invertible k×k Convolutions for Normalizing Flows by Kallapa, Nagar et al.
   propose a k×k convolutional layer and Deep Normalizing Flow architecture which i) has a fast parallel inversion algorithm with running time O(nk^2) (n is height and width of the input image and k is kernel size), ii) masks the minimal amount of learnable parameters in a layer. iii) gives better forward pass and sampling times comparable to other k×k convolution-based models on real-world benchmarks. We provide an implementation of the proposed parallel algorithm for sampling using our invertible convolutions on GPUs. [Code]

3. 2022-10-15 - Invertible Monotone Operators for Normalizing Flows by Ahn, Kim et al.
   This work proposes the monotone formulation to overcome the issue of the Lipschitz constants in previous ResNet-based normalizing flows using monotone operators and provides an in-depth theoretical analysis. Furthermore, this work constructs an activation function called Concatenated Pila (CPila) to improve gradient flow. The resulting model, Monotone Flows, exhibits an excellent performance on multiple density estimation benchmarks (MNIST, CIFAR-10, ImageNet32, ImageNet64). [Code]

4. 2022-08-18 - ManiFlow: Implicitly Representing Manifolds with Normalizing Flows by Postels, Danelljan et al.
   The invertibility constraint of NFs imposes limitations on data distributions that reside on lower dimensional manifolds embedded in higher dimensional space. This is often bypassed by adding noise to the data which impacts generated sample quality. This work generates samples from the original data distribution given full knowledge of perturbed distribution and noise model. They establish NFs trained on perturbed data implicitly represent the manifold in regions of maximum likelihood, then propose an optimization objective that recovers the most likely point on the manifold given a sample from the perturbed distribution.

5. 2022-06-03 - Graphical Normalizing Flows by Wehenkel, Louppe
   This work revisits coupling and autoregressive transformations as probabilistic graphical models showing they reduce to Bayesian networks with a pre-defined topology. From this new perspective, the authors propose the graphical normalizing flow, a new invertible transformation with either a prescribed or a learnable graphical structure. This model provides a promising way to inject domain knowledge into normalizing flows while preserving both the interpretability of Bayesian networks and the representation capacity of normalizing flows. [Code]

6. 2022-05-16 - Multi-scale Attention Flow for Probabilistic Time Series Forecasting by Feng, Xu et al.
   Proposes a novel non-autoregressive deep learning model, called Multi-scale Attention Normalizing Flow(MANF), where one integrates multi-scale attention and relative position information and the multivariate data distribution is represented by the conditioned normalizing flow.

7. 2022-03-02 - Adaptive Monte Carlo augmented with normalizing flows by Gabrié, Rotskoff et al.
   Markov Chain Monte Carlo (MCMC) algorithms struggle with sampling from high-dimensional, multimodal distributions, requiring extensive computational effort or specialized importance sampling strategies. To address this, an adaptive MCMC approach is proposed, combining local updates with nonlocal transitions via normalizing flows. This method blends standard transition kernels with generative model moves, adapting the generative model using generated data to improve sampling efficiency. Theoretical analysis and numerical experiments demonstrate the algorithm's ability to equilibrate quickly between metastable modes, sampling effectively across large free energy barriers and achieving significant accelerations over traditional MCMC methods. [Code]

8. 2022-01-14 - E(n) Equivariant Normalizing Flows by Satorras, Hoogeboom et al.
   Introduces equivariant graph neural networks into the normalizing flow framework which combine to give invertible equivariant functions. Demonstrates their flow beats prior equivariant models and allows sampling of molecular configurations with positions, atom types and charges.

9. 2021-07-16 - Efficient Bayesian Sampling Using Normalizing Flows to Assist Markov Chain Monte Carlo Methods by Gabrié, Rotskoff et al.
   Normalizing flows have potential in Bayesian statistics as a complementary or alternative method to MCMC for sampling posteriors. However, their training via reverse KL divergence may be inadequate for complex posteriors. This research proposes a new training approach utilizing direct KL divergence, which involves augmenting a local MCMC algorithm with a normalizing flow to enhance mixing rate and utilizing the resulting samples to train the flow. This method requires minimal prior knowledge of the posterior and can be applied for model validation and evidence estimation, offering a promising strategy for efficient posterior sampling.

10. 2021-07-03 - CInC Flow: Characterizable Invertible 3x3 Convolution by Nagar, Dufraisse et al.
    Seeks to improve expensive convolutions. They investigate the conditions for when 3x3 convolutions are invertible under which conditions (e.g. padding) and saw successful speedups. Furthermore, they developed a more expressive, invertible Quad coupling layer. [Code]

11. 2021-04-14 - Orthogonalizing Convolutional Layers with the Cayley Transform by Trockman, Kolter
    Parametrizes the multichannel convolution to be orthogonal via the Cayley transform (skew-symmetric convolutions in the Fourier domain). This enables the inverse to be computed efficiently. [Code]

12. 2021-04-14 - Improving Normalizing Flows via Better Orthogonal Parameterizations by Goliński, Lezcano-Casado et al.
    Parametrizes the 1x1 convolution via the exponential map and the Cayley map. They demonstrate an improved optimization for the Sylvester normalizing flows.

13. 2020-09-28 - Multivariate Probabilistic Time Series Forecasting via Conditioned Normalizing Flows by Rasul, Sheikh et al.
    Models the multi-variate temporal dynamics of time series via an autoregressive deep learning model, where the data distribution is represented by a conditioned normalizing flow. [OpenReview.net] [Code]

14. 2020-09-21 - Haar Wavelet based Block Autoregressive Flows for Trajectories by Bhattacharyya, Straehle et al.
    Introduce a Haar wavelet-based block autoregressive model.

15. 2020-07-15 - AdvFlow: Inconspicuous Black-box Adversarial Attacks using Normalizing Flows by Dolatabadi, Erfani et al.
    An adversarial attack method on image classifiers that use normalizing flows. [Code]

16. 2020-07-06 - SurVAE Flows: Surjections to Bridge the Gap between VAEs and Flows by Nielsen, Jaini et al.
    They present a generalized framework that encompasses both Flows (deterministic maps) and VAEs (stochastic maps). By seeing deterministic maps x = f(z) as limiting cases of stochastic maps x ~ p(x|z), the ELBO is reinterpreted as a change of variables formula for the stochastic maps. Moreover, they present a few examples of surjective layers using stochastic maps, which can be composed together with flow layers. [Video] [Code]

17. 2020-06-15 - Why Normalizing Flows Fail to Detect Out-of-Distribution Data by Kirichenko, Izmailov et al.
    This study how traditional normalizing flow models can suffer from out-of-distribution data. They offer a solution to combat this issue by modifying the coupling layers. [Tweet] [Code]

18. 2020-06-03 - Equivariant Flows: exact likelihood generative learning for symmetric densities by Köhler, Klein et al.
    Shows that distributions generated by equivariant NFs faithfully reproduce symmetries in the underlying density. Proposes building blocks for flows which preserve typical symmetries in physical/chemical many-body systems. Shows that symmetry-preserving flows can provide better generalization and sampling efficiency.

19. 2020-06-02 - The Convolution Exponential and Generalized Sylvester Flows by Hoogeboom, Satorras et al.
    Introduces exponential convolution to add the spatial dependencies in linear layers as an improvement of the 1x1 convolutions. It uses matrix exponentials to create cheap and invertible layers. They also use this new architecture to create convolutional Sylvester flows and graph convolutional exponentials. [Code]

20. 2020-05-11 - iUNets: Fully invertible U-Nets with Learnable Upand Downsampling by Etmann, Ke et al.
    Extends the classical UNet to be fully invertible by enabling invertible, orthogonal upsampling and downsampling layers. It is rather efficient so it should be able to enable stable training of deeper and larger networks.

21. 2020-04-08 - Normalizing Flows with Multi-Scale Autoregressive Priors by Mahajan, Bhattacharyya et al.
    Improves the representational power of flow-based models by introducing channel-wise dependencies in their latent space through multi-scale autoregressive priors (mAR). [Code]

22. 2020-03-31 - Flows for simultaneous manifold learning and density estimation by Brehmer, Cranmer
    Normalizing flows that learn the data manifold and probability density function on that manifold. [Tweet] [Code]

23. 2020-03-04 - Gaussianization Flows by Meng, Song et al.
    Uses a repeated composition of trainable kernel layers and orthogonal transformations. Very competitive versus some of the SOTA like Real-NVP, Glow and FFJORD. [Code]

24. 2020-02-27 - Gradient Boosted Normalizing Flows by Giaquinto, Banerjee
    Augment traditional normalizing flows with gradient boosting. They show that training multiple models can achieve good results and it's not necessary to have more complex distributions. [Code]

25. 2020-02-24 - [Modeling Continuous Stochastic Processes with Dynamic Normalizing