Improving Continuous Normalizing Flows using a Multi-Resolution Framework
Abstract
Recent work has shown that Continuous Normalizing Flows (CNFs) can serve as generative models of images with exact likelihood calculation and invertible generation/density estimation.
In this work we introduce a Multi-Resolution variant of such models (MRCNF).
We introduce a transformation between resolutions that allows for no change in the log likelihood.
We show that this approach yields comparable likelihood values for various image datasets, with improved performance at higher resolutions, with fewer parameters, using only 1 GPU.
1. Introduction
Reversible generative models derived through the use of the change of variables technique (Dinh et al., 2017; Kingma & Dhariwal, 2018; Ho et al., 2019; Yu et al., 2020) are growing in interest as generative models, because they enable efﬁcient density estimation, efﬁcient sampling, and computation of exact likelihoods.
A promising variation of the change-of-variable approach is based on the use of a continuous time variant of normalizing ﬂows (Chen et al., 2018; Grathwohl et al., 2019), which uses an integral over continuous time dynamics to transform a base distribution into the model distribution, called Continuous Normalizing
Flows (CNF).
CNFs have been shown to be capable of modelling complex distributions such as those associated with images.
While this new paradigm for the generative modelling of images is not as mature as Generative Adversarial
Networks (GANs) (Goodfellow et al., 2016) or Variational
Autoencoders (VAEs) (Kingma & Welling, 2013) in terms of the generated image quality, it is a promising direction of research.
workshop on Invertible Neural Networks, Normalizing Flows, and Explicit
Models (ICML 2021).
Copyright In this work, we focus on making the training of continuous normalizing ﬂows feasible for higher resolution images, and help reduce computation time.
We thus introduce a novel multi-resolution technique for continuous normalizing ﬂows, by modelling the conditional distribution of highlevel information at each resolution in an autoregressive fashion.
We show that this makes the models perform better at higher resolutions.
A high-level view of our approach is shown in Figure 1.
Our main contributions are:
We introduce Multi-Resolution Continuous Normalizing Flows (MRCF) , through which we achieve stateof-the-art Bits-per-dimension (BPD) (negative log likelihood per pixel) on ImageNet64 using fewer model parameters relative to comparable methods.
We propose a multi-resolution transformation that does not add cost in terms of likelihood.
2. Background 2.1. Normalizing Flows
Normalizing ﬂows (Tabak &
Turner, 2013;
Jimenez Rezende & Mohamed, 2015; Dinh et al., 2017; Papamakarios et al., 2019; Kobyzev et al., 2020) are generative models that map a complex data distribution, such as real images, to a known noise distribution.
They are trained by maximizing the log likelihood of their input images.
Suppose a normalizing ﬂow g produces output z from an input x i.e. z = g ( x ) .
The change-of-variables formula provides the likelihood of the image under this transformation as:
[Equation]
The ﬁrst term on the right (log determinant of the Jacobian) is often intractable, however, previous works on normalizing ﬂows have found ways to estimate this efﬁciently.
The second term, log p ( z ) , is computed as the log probability of z under a known noise distribution, typically the standard
Gaussian N .
2.2. Continuous Normalizing Flows
Continuous Normalizing Flows (CNF) (Chen et al., 2018;
Grathwohl et al., 2019; Finlay et al., 2020) are a variant of normalizing ﬂows that operate in the continuous domain, using the framework of Neural ODEs (Chen et al., 2018).
Suppose CNF g transforms its state v ( t ) using a Neural
ODE with neural network f deﬁning the differential.
Here, v ( t 0 ) =
x is, say, an image, and at the ﬁnal time step v ( t 1 ) =
z is a sample from a known noise distribution.
[Equation]
Chen et al.
Grathwohl et al.
(2019) proposed a more efﬁcient method to compute the change in log-probability in the context of CNFs, called the instantaneous variant of the change-of-variables formula:
[Equation]
An ODE solver solves both differential equations eq.
(2) and eq.
Thus, a CNF provides both the ﬁnal state v ( t 1 ) as well as the change in log probability ∆ log p v ( t 0 )
→ v ( t 1 ) .
Prior works (Grathwohl et al., 2019; Finlay et al., 2020;
Ghosh et al., 2020; Onken et al., 2021; Huang & Yeh, 2021) have trained CNFs as reversible generative models of images, by maximizing the likelihood of images:
[Equation]
where x is an image, z and ∆ log p x
→ z are computed by the CNF using eq.
(2) and eq.
(3), and log p ( z ) is the likelihood of the computed z under a known noise distribution, typically the standard Gaussian N ( 0 , I ) .
CNF g is trained by maximizing E x log p
( x ) .
Novel images are generated by sampling z from the known noise distribution, and running it through the CNF in reverse.
3. Our method
Our method is a reversible generative model of images that builds on top of CNFs.
We introduce the notion of multiple resolutions in images, and connect the different resolutions in an autoregressive fashion.
This helps generate images faster, with better likelihood values at higher resolutions.
Moreover, we used only one GPU in all our experiments.
We call this model Multi-Resolution Continuous Normalizing
Flow (MRCNF).
3.1. Multi-Resolution image representation
Multi-resolution representations of images have been explored in computer vision for decades (Burt, 1981; Witkin, 1987; Burt & Adelson, 1983; Mallat, 1989; Marr, 2010).
We express an image x as a series of high-level information y s not present in the immediate coarser images x s +1 (obtained by averaging every 2 × 2 patch), and a ﬁnal coarse image x S : x → ( y 1 , x 2 ) → · · · → ( y 1 , y 2 ,
[Equation]
Our overall method is to map these S terms to S noise samples using S CNFs.
Deﬁning the high-level information y s
The multi-resolution representation in eq.
(5) needs to be invertible, i.e. it should be possible to deterministically obtain x s from y s and x s +1 , and vice versa.
Further, it is preferable that this transformation incurs minimal additional computational cost, and does not add too much change in log-likelihood.
Hence, we choose to perform a linear transformation taking into account the following properties: 1) volume preserving i.e. determinant is 1, 2) angle preserving, and 3)
range preserving (respecting the maximum principle, studied for some time, under the notion of the maximum principle (Varga, 1966)).
Consider the simplest case of 2 resolutions where x 1 is a 2 × 2 image with pixel values x 1 , x 2 , x 3 , x 4 , and x 2 is a 1 ×image with pixel value ¯ x = 14 ( x 1 + x
+ x 3 + x 4 ) .
We require three values ( y 1 , y 2 , y 3 ) =
y 1 that contain information not present in x 2 , such that when they are combined with x 2 , xis obtained.
This could be viewed as a problem of ﬁnding a matrix M such that: [ x 1 , x 2 , x 3 , x 4 ] (cid:62) =
M [ y 1 , y 2 , y 3 , ¯ x ] (cid:62) .
We ﬁx the last column of M as [1 , 1 , 1 , 1] (cid:62) , since every pixel value in x 1 depends on ¯ x .
Finding the rest of the parameters can be viewed as requiring four 3D vectors that are (ideally) non-trivially equally spaced.
These can be considered as the four corners of a tetrahedron in 3D space, under any rotation in space and scaling of the vectors.
Out of the many possibilities for this tetrahedron, we could choose the matrix that performs the Discrete Haar Wavelet
Transform (DHWT) (Mallat, 1989; Mallat & Peyr´e, 2009).
However, this has log (cid:12)(cid:12) det( M − 1 ) (cid:12)(cid:12) =
log(1 / 2) , and is therefore not volume preserving.
We introduce a variant of the DHWT matrix that is unimodular, i.e. has a determinant of 1 (therefore volume preserving), while also preserving the range of the images for the input and its average:  
[Equation]

