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: This paper proposes a DARTS-like method for searching automated normalization flow models. Instead of directly using the output ensembles, which leads to infeasible flow models, this work proposed distribution mixture to guarantee that the supernet is always a valid flow model. The upper bound of the loss function is optimized jointly with resource constraints. Experiments on small-to-medium scale datasets valid the effectiveness of the proposed method.