Universal Joint Approximation of Manifolds and Densities by Simple Injective Flows

1 INTRODUCTION . In the past several years , invertible flow networks emerged as powerful deep learning models to learn maps between distributions ( Durkan et al. , 2019a ; Grathwohl et al. , 2018 ; Huang et al. , 2018 ; Jaini et al. , 2019 ; Kingma et al. , 2016 ; Kingma & Dhariwal , 2018 ; Kobyzev et al. , 2020 ; Kruse et al. , 2019 ; Papamakarios et al. , 2019 ) . They generate excellent samples ( Kingma & Dhariwal , 2018 ) and facilitate solving scientific inference problems ( Brehmer & Cranmer , 2020 ; Kruse et al. , 2021 ) . By design , invertible flows are bijective and , hence , may not be a natural choice when the target distribution has low-dimensional support . This problem can be overcome by combining bijective flows with expansive , injective layers , which map to higher dimensions ( Brehmer & Cranmer , 2020 ; Cunningham et al. , 2020 ; Kothari et al. , 2021 ) . Despite their empirical success , the theoretical aspects of such globally injective architectures are not well understood . In this work , we present approximation-theoretic properties of injective flows whose architecture combines bijective flows and expansive , injective layers , with an emphasis on approximating measures with low-dimensional support . We state conditions under which these networks are universal approximators and describe how their design enables applications to inference and inverse problems . 1.1 PRIOR WORK . The idea to combine invertible ( coupling ) layers with expansive layers has been explored by Brehmer & Cranmer ( 2020 ) and Kothari et al . ( 2021 ) . Brehmer & Cranmer ( 2020 ) combine two flow networks with a simple expansive element ( in the sense made precise in Section 2.1 ) and obtain a network that paramterizes probability distributions supported on manifolds . They suggest that such constructions may be universal but neither they not Kothari et al . ( 2021 ) derive theoretical results . We discuss in detail the connection between these two empirical works and the approximation results derived here in Appendix B.1 . Kothari et al . ( 2021 ) propose expansive coupling layers and build networks similar to that of Brehmer & Cranmer ( 2020 ) but with an arbitrary number of expressive and expansive elements . They observe that the resulting network trains much faster with a smaller memory footprint , while producing high-quality samples on a variety of benchmark datasets . While to the best of our knowledge , there are no approximation-theoretic results for injective flows , there exists a body of work on universality of invertible flows ; see Kobyzev et al . ( 2020 ) for an overview . Several works show that certain bijective flow architectures are distributionally universal . This was proved for autoregressive flows with sigmoidal activations by Huang et al . ( 2018 ) and for sum-of-squares polynomial flows by Jaini et al . ( 2019 ) . Teshima et al . ( 2020 ) show that several flow networks including those from Huang et al . ( 2018 ) and Jaini et al . ( 2019 ) are also universal approximators of diffeomorphisms . The injective flows considered here have key applications in inference and inverse problems ; for an overview of deep learning approaches to inverse problems , see Arridge et al . ( 2019 ) . Bora et al . ( 2017 ) proposed to regularize compressed sensing problems by constraining the recovery to the range of ( pre-trained ) generative models . Injective flows with efficient inverses as generative models gives an algorithmic projection1 on the range , which facilitates implementation of reconstruction algorithms . An alternative approach is Bayesian , where flows are used to obtain tractable variational approximations of posterior distributions over parameters of interest , via supervised training on labeled input-output data pairs . Ardizzone et al . ( 2018 ) encode the dimension-reducing forward process by an invertible neural network ( INN ) , with additional outputs used to encode posterior variability . Invertibility guarantees that a model of the inverse process is learned implicitly . For a given measurement , the inverse pass of the INN approximates the posterior over parameters . Sun & Bouman ( 2020 ) propose variational approximations of the posterior using an untrained deep generative model . They train a normalizing flow which produces samples from the posterior , with the prior and the noise model given implicitly by the regularized misfit functional . In Kothari et al . ( 2021 ) this procedure is adapted to priors specified by injective flows which yields significant improvements in computational efficiency . 1.2 OUR CONTRIBUTION . We derive new approximation results for neural networks composed of bijective flows and injective expansive layers , including those introduced by Brehmer & Cranmer ( 2020 ) and Kothari et al . ( 2021 ) . We show that these networks universally jointly approximate a large class of manifolds and densities supported on them . We build on the results of Teshima et al . ( 2020 ) and develop a new theoretical device which we refer to as the embedding gap . This gap is a measure of how nearly a mapping from Ro → Rm embeds an n-dimensional manifold in Rm , where n ≤ o . We find a natural relationship between the embedding gap and the problem of approximating probability measures with low-dimensional support . We then relate the embedding gap to a relaxation of universality we call manifold embedding property . We show that this property captures the essential geometric aspects of universality and uncover important topological restrictions on the approximation power of these networks , to our knowledge , heretofore unknown in the literature . We give an example of an absolutely continuous measure µ and embedding f : R2 → R3 such that f # µ can not be approximated with combinations of flow layers and linear expansive layers . This may be surprising since it was previously thought that networks such as those of Brehmer & Cranmer ( 2020 ) can approximate any “ nice ” density supported on a “ nice ” manifold . We establish universality for manifolds with suitable topology , described in terms of extendable embeddings . Our proof shows that optimality of the approximating network can be established in reverse : optimality of a given layer can be established without optimality of preceding layers . This settles a ( generalization of a ) conjecture posed for a two-layer case in ( Brehmer & Cranmer , 2020 ) . Finally , we show that these universal architectures are also practical and admit exact layer-wise projections , as well as other properties discussed in Section 3.4 . 2 DESCRIPTION OF THE ARCHITECTURE . Let C ( X , Y ) denote the space of continuous functions X → Y . We study networks in F ⊂ C ( X , Y ) that are of the form : F = T nLL ◦ R nL−1 , nL L ◦ · · · ◦ T n11 ◦ Rn0 , n11 ◦ T n00 ( 1 ) 1Idempotent but in general not orthogonal . whereRn ` −1 , n `` ⊂ C ( Rn ` −1 , Rn ` ) , T n `` ⊂ C ( Rn ` , Rn ` ) , L ∈ N is the number of networks , n0 = n , nL = m , and n ` ≥ n ` −1 for ` = 1 , . . . , L. We introduce a well-tuned shorthand notation and write H ◦ G : = { h ◦ g : h ∈ H , g ∈ G } throughout the paper . We identify R with the expansive layers and T with the bijective flows . Loosely speaking , the purpose of the expansive layers is to allow the network to parameterize high-dimensional functions by low-dimensional coordinates in an injective way . The flow networks give the network the expressivity necessary for universal approximation of manifold-supported distributions . 2.1 EXPANSIVE LAYERS . The expansive elements transform an n-dimensional manifoldM embedded in Rn ` −1 , and embed it in a higher dimensional space Rn ` . To preserve the topology of the manifold , this must be done injectively . We thus make the following assumptions about the expansive elements : Definition 1 ( Expansive Element ) . Let ` = 1 , . . . , L , and Rn ` −1 , n `` be a family of functions from Rn ` −1 → Rn ` . Rn ` −1 , n `` is a family of expansive elements if every R ∈ R n ` −1 , n ` ` is injective and Lipschitz . Examples of expansive elements include ( R1 ) Zero padding : R ( x ) = [ xT ,0 ( m−n ) ] T where 0 ( m−n ) is the zero vector ( Brehmer & Cran- mer , 2020 ) . ( R2 ) Multiplication by an arbitrary full-rank matrix , or one-by-one convolution : R ( x ) = Wx , or R ( x ) = w ? x ( 2 ) where W ∈ Rm×n and rank ( W ) = n ( Cunningham et al. , 2020 ) , and w is a convolution kernel ? denotes convolution Kingma & Dhariwal ( 2018 ) . ( R3 ) Injective ReLU layers : R = ReLU ( Wx ) , W = [ BT , −DBT , MT ] T , R ( x ) = ReLU ( [ wT , −wT ] ? x ) for matrix B ∈ GLn ( R ) , positive diagonal matrix D ∈ Rn×n , and arbitrary matrix M ∈ R ( m−2n ) ×n ( Puthawala et al. , 2020 ) . ( R4 ) Injective ReLU networks ( Puthawala et al. , 2020 , Theorem 15 ) . These are functions R : Rn → Rm of the formR ( x ) = WL+1 ReLU ( . . .ReLU ( W1x+b1 ) . . . ) +bL whereW ` are n ` +1 × n ` matrices and b ` are the bias vectors in Rn ` +1 . The weight matrices WL satisfy the Directed Spanning Set ( DSS ) condition for ` ≤ L ( that make all layers injective ) and WL+1 is a generic matrix which makes the mapR : Rn → Rm injective . Note that the DSS condition requires that n ` ≥ 2n ` −1 + 1 for ` ≤ L and we have n1 = n and nL+1 = m . 2.2 BIJECTIVE FLOW NETWORKS . The bulk of our theoretical analysis is devoted to expressive elements . The expressive elements bend the range of the expansive elements into the correct shape . We make the following assumptions about the expressive elements : Definition 2 ( Bijective Flow Network ) . Let ` = 0 , . . . , L be given and let T n `` ⊂ C ( Rn ` , Rn ` ) . T n `` is a family of bijective flow networks if every T ∈ T n `` is Lipschitz continuous and bijective . Examples of bijective flow networks include ( T1 ) Coupling flows , introduced by Dinh et al . ( 2014 ) consider R ( x ) = Hk ◦ · · · ◦H1 ( x ) where Hi ( x ) = [ hi ( [ x ] 1 : d , gi ( [ x ] d+1 : n ) ) [ x ] d+1 : n ] . ( 3 ) In Eqn . 3 , hi : Rd × Re → Rd is invertible w.r.t . the first argument given the second , and gi : Rn−d → Re is arbitrary . Typically in practice the operation in Eqn . 3 is combined with additional invertible operations such as permutations , masking or convolutions Dinh et al . ( 2014 ; 2016 ) ; Kingma & Dhariwal ( 2018 ) . ( T2 ) Autoregressive flows , introduced by Kingma et al . ( 2016 ) are generalizations of triangular flows A : Rn → Rn where for i = 1 , . . . , n the i ’ th value of A is given by of the form [ A ] i ( x ) = hi ( [ x ] i , gi ( [ x ] 1 : i−1 ) ) ( 4 ) In Eqn . 4 , hi : R × Rm → R where again hi is invertible w.r.t . the first argument given the second , and gi : Ri−1 → Rm is arbitrary except for g1 = 0 . In Huang et al . ( 2018 ) , the authors choose hi ( x , y ) , where y ∈ Rm , to be a multi-layer perceptron ( MLP ) of the form hi ( x , y ) = φ ◦Wp , y ◦ · · · ◦ φ ◦W1 , y ( x ) ( 5 ) where φ is a sigmoidal increasing non-linear activation function .

Summary: This paper studied the expressive power of models composed of invertible flows and injective embeddings. First, this paper defined the concept of the Embedding Gap as the measure by which a model can approximate an embedding with low-dimensional support. Then, this paper defined the concept of MEP as the ability to approximate the target embedding arbitrarily small in terms of the Embedding Gap. This paper showed that when invertible layers have the MEP properties and the first layer is a distributively universal approximator (and some additional assumptions), the model is a universal approximator in terms of 2-Wasserstein distance (Theorem 1). Finally, this paper gave a method to compute the inverse transformation of an injective layer with a special form of linear transformation and ReLU nonlinearity. (Theorem 2).