Multi-Resolution Continuous Normalizing Flows

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. , 2019a ; Yu et al. , 2020 ) are growing in interest as alternatives to generative models based on Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2016 ) and Variational Autoencoders ( VAEs ) ( Kingma & Welling , 2013 ) . While GANs and VAEs have been able to produce visually impressive samples of images , they have a number of limitations . A change of variables approach facilitates the transformation of a simple base probability distribution into a more complex model distribution . Reversible generative models using this technique are attractive because they enable efficient density estimation , efficient 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 flows ( Chen et al. , 2018a ; Grathwohl et al. , 2019 ; Finlay et al. , 2020 ) , which uses an integral over continuous time dynamics to transform a base distribution into the model distribution , called Continuous Normalizing Flows ( CNF ) . This approach uses ordinary differential equations ( ODEs ) specified by a neural network , or Neural ODEs . 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 GANs or VAEs in terms of the generated image quality , it is a promising direction of research as it does not have some key shortcomings associated with GANs and VAEs . Specifically , GANs are known to suffer from mode-collapse ( Lin et al. , 2018 ) , and are notoriously difficult to train ( Arjovsky & Bottou , 2017 ) compared to feed forward networks because their adversarial loss seeks a saddle point instead of a local minimum ( Berard et al. , 2020 ) . CNFs are trained by mapping images to noise , and their reversible architecture allows images to be generated by going in reverse , from noise to images . This leads to fewer issues related to mode collapse , since any input example in the dataset can be recovered from the flow using the reverse of the transformation learned during training . VAEs only provide a lower bound on the marginal likelihood whereas CNFs provide exact likelihoods . Despite the many advantages of reversible generative models built with CNFs , quantitatively such methods still do not match the widely used Fréchet Inception Distance ( FID ) scores of GANs or VAEs . However their other advantages motivate us to explore them further . Furthermore , state-of-the art GANs and VAEs exploit the multi-resolution properties of images , and recent top-performing methods also inject noise at each resolution ( Brock et al. , 2019 ; Shaham et al. , 2019 ; Karras et al. , 2020 ; Vahdat & Kautz , 2020 ) . While shaping noise is fundamental to normalizing flows , only recently have normalizing flows exploited the multi-resolution properties of images . For example , WaveletFlow ( Yu et al. , 2020 ) splits an image into multiple resolutions using the Discrete Wavelet Transform , and models the average image at each resolution using a normalizing flow . While this method has advantages , it suffers from many issues such as high parameter count and long training time . In this work , we consider a non-trivial multi-resolution approach to continuous normalizing flows , which fixes many of these issues . A high-level view of our approach is shown in Figure 1 . Our main contributions are : 1 . We propose a multi-resolution transformation that does not add cost in terms of likelihood . 2 . We introduce Multi-Resolution Continuous Normalizing Flows ( MRCNF ) . 3 . We achieve comparable Bits-per-dimension ( BPD ) ( negative log likelihood per pixel ) on image datasets using fewer model parameters and significantly less training time with only one GPU . 2 BACKGROUND . 2.1 NORMALIZING FLOWS . Normalizing flows ( 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 flow 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 : log p ( x ) = log ∣∣∣∣det dgdx ∣∣∣∣+ log p ( z ) ( 1 ) The first term on the right ( log determinant of the Jacobian ) is often intractable , however , previous works on normalizing flows have found ways to estimate this efficiently . The second term , log p ( z ) , is computed as the log probability of z under a known noise distribution , typically the standard Gaussian N ( 0 , I ) . 2.2 WAVELET FLOW ( YU ET AL. , 2020 ) . WaveletFlow decomposes an image using the Discrete Wavelet Transformation , and maps the average image at each resolution to noise using a normalizing flow . The WaveletFlow model builds on the Glow ( Kingma & Dhariwal , 2018 ) architecture , with a high parameter count and a long training time . It uses an orthogonal transformation , which does not preserve the range of the image , and adds a constant term to the log likelihood at each resolution . We aim to fix these issues using our MRCNF . 2.3 CONTINUOUS NORMALIZING FLOWS . Continuous Normalizing Flows ( CNF ) ( Chen et al. , 2018a ; Grathwohl et al. , 2019 ; Finlay et al. , 2020 ) are a variant of normalizing flows that operate in the continuous domain . A CNF creates a geometric flow between the input and target ( noise ) distributions , by assuming that the state transition is governed by an Ordinary Differential Equation ( ODE ) . It further assumes that the differential function is parameterized by a neural network , this model is called a Neural ODE ( Chen et al. , 2018a ) . Suppose CNF g transforms its state v ( t ) using a Neural ODE , with neural network f defining the differential . Here , v ( t0 ) = x is , say , an image , and at the final time step v ( t1 ) = z is a sample from a known noise distribution . dv ( t ) dt = f ( v ( t ) , t ) =⇒ v ( t1 ) = g ( v ( t0 ) ) = v ( t0 ) + ∫ t1 t0 f ( v ( t ) , t ) dt ( 2 ) This integration is typically performed by an ODE solver . Since this integration can be run backwards as well to obtain the same v ( t0 ) from v ( t1 ) , a CNF is a reversible model . Equation 1 can be used to compute the change in log-probability induced by the CNF . However , Chen et al . ( 2018a ) and Grathwohl et al . ( 2019 ) proposed a more efficient variant in the context of CNFs , called the instantaneous change-of-variables formula : ∂ log p ( v ( t ) ) ∂t = −Tr ( ∂f ∂v ( t ) ) =⇒ ∆ log pv ( t0 ) →v ( t1 ) = − ∫ t1 t0 Tr ( ∂f ∂v ( t ) ) dt ( 3 ) Hence , the change in log-probability of the state of the Neural ODE i.e . ∆ log pv is expressed as another differential equation . The ODE solver now solves both differential equations eq . ( 2 ) and eq . ( 3 ) by augmenting the original state with the above . Thus , a CNF provides both the final state v ( t1 ) as well as the change in log probability ∆ log pv ( t0 ) →v ( t1 ) together . 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 the images under the model : z = g ( x ) ; log p ( x ) = ∆ log px→z + log p ( z ) ( 4 ) where x is an image , z and ∆ log px→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 ) . 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 , using 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 ; Marr , 2010 ; Witkin , 1987 ; Burt & Adelson , 1983 ; Mallat , 1989 ; Lindeberg , 1990 ) . This implies that much of the content of an image at a resolution is a composition of low-level information captured at coarser resolutions , and some high-level information not present in the coarser images . We take advantage of this property by first decomposing an image in resolution space i.e . by expressing it as a series of S images at decreasing resolutions : x → ( x1 , x2 , . . . , xS ) , where x1 = x is the finest image , xS is the coarsest , and every xs+1 is the average image of xs . This called an image pyramid , or a Gaussian Pyramid if the upsampling-downsampling operations include a Gaussian filter ( Burt , 1981 ; Burt & Adelson , 1983 ; Adelson et al. , 1984 ; Witkin , 1987 ; Lindeberg , 1990 ) . In this work , we obtain a coarser image simply by averaging pixels in every 2×2 patch , thereby halving the width and height . However , this representation is redundant since much of the information in x1 is contained in xs > 1 . Instead , we express x as a series of high-level information ys not present in the immediate coarser images xs+1 , and a final coarse image xS : x→ ( y1 , x2 ) → ( y1 , y2 , x3 ) → · · · → ( y1 , y2 , . . . , yS−1 , xS ) ( 5 ) Our overall method is to map these S terms to S noise samples using S CNFs . 3.2 DEFINING THE HIGH-LEVEL INFORMATION ys We choose to design a linear transformation with the following properties : 1 ) invertible i.e . it should be possible to deterministically obtain xs from ys and xs+1 , and vice versa ; 2 ) volume preserving i.e . determinant is 1 , change in log-likelihood is 0 ; 3 ) angle preserving ; and 4 ) range preserving ( under the notion of the maximum principle ( Varga , 1966 ) ) . Consider the simplest case of 2 resolutions where x1 is a 2×2 image with pixel values x1 , x2 , x3 , x4 , and x2 is a 1×1 image with pixel value x̄ = 14 ( x1 + x2 + x3 + x4 ) . We require three values ( y1 , y2 , y3 ) = y1 that contain information not present in x2 , such that x1 is obtained when y1 and x2 are combined . This could be viewed as a problem of finding a matrix M such that : [ x1 , x2 , x3 , x4 ] > = M [ y1 , y2 , y3 , x̄ ] > . We fix the last column of M as [ 1 , 1 , 1 , 1 ] > , since every pixel value in x1 depends on x̄ . Finding the rest of the parameters can be viewed as requiring four 3D vectors that are spaced such that they do not degenerate the number of dimensions of their span . These can be considered as the four corners of a tetrahedron in 3D space , under any configuration ( rotated in 3D space ) , and any scaling of the vectors ( see Figure 2 ) . Out of the many possibilities for this tetrahedron , we could choose the matrix that performs the Discrete Haar Wavelet Transform ( Mallat , 1989 ; Mallat & Peyré , 2009 ) : x1x2x3 x4 = 1 2 1 2 1 2 1 1 2 − 1 2 − 1 2 1 − 12 1 2 − 1 2 1 − 12 − 1 2 1 2 1 y1y2y3 x̄ ⇐⇒ y1y2y3 x̄ = 1 2 1 2 − 1 2 − 1 2 1 2 − 1 2 1 2 − 1 2 1 2 − 1 2 − 1 2 1 2 1 4 1 4 1 4 1 4 x1x2x3 x4 ( 6 ) However , this has log ∣∣det ( M−1 ) ∣∣ = log ( 1/2 ) ( eq . ( 6 ) ) , and is therefore not volume preserving . Other simple scaling of eq . ( 6 ) has been used in the past , for example multiplying the last row of eq . ( 6 ) by 2 , yielding an orthogonal transformation , such as in WaveletFlow ( Yu et al. , 2020 ) . However , this transformation neither preserves the volume i.e . the log determinant is not 0 , nor the maximum i.e . the range of xs changes . We wish to find a transformation M where : one of the results is the average of the inputs , x̄ ; it is unit determinant ; the columns are orthogonal ; and it preserves the range of x̄ . Fortunately such a matrix exists – although we have not seen it discussed in prior literature . It can be seen as a variant of the Discrete Haar Wavelet Transformation matrix that is unimodular , i.e . has a determinant of 1 ( and is therefore volume preserving ) , while also preserving the range of the images for the input and its average : x1x2x3 x4 = 1 a c c c ac −c −c a−c c −c a −c −c c a y1y2y3 x̄ ⇐⇒ y1y2y3 x̄ = c −1 c−1 −c−1 −c−1 c−1 −c−1 c−1 −c−1 c−1 −c−1 −c−1 c−1 a−1 a−1 a−1 a−1 x1x2x3 x4 ( 7 ) where c = 22/3 , a = 4 . Hence , log ∣∣det ( M−1 ) ∣∣ = 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 xs+1 with the three corresponding pixels in ys to make the 2×2 patch at that location in xs using eq . ( 7 ) : xs = M ( ys , xs+1 ) ⇐⇒ ys , xs+1 = M−1 ( xs ) ( 8 ) Equation 1 can be used to compute the change in log likelihood from this transformation xs → ( ys , xs+1 ) : log p ( xs ) = ∆ log pxs→ ( ys , xs+1 ) + log p ( ys , xs+1 ) = log ∣∣det ( M−1 ) ∣∣+ log p ( ys | xs+1 ) + log p ( xs+1 ) ( 9 ) where log ∣∣det ( M−1 ) ∣∣ = dims ( xs+1 ) log ( 1/2 ) in the case of eq . ( 6 ) , where “ dims ” is the number of pixels times the number of channels ( typically 3 ) in the image , and log ∣∣det ( M−1 ) ∣∣ = 0 for eq . ( 7 ) . Although a model could learn to compensate for the additional log-likelihood due to eq . ( 6 ) , we find that using eq . ( 7 ) improves log likelihood values and decreases training time .

Summary: This paper proposes a new architecture for continuous normalizing flows that explicitly models images at multiple resolutions. The authors achieve this by learning an unconditional distribution of coarse images and then learning conditional distributions at progressively finer resolutions, each with a continuous normalizing flow. In contrast to the typical “squeeze” layers used in normalizing flows, where the noise is only passed at the first layer, the proposed model injects noise at several resolutions (which has already proved successful with other generative models). The authors introduce a new (invertible) transformation to go from a coarse image to a finer image which preserves the range of the image while having unit determinant (implying the transformation leaves the log likelihood unchanged). The whole model is then trained via maximum likelihood.  The authors evaluate their method on standard image datasets and perform several ablations to test the contributions of their model.  The main contributions of the paper in my eyes are then: - The introduction of a new architecture and layer for learning multi resolution normalizing flows - Demonstration that the proposed model can be trained with limited computational resources even on fairly large scale datasets (such as Imagenet128) - Ablation studies evaluating the improvements from each aspect of the proposed model