AutoNF: Automated Architecture Optimization of Normalizing Flows Using a Mixture Distribution Formulation

Although various flow models based on different transformations have been proposed , there still lacks a quantitative analysis of performance-cost trade-offs between different flows as well as a systematic way of constructing the best flow architecture . To tackle this challenge , we present an automated normalizing flow ( NF ) architecture search method . Our method aims to find the optimal sequence of transformation layers from a given set of unique transformations with three folds . First , a mixed distribution is formulated to enable efficient architecture optimization originally on the discrete space without violating the invertibility of the resulting NF architecture . Second , the mixture NF is optimized with an approximate upper bound which has a more preferable global minimum . Third , a block-wise alternating optimization algorithm is proposed to ensure efficient architecture optimization of deep flow models . 1 INTRODUCTION . Normalizing flow ( NF ) is a probabilistic modeling tool that has been widely used in density estimation , generative models , and random sampling . Various flow models have been proposed in recent years to improve their expressive power . Discrete flow models are either built based on elementalwise monotonical functions , named autoregressive flow or coupling layers ( Papamakarios et al. , 2017 ) , or built with transformations where the determinant of the flow can be easily calculated with matrix determinant lemma ( Rezende & Mohamed , 2015 ) . In the continuous flow family , the models are constructed by neural ODE ( Grathwohl et al. , 2019 ) . Despite the variety of flow models , there ’ s yet no perfect flow concerning the expressive power and the computation cost . The flow models with higher expressive power usually have higher computational costs in either forward and inverse pass . In contrast , flows which are fast to compute are not able to model rich distributions and are limited to simple applications . For instance , autoregressive flows ( Papamakarios et al. , 2017 ) are universal probability approximators but are D times slower to invert than forward calculation , where D is the dimension of the modeled random variable x ( Papamakarios et al. , 2021 ) . Flows based on coupling layers ( Dinh et al. , 2015 ; 2017 ; Kingma & Dhariwal , 2018 ) have an analytic one-pass inverse but are less expressive than their autoregressive counterparts . Other highly expressive NF models ( Rezende & Mohamed , 2015 ; Behrmann et al. , 2019 ) can not provide an analytic inverses and relies on numerical optimizations . For different applications , the optimal flow model can be drastically different , especially if the computation cost is taken into consideration . For generative models ( Dinh et al. , 2015 ; Kingma & Dhariwal , 2018 ) , flows with the fast forward pass are preferable since the forward transformations need to be applied to every sample from the base distribution . For density estimation ( Papamakarios et al. , 2017 ; Rippel & Adams , 2013 ) , flows with cheap inverse will prevail . For applications where flow is utilized as a co-trained kernel ( Mazoure et al. , 2020 ) , the computation cost and performance trade-off are more important , i.e. , having a fast model with relatively good performance . However , in the current body of work , the architecture designs of the flow models are all based on manual configuration and tuning . To this date , there is a lack of a systematic way that could automatically construct an optimal flow architecture with a preferred cost . In this paper , we propose AutoNF , an automated method for normalizing flow architecture optimization . AutoNF has a better performance-cost trade-off than hand-tuned SOTA flow models based on a given set of transformations . Our approach employs a mixture distribution formulation that can search a large design space of different transformations while still satisfying the invertibility requirement of normalizing flow . The proposed mixture NF is optimized via approximate upper bound which provides a better optimization landscape for finding the desired flow architecture . Besides , to deal with exponentially growing optimization complexity , we introduce a block-wise optimization method to enable efficient optimization of deep flow models . 2 RELATED WORK . Normalizing Flows : Various normalizing flow models have been proposed since the first concept in ( Tabak & Turner , 2013 ) . Current flow models can be classified into two categories : finite flows based on layer structure , and continuous flow based on neural ODE ( Grathwohl et al. , 2019 ) . The finite flow family includes flows based on elemental-wise transformation ( Papamakarios et al. , 2017 ; Kingma & Dhariwal , 2018 ) and flows whose transformations are restricted to be contractive ( Behrmann et al. , 2019 ) . In elemental-wise transformation flows , autoregressive flow and coupling layers are two major flavors and extensive work has been proposed to improve the expressive power of both flow models . In Huang et al . ( 2018 ) , the dimension-wise scalar transformation is implemented by a sigmoid neural network , which increases the expressive power at the cost of being not analytically invertible . In Durkan et al . ( 2019 ) , piecewise splines are used as drop-in replacement of affine or additive transformations ( Dinh et al. , 2015 ; 2017 ) and is the current SOTA flow model . Consequently many recent research efforts have been devoted to closing the gap of expressive power , albeit at the cost of more complex and expensive transformations . Moreover , there has been no quantitative trade-off analysis between the performance and cost among different flows . Neural Architecture Search : Many algorithms have been proposed or applied for neural architecture search . For instance , reinforcement learning ( Zoph & Le , 2017 ) , genetic algorithm ( Real et al. , 2017 ; Suganuma et al. , 2018 ; Liu et al. , 2018 ) , Monte Carlo tree search ( Negrinho & Gordon , 2017 ) or Bayesian optimization ( Kandasamy et al. , 2018 ) . However , these methods all face the challenge of optimizing on a large discrete space and can take thousand of GPU days to find a good architecture . To address this issue , DARTS ( Liu et al. , 2019 ) proposes to relax the search space from discrete to continuous and allows efficient differentiable architecture search with gradient method which could reduce the search time to a single GPU day while still producing the SOTA architecture . However , all current NAS methods focus on optimizing traditional neural network structures ( CNN , RNN ) and there has yet been any implementation on normalizing flow . Necessity for the Trade-off Between Performance and Cost : Despite various transformations proposed in the literature , there is no perfect transformation with strong expressive power and low computational cost . Autoregressive flows have better expressive power , but the inverse computation cost grows linearly with data dimension . Coupling layers ’ inverse calculation is as fast as the forward pass , but their expressive power is generally worse than autoregressive flow with the same element-wise transformation . Even in the same autoregressive flow or coupling layer family , flows with different element-wise transformations have different performance and computation costs . For instance , additive or affine coupling layers ( Dinh et al. , 2017 ; 2015 ) have very fast forward and inverse calculation with limited expressive power while the flow in ( Durkan et al. , 2019 ) are highly expressive but are more demanding on computation . In most applications , it is necessary to find the best performance while minimizing at least one specific component of the cost . Unfortunately , the current design of flow models is empirical and therefore can not ensure the optimal trade-offs . 3 METHOD . In this work , we aim to tackle the challenge of finding an optimal flow model for a given task via an automated architecture search algorithm . Assumptions : In the remaining part of this paper , without losing generality , we assume that the transformation is properly modeled such that during the training process , only forward computation is needed . Under this assumption , when the flow model is used for density modeling ( Durkan et al. , 2019 ) , the forward calculation is the dominant computation . When the flow model is used for random sampling ( Kingma & Dhariwal , 2018 ) , the inverse calculation is computationally intensive . When the flow model is utilized as a module and trained together with other components , e.g. , policy network in maximum entropy learning ( Mazoure et al. , 2020 ) , the training cost of the flow model is an important consideration . Problem Definition : Given a transformation set with m options { T 1 , T 2 , ... Tm } , the goal is to construct an optimal flow model with n layers of transformations from the set . The flow model pNF ( x ; θ ) = pT1T2 ... Tn ( x ; θ ) should minimize the KL divergence between the target distribution p∗ ( x ) and itself while minimizing its computational cost CNF . Here , θ are the parameters of the transformation in the flow model . In this paper , we use the forward KL divergence as our target loss function ( Papamakarios et al. , 2021 ) : θ∗ =argmin θ { DKL [ p∗ ( x ) || pT1T2 ... Tn ( x ; θ ) ] + λ · CNF } s.t . Ti ∈ { T 1 , T 2 , ... Tm } ( 1 ) While λ is a tuning factor capturing the relative importance of the performance-cost trade-off . Finding this optimal flow model is a discrete optimization problem with exponential complexity . To enable efficient architecture optimization , we use proposed method of relaxing the discrete search space to continuous space as suggested in Liu et al . ( 2019 ) . 3.1 MIXED FLOW ENSEMBLE . For the ith transformation layer with m options , we introduce a corresponding weight w j i for each option T j which reflects how likely the transformation will be selected . The weight is parameterized by a vector α and made continuous via softmax : wji = exp ( αji ) ∑m j=1 exp ( α j i ) ( 2 ) By applying this parameterization for each transformation layer , we can construct a mixed flow ensemble pMix ( x ; θ , α ) , where each layer in this mixed model reflects a weighted combination of the effect of all possible transformations . In this case , the architecture optimization problem is reduced to learning the weight vector for each layer and , at the end of the optimization process , weights will be binarized and the transformation with the highest weight in one layer will be selected as the final transformation . The mixed flow ensemble thus degrades to a normal flow model . The whole procedure is illustrated in Fig . 1 ( left ) . As adopted in ( Liu et al. , 2019 ) , training of the flow ensemble becomes joint optimization of the architecture parameterα and the model parameter θ over the training and validation datasets , which could be written as the following bi-level optimization problem : α∗ =argmin α DvalKL [ p ∗ ( x ) || pMix ( x ; θ∗ , α ) ] + λ · CMix ( α ) s.t . θ∗ = argmin θ DtrainKL [ p ∗ ( x ) || pMix ( x ; θ , α ) ] , ∀ T ∈ pMix , T ∈ { T 1 , T 2 , ... Tm } , ( 3 ) While the optimization problem is well defined , the key challenge is to construct the flow ensemble within the normalizing flow framework . This is different from traditional neural architecture search , which can mix various operations with no additional issue . Normalizing flow has its unique requirement for the invertibility of transformations and a preferred simple Jacobian calculation , which requires careful handling . The mixed flow ensemble pMix ( x ; θ∗ , α ) must satisfy two requirements . First , it must be a legal density function such that it can be optimized by the KL divergence formulation . Second , each transformation layer in pMix ( x ; θ∗ , α ) should represent a weighted combination of all possible transformations . Consider the ith layer in the mixed flow ensemble with input random variable xin and output random variable xout , and pxin ( xin ) and pxout ( xout ) are their corresponding density functions . This layer has m transformation options in { T 1i , T 2i , ... Tmi } and w j i is the corresponding weight for each transformation . As discussed in Assumption , we assume all transformations directly model the inverse transformation , i.e . xin = T j i ( xout ) . Two approaches can be used to construct the mixed flow ensemble . Construction by Mixed Transformations : The straight forward way of building the ith mix flow ensemble layer is to mix all transformations by weighted summation , as shown in Fig . 1 ( right-top ) . The final weighted transformation for this layer can be thus represented as : Ti ( xin ) = m∑ j=1 wji · T j i ( xout ) ( 4 ) There are two drawbacks of this formulation despite its simplicity . First , definition of normalizing flow requires the mixed transformation Ti be invertible and differentiable in order to ensure pxout ( xout ) legal density function . However , this invertibility is not guaranteed even if all candidate transformations are invertible . Second , even if the mixed transformation is invertible , there is no easy way to calculate the Jacobian determinant of this weighted summation of transformations . Meeting the requirement of invertibility and ease of calculating Jacobian determinant brings strict restrictions on the candidate transformations and prevents the optimization of flow architectures on a wider search space . As a result , the construction of the mixed flow ensemble by weighted summation of transformations is not adopted in this paper . Construction by Mixed Distributions : An alternating way is to build the mixed flow ensemble by mixing distributions . For a given transformation T ji in this ith layer , applying the transformation to the input random variable will result in a new distribution : pT ji ( xout ) = pxin ( T j i ( xout ) ) · | detJT ji ( xout ) | ( 5 ) By applying this to every transformation option in { T 1i , T 2i , ... Tmi } , we can obtain k different distributions , and it is possible to mix all the density functions together by their weighted summation , to get a mixture model as shown in eq. ( 6 ) . pTi ( xout ) = m∑ j=1 wji · pT ji ( xout ) ( 6 ) An illustration of this process is shown in Fig . 1 ( right-bottom ) . Different from the previous approach , the mixture model has a legal density function as : pTi ( xout ) . By the definition of normalizing flow , we can assume that there exists an invertible and differentiable transformation Ti , which transforms xin to xout , although the transformation itself can not be explicitly written out . For the next ( i + 1 ) th layer , the density of the mixture model will be used as the input density function pxin ( xin ) as in the previous layer . By applying this formulation for n layers , the final mixed flow ensemble can be written as : pMix ( x ; θ , a ) = mn∑ k=1 Wk · pT1T2 ... Tn ( x , θ ) = mn∑ k=1 Wk · pi ( x ; θi ) where each Wk = n∏ i=1 wi and mn∑ k Wk = 1 ( 7 ) Each wi is defined in eq . ( 2 ) and we use pk ( x ; θk ) to represent a “ normal flow architecture ” with n transformation layers . Clearly , the final mixed flow ensemble is a legal density function which is in fact , a weighted summation of all possible flow models built with n layers of transformations .

Summary: In this paper, the authors present an automated normalizing flow(NF) architecture search method. The method employs a mixture distribution formulation that can construct an optimal flow model with n layers of transformations from the transformation set. Besides, the authors introduce a block-wise optimization method to deal with exponentially growing optimization complexity. In the experiment, the authors proved the effectiveness of the optimization method which via approximate upper bound. And AutoNF has a better performance-cost trade-off than hand-tuned SOTA flow models.