Tractable Density Estimation on Learned Manifolds with Conformal Embedding Flows

1 Introduction . Deep generative modelling is the task of modelling a complex , high-dimensional data distribution from a sample set . Research has encompassed major approaches such as normalizing flows ( NFs ) [ 16 , 59 ] , generative adversarial networks ( GANs ) [ 23 ] , variational autoencoders ( VAEs ) [ 36 ] , autoregressive models [ 52 ] , energy-based models [ 18 ] , score-based models [ 64 ] , and diffusion models [ 29 , 63 ] . NFs in particular describe a distribution by modelling a change-of-variables mapping to a known base density . This approach provides the unique combination of efficient inference , efficient sampling , and exact density estimation , but in practice generated images have not been as detailed or realistic as those of those of other methods [ 7 , 10 , 29 , 32 , 68 ] . One limitation of traditional NFs is the use of a base density with the same dimensionality as the data . This stands in contrast to models such as GANs and VAEs , which generate data by sampling from a low-dimensional latent prior and mapping the sample to data space . In many application domains , it is known or commonly assumed that the data of interest lives on a lower-dimensional manifold embedded in the higher-dimensional data space [ 21 ] . For example , when modelling images , data samples belong to [ 0 , 1 ] n , where n is the number of pixels in each image and each pixel has a brightness in the domain [ 0 , 1 ] . However , most points in this data space correspond to meaningless noise , whereas meaningful images of objects lie on a submanifold of dimension m ⌧ n. A traditional NF can not take advantage of the lower-dimensional nature of realistic images . There is growing research interest in injective flows , which account for unknown manifold structure by incorporating a base density of lower dimensionality than the data space [ 6 , 12 , 13 , 39 , 41 ] . Flows with low-dimensional latent spaces could benefit from making better use of fewer parameters , being more memory efficient , and could reveal information about the intrinsic structure of the data . Properties of the data manifold , such as its dimensionality or the semantic meaning of latent directions , 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . can be of interest as well [ 36 , 58 ] . However , leading injective flow models still suffer from drawbacks including intractable density estimation [ 6 ] or reliance on stochastic inverses [ 13 ] . In this paper we propose Conformal Embedding Flows ( CEFs ) , a class of flows that use conformal embeddings to transform from low to high dimensions while maintaining invertibility and an efficiently computable density . We show how conformal embeddings can be used to learn a lower dimensional data manifold , and we combine them with powerful NF architectures for learning densities . The overall CEF paradigm permits efficient density estimation , sampling , and inference . We propose several types of conformal embedding that can be implemented as composable layers of a flow , including three new invertible layers : the orthogonal k ⇥ k convolution , the conditional orthogonal transformation , and the special conformal transformation . Lastly , we demonstrate their efficacy on synthetic and real-world data . 2 Background . 2.1 Normalizing Flows . In the traditional setting of a normalizing flow [ 16 , 59 ] , an independent and identically distributed sample { xi } ⇢ X = Rn from an unknown ground-truth distribution with density p⇤x ( x ) is used to learn an approximate density px ( x ) via maximum likelihood estimation . The approximate density is modelled using a diffeomorphism f : Z ! X which maps a base density pz ( z ) over the space Z = Rn , typically taken to be a multivariate normal , to px ( x ) via the change of variables formula px ( x ) = pz f 1 ( x ) detJf f 1 ( x ) 1 , ( 1 ) where Jf ( z ) is the Jacobian matrix of f at the point z . In geometric terms , the probability mass in an infinitesimal volume dz of Z must be preserved in the volume of X corresponding to the image f ( dz ) , and the magnitude of the Jacobian determinant is exactly what accounts for changes in the coordinate volume induced by f . By parameterizing classes of diffeomorphisms f✓ , the flow model can be fitted via maximum likelihood on the training data { xi } ⇢ X . Altogether , the three following operations must be tractable : sampling with f ( z ) , inference with f 1 ( x ) , and density estimation with the |detJf ( z ) | factor . To scale the model , one can compose many such layers f = fk · · · f1 , and the |detJfi ( z ) | factors multiply in Eq . ( 1 ) . Generally , there is no unifying way to parameterize an arbitrary bijection satisfying these constraints . Instead , normalizing flow research has progressed by designing new and more expressive component bijections which can be parameterized , learned , and composed . In particular , progress has been made by designing invertible layers whose Jacobian determinants are tractable by construction . A significant theme has been to structure flows to have a triangular Jacobian [ 16 , 59 , 37 , 17 ] . Kingma and Dhariwal [ 35 ] introduced invertible 1⇥ 1 convolution layers for image modelling ; these produce block-diagonal Jacobians whose blocks are parameterized in a PLU -decomposition , so that the determinant can be computed in O ( c ) , the number of input channels . See Papamakarios et al . [ 54 ] for a thorough survey of normalizing flows . 2.2 Injective Flows . The requirement that f be a diffeomorphism fixes the dimensionality of the latent space . In turn , px ( x ) must have full support over X , which is problematic when the data lies on a submanifold M ⇢ X with dimension m < n. Dai and Wipf [ 14 ] observed that if a probability model with full support is fitted via maximum likelihood to such data , the estimated density can converge towards infinity on M while ignoring the true data density p⇤ ( x ) entirely . Behrmann et al . [ 5 ] point out that analytically invertible neural networks can become numerically non-invertible , especially when the effective dimensionality of data and latents are mismatched . Correctly learning the data manifold along with its density may circumvent these pathologies . Injective flows seek to learn an explicitly low-dimensional support by reducing the dimensionality of the latent space and modelling the flow as a smooth embedding , or an injective function which is diffeomorphic to its image1 . This case can be accommodated with a generalized change of variables formula for densities as follows [ 22 ] . Let g : U ! X be a smooth embedding from a latent space U onto the data manifold M ⇢ X . That is , M = g ( U ) is the range of g. Accordingly , g has a left-inverse2 g† : X ! U which is smooth on M and satisfies g† ( g ( u ) ) = u for all u 2 U . Suppose pu ( u ) is a density on U described using coordinates u . The same density can be described in the ambient space X using x coordinates by pushing it through g. The quantity that accounts for changes in the coordinate volume at each point u isq det ⇥ JTg ( u ) Jg ( u ) ⇤ , where the Jacobian Jg is now a n⇥m matrix [ 43 ] . Hence , using the shorthand u = g† ( x ) , the generalized change of variables formula defined for x 2 M can be written px ( x ) = pu ( u ) det ⇥ JTg ( u ) Jg ( u ) ⇤ 1 2 . ( 2 ) While g describes how the data manifold M is embedded in the larger ambient space , the mapping g alone may be insufficient to represent a normalized base density . As before , it is helpful to introduce a latent space Z of dimension m along with a diffeomorphism h : Z ! U representing a bijective NF between Z and U [ 6 ] . Taking the overall injective transformation g h and applying the chain rule Jg h = JgJh simplifies the determinant in Eq . ( 2 ) since the outer Jacobian Jh is square , det ⇥ JThJ T gJgJh ⇤ = ( detJh ) 2 det ⇥ JTgJg ⇤ . ( 3 ) Finally , writing z = h 1 ( u ) , the data density is modelled by px ( x ) = pz ( z ) |detJh ( z ) | 1 det ⇥ JTg ( u ) Jg ( u ) ⇤ 1 2 , ( 4 ) with the entire process depicted in Fig . 1 . Generating samples from px ( x ) is simple ; a sample z ⇠ pz ( z ) is drawn from the base density and passed through g h. Inference on a data sample x ⇠ px ( x ) is achieved by passing it through h 1 g† , evaluating the density according to pz ( z ) , and computing both determinant factors . Notably , the learned density px ( x ) only has support on a low-dimensional subset M of X , as per the manifold hypothesis . This formulation leads the learned manifold M to be diffeomorphic to Euclidean space , which can cause numerical instability when the data ’ s support differs in topology [ 11 ] , but we leave this issue to future work . In practice , there will be off-manifold points during training or if g ( u ) can not perfectly fit the data , in which case the model ’ s log-likelihood will be 1 . Cunningham et al . [ 13 ] remedy this by adding 1Throughout this work we use “ embedding ” in the topological sense : a function which describes how a low-dimensional space can sit inside a high-dimensional space . This is not to be confused with other uses for the term in machine learning , namely a low-dimensional representation of high-dimensional or discrete data . 2† denotes a left-inverse function , not necessarily the matrix pseudoinverse . an off-manifold noise term to the model , but inference requires a stochastic inverse , and the model must be optimized using an ELBO-like objective . Other work [ 6 , 9 , 39 ] has projected data to the manifold via g g† prior to computing log-likelihoods and optimized g using the reconstruction loss Ex⇠p⇤x ||x g ( g † ( x ) ) ||2 . We prove in App . A that minimizing the reconstruction loss brings the learned manifold into alignment with the data manifold . When computing log-likelihoods , the determinant term log det ⇥ JTgJg ⇤ presents a computational challenge . Kumar et al . [ 41 ] maximize it using an approximate lower bound , while Brehmer and Cranmer [ 6 ] and Kothari et al . [ 39 ] circumvent its computation altogether by only maximizing the other terms in the log-likelihood . In concurrent work , Caterini et al . [ 9 ] optimize injective flows using a stochastic estimate of the log-determinant ’ s gradient . They are also able to optimize log det ⇥ JTgJg ⇤ exactly for smaller datasets , but this procedure involves the explicit construction of Jg , which would be memory-intensive to scale to larger data such as CelebA . In line with research to build expressive bijective flows where detJf is tractable , our work focuses on designing and parameterizing injective flows where log det ⇥ JTgJg ⇤ as a whole is efficiently computable . In contrast to past injective flow models , our approach allows for straightforward evaluation and optimization of log det ⇥ JTgJg ⇤ in the same way standard NFs do for log |detJf | . As far as we can find , ours is the first approach to make this task tractable at scale .

