Densely connected normalizing flows

1 Introduction . One of the main tasks of modern artificial intelligence is to generate images , audio waveforms , and natural-language symbols . To achieve the desired goal , the current state of the art uses deep compositions of non-linear transformations [ 1 , 2 ] known as deep generative models [ 3 , 4 , 5 , 6 , 7 ] . Formally , deep generative models estimate an unknown data distribution pD given by a set of i.i.d . samples D = { x1 , ... , xn } . The data distribution is approximated with a model distribution pθ defined by the architecture of the model and a set of parameters θ . While the architecture is usually handcrafted , the set of parameters θ is obtained by optimizing the likelihood across the training distribution pD : θ∗ = argmin θ∈Θ Ex∼pD [ − ln pθ ( x ) ] . ( 1 ) Properties of the model ( e.g . efficient sampling , ability to evaluate likelihood etc . ) directly depend on the definition of pθ ( x ) , or decision to avoid it . Early approaches consider unnormalized distribution [ 3 ] which usually requires MCMC-based sample generation [ 8 , 9 , 10 ] with long mixing times . Alternatively , the distribution can be autoregressively factorized [ 7 , 11 ] , which allows likelihood estimation and powerful but slow sample generation . VAEs [ 4 ] use a factorized variational approximation of the latent representation , which allows to learn an autoencoder by optimizing a lower bound of the likelihood . Diffussion models [ 12 , 13 , 14 ] learn to reverse a diffusion process , which is a fixed Markov chain that gradually adds noise to the data in the opposite direction of sampling until the signal is destroyed . Generative adversarial networks [ 5 ] mimic the dataset samples by competing in a minimax game . This allows to efficiently produce high quality samples [ 15 ] , which however often do not span the entire training distribution support [ 16 ] . Additionally , the inability to `` invert '' the generation process in any meaningful way implies inability to evaluate the likelihood . Contrary to previous approaches , normalizing flows [ 6 , 17 , 18 ] model the likelihood using a bijective mapping to a predefined latent distribution p ( z ) , typically a multivariate Gaussian . Given the bijection 1Code available at : https : //github.com/matejgrcic/DenseFlow 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . fθ , the likelihood is defined using the change of variables formula : pθ ( x ) = p ( z ) ∣∣∣∣det ∂z∂x ∣∣∣∣ , z = fθ ( x ) . ( 2 ) This approach requires computation of the Jacobian determinant ( det ∂z∂x ) . Therefore , during the construction of bijective transformations , a great emphasis is placed on tractable determinant computation and efficient inverse computation [ 18 , 19 ] . Due to these constraints , invertible transformations require more parameters to achieve a similar capacity compared to standard NN building blocks [ 20 ] . Still , modeling pθ ( x ) using bijective formulation enables exact likelihood evaluation and efficient sample generation , which makes this approach convenient for various downstream tasks [ 21 , 22 , 23 ] . The bijective formulation ( 2 ) implies that the input and the latent representation have the same dimensionality . Typically , convolutional units of normalizing-flow approaches [ 18 ] internally inflate the dimensionality of the input , extract useful features , and then compress them back to the original dimensionality . Unfortunately , the capacity of such transformations is limited by input dimensionality [ 24 ] . This issue can be addressed by expressing the model as a sequence of bijective transformations [ 18 ] . However , increasing the depth alone is a suboptimal approach to improve capacity of a deep model [ 25 ] . Recent works propose to widen the flow by increasing the input dimensionality [ 24 , 26 ] . We propose an effective development of that idea which further improves the performance while relaxing computational requirements . We increase the expressiveness of normalizing flows by incremental augmentation of intermediate latent representations with Gaussian noise . The proposed cross-unit coupling applies an affine transformation to the noise , where the scaling and translation are computed from a set of previous intermediate representations . In addition , we improve intra-module coupling by proposing a transformation which fuses the global spatial context with local correlations . The proposed imageoriented architecture improves expressiveness and computational efficiency . Our models set the new state-of-the-art result in likelihood evaluation on ImageNet32 and ImageNet64 . 2 Densely connected normalizing flows . We present a recursive view on normalizing flows and propose improvements based on incremental augmentation of latent representations , and densely connected coupling modules paired with selfattention . The improved framework is then used to develop an image-oriented architecture , which we evaluate in the experimental section . 2.1 Normalizing flows with cross-unit coupling . Normalizing flows ( NF ) achieve their expressiveness by stacking multiple invertible transformations [ 18 ] . We illustrate this with the scheme ( 3 ) where each two consecutive latent variables zi−1 and zi are connected via a dedicated flow unit fi . Each flow unit fi is a bijective transformation with parameters θi which we omit to keep notation uncluttered . The variable z0 is typically the input x drawn from the data distribution pD ( x ) . z0 f1←→ z1 f2←→ z2 f3←→ · · · fi−1←→ zi fi←→ · · · fK←→ zK , zK ∼ N ( 0 , I ) . ( 3 ) Following the change of variables formula , log likelihoods of consecutive random variables zi and zi+1 can be related through the Jacobian of the corresponding transformation Jfi+1 [ 18 ] : ln p ( zi ) = ln p ( zi+1 ) + ln |detJfi+1 | . ( 4 ) This relation can be seen as a recursion . The term ln p ( zi+1 ) can be recursively replaced either with another instance of ( 4 ) or evaluated under the latent distribution , which marks the termination step . This setup is characteristic for most contemporary architectures [ 17 , 18 , 19 , 27 ] . The standard NF formulation can be expanded by augmenting the input by a noise variable ei [ 24 , 26 ] . The noise ei subjects to some known distribution p∗ ( ei ) , e.g . a multivariate Gaussian . We further improve this approach by incrementally concatenating noise to each intermediate latent representation zi . A tractable formulation of this idea can be obtained by computing the lower bound of the likelihood p ( zi ) through Monte Carlo sampling of ei : ln p ( zi ) ≥ Eei∼p∗ ( e ) [ ln p ( zi , ei ) − ln p ∗ ( ei ) ] . ( 5 ) The learned joint distribution p ( zi , ei ) approximates the product of the target distributions p∗ ( zi ) and p∗ ( ei ) , which is explained in more detail in Appendix D. We transform the introduced noise ei with element-wise affine transformation . Parameters of this transformation are computed by a learned non-linear transformation gi ( z < i ) of previous representations z < i = [ z0 , ... , zi−1 ] . The resulting layer hi can be defined as : z ( aug ) i = hi ( zi , ei , z < i ) = [ zi , σ ei + µ ] , ( µ , σ ) = gi ( z < i ) . ( 6 ) Square brackets [ · , · ] denote concatenation along the features dimension . In order to compute the likelihood for ( zi , ei ) , we need the determinant of the jacobian ∂z ( aug ) i ∂ [ zi , ei ] = [ I 0 0 diag ( σ ) ] . ( 7 ) Now we can express p ( zi , ei ) in terms of p ( z ( aug ) i ) according to ( 4 ) : ln p ( zi , ei ) = ln p ( z ( aug ) i ) + ln |det diag ( σ ) | . ( 8 ) We join equations ( 5 ) and ( 8 ) into a single step : ln p ( zi ) ≥ Eei∼p∗ ( ei ) [ ln p ( z ( aug ) i ) − ln p ∗ ( ei ) + ln |det diag ( σ ) | ] . ( 9 ) We refer the transformation hi as cross-unit coupling since it acts as an affine coupling layer [ 17 ] over a group of previous invertible units . The latent part of the input tensor is propagated without change , while the noise part is linearly transformed . The noise transformation can be viewed as reparametrization of the distribution from which we sample the noise [ 4 ] . Note that we can conveniently recover zi from z ( aug ) i by removing the noise dimensions . This step is performed during model sampling . Fig . 1 compares the standard normalizing flow ( a ) normalizing flow with input augmentation [ 24 ] ( b ) and the proposed densely connected incremental augmentation with cross-unit coupling ( c ) . Each flow unit fDFi consists of several invertible modules mi , j and cross-unit coupling hi . The main novelty of our architecture is that each flow unit fDFi+1 increases the dimensionality with respect to its predecessor fDFi . Cross-unit coupling hi augments the latent variable zi with affinely transformed noise ei . Parameters of the affine noise transformation are obtained by an arbitrary function gi which accepts all previous variables z < i . Note that reversing the direction does not require evaluating gi since we are only interested in the value of zi . For further clarification , we show the likelihood computation for the extended framework . Example 1 ( Likelihood computation ) Let m1 and m2 be the bijective mappings from z0 to z1 and z ( aug ) 1 to z2 , respectively . Let h1 be the cross-unit coupling from z1 to z ( aug ) 1 , z ( aug ) 1 = [ z1 , σ e1+µ ] . Assume σ and µ are computed by any non-invertible neural network g1 . The network accepts z0 as the input . We calculate log likelihood of the input z0 according to the following sequence of equations : [ transformation , cross-unit coupling , transformation , termination ] . ln p ( z0 ) = ln p ( z1 ) + ln |det Jf1 | , ( 10 ) ln p ( z1 ) ≥ Ee1∼p∗ ( e1 ) [ ln p ( z ( aug ) 1 ) − ln p ( e1 ) + ln |det diag ( σ ) | ] , ( σ , µ ) = g1 ( z0 ) , ( 11 ) ln p ( z ( aug ) 1 ) = ln p ( z2 ) + ln |det Jf2 | , ( 12 ) ln p ( z2 ) = lnN ( z2 ; 0 , I ) . ( 13 ) We approximate the expectation using MC sampling with a single sample during training and a few hundreds of samples during evaluation to reduce the variance of the likelihood . Note however that our architecture generates samples with a single pass since the inverse does not require MC sampling . We repeatedly apply the cross-unit coupling hi throughout the architecture to achieve incremental augmentation of intermediate latent representations . Consequently , the data distribution is modeled in a latent space of higher dimensionality than the input space [ 24 , 26 ] . This enables better alignment of the final latent representation with the NF prior . We materialize the proposed expansion of the normalizing flow framework by developing an image-oriented architecture which we call DenseFlow . 2.2 Image-oriented invertible module We propose a glow-like invertible module ( also known as step of flow [ 19 ] ) consisting of activation normalization , 1 × 1 convolution and intra-module affine coupling layer . The attribute `` intra-module '' emphasizes distinction with respect to cross-unit coupling . Different than in the original glow design , our coupling network leverages advanced transformations based on dense connectivity and fast selfattention . All three layers are designed to capture complex data dependencies while keeping tractable Jacobians and efficient inverse computation . For completeness , we start by reviewing elements of the original glow module [ 19 ] . ActNorm [ 19 ] is an invertible substitute for batch normalization [ 30 ] . It performs affine transformation with per-channel scale and bias parameters : yi , j = s xi , j + b . ( 14 ) Scale and bias are calculated as the variance and mean of the initial minibatch . Invertible 1 × 1 Convolution is a generalization of channel permutation [ 19 ] . Convolutions with 1 × 1 kernel are not invertible by construction . Instead , a combination of orthogonal initialization and the loss function keeps the kernel inverse numerically stable . The normalizing flow loss maximizes ln |detJf | which is equivalent to maximizing∑ i ln |λi| , where λi are eigenvalues of the Jacobian . Maintaining a relatively large amplitude of the eigenvalues ensures a stable inversion . The Jacobian of this transformation can be efficiently computed by LU-decomposition [ 19 ] . Affine Coupling [ 18 ] splits the input tensor x channel-wise into two halves x1 and x2 . The first half is propagated without changes , while the second half is linearly transformed ( 15 ) . The parameters of the linear transformation are calculated from the first half . Finally , the two results are concatenated as shown in Fig . 2. y1 = x1 , y2 = s x2 + t , ( s , t ) = coupling_net ( x1 ) . ( 15 ) Parameters s and t are calculated using a trainable network which is typically implemented as a residual block [ 18 ] . However , this setting can only capture local correlations . Motivated by recent advances in discriminative architectures [ 29 , 31 , 32 ] , we design our coupling network to fuse both global context and local correlations as shown in Fig . 2 : First , we project the input into a low-dimensional manifold . Next , we feed the projected tensor to a densely-connected block [ 29 ] and self-attention module [ 31 , 33 ] . The densely connected block captures the local correlations [ 34 ] , while the self-attention module captures the global spatial context . Outputs of these two branches are concatenated and blended through a BN-ReLU-Conv unit . As usual , the obtained output parameterizes the affine coupling transformation ( 15 ) . Differences between the proposed coupling network and other network designs are detailed in related work . It is well known that full-fledged self-attention layers have a very large computational complexity . This is especially true in the case of normalizing flows which require many coupling layers and large latent dimensionalities . We alleviate this issue by approximating the keys and queries with their low-rank approximations according to the Nystrom method [ 28 ] .

Summary: The paper proposes to use incremental augmentation, gradually increasing dimensionality towards the latent space, in normalizing flows. The augmentation distribution at each level is conditioned on previous feature maps. This is referred to as "cross-unit coupling". Additionally, the authors introduce a new neural architecture for use in the usual coupling layers, referred to as "intra-unit coupling". The result is an efficiently trained flow with state-of-the-art likelihoods.