Invertible DenseNets with Concatenated LipSwish

We introduce Invertible Dense Networks ( i-DenseNets ) , a more parameter efficient extension of Residual Flows . The method relies on an analysis of the Lipschitz continuity of the concatenation in DenseNets , where we enforce invertibility of the network by satisfying the Lipschitz constant . Furthermore , we propose a learnable weighted concatenation , which not only improves the model performance but also indicates the importance of the concatenated weighted representation . Additionally , we introduce the Concatenated LipSwish as activation function , for which we show how to enforce the Lipschitz condition and which boosts performance . The new architecture , i-DenseNet , out-performs Residual Flow and other flow-based models on density estimation evaluated in bits per dimension , where we utilize an equal parameter budget . Moreover , we show that the proposed model outperforms Residual Flows when trained as a hybrid model where the model is both a generative and a discriminative model . 1 Introduction . Neural networks are widely used to parameterize non-linear models in supervised learning tasks such as classification . In addition , they are also utilized to build flexible density estimators of the true distribution of the observed data [ 25 , 33 ] . The resulting deep density estimators , also called deep generative models , can be further used to generate realistic-looking images that are hard to separate from real ones , detection of adversarial attacks [ 9 , 17 ] , and for hybrid modeling [ 27 ] which have the property to both predict a label ( classify ) and generate . Many deep generative models are trained by maximizing the ( log- ) likelihood function and their architectures come in different designs . For instance , causal convolutional neural networks are used to parameterize autoregressive models [ 28 , 29 ] or various neural networks can be utilized in Variational Auto-Encoders [ 19 , 32 ] . The other group of likelihood-based deep density estimators , flow-based models ( or flows ) , consist of invertible neural networks since they are used to compute the likelihood through the change of variable formula [ 31 , 37 , 36 ] . The main difference that determines an exact computation or approximation of the likelihood function for a flow-based model lies in the design of the transformation layer and tractability of the Jacobian-determinant . Many flow-based models formulate the transformation that is invertible and its Jacobian is tractable [ 3 , 6–8 , 21 , 30 , 31 , 38 ] . Recently , Behrmann et al . [ 2 ] proposed a different approach , namely , deep-residual blocks as a transformation layer . The deep-residual networks ( ResNets ) of [ 12 ] are known for their successes in supervised learning approaches . In a ResNet block , each input of the block is added to the output , which forms the input for the next block . Since ResNets are not necessarily invertible , Behrmann et al . [ 2 ] enforce the Lipschitz constant of the transformation to be smaller than 1 ( i.e. , it becomes a contraction ) that allows applying an iterative procedure to invert the network . Furthermore , Chen et al . [ 4 ] proposed Residual Flows , an improvement of i-ResNets , that uses an unbiased estimator for the logarithm of the Jacobian-determinant . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) In supervised learning , an architecture that uses fewer parameters and is even more powerful than the deep-residual network is the Densely Connected Convolution Network ( DenseNet ) , which was first presented in [ 15 ] . Contrary to a ResNet block , a DenseNet layer consists of a concatenation of the input with the output . The network showed to improve significantly in recognition tasks on benchmark datasets such as CIFAR10 , SVHN , and ImageNet , by using fewer computations and having fewer parameters than ResNets while performing at a similar level . In this work , we extend Residual Flows [ 2 , 4 ] , and use densely connected blocks ( DenseBlocks ) as a residual layer . First , we introduce invertible Dense Networks ( i-DenseNets ) , and we show that we can derive a bound on the Lipschitz constant to create an invertible flow-based model . Furthermore , we propose the Concatenated LipSwish ( CLipSwish ) as an activation function , and derive a stronger Lipschitz bound . The CLipSwish function preserves more signal than LipSwish activation functions . Finally , we demonstrate how i-DenseNets can be efficiently trained as a generative model , outperforming Residual Flows and other flow-based models under an equal parameter budget . 2 Background . Flow-based models Let us consider a vector of observable variables x 2 Rd and a vector of latent variables z 2 Rd . We define a bijective function f : Rd ! Rd that maps a latent variable to a datapoint x = f ( z ) . Since f is invertible , we define its inverse as F = f 1 . We use the change of variables formula to compute the likelihood of a datapoint x after taking the logarithm , that is : ln pX ( x ) = ln pZ ( z ) + ln | det JF ( x ) | , ( 1 ) where pZ ( z ) is a base distribution ( e.g. , the standard Gaussian ) and JF ( x ) is the Jacobian of F at x . The bijective transformation is typically constructed as a sequence of K invertible transformations , x = fK · · · f1 ( z ) , and a single transformation fk is referred to as a flow [ 31 ] . The change of variables formula allows evaluating the data in a tractable manner . Moreover , the flows are trained using the log-likelihood objective where the Jacobian-determinant compensates the change of volume of the invertible transformations . Residual flows Behrmann et al . [ 2 ] construct an invertible ResNet layer which is only constrained in Lipschitz continuity . A ResNet is defined as : F ( x ) = x + g ( x ) , where g is modeled by a ( convolutional ) neural network and F represents a ResNet layer ( see Figure 1 ( a ) ) which is in general not invertible . However , g is constructed in such a way that it satisfies the Lipschitz constant being strictly lower than 1 , Lip ( g ) < 1 , by using spectral normalization of [ 10 , 26 ] : Lip ( g ) < 1 , if ||Wi||2 < 1 , ( 2 ) where || · ||2 is the ` 2 matrix norm . Then Lip ( g ) = K < 1 and Lip ( F ) < 1+K . Only in this specific case the Banach fixed-point theorem holds and ResNet layer F has a unique inverse . As a result , the inverse can be approximated by fixed-point iterations . To estimate the log-determinant is , especially for high-dimensional spaces , computationally intractable due to expensive computations . Since ResNet blocks have a constrained Lipschitz constant , the log-likelihood estimation of Equation ( 1 ) can be transformed to a version where the logarithm of the Jacobian-determinant is cheaper to compute , tractable , and approximated with guaranteed convergence [ 2 ] : ln p ( x ) = ln p ( f ( x ) ) + tr 1X k=1 ( 1 ) k+1 k [ Jg ( x ) ] k ! , ( 3 ) where Jg ( x ) is the Jacobian of g at x that satisfies ||Jg||2 < 1 . The Skilling-Hutchinson trace estimator [ 35 , 16 ] is used to compute the trace at a lower cost than to fully compute the trace of the Jacobian . Residual Flows [ 4 ] use an improved method to estimate the power series at an even lower cost with an unbiased estimator based on `` Russian roulette '' of [ 18 ] . Intuitively , the method estimates the infinite sum of the power series by evaluating a finite amount of terms . In return , this leads to less computation of terms compared to invertible residual networks . To avoid derivative saturation , which occurs when the second derivative is zero in large regions , the LipSwish activation is proposed . 3 Invertible Dense Networks . In this section , we propose Invertible Dense Networks by using a DenseBlock as a residual layer . We show how the network can be parameterized as a flow-based model and refer to the resulting model as i-DenseNets . The code can be retrieved from : https : //github.com/yperugachidiaz/ invertible_densenets . 3.1 Dense blocks . The main component of the proposed flow-based model is a DenseBlock that is defined as a function F : Rd ! Rd with F ( x ) = x + g ( x ) , where g consists of dense layers { hi } ni=1 . Note that an important modification to make the model invertible is to output x + g ( x ) whereas a standard DenseBlock would only output g ( x ) . The function g is expressed as follows : g ( x ) = Wn+1 hn · · · h1 ( x ) , ( 4 ) where Wn+1 represents a 1⇥ 1 convolution to match the output size of Rd . A layer hi consists of two parts concatenated to each other . The upper part is a copy of the input signal . The lower part consists of the transformed input , where the transformation is a multiplication of ( convolutional ) weights Wi with the input signal , followed by a non-linearity having Lip ( ) 1 , such as ReLU , ELU , LipSwish , or tanh . As an example , a dense layer h2 can be composed as follows : h1 ( x ) = x ( W1x ) , h2 ( h1 ( x ) ) = h1 ( x ) ( W2h1 ( x ) ) . ( 5 ) In Figure 1 , we schematically outline a residual block ( Figure 1 ( a ) ) and a dense block ( Figure 1 ( b ) ) . We refer to concatenation depth as the number of dense layers in a DenseBlock and growth as the channel growth size of the transformation in the lower part . 3.2 Constraining the Lipschitz constant . If we enforce function g to satisfy Lip ( g ) < 1 , then DenseBlock F is invertible since the Banach fixed point theorem holds . As a result , the inverse can be approximated in the same manner as in [ 2 ] . To satisfy Lip ( g ) < 1 , we need to enforce Lip ( hi ) < 1 for all n layers , since Lip ( g ) Lip ( hn+1 ) · . . . · Lip ( h1 ) . Therefore , we first need to determine the Lipschitz constant for a dense layer hi . For the full derivation , see Appendix A . We know that a function f is K-Lipschitz if for all points v and w the following holds : dY ( f ( v ) , f ( w ) ) KdX ( v , w ) , ( 6 ) where we assume that the distance metrics dX = dY = d are chosen to be the ` 2-norm . Further , let two functions f1 and f2 be concatenated in h : hv = f1 ( v ) f2 ( v ) , hw = f1 ( w ) f2 ( w ) , ( 7 ) where function f1 is the upper part and f2 is the lower part . We can now find an analytical form to express a limit on K for the dense layer in the form of Equation ( 6 ) : d ( hv , hw ) 2 = d ( f1 ( v ) , f1 ( w ) ) 2 + d ( f2 ( v ) , f2 ( w ) ) 2 , d ( hv , hw ) 2 ( K21 +K22 ) d ( v , w ) 2 , ( 8 ) where we know that the Lipschitz constant of h consist of two parts , namely , Lip ( f1 ) = K1 and Lip ( f2 ) = K2 . Therefore , the Lipschitz constant of layer h can be expressed as : Lip ( h ) = q ( K21 +K 2 2 ) . ( 9 ) With spectral normalization of Equation ( 2 ) , we know that we can enforce ( convolutional ) weights Wi to be at most 1-Lipschitz . Hence , for all n dense layers we apply the spectral normalization on the lower part which locally enforces Lip ( f2 ) = K2 < 1 . Further , since we enforce each layer hi to be at most 1-Lipschitz and we start with h1 , where f1 ( x ) = x , we know that Lip ( f1 ) = 1 . Therefore , the Lipschitz constant of an entire layer can be at most Lip ( h ) < p 12 + 12 = p 2 , thus dividing by this limit enforces each layer to be at most 1-Lipschitz . 3.3 Learnable weighted concatenation We have shown that we can enforce an entire dense layer to have Lip ( hi ) < 1 by applying a spectral norm on the ( convolutional ) weights Wi and then divide the layer hi by p 2 . Although learning a weighting between the upper and lower part would barely affect a standard dense layer , it matters in this case because the layers are regularized to be 1-Lipschitz . To optimize and learn the importance of the concatenated representations , we introduce learnable parameters ⌘1 and ⌘2 for , respectively , the upper and lower part of each layer hi . Since the upper and lower part of the layer can be at most 1-Lipschitz , multiplication by these factors results in functions that are at most ⌘1-Lipschitz and ⌘2-Lipschitz . As indicated by Equation ( 9 ) , the layer is then at mostp ⌘21 + ⌘ 2 2 Lipschitz . Dividing by this factor results in a bound that is at most 1-Lipschitz . In practice , we initialize ⌘1 and ⌘2 at value 1 and during training use a softplus function to avoid them being negative . The range of the normalized parameters is between ⌘̂1 , ⌘̂2 2 [ 0 , 1 ] and can be expressed on the unit circle as shown in Figure 2 . In the special case where ⌘1 = ⌘2 , the normalized parameters are ⌘̂1 = ⌘̂2 = 12 p 2 . This case corresponds to the situation in Section 3.2 where the concatenation is not learned . An additional advantage is that the normalized ⌘̂1 and ⌘̂2 express the importance of the upper and lower signal . For example , when ⌘̂1 > ⌘̂2 , the input signal is of more importance than the transformed signal .

Summary: In the article, a new normalizing flow architecture, called invertible DenseNets (i-DenseNets), is introduced. It is an extension of residual flows with two new contributions. First, the invertible residual blocks are replaced by invertible dense blocks. Similar to residual blocks, the authors proof invertibility by computing the Lipschitz constant and guarantee that it is smaller than one by a suitable rescaling operation. Second, the LipSwish activation function is replaced by its concatenated cousin, which is rescaled as well to ensure invertibility. The method is applied to three illustrative toy datasets as well as CIFAR10 and ImageNet32, where it outperforms competing methods and sets a new SOTA for flow models with uniform dequantization. They also use an i-DenseNet to perform hybrid modelling, where is outperform the competition as well.