Rectangular Flows for Manifold Learning
Authors: Anthony L. Caterini
Gabriel Loaiza-Ganem
Geoff Pleiss
John Patrick Cunningham
Venue: NeurIPS 2021 Poster (2021)
Keywords: Normalizing Flows
Injective Flows
Manifold Learning
Density Estimation
Generative Modelling

Abstract: Normalizing flows are invertible neural networks with tractable change-of-volume terms, which allow optimization of their parameters to be efficiently performed via maximum likelihood. However, data of interest are typically assumed to live in some (often unknown) low-dimensional manifold embedded in a high-dimensional ambient space. The result is a modelling mismatch since -- by construction -- the invertibility requirement implies high-dimensional support of the learned distribution. Injective flows, mappings from low- to high-dimensional spaces, aim to fix this discrepancy by learning distributions on manifolds, but the resulting volume-change term becomes more challenging to evaluate. Current approaches either avoid computing this term entirely using various heuristics, or assume the manifold is known beforehand and therefore are not widely applicable. Instead, we propose two methods to tractably calculate the gradient of this term with respect to the parameters of the model, relying on careful use of automatic differentiation and techniques from numerical linear algebra. Both approaches perform end-to-end nonlinear manifold learning and density estimation for data projected onto this manifold. We study the trade-offs between our proposed methods, empirically verify that we outperform approaches ignoring the volume-change term by more accurately learning manifolds and the corresponding distributions on them, and show promising results on out-of-distribution detection. Our code is available at https://github.com/layer6ai-labs/rectangular-flows.