[Equation]
[Equation]
 
[Equation]

[Equation]
  

[Equation]   [Equation] (7) where c = 2 2 / 3 , a = 4 , and log (cid:12)(cid:12) det( M − 1 ) (cid:12)(cid:12)
= log(1)
= 0 .
This can be scaled up to larger spatial regions by performing the same calculation for each 2 × 2 patch.
Let M be the function that uses matrix M from above and combines every pixel in x s +1 with 3 corresponding pixels in y s to make the 2 × 2 patch at that location in x s using eq.
[Equation]
3.3. Multi-Resolution Continuous Normalizing Flows
Using the multi-resolution image representation in eq.
(5), we characterize the conditional distribution over the additional degrees of freedom ( y s ) required to generate a higher resolution image ( x s ) that is consistent with the average ( x s +1 ) over the equivalent pixel space.
At each resolution s , we use a CNF to reversibly map between y s (or x S when s = S ) and a sample z s from a known noise distribution.
At generation, y s only adds information missing in x s +1 , but conditional on it.
This framework ensures that one coarse image could generate several potential ﬁne images, but these ﬁne images have the same coarse image as their average.
This fact is preserved across resolutions.
In principle, any generative model could be used to map between the multi-resolution image and noise.
Normalizing ﬂows are good candidates for this as they are probabilistic generative models that perform exact likelihood estimates, and can be run in reverse to generate novel data from the model’s distribution.
This allows model comparison and measurement of generalization to unseen data.
We choose to use the CNF variant of normalizing ﬂows at each resolution, since CNFs have recently been shown to be effective in modeling image distributions using a fraction of the number of parameters typically used in normalizing ﬂows (and non ﬂow-based approaches), and their underlying framework of Neural ODEs have been shown to be more robust than convolutional layers (Yan et al., 2020).
Training : We train an MRCNF by maximizing the average log-likelihood of the images in the training dataset under the model, i.e. max E x log p ( x ) .
The log probability of each image log p ( x ) can be estimated recursively as:
[Equation]
where log p ( x S ) is computed by CNF g S using eq.
[Equation]
[Equation]
and log p ( y s | x s +1 ) is also computed by CNFs g s similarly, conditioning on the coarser image:
[Equation]
This model could be seen as a stack of CNFs connected in an autoregressive fashion.
Typically, likelihood-based generative models are compared using the metric of bitsper-dimension (BPD), i.e. the negative log likelihood per pixel in the image:
[Equation]
Hence, we train our MRCNF to minimize the average
BPD of the images in the training dataset, computed using eq.
Although the ﬁnal log likelihood log p ( x ) involves sequentially summing over values returned by all S CNFs, each CNF can be trained independently, in parallel.
We use FFJORD (Grathwohl et al., 2019) as the baseline model for our CNFs.
In addition, we use to two regularization terms introduced by RNODE (Finlay et al., 2020) to speed up the training of FFJORD models.
Generation : First, z s , s = 1 , . . .
, S are sampled from the latent noise distributions.
Given an S -resolution model,
CNF g s at resolution s transforms the noise sample z
s to high-level information y s conditioned on the immediate coarse image x
s +1 (except g S which is unconditioned).
y s and x s +1 are then combined to form x s as described in section 3.2 (see ﬁg. 1).
This process is repeated progressively from coarser to ﬁner resolutions:
[Equation]
4. Related work
Several prior works on normalizing ﬂows — Glow (Kingma & Dhariwal, 2018), Hoogeboom et al. (2019a;b), MintNet (Song et al., 2019), MaCow (Ma et al., 2019), Durkan et al.
; Chen et al. (2020), Flow++ (Ho et al., 2019), NanoFlow (Lee et al., 2020), Wavelet Flow (Yu et al., 2020), DenseFlow (Grci´c et al., 2021) — build on RealNVP (Dinh et al., 2017).
Although they achieve great results in terms of BPD and image quality, they nonetheless report results from signiﬁcantly higher parameter (some with 100x!), and several times GPU hours of training.
Although our MRCNF model is similar to the recently published WaveletFlow (Yu et al., 2020), we generalize the notion of a multi-resolution image representation.
WaveletFlow builds on the Glow (Kingma & Dhariwal, 2018) architecture, while ours builds on CNFs.
WaveletFlow claims to have orthonormal transformation, our eq.
(6) involves a unimodular transformation.
Finally, WaveletFlow applies special sampling techniques to obtain better samples from its model.
We have so far not used such techniques for generation, but we believe they can potentially help our models as well.
“Multiple scales” in prior normalizing ﬂows : Normalizing ﬂows (Dinh et al., 2017; Kingma & Dhariwal, 2018;
Grathwohl et al., 2019) try to be “multi-scale” by transforming the input at one resolution in a smart way (squeezing operation) such that the width of the features progressively reduces.
In contrast, our model stacks normalizing ﬂows at multiple resolutions in an autoregressive fashion.
5. Experimental results
We train Multi-Resolution Continuous Normalizing Flow (MRCNF) and Multi-Resolution Continuous Normalizing
Flow - Wavelet (MRCNF-Wavelet) models on the ImageNet (Deng et al., 2009) dataset at 32x32, 64x64, 128x128.
We build on the code provided in (Finlay et al., 2020) (https://github.com/cﬁnlay/ffjord-rnode).
In all cases, we train using only one NVIDIA RTX 2080 Ti GPU with 11GB.
At lower resolution spaces, we achieve comparable BPDs in lesser time with far fewer parameters than previous normalizing ﬂows (and non ﬂow-based approaches).
However, the power of the multi-resolution formulation is more evident at higher resolutions: we achieve state-of-the-art BPD for
ImageNet64 with signiﬁcantly fewer parameters and lower time using only one GPU.
Progressive training : Since each resolution can be trained independently, we train an MRCNF model on ImageNet by training only the ﬁnest resolution ( 128 × 128 ) conditioned on the immediate coarser ( 64 × 64 ) images, and attach that to a 3-resolution 64 × 64 model.
The resulting 4-resolution
ImageNet128 model gives a BPD of 3.31 (Table 2) with just 2.74M parameters and 59 GPU hours of total training time.
We also report the number of parameters in the models, and the time taken to train (in GPU hours).
Most previous models use multiple GPUs for training, all our models were trained on only one GPU: NVIDIA RTX 2080 Ti 11GB.
As reported in (Ghosh et al., 2020).
§ Re-implemented by us.
‘x’: Fails to train.
Blank spaces indicate unreported values.
RNODE (Finlay et al., 2020) used 4 GPUs to train on ImageNet64. [Table] [Figure]
[Table]
Parallel Multiscale
[Table]
[Table] 6. Conclusion We presented a Multi-Resolution approach to CNFs, which provides an efﬁcient framework for likelihood calculations by training on a single GPU in lesser time with a signiﬁcantly fewer parameters.
We see a marked improvement in
BPD for ImageNet64 and above.