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: The paper proposes a multi-resolution variant of continuous normalizing flows for images. They show empirical results for image sizes upto 64 x 64 where they seem to be better than regular continuous normalizing flows. The key proposed benefit seems to be in number of parameters and training times.