InvertibleNetworks.jl: A Julia package for scalable normalizing flows

Keywords Julia  $\cdot$ inverse problems  $\cdot$ Bayesian inference  $\cdot$ imaging  $\cdot$ normalizing flows

1 Summary

Normalizing flows is a density estimation method that provides efficient exact likelihood estimation and sampling [1] from high dimensional distributions. This method depends on the use of the change of variables formula which requires an invertible transform. Thus normalizing flow architectures are built to be invertible by design [1]. In theory, the invertibility of architectures constrains the expressiveness but the use of coupling layers allows normalizing flows to exploit the power of arbitrary neural networks that need not be invertible [2] and layer invertibility means if properly implemented many layers can be stacked to increase expressiveness without creating a training memory bottleneck.

The package we present, InvertibleNetworks.jl, is a pure Julia [3] implementation of normalizing flows. We have implemented many relevant neural network layers, including GLOW 1x1 invertible convolutions [4], affine/additive coupling layers [1], Haar wavelet multiscale transforms [5] and Hierarchical invertible neural transport (HINT) [6] among others. These modular layers are easily composed and modified to create different types of normalizing flows. As starting points, we have implemented RealNVP, GLOW, HINT, Hyperbolic networks [7] and their conditional counterparts for users to quickly implement their individual applications.

2 Statement of need

This software package focuses on memory efficiency. The promise of neural networks is in learning high-dimensional distributions from examples thus normalizing flow packages should allow easy application to large dimensional inputs such as images or 3D volumes. Interestingly, the invertibility of normalizing flows naturally alleviates memory concerns since intermediate networks activations can be recomputed instead of saved in memory, greatly reducing the memory needed during backpropagation. The problem is that directly implementing normalizing flows in automatic differentiation frameworks such as PyTorch [8] will not automatically exploit this invertibility. The available packages for normalizing flows such as: nflows [9], normflows [10] and FrEIA [11] are built depending on automatic differentiation frameworks thus do not exploit invertibility for memory efficiently.

3 Memory efficiency

By implementing gradients by hand instead of depending completely on automatic differentiation, our layers are capable of scaling to large inputs. By scaling, we mean that these codes are not prone to out-of-memory errors when training on GPU accelerators. Indeed, previous literature has described memory problems when using normalizing flows as their invertibility requires the latent code to maintain the same dimensionality as the input [12].

Figure 1: Our package InvertibleNetworks.jl provides memory frugal implementations of normalizing flows. Here, we compare our implementation of GLOW with an equivalent implementation in a PyTorch package. Using a 40GB A100 GPU, the PyTorch package can not train on image sizes larger than 480x480, while our package can handle sizes larger than 1024x1024.

In Figure 1, we show the relation between input size and the memory required for a gradient calculation in a PyTorch normalizing flow package (normflows [10]) as compared to our package. The two tests were run with identical normalizing flow architectures. We note that the PyTorch implementation quickly increases the memory load and throws an out of memory error on the 40GB A100 GPU at the spatial image size of 480x480 while our InvertibleNetworks.jl implementation still has not run out of memory at spatial size 1024x1024. Note that this is in the context of a typical learning routine, so the images include 3 channels (RGB) and we used a batchsize of 8.

Figure 2: Due to the invertibility of the normalizing flow architecture, the memory consumption does not increase as we increase the depth of the network. Our package properly exploits this property thus shows constant memory consumption whereas the PyTorch package does not.

Since traditional normalizing flow architectures need to be invertible these might be less expressive as compared to non-invertible counterparts. In order to increase their expressiveness, practitioners stack many invertible layers to increase the overall expressive power. Increasing the depth of a neural network would in most cases increase the memory consumption of the network but in this case since normalizing flows are invertible, the memory consumption does not increase. Our package displays this phenomena as shown in Figure 2 while the PyTorch (normflows) package that has been implemented with automatic differentiation does not display this constant memory phenomena.

4 Ease of use

Although the normalizing flow layers gradients are hand-written, the package is fully compatibly with ChainRules [13] in order to integrate with automatic differentiation frameworks in Julia such as Zygote [14]. This integration allows users to add arbitrary neural networks which will be differentiated by automatic differentiation while the memory bottleneck created by normalizing flow gradients will be dealt with InvertibleNetworks.jl. The typical use case for this combination is the summary networks used in amortized variational inference such as BayesFlow [15] which has been implemented in our package.

All implemented layers are tested for invertibility and correctness of their gradients with continuous integration testing via GitHub actions. There are many example for layers, networks and application workflows allowing new users to quickly build networks for a variety of applications. The ease of use is demonstrated by the publications that made use of the package.

Many publications have used InvertibleNetworks.jl for diverse applications including: change point detection, [16], acoustic data denoising [17], seismic imaging [18, 19, 20, 21, 22, 23], fluid flow dynamics [24], medical imaging [25, 26, 27, 28] and monitoring CO2 for combating climate change [29].

5 Future work

The neural network primitives (convolutions, non-linearities, pooling etc) are implemented in NNlib.jl abstractions thus support for AMD, Intel and Apple GPU can be trivially extended. Also, while our package currently can handle 3D inputs and has been used on large volume-based medical imaging [30] there are interesting avenues of research regarding the "channel explosion" seen in invertible down and upsampling used in invertible networks [31].