Rectangular Flows for Manifold Learning
Anthony L. Caterini" University of Oxford Layer Al anthony@layer6. ai
Gabriel Loaiza-Ganem"
Layer 6 gabriel@layer..i
Geoff' Pleiss
John P. Cunningham Columbia University jpc2181@columbie edu
Columbia University gmp21628columbia edu
Abstract
Normalizing flows are invertible neural networks with tractable change-of-volume volluumee terms, which allow optimization of their parameters be efficiently performed via maximum likelihood. However, data of interest are typically assumed to live some (often unknown) low-dimensional manifold embedded in high-dimensional ambient space. The result a modelling mismatch since - by construction - the invertibility requirement implies high-dimensional support of the learned distri- bution. Injective flows, mappings from low- to high-dimensional spaces, aim fix this discrepancy by learning distributions on manifolds, but the resulting volume-change term becomes more challenging to evaluate. Current approaches either avoid computing this term entirely using various heuristics, or assume the manifold is known beforehand and therefore are not widely applicable. Instead, we propose two methods to tractably calculate the gradient of this term with respect the parameters of the model, relying on careful use of automatic differentiation and techniques from numerical linear algebra. Both approaches perform end-to-end nonlinear manifold learning and density estimation for data projected onto this manifold. We study the trade-offs between our proposed methods, empirically verify that outperform approaches ignoring the volume-change term more accurately learning manifolds and the corresponding distributions on them, and show promising results -ittt---iittibuution detection. Our code is available at htttps://githhi
-abs http://ayveraaaannneeeeeccaaarua
Introduction
In recent years, Normalizing Flows (NFs) have become a staple of generative modelling, being widely used for density estimation [
28 [6], variational inference [52 30), maximum entropy modelling [37], and more 44 In density estimation, we typically have access to set of points living in some sigh-dimessionll space E.,, NFs model the corresponding cata-gennrating distribution as the pushforward of simple distribution on Rº often Gaussian through smooth bijective mapping. Clever construction of these bijections allows for tractable density evaluation and thus maximum likelihood estimation of the parameters. However, as immediate consequence of this choice, the learned distribution has support homeomorphic to RPD; particular, the resulting distribution is supported on set of dimension D. This is not realistic assumption in practice especially for density estimation as directly contradicts the manifold hypothesis (6] which states that high-dimensional data líves on a lower-dimensiooal manifold embedded in ambient space.
0
45
Authors contributed equally.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

A natural idea to cireumvent this misspecification is to consider injective instead of bijective flows, which now push forward a random variable on IR with d to obtain distribution on some d-dimensional manifold embedded in RRP These mappings admit change-of- vaariable 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 [8 533 8 35 40 12,, these approaches either assume that the manifold is computation. Both of these undesirable, because, while we should expect most high-dimensional data of interest to exhibit oow-dimensional structure, this structure is almost always unknown. On the other hand, we argue that avoiding the volume-change term may result in learning 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.
known beforehand or propose various heuristics to avoid the change-of-variale
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-varian formula. The first method involves exact evaluation of this term and its gradient which incurs a higher memory cost; the second uses conjugate gradients [4 and Hutchinson's 's trace estimator [23] to obtain unbiased stochastic gradient estimates. Unlike previous work, methods do not need the data manifold to be specified beforchand, but instead simultaneously estimate this manifold along with the distribution 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 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 lkeelihood-base Out-of-Dirn -iistr--Ditiittiono (OoD) detection, assigning higher likelihoods to Fashion-MNIST (FMNIST) [57 than MNIST 336] with model trained on the former.
2 Background
2.1 Square Normalizing Flows
A normalizing flow 52 [5] is a diffeomorphism So IR - IR parametrized by 0, that is, a differentiable bijection with differentiable inverse. Starting with random variable Z pz for simple density pz supported on IR", e.g. standard Gaussian, the change-of- formula states that the random variable X == f2) has density on IR given by:
(x = pz (Fo (z)) det J
[ (Fo (x))
(x)) where [[.] is the differentiation operator, so that ( E DDD is the Jacobian of fo (with respect to the inputs and not 6) evaluated at (z). We refer to this now standard setup square flows since the Jacobian is square matrix. The change-of-variae
formula is often written in terms of the Jacobian of fo but we use the form of [I as it is more applicable for the next section. NFs are typically constructed in such way that not only ensures bijectivity, but also so that the Jacobian determinant in (I) can be efficiently evaluated. When provided with dataset {x; [==)) C RR,, an NE models its generating distribution as the pushforward of pz through fo, and thus the parameters can be estimated via maximum likelhood as 0 = arg max, Σn= log Px (z:).
2.2 Rectangular Normalizing Flows
As previously mentioned, square NFs unrealistically result in the learned density px having D- dimensional support. We follow the injective flow construction of Brehmer and Cranmer (8) where smooth and injective mapping 96 R² + Rº with d < D is constructed. In this setting, Z E Rd the fow-dimensional variable used model the data as X := (Z). ). well-known result from differential geometry [32] provides an applicable of formula:
1/2 1(x E Ma),
Px (x) = Pz (15'(x)) ddet J[gol (95" (o)((a)) Jlgol (5 (x)
(2)
2

---

where Mg. == (g<(2) E R. The Jacobian aerobiantrrassosee-eee transpose-Jacobian determinant now characterizes the change in volume from Z to X. We make several relevant observations: (i) The Jacobian matrix
1o (x)) E DDx is longer square matrix, and we thus refer to these flows as rectangular. (ii) Note that Ma + TR is only properly defined on Ma and not TR: and px is now supported on the d-dimensional manifold Ma. (tei) We write the indicator ((r E Ma) explicitly to highlight the fact that this density is not a density with respect to the Lebesgue measure; rather, the dominating measure is Riemannian measure the manifold Ma. [4] (iv) One can clearly verify sanity check that when d = D, equation
2
reduces to (T).
Since data points x will almost surely not lie exactly on Ma. we use a left inverse gi. RD - IR in place of 9 such that (=)) = 2 for all E RV, which exists because Pø is injective. This properly defined on R2,, unlike dg -1 which only exists over Ma. Equation (2) then becomes:
px(x) = pz (o)( det Jigol' (st(x)) Jggøl (ob(c)) 1/2
(3)
Note that 3 is equivalent to projecting onto Ma as t (x)), and then evaluating the density from 2 at the projected point.
Now, 9o is injectively constructed follows:
9 = fo pad by and
o/b = hn pad o fo
(4)
where fo: DR. - R and hy Rd Rdd are both square flows, & = (0, n), and pad R2 R D and pad!: IR' + Fdd are defined as pad(z) = (z. 0) and pad
= 2, where 0. 2 E R-d
Now, MM depends only on f and not 17, so we write it as Mg from now on. Applying yields:
(x(x)= = pz (sb(x)) det [hnl (od(x)) ( det [[aa] (f)(#) J[fel (r)(#))
1/2 ,
(5)
where fo Fo pad and f = padt fo- We include a derivation of (5) in Appendix A along with note on why injective transformations cannot be stacked as naturally as bijective ones.
Evaluating likelihoods is seemingly intractable since constructing flows with closed-form volume- change term is significantly more challenging than in the square case, even if the relevant matrix is now d x d instead of D x D. Brehmer and Cranmer [ thus propose two-step training procedure to promote tractability wherein So and hy are trained separately. After observing that there is no term
encouraging Me. and that E MM = simply train fe by minimizing the reconstruction error to encourage the observed data to lie on Mg:
(0())(z)) <
fo(ff (æ)), they decide to
0" = arg min Σ1 x. -1
- So (sa(at))
(6)
Note that the above requires computing both fo and f/. so that fo should be chosen as flow allowing
fast evaluation of both fo and fo Architectures such as the Real NVP [15] or follow-up work
[28] 16 are thus natural choices for fo. while architectures with an autoregressive component [45] 30) should be avoided. Then, since h, does not appear in the challenging determinant term in ( h, can
be chosen as any normalizing flow, and optimization - for fixed 0 - can be tractably achieved by maximum likelihood over the lower-dimensional space:
Σ logpz (oo( log /det [[mm (1
= arg max 1j
In practice, gradient steps and 1) are alternated. This entire procedure circumvents evaluation of determinant term in (5), but as we show in sessettton 3, avoiding this term by separately learning the manifold and the density on comes with its downsides. We then show how tractably estimate this term in section 4]
the
3

---

3 Related Work and Motivation
Low-dimensional and topological pathologies The mismatch between the dimension of the mod- elled 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 I D to model data living in a low- low-dimensional manifold results in pathological behavior where the manifold itself is learned, but not the distribution on it. Cornish et al. [T] demonstrate the drawbacks of normalizing flows for estimating the density of poologically-comples data, and provide new numerically stable method for learning NFss when the support is not homeomorphic to R.P However, this approach still models the support as being D- D-dimensional. Behrmann et al. [5] show instabilities associated with NFs particularly lack of numerical invertibility, as also explained theoretically by Cornish et al. M. This is not too surprising, as attempting to learn a smooth invertible function mapping I" 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 manifolds is the Jacobian- Jacobian-transposs transpose- Jacobian 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 fow-dimensional data where the true manifold can realistically be known. entirely through
The second group circumvents the computation of the various heuristics. Kumar et al. (35] use a potentially loose lower bound of the log-likelihood, and do not explicitly enforce injectivity, resulting in method for which the change-of-variable almost surely does not hold. Cunningham et al. [2] propose to convolve the manifold distribution with Gaussian noise, which results in the model having -igeesimensional support. Finally, Brehmer and Cranmer [8] propose the method we described in (subsecction 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 fo and h,, separately without the Jacobian of fe is concerning: even if fo 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, FFigure 1 exemplifies this issue: the top-middle panel shows the ground truth density on a 1-dimensional circle in TR2 and the top-right panel shows the distribution recovered by the two-step method of Brehmer and Cranmer [. 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 fo- maps IR to Ma-; the top of the circle, which should have large densities, also has high speeds. Indeed, there is nothing in the objective discouraging fo to learn this behaviour, which implies that the corresponding dineeninniooal distribution must be concentrated in 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)- (a:)) ) in light red). This failure could have been avoided by learning the manifold in density-awand fashion by including the Jacobian macobian-rranspose prrasppss--aaobiiie determinant in the objective.
4 Maximum Likelihood for Rectangular Flows: Taming the Gradient
4.1 Our Optimization Objective
We have argued that including the sensible. However, as we previously mentioned, (5) corresponds the density of the projection of onto Mg. Thus, simply optimizing the likelihood would not result in learning Ma in such way that observed data lies on it, only encouraging projected data points to have high likelihood. We thus maximize the log- likelihood subject to the constraint that the reconstruction error should be smaller
in the optimization objective is
than some threshold, i.c. o" arg max, Σn= logpx (x:) subject to E ||,, - (f
<
4

---

In practice, we use the KKT conditions [26 34 and maximize the Lagrangian 7
instead:
log pz (ob(x:)) -log det J J|hn (6) (Zx)
1
log det 10 (a,))a((;) (8)
arg max
-3 xs fo (fd (fo(
where we treat B > as a hyperparameter rather than K. and denote (x:)) as Jo(z;) for simplicity. We have dropped the absolute value since Jo (n(r)) always symmetrie positive definite, since Jo( has full rank by injectivity of fo. We now make a technical but relevant observation about our objective: since our likelihoods are Radon-Nikodym derivatives with respect to Riemannian measure on Ma, different values of / will result in different supports and dominating measures. One should thus be careful to compare likelihoods for models with different values of 0. However, thanks to the smoothness of the objective over 0, we should expect likelihoods for values of which are "close enough" to be comparable for practical purposes. In other words, comparisons remain reasonable locally, and the gradient of the volume-change term should contain relevant information to learn Mg in such a way that also facilitates learning h, on the pulled-back dataset ( (z:))
4.2
Optimizing our Objective: Stochastic Gradients
Note that all the terms in (8) are straightforward to evaluate and backpropagate through except for the third one; in this section we show how to obtain unbiased stochastic estimates of its gradient. In what follows drop the dependence of the Jacobian on z; from our notation and write Jg. with the understanding that the end computation will be parallelized over a batch of æys. We assume access to efficient matrix e Rdd We elaborate how we obtain these matrix- products in the next section. It is well known fact from matrix calculus [49] that:
product routine, i.e. computing Jg Joc can be quickly achieved for any
log det J Jo
Jo= = tr (JJ Jo)
a0;
003
where tr denotes the trace operator and 0, is the j-th element of 0. Next, we can use Hutchinson's trace estimator (23), which states that for any matrix M E Iaddd tr(M) = E.[e Me] for any
Rdd-valued random variable e with zero mean and identity covariance matrix. We can thus obtain an unbiased stochastic estimate of our gradient as:
K 00, log det Jo Jo = Σ (JT
Ō J Jack.
(10)
K 1m
00;"
where EK typically sampled either from standard Gaussian or Rademacher distributions. Naïve computation of the above estimate remains intractable without explicitly constructing J Ja- Fortunately, the Jg Joe terms can be trivially obtained using the given matrix-vector product routine, avoiding the construction of JJ Jg. and then , J Joe follows by taking the gradient w.r.t. 0.
There is however still the issue of computing (Jg Jo)-1 [[J Jo)-T We use conjugate gradients (CG) 43 in order to achieve this. CG is an iterative method to solve problems of the form Au= c for given A E dd (in our case = J Jo) and e Rd. we include the CG algorithm in Appendix
B for completeness. CG has several important properties. First, it is known to recover the solution (assuming exact arithmetic) after at most d steps, which means we can evaluate A A-'e. The solution converges exponentially (in the number of iterations r) to the true value [55], so often < d iterations are sufficient for accuracy to many decimal places. practice, we can tolerate a certain amount of bias, we can further increase computational speed by stopping iterations early. Second, CG only requires method to compute matrix-vector products against A, and does not require access to A itself. One such product is performed at each iteration, and CG thus requires at most d matrix-vector products, though again in practice т & d products usually suffice. This results in O(rd²) solve complexity less than the ((d²) required by direct inversion methods. We denote A A-1 computed with conjugate gradients as CG(A; e). We can then compute the estimator from (TO) as:
K
② 00,
log det J/ Jo ag,
ca (JH Jo:
Jj Jock
(11)
K
5

---

In practice we implement this tenn by noting that CG(Jg Jo: 0/00,J0 0/:0)) Joc
J Joye) (Ct(( " Joe, thereby taking advantage of the soop-gradient operation from Automatic Differentiation (AD) libraries and allowing us to avoid implementing a custom backward pass. We thus compute the contribution of point to the training objective as:
log pz (Bb(or) - -log det det. [hn] (o)())) - 3 1|0 -f (ro()))
(12)
K
stop-gradient (cG (Jo Jo: Jo:ex) )) J0 Jock
2K k=1
which gives the correct gradient estimate when taking the derivative with respect to Φ.
Linear solvers for Jacobian terms We note that linear solvers like CG have been used before to backpropagate through log determinant computations in the context of Gaussian processes [[7), and more recently for square NFs with flexible architectures which do not allow for straightforward Jacobian determinant computations [22 39. However, none of these methods require the Jacobian- casssoss--acobian---e vector product routine that we derive in the next section, and to the best of our knowledge, these techniques have not been previously applied for training rectangular NFs. We also point out that recently Oktay al. -A proposed method to efficiently obtain stochastic estimates of Jge. While their method cannot be used as a drop-in replacement within our framework as it would result in biased CG output, we believe this could be an interesting direction for future work. Finally, we note that CG has recently been combined with the Russian roulette estimator [25) to avoid having to always iterate d times while maintaining unbiasedness, again in the context of Gaussian processes [50]. We also leave the exploration of this estimator within our method for future work.
4.3 AD Considerations: The Exact Method and the
AD Trick
In this section we derive the aforementioned routine for vector products against J Jo, well as an exact method that avoids the need for stochastic gradients (for given x)) at the price of increased memory requirements. But first, let us ask: why are these methods needed in the first place? There is work using power series to obtain stochastic estimates of log determinants [20] [), and one might consider using them in our setting. However, these series require knowledge of the singular values of Jg. Jg. to which we do not have access (constructing Jg Jg to obtain its singular values would defeat the purpose of using the power series in the first place), and we would thus not have a guarantee that the series are valid. Additionally, they have to be truncated and thus result in biased estimators, and using Russian roulette estimators avoid bias [9] can result in infinite variance [TI. Finally, these series compute and backpropagate (w...t& 8) through products of the form c (Jo Ja) for different values of m, which can easily require more matrix-vector products than our methods. Behrmann et al.
( address some of the issues with power series approximations as the result of controlling Lipschitz constants, although their estimates remain biased and potentially expensive.
Having motivated our approach, we now use commonly-knowe properties of AD to derive it; we briefly review these properties in infpppedix C] referring the reader to Baydin et al. [3] for more detail. First, we consider the problem of explicitly constructing Jg. This construction can then be used to evaluate J Jo and exactly compute its log determinant either for log density evaluation of a trained model, or to backpropagate (with respect to 8) through both the log determinant computation and the matrix construction, thus avoiding having use stochastic gradients in the previous section. We refer to this procedure as the exact method. Naïvely, one might try to explicitly construct Jo using only backward-mode AD, which would require D vector-Jacobian products (vjps) of the form v Jo - one one per basis vector v E RDD (and then stacking the resulting row vectors vertically). A better way to explicitly construct Ja is with forward-mode AD, which only requires d Jacobian-vector products (jvps) Jęc, again one per basis vector E Ed (and then stacking the resulting column vectors horizontally). We use a custom implementation of norward-mode AD the popular PyTorch 47 library for the exact method, as well as for the
AD trick described below.
We now explain how to combine forward- and backward-mods AD to obtain efficient matrix-vectes products against Jg Jo in order to obtain the tractable gradient estimates from the previous section.
Ppy Torch has a forward-moot AD implementation which relies on the "double backward" trick, which is known to be memory-inefficient inefficient. Seelhttps: //j- wttpj/tttwsssoooooooo www.o...ttuu............. 1t//////////// 0t/////0/12//// new- trick. html for a description.

---

Table Number of jups and vjps (with respect to inputs) needed for forward and backward passes (with respect to 8), along with the corresponding variance of gradient entries.
Method
FORWARD
BACKWARD
VARIANCE
Exact (naíve)
D vips
vips
Exact
jvps
d jvps
0
Stochastic
K(T + 1)jvps +K(r + vips
K Kjvps +Kvips
X 1/K
Note that u := J&c can be computed with a single jvp call, and then J JDc = Jol can be efficiently computed using only a vjp call. We refer this way of computing matrix-vector products against JØJ Jg as the forward-buckward AD trick. We summarize both of our gradient estimators in Appendix D Note that (12) requires [(T + 1) such matrix-vector products, which is seemingly less efficient as it is potentially greater than the d jvps required by the exact method. However, the stochastic method is much more memory-efficiert than its exact counterpart when optimizing over 0: of the K (T + 1) matrix-vector products needed to evaluate (12). only K require gradients with respect to 0. Thus only K jvps and K vjps, along with their intermediate steps, must be stored in memory over a training step. In contrast, the exact method requires gradients (w.r.t. /) for every one of its d jvp computations, which requires storing these computations along with their intermediate steps in memory.
Our proposed methods thus offer memory vs. variance trade-off. Increasing K in the stochastic method results in larger memory requirements which imply longer training times, as the batch size must be set smaller value. On the other hand, the larger the memory cost, the smaller the variance of the gradient. This still holds true for the exact method, which results in exact gradients, at the cost of increased memory requirements (as long as K < d; if K is large enough the stochastic method should never be used over the exact one). Table I summarizes this trade-off.
Experiments
We now compare our methods against the two-step baseline of Brehmer and Cranmer [8] and also study the memory vs. variance trade-off. We use the real NVP [15] architecture for all flows, except we do not use batch normalization [24] as it causes issues with vjp computations. We point out that
all comparisons remain fair, and we include detailed explanation of this phenomenon in Appendix [ along with all experimental details in Appendix G Throughout, we use the abbreviations RNFs-ML for our maximum likelihood training method, RNFs-TS for the two-step method, and RNFs for rectangular NFs in general. For most runs, we found it useful to anneal the likelihood term(s). That is, at the beginning of training we optimize only the reconstruction term, and then slowly incorporate the other terms. This likelihood annealing procedure helped avoid local optima where the manifold is not recovered (large reconstruction error) but the likelihood of projected data is high.
5.1 Simulated Data
We consider a simulated dataset where we have access to ground truth, which allows us to empirically verify the deficiencies of RNFs-TS. We use one-dimensional unit circle in 2R:2 Figure 1 shows this distribution, along with its estimates from RNFs-ML (exact) and RNFs-TS. As previously observed, RNFs-TS correctly approximate the manifold, but fail to learn the right distribution on it. In contrast we can see that RNFs-ML, by virtue of including the Jacobian-transpe transpos---aaooiaaan term in the optimization, manage to recover both the manifold and the distribution on it (top left panel), while also resulting in an casier-to-leane low- dimensional distribution (bottom middle panel) thanks to fe- mapping to Ma- at a more consistent speed (bottom left panel). We do point out that, while the results presented here are representative of usual runs for both methods, we did have runs with different results which we include in Appendix G for completeness. We finish with the observation that even though the line and the circle are not homeomorphic and thus RNFs are not perfectly able to recover the support, they manage to adequately approximate it.
Mises distribution, which is supported the
7

---

(exact) density
ground
TS density
-0.55
-150
-.3
1.7
-275
-.35
-0.5
RNFs-MI. (exace) speed
Distribution of fg= (X)
RNFs-TS spead
RNPs lesacti target
0.00
target
0000
RNNITT learned
0025
6.
(ggnt)
9000
00000
02
0.000
Figure 1: Top row: RNFs-ML (exact) (left), von Mises ground truth (middle), and RNF-TS (right). Bottom row: Speed which fo- maps to Ma- (measured distance between uniformly spaced consecutive points in R mapped through fo-) for RNFs-ML (exact) (left), RNFs-TS (right), and distribution by has to learn order to recover the ground truth, fixing (middle). See text for discussion.
Table 2: FID-like metric for tabular data (lower is better). Bolded runs are the best or overlap with it.
Method
POWER
GAS
HEPMASS
MINIBONNE
RNFs-ML (exact)
0.067 0.016
0.138 + 0.023
0.486 = 0.032
0.978 + 0. 0.082
RNFs-ML (K-
0.083 0.015
0.110 + 0.021
0.779 + 0.191
1.001 0.051
RNFs-ML (K- = 10)
0.113 0.037
0.140 + 0.013
0.495 + 0.055
0.878 + 0.083
RNFs-TS
0.178 + 0.021
0.161 + 0.016
0.649 + 0.081
1.085 + 0.062
5.2
Tabular Data
We now turn our attention to the tabular datasets used by Papamakarios et al. 45), now common benchmark for NFs as well. As previously mentioned, one should careful when comparing models cannot rely test likelihoods FID score [2I]. which is commonly used to evaluate quality generated images when likelihoods not available. The FID score compares the first and second moments of well-chosen statistic taken in practice to be the values of the last hidden layer of a pre-trained inception network [56] from the model and data distributions using the squared Wasserstein-2 metric (between Gaussians). Here, we take the statistic to be the data itself instead of the final hidden units of a pre-trained classifier: in other words, our metric compares the mean and covariance of generated data against those of observed data with the same squared Wasserstein-22 metric. We include the mathematical formulas for computing both FID and modified version for tabular data in Appendix [F] We early stopping with our FID-like score across all models. Our results are summarized in Table 2 where we can that RNFs-ML consistently do a better job recovering the underlying distribution. in the
with different supports,
metric. We take inspiration from the
Once again, these results emphasize the benefits of including the traabiiantttassss---aaeroo objective. Interestingly, except for HEPMASS, the results from our stochastic version with K 1 are not significantly exceeded by the exact version or using larger value of K, suggesting that the added variance does not result in decreased empirical performance. We highlight that no tuning was done (except on GAS for which we changed d from 4 to 2), RNFs-ML outperformed RNFs-TS out-of-the-bbx here (details are Appendix G). We report training times in Appendix G and observe that RNFs-ML. take similar amount of time as RNFs-TS to train for datasets with lower values of D, and while we do take longer to train for the other datasets, our training times remain reasonable and often require fewer epochs to converge.
8

---

Table 3: FID scores (lower better) and decision stump OoD accuracy (higher is better).
FID
OoD ACCURACY
Method
CIFAR-10
MNIST
FMNIST
MNIST + FMNIST
FMNIST - MNIST
RNF (exact)
643.31
36.09
296.01
92%
91%
RNFs-ML (K = 1)
830.94
33 33.98
288.39
97%
78%
RNFs-ML (K 4)
42.90
342.91
77%
89%
RNFs-TS
731.46
35.52
318.59
98%
96%
5.3 Image Data and Out-of-Distribution Detection
We also compare RNFs-ML to RNFs-TS for image modelling on MNIST and FMNIST. We point out that these datasets have ambient dimension D 784, and being able to fit RNFs-ML is in itself noteworthy: to the best of our knowledge no previous method has scaled optimizing the Jacobian- transpose-Jacobian term to these dimensions. We use FID scores both for comparing models and for early stopping during training. We also used likelihood annealing, with all experimental details again given in Appendix G We report FID scores in Table 3 where we can see that we outperform RNFs-TS. Our RNFs-ML (K = 1) variant also outperforms its decreased-variancce This initially puzzling behaviour is partially explained by the fact that we used the K = 1 variant to tune the model (being the cheapest one to train), and then used the tuned hyperparameters for single run of the other two variants. Nonetheless, once again these results suggest that the variance induced by stochastic method is not empirically harmful, and that while using the exact method should be the default whenever feasible, using K 1 otherwise is sensible. We also report training times where can see the computational benefits of our stochastic method, as well as visualizations of samples, in Appendix G
variance counterparts.
We also compare performance on the CIFAR-10 dataset [33]. for which D = 3.072. Once again, being able to fit RNFs-ML in this setting is in itself remarkable. We do not include RNFs-ML (K = 4) results because of limited experimentation details, including hyperparameters which we tried, are given in Appendix G We can see that, while RNFs-TS outperformed RNFs-ML (K 1) which we hypothesize might be reversed given more tuning our RNFF-ML (exact) version is the best performing model, yet again highlighting the importance of including the change-of- volume term in the objective.
CIFAR- 10 due to computational cost (experimental
We further evaluate the performance of RNFs for OoD detection. Nalisnick et al. [41] pointed out that square NFs trained on FMNIST assign higher likelihoods to MNIST than they do to FMNIST. While there has been research attempting to fix this puzzling behaviour 10 to the best of our knowledge no method has managed to correct it using only likelihoods of trained models. Figure 2 shows that RNFs remedy this phenomenon, and that models trained on FMNIST assign higher test likelihoods to FMNIST than to MNIST. This correction does not come at the cost of strange behaviour now emerging in the opposite direction (i.e. when training on MNIST, see Appendix G for a histogram). Table 3 quantifies these results (arrows point from in-distribution datasets OoD ones) with the accuracy of a decision stump using only log-likelihood, and we can see that the berf-perfo RNFs models essentially solve this OoD task. While we leave a formal explanation of this result for future work, we believe this discovery highlights the importance of properly specifying models and of ensuring the use of appropriate inductive biases, in this case low intrinsic dimensionality of the observed data. The strong performance of RNFs-TS here seems to indicate that this is a property of RNFs rather than of our ML training method specifically, although our exact approach is still used to compute these log- liglliiessooods at test time. We include additional results on OoD detection using reconstruction errors along with discussion in Appendix G where we found the opposite unexpected behaviour: FMNIST always has smaller reconstruction errors, regardless of which dataset was used for training.
Trained FMNIST
FMNIST
620
MNIST
decision boundary
500
420-
200
220
300
200
-175
-250
-125
-200
-15
50
log
Figure 2: OoD detection with RNFs-ML. (exact).

---

6 Scope and Limitations
In this paper we address the dimensionality-bass misspecification of square NFs while properly using maximum likelihood as the training objective, thus providing an advancement in the training of RNFs. Our methods however remain topologically misspecified: even though we can better address dimensionality,
currently only learn manifolds homeomorphic to R4, For example, one could conceive of the MNIST manifold as consisting 10 connected components (one per digit), which cannot be learned by fa- It is nonetheless worth noting that this limitation is shared by other deep generative modelling approaches, for example GANs [19] result in connected supports (since the image of connected set under continuous function is connected). We observed during training in image data that the residuals of CG were not close to O numerically, even after d steps, indicating poor conditioning and thus possible numerical -on-invertibiliiyy of the matrix Jg Jo- We hypothesize that this phenomenon is caused by topological mismatch, which we also conjecture affects us more than the baseline as CG-obtained (or from the exact method) gradients might point in inaccurate direction. We thus expect our methods in particular to benefit from improved research on making flows match the target topology, for example via continuous indexing [
Additionally, while we have successfully scaled likelihood based training of RNFs far beyond current capabilities, methods- dimensions, and further computational gains remain an open problem. We also attempted OoD detection CIFAR-10 10 against the SVHN dataset [42]. and found that neither RNFs-ML nor RNFs- TS has good performance, although anecdotally we may have at least improved on the situation outlined by Nalisnick et al. [41). We hypothesize these results might be either caused by topological mismatch, or corrected given more tuning.
the stochastic
remain computationally expensive for higher
7 Conclusions and Broader Impact
In this paper we argue for the importance of likelihood-basse training rectangular flows, and introduce two methods allowing to do We study the benefits of our methods, and empirically show that they are preferable to current alternatives. Given the methodological nature of our contributions, we do not foresee our work having any negative ethical implications societal consequences.
Acknowledgements
We thank Brendan Ross, Jesse Cresswell, and Maksims Volkovs for useful comments and feedback. We would also like to thank Rob Comish for the excellent CIFs codebase upon which our code built, and Emile Mathieu for plotting suggestions. GP and JPC are supported by the Simons Foundation, McKnight Foundation, the Grossman Center, and the Gatsby Charitable Trust.
References
[1] A. A. Alemi, I. Fischer, J. V. Dillon, and K. Murphy. Deep variational information bottleneck. ICLR, 2017.
[2] A. A. Alemi, I. Fischer, and J. V. Dillon. Uncertainty in the variational information bottleneck. 2018.
arXiv preprint arXiv:
[3] A. G. Baydin, B. A. Pearlmutter, A. A. Radul, and J. M. Siskind. Automatic differentiation in machine learning: survey. Journal of machine learning research, 18, 2018.
[4] J. Behrmann, W. Grathwohl, R. T. Chen, D. Duvenaud, and J.-H. Jacobsen. Invertible residual networks. In International Conference on Machine Learning, pages 573-582. PMLR, 2019.
[5] J. Behrmann, P. Vicol, K.-C Wang, R. Grosse, and J.-H. Jacobsen. Understanding and mitigating exploding inverses invertible neural networks. In International Conference on Artificial Intelligence and Statisties, pages 1792-1800, PMLR, 2021.
[6] Y. Bengio, A. Courville, and P. Vincent. Representation learning: review and new perspectives. IEEE transactions pattern analysis and machine intelligence, 35(8); 2013.
10

---

[7] D. P. Bertsekas. Constrained optimization and Lagrange multiplier methods. Academic press, 2014.
[8] J. Brehmer and K. Cranmer. Flows for simultaneous manifold learning and density estimation. In Advances in Neural Information Processing Systems, volume 33, 2020.
[9] R. T. Q. Chen, J. Behrmann, D. K. Duvenaud, and J.-H. Jacobsen. Residual flows for invertible generative modeling. In Advances in Neural Information Processing Systems, volume 32, 2019.
[10] H. Choi, E. Jang, and A. A. Alemi. Waic, but why? generative ensembles for robust anomaly detection, arXiv preprint arXiv: 1810000323 2018.
[11] R. Cornish, A. Caterini, G. Deligiannidis, and A. Doucet. Relaxing bijectivity constraints with continuously indexed normalising flows. In International Conference on Machine Learning. pages 2133-2143 PMLR, 2020.
[12] E. Cunningham, R. Zabounidis, A. Agrawal, I. Fiterau, and D. Sheldon. Normalizing flows across dimensions. arXiv preprint arXiv:2006.1 2020.
[13] B. Dai and D. Wipf. Diagnosing and enhancing vae models. ICLR, 2019.
[14] L. Dinh, D. Krueger, and Y. Bengio. Nice: Non-linear independent components estimation. arXiv preprint arXiv:1410.8516. 2014.
[15] L. Dinh, J. Sohl-Dickstein, and S. Bengio. Density estimation using real nvp. ICLR, 2017.
[16] C. Durkan, A. Bekasov, I. Murray, and G. Papamakarios. Neural spline flows. In Advances in Neural Information Processing Systems, volume 32, 2019.
[17] Gardner, G. Pleiss, K. Q Weinberger, D. Bindel, and A. G. Wilson. Gpytorch: Blackbox matrix-matrix gaussian process inference with gpu acceleration. Advances in Neural Information Processing Systems, 31, 2018.
[18] M. C. Gemici, D. Rezende, and Mohamed. Normalizing flows on riemannian manifolds. arXiv preprint arXhv:1611.02304 2016.
[19] I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley S. Ozair, A. Courville, and Y. Bengio. Generative adversarial nets. In Proceedings of the 27th International Conference on Neural Information Processing Systems-Volume 2. pages
2014.
[20] I. Han, D. Malioutov, and J. Shin. Large-scale log-determinant computation through stochastic chebyshev expansions. In International Conference on Machine Learning, pages 908-917. PMLR, 2015.
[21] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. In Advances in Neural Information Processing Systems, volume 30, 2017.
[22] C.WW. Huang, R. Chen, Tsirigotis, and A. Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. ICLR, 2021.
[23] M. F. Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Pratistics-Simllatie and Computation, 18(3):1059- 1076, 1989.
[24] S. loffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456. PMLR, 2015.
[25] H. Kahn. Use of different Monte Carlo sampling techniques. Rand Corporation, 1955.
[26] W. Karush. Minima of functions of several variables with inequalities as side constraints. Sc. Dissertation. Dept. of Mathematics, Univ. of Chicago, 1939.
[27] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. ICLR, 2015.

---

[28] D. P. Kingma and P. Dhariwal. Glow: generative flow with invertible 1x 1 convolutions. In Advances Neural Information Processing Systems, volume 31, 2018.
[29] D. P. Kingma and M. Welling. Auto-encoding variational bayes. ICLR, 2014.
[30] D. P. Kingma, T. Salimans, R. Jazefowicz, variational inference with inverse autoregressive flow. In Advances in Neural Information Processing Systems, volume 30, 2016.
X. Chen, I. Sutskever, and M. Welling, Improved
[31] I. Kobyzev, S. Prince, and M. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020.
[32] S. G. Krantz and H. R. Parks. Geometric integration theory. Springer Science & Business Media, 2008.
[33] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[34] H. W. Kuhn and A. Tucker. W., 1951," nonlinear programming.". In Proceedings of the Second Berkeley Symposium on Mathematical Statistics and Probabilit(Unniversty of Califormia Press, Berkeley, CA), volume 48 1492, 1951.
[35] A. Kumar, B. Poole, and K. Murphy. Regularized autoencoders via relaxed injective probability flow. In International Conference on Artificial Intelligence and Statistics, pages 4292-4300 PMLR, 2020.
[36] Y. LeChun. The mnist database of handwritten digits, 1998. URL http://yann. lecun. com/exdb/nnist, 1998.
[37] G. Loaiza-Ganent Y. Gao, and J. Cunningham. Maximum entropy flow networks. ICLR, 2017.
[38] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. ICLR, 2019.
[39] C. Lu, Chen, C. Li, Q. Wang, and J. Zhu. Implicit normalizing flows. ICLR, 2021.
[40] E. Mathieu and M. Nickel. Riemannian continuous normalizing flows. In Advances in Neural Information Processing Systems, volume 33, 2020.
[41] E. Nalisnick, A. Matsukawa, W. Teh, D. Gorur, and B. Lakshminarayanan. Do deep generative models know what they don't know? ICLR, 2019.
[42] Y. Netzer, T. Wang, A. Coates, A. Bissacco, B. Wu, and A. Y. Ng. Reading digits in natural images with unsupervised feature learning. 2011.
[43] J. Nocedal and S. Wright. Numerical optimization. Springer Science & Business Media, 2006.
[44] D. Oktay, N. McGreivy, J. Aduol, A. Beatson, and R. P. Adams. Randomized automatic differentiation. ICLR, 2021.
[45] G. Pupamakarios, T. Pavlakou, and I. Murray. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems, volume 30, 2017.
[46] G. Papamakarios, E. Nalisnick, D. Rezende, S. Mohamed, and B. Lakshminarayanan. Nor- malizing flows for probabilistic modeling and inference. arXiv preprint 1rXwv:191.002762. 2019.
[47] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L.. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, I.. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, volume 32. 2019.
[48] X. Pennec, Intrinsic statistics on riemannian manifolds: Basic tools for geometric measurements. Journal of Mathematical Imaging and Vision, 25(1):127-154, 2006.
12

---

[49] K. B. Petersen and M. S. Pedersen. The matrix cookbook, Oct. 2008. URL http: wwww. imm åtu.dk; wttp:///pppppppppopobo php?3274 Version 20081 1100.
[50] A. Potapezynski, L. Wu, D. Biderman, G. Pleiss, and J. P. Cunningham. Bias-free scalable gaussian processes via randomized truncations. International Conference on Machine Learning, to appear, 2021.
[51] J. Ren, P. Liu, E. Fertig, J. Snoek, R. Poplin, M. Depristo, J. Dillon, and B. Lakshminarayanan. Likelihood ratios for out-of-distribution detection. In Advances in Neural Information Process- ing Systems, volume 32, 2019.
[52] D. Rezende and S. Mohamed. Variational inference with normalizing flows, In International Conference on Machine Learning, pages 1530-1538. PMLR, 2015.
[53] D. Rezende, G. Papamakarios, S. Racaniere, M. Albergo, Kanwar, Shanahan, and K. Cranmer. Normalizing flows on tori and spheres. In International Conference on Machine Learning, pages 8083-8092. PMLR, 2020.
[54] B. L. Ross and J. C. Cresswell. Tractable density estimation on learned manifolds with conformal embedding flows. arXiv preprint vrXiv:2106.05275 2021. 2021.
[55] R. Shewchuk al. An introduction the conjugate gradient method without the agonizing pain, 1994.
[56] C. Szegedy, W. Liu, Y. Jia, P. Sermanet, S. Reed, D. Anguelov, D. Erhan, V. Vanhoucke, and A. Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1 1-9, 2015.
[57] H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint 17X1v:1708.07747, 2017.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 6
(c) Did you discuss any potential negative societal impacts of your work? [Yes] Although we do not expect any. See Section
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes] As supplemental material, and we will release the code upon publication.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix G
(c) Did you report error bars (e-g., with respect the random seed after running experi-
ments multiple times)? [Yes] When computationally tractable to repeat experiments.
See Section 5
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] See Appendix G
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] See the README in our code and the references in our paper.
(b) Did you mention the license of the assets? [No] Because the code we based ours on does not seem to have had a license. See the README in our code.
(c) Did you include any assets either in the supplemental material
URL? [Yes]
Just code.
(d) Did you discuss whether and how consent was obtained from people whose data you'l 're using/curating? [No] But we are using publicly available and widely used datasets only.
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [No] But only because we do not believe does.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "vyRYCAllrx", "reviewer": "Reviewer_b6vS", "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.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality. The presented method in Section 4.2, for calculating the volume-change term, is believed new; it's a novel combination of existing techniques. \n\nQuality. The presented method is technically sound. Some empirical observations are not discussed/analyzed in detail. For example, the inconsistent results on simulated data, and the unexpected behavior on MNIST/FMNIST.\n\nClarity. The writing is OK, but can be certainly improved; for example, to highlight the underlying logic, the main contributions over existing methods, algorithms that show the big picture and/or the details (like VJP/JVP), and so on.\n\nSignificance. The contributions might be incremental; the presented method enables joint training for existing work, but the corresponding importance is not well demonstrated with serious experiments and/or analysis.\n\nOther comments:\n\nLine 48, the exact method is quite straightforward. I don't think it is a contribution, so no credit for it.\n\nEq (8). Since beta is just treated as a hyperparameter, the KKT conditions are not met or used. Please clarify.\n\nLines 164-165. What's the meaning of the statements here?\n\nHow to calculate a VJP or a JVP? It's not clear in the manuscript.\n\nEq (12). Since the (J^T J eps) term of the loss is calculated via VJP/JVP, how do you guarantee the correct gradient wrt theta?\n\nLine 323. Why RNFs can remedy the phenomenon? Also, too many conjectures are made in this paragraph without solid analysis.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "YhdwG8XKNoS", "reviewer": "Reviewer_1ptY", "summary": "This paper is motivated by the desire to circumvent a problematic requirement in Normalizing Flows (NFs) on the function $f: \\mathbb{R}^d \\rightarrow \\mathbb{R}^D$ that produces the pushforward measure --- namely, that $f$ must be strictly invertible. This is difficult for generative models, since the dimension of the latent space $d$ is usually much smaller than that of the target data space $D$, i.e. $d << D$. The invertibility requirement centers around the change-of-variables formula, which is crucial for computing the cost function in NFs, and which requires computing the Jacobian of $f^{-1}$. To enable the use of NFs on latent/data spaces where $d << D$, the authors propose \"Rectangular\" NFs (RNFs). The name comes from the fact that the Jacobian of $f$ is now rectangular in the $d << D$ case. This involves the use of a differential geometric variant of the change-of-variables formula, and is much more computationally taxing, but provides superior density estimation performance to previous work.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors explore a very natural and compelling question that arises from the invertibility requirement (of the aforementioned $f$) in NFs, when applying NFs to generative modeling. \n\n$f$, which can be thought of as the \"decoder\" or \"generator\" network, often maps from some lower dimensional latent space to a higher dimensional data space. If $f$ and the latent space are suitably well behaved, it would be reasonable to assume that, even though $f$ is not invertible, its restriction to the data manifold $\\mathcal{M}$ could be. Intuitively, if $f$ is injective, this suggests that there does still exist a tractable way to compute the pushforward density, and one wonders if there is a more general case of the change of variables formula that can be used.\n\nTo this end the authors demonstrate that, while this general form does exist, the computation of its terms (namely, the inverse of the Jacobian) is unfortunately quite heavy and numerically unstable, and moreover, stochastic estimations of this Jacobian via random projections is noisy. \n\nThough these are not strictly positive results, they highlight important obstacles to further progress in an area of very active research. Moreover, the work is presented in a lucid and understandable manner, and the limitations of the current method are clearly delineated. For these reasons, I am inclined to accept the paper.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There is limited experimental validation of the proposed technique. The authors compare their method to only one other work, on three datasets: a synthetic circle dataset, MNIST, and FMNIST (they mention having tried CIFAR10 and SVHN, but were unable to tune their model in a reasonable amount of time). Even with MNIST and FMNIST, it could be informative to train an NF with a 2-dimensional latent space and visualize manifold to see if some structure of the dataset were preserved.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "Q64xxvujVjV", "reviewer": "Reviewer_mod8", "summary": "The paper addresses the challenge of learning a normalizing flow, under the assumption that the data probability function is concentrated on a low dimensional manifold. Utilizing the manifold assumption introduces a challenge in computing the volume change term. This paper suggests addressing this computational challenge in two ways: i) If the manifold dimension is low enough, it is reasonable to calculate the map differential exactly based on AD forward mode; ii) alternatively, an unbiased stochastic estimator for the log det can be used.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper addresses the challenge of learning normalizing flows. Previous work has recently suggested incorporating the manifold assumption into the normalizing flow [1].  The loss that follows from the manifold normalizing flows in [1] depends on a volume change term, a term that [1] did not compute directly. The current paper claims that there is an advantage in optimizing directly the loss involving the volume change term, and suggests two options to compute it. Even though there is not a lot of novelty in the techniques related to the suggested computations, I mark this paper as a solid contribution to the ongoing effort of learning normalizing flows and rate it with a positive rating. Detailed comments are provided next. \n\nGenerally, the paper is well written and easy to follow. The introduction and the method presentation are of good quality. I appreciate the author's effort to include all the relevant details to the method, in addition to some issues that were considered in devising the method. This is also reflected in the experiments section, where details about failed experiments are included.\n\nThe main contribution of the paper is that in fact the loss suggested in [1] can be computed directly in some cases or stochastically estimated. It seems to me that the stochastic estimation of the logdet is not novel, see [2]. Also, the exact computation is based on the forward mode AD, which is a known technique. \n\nI am concerned about the inconsistency in the results of figure 1. Do the authors have some thoughts on what can be done? This should be in the main text. \n\nThe experiments show an advantage over [1], except for the OOD experiment. However, why other alternatives have not been considered? In table 3 it would be important to see results from other square flow models (without the manifold assumption), or contractive auto-encoders (that are known to approximate the derivative of the data log probability function [3]). As the current method cannot run on more complex datasets like CIFAR, these additional comparisons can some light if this line of works is worth pursuing. \n\n[1] : Flows for simultaneous manifold learning and density estimation, Brehmer and  Cranmer.\n[2] : Invertible Residual Networks, Behrmann et al.\n[3] : What Regularized Auto-Encoders Learn from the Data Generating Distribution, Alain and Bengio,", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "Pb7zJR6QA5U", "reviewer": "Reviewer_zGxK", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- Paper is clearly written, exposition is coherent with all the  preliminaries and motivations for the technical decisions in place. Derivations are plain and comprehensible. \n\n- While injective flows as a class were described before, previous heuristics for implementing them were oversimplified and in the best case suboptimal. Thus making volume-change tractable for the injective flows seems to be an important step forward. \n\n- While (as stated in the paper) conjugate gradients were previously applied to the calculation of the volume change for the square flows, authors are the first to adopt them to the rectangular flows and Jacobian-Transpose Jacobian product calculation. Thus the novelty in this sense is moderate. Any comments?\n\n- Experiments on the synthetic and tabular data are convincing and clear demonstrate, that two-step method leads to learning incorrect probability distribution. On the Fashion MNIST FID improvements are quite notable. \n\n- Though, for me, the fact that increasing the number of steps and thus decreasing variance leads to worse results is a bit puzzling, and thus I would like to see more experiments in this direction, which would clarify this question.  Any comments?\n\n- Out-of-distribution experiments also are promising and hint that low-dimensional manifold is a important inductive bias. \n\n- Overall the result on the toy data and simple image data are quite convincing, though until it is verified on the real-life level of complexity data the broader significance is under question. Thus I would also l like to see some experiments on CIFAR to get sure.  Any comments?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations.", "review_text": "", "extra_scores": {"ethical_concerns": "not applicable", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.0"}, "extra_text": {}}, {"review_id": "2yrvRU_rGC", "reviewer": "Reviewer_Jpwi", "summary": "The paper presents new methods for estimating densities using rectangular normalizing flows. Rectangular normalizing flows, unlike squared normalizing flows, allow considering injective instead of bijective flows. This property is exploited for estimating the density of data residing on a low (d) dimensional manifold in R^D, for D>d. Existing works using rectangular flows to estimate densities supported on low dimensional manifolds assume that the manifold is known, and therefore are restricted to tori, spheres, etc. In contrast, this work does not make this assumption and shows that density estimation for injective flows based on ML is tractable. Two methods are proposed, which are based on incorporating and estimating the “Jacobian-transpose-Jacobian” in the objective. One is based on an exact derivation, which is computationally demanding. The other is based on stochastic gradient estimates. Experimental results on a simulation and a couple of datasets are shown, demonstrating the advantage of the proposed methods compared to an existing method as well as the tradeoff between memory and variance incurred by the two proposed methods.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This is a well-written, timely, and interesting paper. It addresses an important problem and presents a solid contribution. \nThe work has several notable limitations:\n- It can only learn densities supported on manifolds homeomorphic to R^d. \n- It is computationally expensive and feasible only when the manifold dimension is low.\n- Partly as a result of the limitation above, the experimental study is not impressive as it includes only relatively simple examples. \nI commend the authors for clearly putting these limitations forward in a nice discussion in Section 6, which also highlights some intriguing results and the gained insights. \nDespite its shortcomings, I believe the paper improves existing methods and advances this line of work, and I would be happy to see it published.\n\nMinor comments:\n- A clear description of the two algorithms would help (even if put in the appendix due to length limitations).\n- More details on the choice of the number of gradients to estimate K (in the stochastic variant) would help practitioners.\n- In Appendix F, failed runs on the simulated data are presented, which raises two questions: is there an explanation for the failures? can the results be quantified (over many runs)?  \n- The FID score, which is used to evaluate the results in Sec 5.2, is defined only in the appendix. I believe that some details, or a basic definition, need to appear in the main paper.\n- In the experimental study on MNIST (Sec 5.3), the dimension d is set to 20. Why? Was it obtained empirically? Is it based on previous works? Is this a result of the computational limitation? Please elaborate.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors very nicely discuss the limitations of their work. \nI do not foresee any potential negative societal impact", "review_text": "", "extra_scores": {"ethical_concerns": "I do not have any ethical concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Congratulations on the acceptance of your paper! Please incorporate changes, edits and additional promised experiments from "Author Discussion" in the final paper/appendix.
Decision: Accept (Poster)