Rectangular Flows for Manifold Learning

1 Introduction . In recent years , Normalizing Flows ( NFs ) have become a staple of generative modelling , being widely used for density estimation [ 14 , 15 , 45 , 28 , 16 ] , variational inference [ 52 , 30 ] , maximum entropy modelling [ 37 ] , and more [ 46 , 31 ] . In density estimation , we typically have access to a set of points living in some high-dimensional space RD . NFs model the corresponding data-generating distribution as the pushforward of a simple distribution on RD – often a Gaussian – through a smooth bijective mapping . Clever construction of these bijections allows for tractable density evaluation and thus maximum likelihood estimation of the parameters . However , as an immediate consequence of this choice , the learned distribution has support homeomorphic to RD ; in particular , the resulting distribution is supported on a set of dimension D. This is not a realistic assumption in practice – especially for density estimation – as it directly contradicts the manifold hypothesis [ 6 ] which states that high-dimensional data lives on a lower-dimensional manifold embedded in ambient space . ⇤Authors contributed equally . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . A natural idea to circumvent this misspecification is to consider injective instead of bijective flows , which now push forward a random variable on Rd with d < D to obtain a distribution on some d-dimensional manifold embedded in RD . These mappings admit a change-of-variable formula bearing resemblance to that of bijective flows , but unfortunately the volume-change term becomes computationally prohibitive , which then impacts the tractability of maximum likelihood . While there have been recent efforts towards training flows where the resulting distribution is supported on a low-dimensional manifold [ 18 , 53 , 8 , 35 , 40 , 12 ] , these approaches either assume that the manifold is known beforehand or propose various heuristics to avoid the change-of-variable computation . Both of these are undesirable , because , while we should expect most high-dimensional data of interest to exhibit low-dimensional structure , this structure is almost always unknown . On the other hand , we argue that avoiding the volume-change term may result in learning a manifold to which it is difficult to properly assign density , and this approach further results in methods which do not take advantage of density evaluation , undermining the main motivation for using NFs in the first place . We show that density estimation for injective flows based on maximum likelihood can be made tractable . By carefully leveraging forward- and backward-mode automatic differentiation [ 3 ] , we propose two methods that allow backpropagating through the volume term arising from the injective change-of-variable formula . The first method involves exact evaluation of this term and its gradient which incurs a higher memory cost ; the second uses conjugate gradients [ 43 ] and Hutchinson ’ s trace estimator [ 23 ] to obtain unbiased stochastic gradient estimates . Unlike previous work , our methods do not need the data manifold to be specified beforehand , but instead simultaneously estimate this manifold along with the distribution on it end-to-end , thus enabling maximum likelihood training to occur . To the best of our knowledge , ours are the first methods to scale backpropagation through the injective volume-change term to ambient dimensions D close to 3,000 . We study the trade-off between memory and variance introduced by our methods and show empirical improvements over injective flow baselines for density estimation . We also show that injective flows obtain state-of-the-art performance for likelihood-based Out-of-Distribution ( OoD ) detection , assigning higher likelihoods to Fashion-MNIST ( FMNIST ) [ 57 ] than to MNIST [ 36 ] with a model trained on the former . 2 Background . 2.1 Square Normalizing Flows . A normalizing flow [ 52 , 15 ] is a diffeomorphism f̃✓ : RD ! RD parametrized by ✓ , that is , a differentiable bijection with differentiable inverse . Starting with a random variable Z ⇠ pZ for a simple density pZ supported on RD , e.g . a standard Gaussian , the change-of-variable formula states that the random variable X : = f̃✓ ( Z ) has density pX on RD given by : pX ( x ) = pZ ⇣ f̃ 1✓ ( x ) ⌘ detJ h f̃✓ i ⇣ f̃ 1✓ ( x ) ⌘ 1 , ( 1 ) where J [ · ] is the differentiation operator , so that J [ f̃✓ ] ( f̃ 1✓ ( x ) ) 2 RD⇥D is the Jacobian of f̃✓ ( with respect to the inputs and not ✓ ) evaluated at f̃ 1✓ ( x ) . We refer to this now standard setup as square flows since the Jacobian is a square matrix . The change-of-variable formula is often written in terms of the Jacobian of f̃ 1✓ , but we use the form of ( 1 ) as it is more applicable for the next section . NFs are typically constructed in such a way that not only ensures bijectivity , but also so that the Jacobian determinant in ( 1 ) can be efficiently evaluated . When provided with a dataset { xi } ni=1 ⇢ RD , an NF models its generating distribution as the pushforward of pZ through f̃✓ , and thus the parameters can be estimated via maximum likelhood as ✓⇤ : = argmax✓ Pn i=1 log pX ( xi ) . 2.2 Rectangular Normalizing Flows . As previously mentioned , square NFs unrealistically result in the learned density pX having Ddimensional support . We follow the injective flow construction of Brehmer and Cranmer [ 8 ] , where a smooth and injective mapping g : Rd ! RD with d < D is constructed . In this setting , Z 2 Rd is the low-dimensional variable used to model the data as X : = g ( Z ) . A well-known result from differential geometry [ 32 ] provides an applicable change-of-variable formula : pX ( x ) = pZ ⇣ g 1 ( x ) ⌘ detJ [ g ] > ⇣ g 1 ( x ) ⌘ J [ g ] ⇣ g 1 ( x ) ⌘ 1/2 1 ( x 2M ) , ( 2 ) where M : = { g ( z ) : z 2 Rd } . The Jacobian-transpose-Jacobian determinant now characterizes the change in volume from Z to X . We make several relevant observations : ( i ) The Jacobian matrix J [ g ] ( g 1 ( x ) ) 2 RD⇥d is no longer a square matrix , and we thus refer to these flows as rectangular . ( ii ) Note that g 1 : M ! Rd is only properly defined on M and not RD , and pX is now supported on the d-dimensional manifold M . ( iii ) We write the indicator 1 ( x 2M ) explicitly to highlight the fact that this density is not a density with respect to the Lebesgue measure ; rather , the dominating measure is a Riemannian measure on the manifold M [ 48 ] . ( iv ) One can clearly verify as a sanity check that when d = D , equation ( 2 ) reduces to ( 1 ) . Since data points x will almost surely not lie exactly on M , we use a left inverse g† : RD ! Rd in place of g 1 such that g † ( g ( z ) ) = z for all z 2 Rd , which exists because g is injective . This is properly defined on RD , unlike g 1 which only exists over M . Equation ( 2 ) then becomes : pX ( x ) = pZ ⇣ g† ( x ) ⌘ detJ [ g ] > ⇣ g† ( x ) ⌘ J [ g ] ⇣ g† ( x ) ⌘ 1/2 . ( 3 ) Note that ( 3 ) is equivalent to projecting x onto M as x g ( g† ( x ) ) , and then evaluating the density from ( 2 ) at the projected point . Now , g is injectively constructed as follows : g = f̃✓ pad h⌘ and g† = h 1 ⌘ pad † f̃ 1✓ , ( 4 ) where f̃✓ : RD ! RD and h⌘ : Rd ! Rd are both square flows , : = ( ✓ , ⌘ ) , and pad : Rd ! RD and pad† : RD ! Rd are defined as pad ( z ) = ( z,0 ) and pad† ( z , z0 ) = z , where 0 , z0 2 RD d. Now , M depends only on ✓ and not ⌘ , so we write it as M✓ from now on . Applying ( 3 ) yields : pX ( x ) = pZ ⇣ g† ( x ) ⌘ detJ [ h⌘ ] ⇣ g† ( x ) ⌘ 1 detJ [ f✓ ] > ⇣ f†✓ ( x ) ⌘ J [ f✓ ] ⇣ f†✓ ( x ) ⌘ 1/2 , ( 5 ) where f✓ = f̃✓ pad and f†✓ = pad † f̃ 1✓ . We include a derivation of ( 5 ) in Appendix A , along with a note on why injective transformations can not be stacked as naturally as bijective ones . Evaluating likelihoods is seemingly intractable since constructing flows with a closed-form volumechange term is significantly more challenging than in the square case , even if the relevant matrix is now d⇥ d instead of D ⇥D . Brehmer and Cranmer [ 8 ] thus propose a two-step training procedure to promote tractability wherein f✓ and h⌘ are trained separately . After observing that there is no term encouraging x 2M✓ , and that x 2M✓ ( ) x = g ( g† ( x ) ) ( ) x = f✓ ( f † ✓ ( x ) ) , they decide to simply train f✓ by minimizing the reconstruction error to encourage the observed data to lie on M✓ : ✓⇤ = argmin ✓ nX i=1 xi f✓ ⇣ f†✓ ( xi ) ⌘ 2 2 . ( 6 ) Note that the above requires computing both f✓ and f†✓ , so that f̃✓ should be chosen as a flow allowing fast evaluation of both f̃✓ and f̃ 1✓ . Architectures such as the Real NVP [ 15 ] or follow-up work [ 28 , 16 ] are thus natural choices for f̃✓ , while architectures with an autoregressive component [ 45 , 30 ] should be avoided . Then , since h⌘ does not appear in the challenging determinant term in ( 5 ) , h⌘ can be chosen as any normalizing flow , and optimization – for a fixed ✓ – can be tractably achieved by maximum likelihood over the lower-dimensional space : ⌘⇤ = argmax ⌘ nX i=1 n log pZ ⇣ g† ( xi ) ⌘ log detJ [ h⌘ ] ⇣ g† ( xi ) ⌘ o . ( 7 ) In practice , gradient steps in ✓ and ⌘ are alternated . This entire procedure circumvents evaluation of the Jacobian-transpose-Jacobian determinant term in ( 5 ) , but as we show in section 3 , avoiding this term by separately learning the manifold and the density on it comes with its downsides . We then show how to tractably estimate this term in section 4 . 3 Related Work and Motivation . Low-dimensional and topological pathologies The mismatch between the dimension of the modelled support and that of the data-generating distribution has been observed throughout the literature in various ways . Dai and Wipf [ 13 ] show , in the context of variational autoencoders [ 29 ] , that using flexible distributional approximators supported on RD to model data living in a low-dimensional manifold results in pathological behavior where the manifold itself is learned , but not the distribution on it . Cornish et al . [ 11 ] demonstrate the drawbacks of normalizing flows for estimating the density of topologically-complex data , and provide a new numerically stable method for learning NFs when the support is not homeomorphic to RD . However , this approach still models the support as being D-dimensional . Behrmann et al . [ 5 ] show instabilities associated with NFs – particularly a lack of numerical invertibility , as also explained theoretically by Cornish et al . [ 11 ] . This is not too surprising , as attempting to learn a smooth invertible function mapping RD to some low-dimensional manifold is an intrinsically ill-posed problem . This body of work strongly motivates the development of models whose support has matching topology – including dimension – to that of the true data distribution . Manifold flows A challenge to overcome for obtaining NFs on manifolds is the Jacobian-transposeJacobian determinant computation . Current approaches for NFs on manifolds approach this challenge in one of two ways . The first assumes the manifold is known beforehand [ 18 , 53 , 40 ] , severely limiting applicability to low-dimensional data where the true manifold can realistically be known . The second group circumvents the computation of the Jacobian-transpose-Jacobian entirely through various heuristics . Kumar et al . [ 35 ] use a potentially loose lower bound of the log-likelihood , and do not explicitly enforce injectivity , resulting in a method for which the change-of-variables almost surely does not hold . Cunningham et al . [ 12 ] propose to convolve the manifold distribution with Gaussian noise , which results in the model having high-dimensional support . Finally , Brehmer and Cranmer [ 8 ] propose the method we described in subsection 2.2 , where manifold learning and density estimation are done separately in order to avoid the log determinant computation . Concurrently to our work , Ross and Cresswell [ 54 ] proposed a rectangular flow construction which sacrifices some expressiveness but allows for exact likelihood evaluation . Why optimize the volume-change term ? Learning f✓ and h⌘ separately without the Jacobian of f✓ is concerning : even if f✓ maps to the correct manifold , it might unnecessarily expand and contract volume in such a way that makes correctly learning h⌘ much more difficult than it needs to be . Looking ahead to our experiments , Figure 1 exemplifies this issue : the top-middle panel shows the ground truth density on a 1-dimensional circle in R2 , and the top-right panel shows the distribution recovered by the two-step method of Brehmer and Cranmer [ 8 ] . We can see that , while the manifold is correctly recovered , the distribution on it is not . The bottom-right panel shows the speed at which f✓⇤ maps R to M✓⇤ : the top of the circle , which should have large densities , also has high speeds . Indeed , there is nothing in the objective discouraging f✓ to learn this behaviour , which implies that the corresponding low-dimensional distribution must be concentrated in a small region and thus making it harder to learn . The bottom-middle panel confirms this explanation : the learned low-dimensional distribution ( dark red ) does not match what it should ( i.e . the distribution of { f†✓⇤ ( xi ) } n i=1 , in light red ) . This failure could have been avoided by learning the manifold in a density-aware fashion by including the Jacobian-transpose-Jacobian determinant in the objective .

Summary: The task of interest is to simultaneously learn a manifold (with a fixed predefined dimensionality) and the density on it, from training data. The proposed method, rectangular flows, is developed based on an existing work (Brehmer and Cranmer [6]) that used a two-step training procedure (for manifold matching and density matching, respectively). The main contributions include: (1) revealing that the two-step training can be improved by joint training; (2) a new technique to enable computing the Jacobian-transpose-Jacobian term of the joint training objective with better efficiency. Experiments on simulated data and MNIST/FMNIST are conducted to demonstrate the proposed method.