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: This paper addresses the problem of constructing injective normalizing flows, which connect some low-dimensional space with the data manifold of interest in the high-dimensional space. Typically injective (or rectangular) flows are composed of two square flows with the “upsampling” padding layer between them. Because the volume change term in the case of injective flows is seemingly intractable, current approaches find workarounds by training these two square flows separately step-by-step. Though this decision has certain drawbacks and may result in suboptimal performance. Current work suggests to employ numerical linear algebra methods to make volume change computation tractable. Experiments demonstrate that taking into account volume сhanging term improves the generation results both for synthetic and real-life data. Interestingly, it also improves out of distribution detection score, allowing the model trained on FashionMNItST to assign lower probability for the MNIST data samples.