Likelihood Contribution based Multi-scale Architecture for Generative Flows

1 INTRODUCTION . Deep Generative Modeling aims to learn the embedded distributions and representations in input ( especially unlabelled ) data , requiring no/minimal human labelling effort . Learning without knowledge of labels ( unsupervised learning ) is of increasing importance because of the abundance of unlabelled data and the rich inherent patterns they posses . The representations learnt can then be utilized in a number of downstream tasks such as semi-supervised learning ( Kingma et al. , 2014 ; Odena , 2016 ) , synthetic data augmentation and adversarial training ( Cisse et al. , 2017 ) , text analysis and model based control etc . The repository of deep generative modeling majorly includes Likelihood based models such as autoregressive models ( Oord et al. , 2016b ; Graves , 2013 ) , latent variable models ( Kingma & Welling , 2013 ) , flow based models ( Dinh et al. , 2014 ; 2016 ; Kingma & Dhariwal , 2018 ) and implicit models such as generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) . Autoregressive models ( Salimans et al. , 2017 ; Oord et al. , 2016b ; a ; Chen et al. , 2017 ) achieve exceptional log-likelihood score on many standard datasets , indicative of their power to model the inherent distribution . But , they suffer from slow sampling process , making them unacceptable to adopt in real world applications . Latent variable models such as variational autoencoders ( Kingma & Welling , 2013 ) tend to better capture the global feature representation in data , but do not offer an exact density estimate . Implicit generative models such as GANs which optimize a generator and a discriminator in a min-max fashion have recently become popular for their ability to synthesize realistic data ( Karras et al. , 2018 ; Engel et al. , 2019 ) . But , GANs do not offer a latent space suitable for further downstream tasks , nor do they perform density estimation . Flow based generative models ( Dinh et al. , 2016 ; Kingma & Dhariwal , 2018 ) perform exact density estimation with fast inference and sampling , due to their parallelizability . They also provide an information rich latent space suitable for many applications . However , the dimension of latent space for flow based generative models is same as the high-dimensional input space , by virtue of bijectivity nature of flows . This poses a bottleneck for flow models to scale with increasing input dimensions due to computational complexity . An effective solution to the above challenge is a multi-scale architecture , introduced by Dinh et al . ( 2016 ) , which performs iterative early gaussianization of a part of the total dimensions at regular intervals of flow layers . This not only makes the model computational and memory efficient but also aids in distributing the loss function throughout the network for better training . Many prior works including Kingma & Dhariwal ( 2018 ) ; Atanov et al . ( 2019 ) ; Durkan et al . ( 2019 ) ; Kumar et al . ( 2019 ) implement multi-scale architecture in their flow models , but use static masking methods for factorization of dimensions . We propose a multi-scale architecture which performs data dependent factorization to decide which dimensions should pass through more flow layers . For the decision making , we introduce a heuristic based on the amount of total log-likelihood contributed by each dimension , which in turn signifies their individual importance . We lay the ground rules for quantitative estimation and qualitative sampling to be satisfied by an ideal factorization method for a multi-scale architecture . Since in the proposed architecture , the heuristic is obtained as part of the flow training process , it can be universally applied to generic flow models . We present such implementations for flow models based on affine/additive coupling and ordinary differential equation ( ODE ) and achieve quantitative and qualitative improvements . We also perform ablation studies to confirm the novelty of our method . Summing up , the contributions of our research are , Contributions : 1 . A log-determinant based heuristic which entails the contribution by each dimensions towards the total log-likelihood in a multi-scale architecture . 2 . A multi-scale architecture based on the above heuristic performing data-dependent splitting of dimensions , implemented for several classes of flow models . 3 . Quantitative and qualitative analysis of above implementations and an ablation study To the best of our knowledge , we are the first to propose a data-dependent splitting of dimensions in a multi-scale architecture . 2 BACKGROUND . In this section , we illustrate the functioning of flow based generative models and the multiscale architecture as introduced by Dinh et al . ( 2016 ) . 2.1 FLOW-BASED GENERATIVE MODELS . Let x be a high-dimensional random vector with unknown true distribution p ( x ) . The following formulation is directly applicable to continous data , and with some pre-processing steps such as dequantization ( Uria et al. , 2013 ; Salimans et al. , 2017 ; Ho et al. , 2019 ) to discrete data . Let z be the latent variable with a known standard distribution p ( z ) , such as a standard multivariate gaussian . Using an i.i.d . dataset D , the target is to model pθ ( x ) with parameters θ . A flow , fθ is defined to be an invertible transformation that maps observed data x to the latent variable z . A flow is invertible , so the inverse function T maps z to x , i.e . z = fθ ( x ) = T −1 ( x ) and x = T ( z ) = f−1θ ( z ) ( 1 ) The log-likelihood can be expressed as , log ( pθ ( x ) ) = log ( p ( z ) ) + log ( ∣∣∣∣det ( ∂fθ ( x ) ∂xT ) ∣∣∣∣ ) ( 2 ) where ∂fθ ( x ) ∂xT is the Jacobian of fθ at x . The invertibile nature of flow allows it to be capable of being composed of other flows of compatible dimensions . In practice , flows are constructed by composing a series of component flows . Let the flow fθ be composed of K component flows , i.e . fθ = fθK ◦ fθK−1 ◦ · · · ◦ fθ1 and the intermediate variables be denoted by yK , yK−1 , · · · , y0 = x . Then the log-likelihood of the composed flow is , log ( pθ ( x ) ) = log ( p ( z ) ) + log ( ∣∣∣∣det ( ∂ ( fθK ◦ fθK−1 ◦ · · · ◦ fθ1 ( x ) ) ∂xT ) ∣∣∣∣ ) ( 3 ) = log ( p ( z ) ) ︸ ︷︷ ︸ Log-latent density + K∑ i=1 log |det ( ∂yi/∂yTi−1 ) |︸ ︷︷ ︸ Log-det ( 4 ) which follows from the fact that det ( A ·B ) = det ( A ) ·det ( B ) . In our work , we refer the first term in Equation 4 as log-latent-density and the second term as log-determinant ( log-det ) . The reverse path , from z to x can be written as a composition of inverse flows , x = f−1θ ( z ) = f −1 θ1 ◦ f−1θ2 ◦ · · · ◦ f −1 θK ( z ) . Confirming with the properties of a flow as mentioned above , different types of flows can be constructed ( Kingma & Dhariwal , 2018 ; Dinh et al. , 2016 ; 2014 ; Behrmann et al. , 2018 ) . 2.2 MULTI-SCALE ARCHITECTURE . Multi-scale architecture is a design choice for latent space dimensionality reduction of flow models , in which part of the dimensions are factored out/early gaussianized at regular intervals , and the other part is exposed to more flow layers . The process is called dimension factorization . In the problem setting as introduced in Section 2.1 , the factoring operation can be mathematically expressed as , y0 = x ( 5 ) zl+1 , yl+1 = fθl+1 ( yl ) , l ∈ { 0 , 1 , · · · , K − 2 } ( 6 ) zK = fθK ( yK−1 ) ( 7 ) z = ( z1 , z2 , · · · , zK ) ( 8 ) The factoring of dimensions at early layers has the benefit of distributing the loss function throughout the network ( Dinh et al. , 2016 ) and optimizing the amount of computation and memory used by the model . We consider the multi-scale architecture for flow based generative models as introduced by Dinh et al . ( 2016 ) ( and later used by state-of-the-art flow models such as Glow ( Kingma & Dhariwal , 2018 ) ) as the base of our research work . 3 LIKELIHOOD CONTRIBUTION BASED MULTISCALE ARCHITECTURE . In a multi-scale architecture , it is apparent that the network will better learn the distribution of variables getting exposed to more layers of flow as compared to the ones which get factored at a finer scale ( earlier layer ) . The method of dimension splitting proposed by prior works such as ( Dinh et al. , 2016 ; Kingma & Dhariwal , 2018 ; Behrmann et al. , 2018 ) are static in nature and do not distinguish between importance of different dimensions . In this section , we introduce a heuristic to estimate the contribution of each dimension towards the total log-likelihood , and introduce a method which can use the heuristic to decide the dimensions to be factored at an earlier layer , eventually achieving preferrential splitting in multiscale architecture . Our approach builds an efficient multiscale architecture which factors the dimensions at each flow layer in a way such that the local variance in the input space is well captured as the flow progresses and the log-likelihood is maximized . We also describe how our multi-scale architecture can be implemented over several standard flow models . Recall from Equation 4 that the log-likelihood is composed of two terms , the log-latent-density term and the log-det term . The log-latent-density term depends on the choice of latent distribution whereas the log-det term depends on the modeling of the flow layers . So , careful design of flow layers can lead to maximized log-determinant , eventually maximizing the likelihood . The total log-det term is nothing but the sum of log-det terms contributed by each dimension . Let the dimension of the input space x be s× s× c , where s is the image height/width and c is the number of channels for image inputs . For the following formulation , let us assume no dimensions were gaussianized early so that we have access to log-det term for all dimensions at each flow layer , and the dimension at all intermediate layer remains same ( i.e . s× s× c ) . We apply a flow ( fθ ) with K component flows to x , z pair , so that z = fθ ( x ) = fθK ◦ fθK−1 ◦ · · · ◦ fθ1 ( x ) . The intermediate variables are denoted by yK , yK−1 , · · · , y0 with yK = z ( since no early gaussianization was done ) and y0 = x . The log-det term at layer l , L ( l ) d , is given by , [ L ( l ) d ] scaler = l∑ i=1 log |det ( ∂yi/∂yTi−1 ) | ( 9 ) The log-det of the jacobian term encompasses contribution by all the s × s × c dimensions . We decompose it to obtain the individual contribution by variables ( dimensions ) towards the total log-det ( ∼ total log-likelihood ) . The log-det term can be viewed ( with slight abuse of notations ) as a s×s× c [ L ( l ) d ] s×s×c representing log-det of variables in a flow layer . ( On right ) It is squeezed to s 2 × s 2 × 4c with local max and min pooling operation . The black ( respectively white ) marked pixels represent dimensions having more ( respectively less ) log-det locally . tensor corresponding to each of the dimensions , summed over the flow layers till l , [ L ( l ) d ] s×s×c = l∑ i=1 [ d ( α , β , γ ) i−1 ] s×s×c , where α , β ∈ { 0 , · · · , s } and γ ∈ { 0 , · · · , c } ( 10 ) s.t . ∑ α , β , γ d ( α , β , γ ) i−1 = log |det ( ∂yi/∂y T i−1 ) | ( 11 ) The entries in [ L ( l ) d ] s×s×c having higher value correspond to the variables which contribute more towards the total log-likelihood , hence are more valuable for better flow formulation . So , we can use the likelihood contribution ( in the form of log-det term ) by each dimension as a heuristic for deciding which variables should be gaussianized early in a multi-scale architecture . Ideally , at each flow layer , the variables with more log-det term should be exposed to more layer of flow and the ones having less log-det term should be factored out . In this manner , selectively more power can be provided to variables which capture meaningful representation ( and are more valuable from log-det perspective ) to be expressive by being exposed to multiple flow layers . This formulation leads to enhanced density estimation performance . Additionally , for many datasets such as images , the spatial nature should be taken into account while deciding dimensions for early gaussianization . Summarily , at every flow layer , an ideal factorization method should , 1 . ( Quantitative ) For efficient density estimation : Early gaussianize the variables having less log-det and expose the ones having more log-det to more flow layers 2 . ( Qualitative ) For qualitative reconstruction : Capture the local variance over the flow layers , i.e . the dimensions being exposed to more flow layers should contain representative pixel variables from throughout the whole image . Keeping the above requirements in mind , variants of hybrid techniques for factorization can be implemented for different types of flow models which involve a multi-scale architecture , to improve their density estimation and qualitative performance . The key requirement is availability of log-det contributions per dimension , which can be fulfilled by decomposition of the log-det of the jacobian . We refer to the method as Likelihood Contribution based Multi-scale Architecture ( LCMA ) . The steps of LCMA implementation for flow models is summarized in Algorithm 1 . Note that in step 2 of dimension factorization phase in algorithm 1 , we group the dimensions having more/less log-det locally and then perform splitting . This preserves the local spatial variation of the image in both parts of the factorization , leveraging both enhanced density estimation as well as qualitative reconstruction . Another important observation is since the factorization of dimensions does not occur during the training time , and before the actual training starts , the decision of dimensions which get factored at each flow layer is fixed , the change of variables formula can be applied . This allows the use of non-invertible operations ( e.g . max and min pooling ) for efficient factorization with log-det heuristic . Step 1 of dimension factorization phase requires computation of individual contribution of dimensions ( [ L ( l ) d ] s×s×c ) towards the total log-likelihood , which can vary depending on the original design of flow Algorithm 1 : LCMA implementation for generative flow models Pre-Training Phase : Pre-train a network with no multiscale architecture ( no dimensionality reduction ) to obtain the log-det term at every flow layer . Dimension Factorization : In this phase , the dimensions to be factored at each flow layer is decided based on the log-det term at that layer 1 . The individual contribution of dimensions towards likelihood ( [ L ( l ) d ] s×s×c ) is computed specifically for corresponding flow models ( Refer Section 3.1 and Section 3.2 ) . 2 . Convert [ L ( l ) d ] s×s×c into a s 2 × s 2 × 4c shaped tensor using local max and min-pooling ( = −max-pooling ( −input ) ) operations ( Figure 1 ) at each flow layer . 3 . Among the 4c channels , one half contains the dimensions having more log-det term compared with its neighbourhood pixel ( Black marked in Fig . 1 ) , while the other half contains the dimensions having less log-det ( White marked in Fig . 1 ) . 4 . Split the tensor along the channel dimension to two parts . 5 . Forward the corresponding dimensions contributing more towards likelihood into more flow layers and early gaussianize the ones contributing less . 6 . Repeat steps 1-5 for all the layers with dimensions passed to that layer till the latent space . Training Phase : The decision of dimensions to be factored at each layer as performed in previous step remains fixed . Finally , the flow model with proposed LCMA is trained . models . Some flow models offer direct decomposition of jacobian into per-dimension components , whereas for others , an indirect estimation method has to be adopted . We now describe such methods to obtain such individual likelihood contribution of dimensions for flow models based on affine coupling ( RealNVP ( Dinh et al. , 2016 ) and Glow ( Kingma & Dhariwal , 2018 ) ) , and flow models involving ordinary differential equation ( ODE ) based density estimators ( i-ResNet ( Behrmann et al. , 2018 ) ) , all of which involve a multiscale architecture .

Summary: This paper propose a heuristic algorithm for deciding which random variables to be Gaussianized early in flow-based generative models. The proposed algorithm involves first training a flow without multi-scale training, for example, 32*32*c  - 32*32*c - 32*32*c. Then, it computes the logdet term for each variable at each layer. It then spatially partition the first flow block by two halves of shape 16*16*2c based on max-pooling the logdet term. Then it recursively Gaussianize one half, and partition the other half as 8*8*4c, still using the pre-computed logdet tensors (Ld in the paper). After partitioning, they train a multi-scale model with the learned partition.