Representational aspects of depth and conditioning in normalizing flows

1 INTRODUCTION . Deep generative models are one of the lynchpins of unsupervised learning , underlying tasks spanning distribution learning , feature extraction and transfer learning . Parametric families of neural-network based models have been improved to the point of being able to model complex distributions like images of human faces . One paradigm that has received a lot attention is normalizing flows , which model distributions as pushforwards of a standard Gaussian ( or other simple distribution ) through an invertible neural network G. Thus , the likelihood has an explicit form via the change of variables formula using the Jacobian of G. Training normalizing flows is challenging due to a couple of main issues . Empirically , these models seem to require a much larger size than other generative models ( e.g . GANs ) and most notably , a much larger depth . This makes training challenging due to vanishing/exploding gradients . A very related problem is conditioning , more precisely the smallest singular value of the forward map G. It ’ s intuitively clear that natural images will have a low-dimensional structure , thus a close-to-singular G might be needed . On the other hand , the change-of-variables formula involves the determinant of the Jacobian of G−1 , which grows larger the more singular G is . While recently , the universal approximation power of various types of invertible architectures has been studied ( Dupont et al. , 2019 ; Huang et al. , 2020 ) if the input is padded with a sufficiently large number of all-0 coordinates , precise quantification of the cost of invertibility in terms of the depth required and the conditioning of the model has not been fleshed out . In this paper , we study both mathematically and empirically representational aspects of depth and conditioning in normalizing flows and answer several fundamental questions . 2 OVERVIEW OF RESULTS . 2.1 RESULTS ABOUT GENERAL ARCHITECTURES . In order to guarantee that the network is invertible , normalizing flow models place significant restrictions on the architecture of the model . The most basic question we can ask is how this restriction affects the expressive power of the model — in particular , how much the depth must increase to compensate . More precisely , we ask : Question 1 : is there a distribution over Rd which can be written as the pushforward of a Gaussian through a small , shallow generator , which can not be approximated by the pushforward of a Gaussian through a small , shallow layerwise invertible neural network ? Given that there is great latitude in terms of the choice of layer architecture , while keeping the network invertible , the most general way to pose this question is to require each layer to be a function of p parameters – i.e . f = f1◦f2◦· · ·◦f ` where ◦ denotes function composition and each fi : Rd → Rd is an invertible function specified by a vector θi ∈ Rp of parameters . This framing is extremely general : for instance it includes layerwise invertible feedforward networks in which fi ( x ) = σ⊗d ( Aix+ bi ) , σ is invertible , Ai ∈ Rd×d is invertible , θi = ( Ai , bi ) and p = d ( d + 1 ) . It also includes popular architectures based on affine coupling blocks ( e.g . Dinh et al . ( 2014 ; 2016 ) ; Kingma & Dhariwal ( 2018 ) ) where each fi has the form fi ( xSi , x [ d ] \Si ) = ( xSi , x [ d ] \Si gi ( xSi ) + hi ( xSi ) ) for some S ⊂ [ d ] which we revisit in more detail in the following subsection . We answer this question in the affirmative : namely , we show for any k that there is a distribution over Rd which can be expressed as the pushforward of a network with depth O ( 1 ) and size O ( k ) that can not be ( even very approximately ) expressed as the pushforward of a Gaussian through a Lipschitz layerwise invertible network of depth smaller than k/p . Towards formally stating the result , let θ = ( θ1 , . . . , θ ` ) ∈ Θ ⊂ Rd ′ be the vector of all parameters ( e.g . weights , biases ) in the network , where θi ∈ Rp are the parameters that correspond to layer i , and let fθ : Rd → Rd denote the resulting function . Define R so that Θ is contained in the Euclidean ball of radius R. We say the family fθ is L-Lipschitz with respect to its parameters and inputs , if ∀θ , θ′ ∈ Θ : Ex∼N ( 0 , Id×d ) ‖fθ ( x ) − fθ′ ( x ) ‖ ≤ L‖θ − θ ′‖ and ∀x , y ∈ Rd , ‖fθ ( x ) − fθ ( y ) ‖ ≤ L‖x− y‖ . 1 We will discuss the reasonable range for L in terms of the weights after the Theorem statement . We show2 : Theorem 1 . For any k = exp ( o ( d ) ) , L = exp ( o ( d ) ) , R = exp ( o ( d ) ) , we have that for d sufficiently large and any γ > 0 there exists a neural network g : Rd+1 → Rd with O ( k ) parameters and depth O ( 1 ) , s.t . for any family { fθ , θ ∈ Θ } of layerwise invertible networks that are L-Lipschitz with respect to its parameters and inputs , have p parameters per layer and depth at most k/p we have ∀θ ∈ Θ , W1 ( ( fθ ) # N , g # N ) ≥ 10γ2d Furthermore , for all θ ∈ Θ , KL ( ( fθ ) # N , g # N ) ≥ 1/10 and KL ( g # N , ( fθ ) # N ) ≥ 10γ 2d L2 . Remark 1 : First , note that while the number of parameters in both networks is comparable ( i.e . it ’ s O ( k ) ) , the invertible network is deeper , which usually is accompanied with algorithmic difficulties for training , due to vanishing and exploding gradients . For layerwise invertible generators , if we assume that the nonlinearity σ is 1-Lipschitz and each matrix in the network has operator norm at most M , 1Note for architectures having trainable biases in the input layer , these two notions of Lipschitzness should be expected to behave similarly . 2In this Theorem and throughout , we use the standard asymptotic notation f ( d ) = o ( g ( d ) ) to indicate that lim supd→∞ f ( d ) g ( d ) = 0 . For example , the assumption k , L , R = exp ( o ( d ) ) means that for any sequence ( kd , Ld , Rd ) ∞ d=1 such that lim supd→∞ max ( log kd , logLd , logRd ) d = 0 the result holds true . then a depth ` network will have L = O ( M ` ) 3 and p = O ( d2 ) . For an affine coupling network with g , h parameterized by H-layer networks with p/2 parameters each , 1-Lipschitz activations and weights bounded by M as above , we would similarly have L = O ( M ` H ) . Remark 2 : We make a couple of comments on the “ hard ” distribution g we construct , as well as the meaning of the parameter γ and how to interpret the various lower bounds in the different metrics . The distribution g for a given γ will in fact be close to a mixture of k Gaussians , each with mean on the sphere of radius 10γ2d and covariance matrix γ2Id . Thus this distribution has most of it ’ s mass in a sphere of radius O ( γ2d ) — so the Wasserstein guarantee gives close to a trivial approximation for g. The KL divergence bounds are derived by so-called transport inequalities between KL and Wasserstein for subgaussian distributions Bobkov & Götze ( 1999 ) . The discrepancy between the two KL divergences comes from the fact that the functions g , fθ may have different Lipschitz constants , hence the tails of g # N and f # N behave differently . In fact , if the function fθ had the same Lipschitz constant as g , both KL lower bounds would be on the order of a constant . 2.2 RESULTS ABOUT AFFINE COUPLING ARCHITECTURES . Next , we prove several results for a particularly common normalizing flow architectures : those based on affine coupling layers ( Dinh et al. , 2014 ; 2016 ; Kingma & Dhariwal , 2018 ) . The appeal of these architecture comes from training efficiency . Although layerwise invertible neural networks ( i.e . networks for which each layer consists of an invertible matrix and invertible pointwise nonlinearity ) seem like a natural choice , in practice these models have several disadvantages : for example , computing the determinant of the Jacobian is expensive unless the weight matrices are restricted . Consequently , it ’ s typical for the transformations in a flow network to be constrained in a manner that allows for efficient computation of the Jacobian determinant . The most common building block is an affine coupling block , originally proposed by Dinh et al . ( 2014 ; 2016 ) . A coupling block partitions the coordinates [ d ] into two parts : S and [ d ] \ S , for a subset S with |S| containing around half the coordinates of d. The transformation then has the form : Definition 1 . An affine coupling block is a map f : Rd → Rd , s.t . f ( xS , x [ d ] \S ) = ( xS , x [ d ] \S s ( xS ) + t ( xS ) ) Of course , the modeling power will be severely constrained if the coordinates in S never change : so typically , flow models either change the set S in a fixed or learned way ( e.g . alternating between different partitions of the channel in Dinh et al . ( 2016 ) or applying a learned permutation in Kingma & Dhariwal ( 2018 ) ) . Of course , a permutation is a discrete object , so difficult to learn in a differentiable manner – so Kingma & Dhariwal ( 2018 ) simply learns an invertible linear function ( i.e . a 1x1 convolution ) as a differentiation-friendly relaxation thereof . 2.2.1 THE EFFECT OF CHOICE OF PARTITION ON DEPTH . The first question about affine couplings we ask is how much of a saving in terms of the depth of the network can one hope to gain from using learned partitions ( ala GLOW ) as compared to a fixed partition . More precisely : Question 2 : Can models like Glow ( Kingma & Dhariwal , 2018 ) be simulated by a sequence of affine blocks with a fixed partition without increasing the depth by much ? We answer this question in the affirmative at least for equally sized partitions ( which is what is typically used in practice ) . We show the following surprising fact : consider an arbitrary partition ( S , [ 2d ] \ S ) of [ 2d ] , such that S satisfies |S| = d , for d ∈ N. Then for any invertible matrix T ∈ R2d×2d , the linear map T : R2d → R2d can be exactly represented by a composition of O ( 1 ) affine coupling layers that are linear , namely have the formLi ( xS , x [ 2d ] \S ) = ( xS , Bix [ 2d ] \S+AixS ) or Li ( xS , x [ 2d ] \S ) = ( CixS + Dix [ 2d ] \S , x [ 2d ] \S ) for matrices Ai , Bi , Ci , Di ∈ Rd×d , s.t . each Bi , Ci is diagonal . For convenience of notation , without loss of generality let S = [ d ] . Then , each of the layers Li is a matrix of the form [ I 0 Ai Bi ] or [ Ci Di 0 I ] , where the rows and columns are partitioned into blocks of size d. With this notation in place , we show the following theorem : 3Note , our theorem applies to exponentially large Lipschitz constants . Theorem 2 . For all d ≥ 4 , there exists a k ≤ 24 such that for any invertible T ∈ R2d×2d with det ( T ) > 0 , there exist matrices Ai , Di ∈ Rd×d and diagonal matrices Bi , Ci ∈ Rd×d≥0 for all i ∈ [ k ] such that T = k∏ i=1 [ I 0 Ai Bi ] [ Ci Di 0 I ] Note that the condition det ( T ) > 0 is required , since affine coupling networks are always orientationpreserving . Adding one diagonal layer with negative signs suffices to model general matrices . In particular , since permutation matrices are invertible , this means that any applications of permutations to achieve a different partition of the inputs ( e.g . like in Glow ( Kingma & Dhariwal , 2018 ) ) can in principle be represented as a composition of not-too-many affine coupling layers , indicating that the flexibility in the choice of partition is not the representational bottleneck . It ’ s a reasonable to ask how optimal the k ≤ 24 bound is – we supplement our upper bound with a lower bound , namely that k ≥ 3 . This is surprising , as naive parameter counting would suggest k = 2 might work . Namely , we show : Theorem 3 . For all d ≥ 4 and k ≤ 2 , there exists an invertible T ∈ R2d×2d with det ( T ) > 0 , s.t . for all Ai , Di ∈ Rd×d and for all diagonal matrices Bi , Ci ∈ Rd×d≥0 , i ∈ [ k ] it holds that T 6= k∏ i=1 [ I 0 Ai Bi ] [ Ci Di 0 I ] Beyond the relevance of this result in the context of how important the choice of partitions is , it also shows a lower bound on the depth for an equal number of nonlinear affine coupling layers ( even with quite complex functions s and t in each layer ) – since a nonlinear network can always be linearized about a ( smooth ) point to give a linear network with the same number of layers . In other words , studying linear affine coupling networks lets us prove a depth lower bound/depth separation for nonlinear networks for free . Finally , in Section 5.3 , we include an empirical investigation of our theoretical results on synthetic data , by fitting random linear functions of varying dimensionality with linear affine networks of varying depths in order to see the required number of layers . The results there suggest that the constant in the upper bound is quite loose – and the correct value for k is likely closer to the lower bound – at least for random matrices .

Summary: The normalizing flows (NF) are among popular generative models, as they have the capability of converting a simple base distribution to complex distributions by successively applying change of variable formula. More importantly, NFs let us do inference by maximizing exact likelihood functions instead of other approximate functions like ELBO in VAEs. NFs have the invertibility constraints which make the calculations of their Jacobean determinants computationally prohibitive and require some tricks that make them easier to compute. All these tricks, among them affine coupling layers are of great popularity, come at the price of losing expressiveness. Therefore, we might need more layers (deep) to compensate for that. However, we lack a theoretical understanding of how much deep is enough to guarantee best results.