Summary: Many datasets of interest that are modelled with normalizing flows are assumed to lie on a lower-dimensional manifold embedded in a higher-dimensional ambient space. Standard normalizing flows that are restricted to learn a *bijective* transformation are not well suited for modelling such data, as the target transformation between the lower-dimensional manifold space and the ambient space (in which said manifold is embedded) is *injective*.  In theory, an *injective flow* can be defined, where we'd parametrize an injective transformation that would only be invertible *for points on the manifold*. In practice there are two problems with this, however. First is that historically there's been much less work on designing flexible *injective* transformations, hence we don't have the same arsenal of invertible layers as we do for bijective flows. Second, the change-of-variable formula for an injective transformation includes a $\det J_f^T J_f$ term, where $J_f$ is the (rectangular) Jacobian matrix of the injective transformation $f$. This determinant, even if tractable, is typically overly expensive to compute exactly.  Authors of this work make an observation that in the bijective flow literature the game has always been to design parametric transformations $f$ that are restricted to be invertible, but *also* ones that have and easy-to-compute $\det J_f$. In line with this, authors propose a class of transformations that would have an easy-to-compute $\det J_f^T J_f$, specifically *conformal embeddings*. A conformal embedding $f$ is a transformation that preserves *angles*, but more importantly has a property that $J_f^T J_f$ is a scalar multiple of the identity, the determinant of which is simply the product of diagonal terms (i.e. an exponent of said scalar).   Authors consider a few examples of conformal embeddings (translation, scaling, an orthogonal transformation, zero-padding, etc.), which could be stacked in the same way as invertible layers are in a bijective flow. In addition, authors consider *piecewise* conformal embeddings, which are conformal *almost everywhere* (i.e. for all points except for a subset of measure zero), and can increase the expressivity of the otherwise simple conformal embeddings.  By using the conformal embeddings as injective transformations authors come up with the first injective flow with a tractable and *cheap-to-compute* density on the learned manifold. This allows said density to be used as a part of the training objective, even though an additional regularization term still needs to be included in the objective to prevent degenerate solutions. The method demonstrates excellent results on a synthetic 3D problem where the target density lies on a sphere, as well as the CelebA image dataset, where the visual fidelity of the samples is improved in comparison to the more expressive baseline flow.