SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\gamma$  is a hyperparameter. In summary, for the transformation parameters  $\phi$ , the log-likelihood of NFs is maximized, while taking into consideration the rectangular nature of the transformation. Additionally, a reconstruction loss with a regularization parameter  $\beta$  is added to ensure that the learned manifold,  $\mathcal{M} \subset \mathbb{R}^d$ , is aligned to the data manifold,  $\mathcal{M} \subset \mathbb{R}^D$ . The  $\ell_1$  loss of the off-diagonal metric tensor ensures that a canonical manifold  $\mathfrak{M}$  is learned.
We compare our method with the rectangular normalizing flow (RNF) and the original Brehmer and Cranmer manifold learning flow (MFlow). No meaningful comparison can be made to linear PCA/ICA as they are not suitable for data on non-linear manifolds. While non-linear PCA/ICA can be used, they require a specific feature extractor or kernel. Our method, a type of manifold learning, directly learns this non-linear transformation from data. We use the same architectures as [7], namely real NVP [7] without batch normalization, for the same reasons explained in the nominal paper<sup>2</sup>. Further experimental details can be found in the appendix. In short, we use the same hyperparameters for all models to ensure fairness. For example, we use a learning rate of  $1 \times 10^{-4}$  across the board, and we have only carried out simple initial searches for the novel parameter  $\gamma$  from Equation (8). Namely, from trying  $\gamma = [0.001, 0.01, 0.1, 1, 10, 50]$ , we concluded in using  $\gamma = \beta = 1$  for the low-dimensional datasets,  $\beta = 5$  and  $\gamma = 0.1$  for the tabular datasets and  $\beta = 5$  and  $\gamma = 0.01$  for the image datasets to ensure stable training.  $\gamma = 1$  also produced good results in our experiments with different datasets, but it was not always stable and long training was required. Overall, the method is not overly sensitive to the hyperparameter  $\gamma$ . Likelihood annealing was implemented for all image training runs. Approximate training times can be found in the appendix.
5.1 Simulated data
We consider some simulated datasets. Specifically, we implement uniform distributions over a two-dimensional sphere and a Möbius band embedded in  $\mathbb{R}^3$ . Then we generate 1000 samples from each of these distributions and use the samples as the simulated data to fit both RNF and CMF. Such datasets have a straightforward and visualizable canonical manifold representation. The samples generated by the trained networks are shown in Figure 2(a) and Figure 2(b) for the RNF and CMF methods, respectively. From the fully sampled density plots Figure 2(a)(i) and Figure 2(b)(i), both methods learn the manifold and the distribution on it. However, the manifold learning of CMF is superior to RNF as the sphere is fully encapsulated, and the density is more uniform. Furthermore, when the three  $z_i$  dimensions are sampled individually in Figure 2(b) (ii),(iii) and (iv), the canonical
2Batch normalization causes issues with vector-Jacobian product computations.
intrinsic basis learned by CMF can be seen. Specifically, two distinct perpendicular components are obtained while the third component approaches zero, i.e. reduced to something small as it is not necessary. In contrast, from Figure 2(a) (ii),(iii) and (iv) it can be seen that in RNF all three components attempt to wrap around the sphere without apparent perpendicularity, and the second component shrinks partially. Even more striking results can be seen for the Möbius band, Figure 2(c) and Figure 2(d). Additionally, the non-trivial topology of the Möbius manifold, the twist and hole, is proving harder to learn for the RNF as compared to CMF. These findings are in direct alignment with the expected canonical manifold learning, where the network is motivated to learn an orthogonal basis and/or implement sparse learning, as explained in Section 4.2. For CMF and RNF methods, the log likelihoods are 1.6553 vs. 1.6517 and KS p-values are 0.17 vs. 0.26. On the sphere, the log likelihoods are 1.97 vs. 1.16. CMF demonstrates better quality than RNF, both quantitatively and qualitatively.
(a) Hallow sphere learned with RNF.
(c) Moebius band learned with RNF.
(b) Hallow sphere learned with CMF.
(d) Moebius band learned with CMF.
Figure 2: Density plot for a uniform distribution on a sphere (left) and uniform distribution on a Möbius (right) band learned with RNF (top) and CMF (bottom), the black dots are the data samples, similar to Figure 1. RNF and CMF are learned via Equations 4 and Equation (8) (i) samples from the model using the full latent space where colors correspond to  $\log p(x)$ . (ii), (iii), (iv) Samples from the model using only the first, the second and third components  $z_{1/2/3} \in \mathcal{Z}$  while setting the others to zero, respectively. Colors correspond to  $\log p(x)$  once more. All densities have been normalized to  $[-1,1]$  for visualization purposes.
Notably, for these simulated datasets we employ full dimensional latent representations,  $D = d$ . First, it allows us to visualize what all the latent dimensions are doing in relation to what we expect them to do. In particular, CMF correctly uses only 2 latent dimensions for representing 2D surfaces even though they have 3 latent dimensions. Second, it clearly shows the advantage to RNFs for the case where dimensions of the embedded manifold are known. Such scenario is equivalent to a standard NF. RNF is based on [33], which is an NF that already tries to solve the learning pathologies of complicated manifolds. Therefore, there is no practical limitation in choosing  $D = d$ , as also seen empirically. Considering the above, the CMF method is showcased to outperform these previous methods, which also aligns with the theoretical intuition.
5.2 Image Data
First, Mflow, RNF and CMF are also compared using image datasets MNIST, Fashion-MNIST and Omniglot  $\in \mathbb{R}^{784}$ . As mentioned, the hyperparameters were set to  $\beta = 5$  and  $\gamma = 0.01$  for all CMF runs,  $\beta = 5$  and  $\gamma = 0$  for RNF and Mflow methods, and the network architectures are all equivalent. Furthermore, to emphasize the ability and importance of CMF, we train for different latent dimensions  $d$ . The test FID scores of the trained models are shown in Table 1. We use  $\mathbb{R}^{d=20}$  initially as in [7] for consistency but also run for lower latent dimensions,  $\mathbb{R}^{d=10}$ , in order to showcase the advantages of CMF relative to the established methods. Namely, a learned canonical intrinsic basis can lead to a more efficient latent space representation and therefore advantageous performance when the latent space is limited. From Table 1, CMF appears to outperform the other methods.
Then, we also experimented using the CIFAR-10 and SVHN datasets with  $\mathbb{R}^{3072}$ , and CelebA with  $\mathbb{R}^{12288}$ . Again we set  $\beta = 5$  and  $\gamma = 0.01$ , however, this time with latent dimensions,  $\mathbb{R}^{d=30}$ , and  $\mathbb{R}^{d=40}$ . The extra dimensions were necessary for the higher dimensional data. The table Table 1
presents the FID scores. CMF outperforms the other alternatives, consistent with the results obtained on lower dimensional imaging data sets.
In order to visualize the sparsity and orthogonalization encouraged by the CMF, the pair-wise absolute cosine similarity between the basis vectors, i.e.,  $\{\partial x / \partial z^i\}_{i = 1}^d$  (bottom) and the diagonal of the metric tensor (top) are visualized in Figure 3. The CMF clearly induces more sparsity as compared to the RNF, seen from the sparse diagonal elements of the metric tensors, see Figure 3(top), for both the Fashion-MNIST and Omniglot examples. Similarly, the cosine similarity for the CMF is evidently lower for both data sets, see Figure 3(bottom); the MACS stands for the mean absolute cosine similarity. Similar results were observed for the tabular datasets, see appendix.
Additionally, in order to demonstrate the more efficient use of latent space by CMF, we trained both CMF and RNF using latent-dimension  $d = 40$  on both the MNIST and the Fashion-MNIST datasets. Then we choose to vary the number of prominent latent dimensions from  $\mathcal{Z} \in \mathbb{R}^{40}$  as those with the largest  $|G_{kk}|$ . We first generate samples from the models using only the prominent dimensions by sampling from the prominent dimensions and setting all others to zero. We compute FID scores for the generated samples. Later, we also reconstruct real samples using only the prominent dimensions and compute reconstruction loss using mean squared error. Here, we take the dimensions, the larger  $G_{ii}$  as the ones with the highest weights; this simple interpretation is used for analysis. We repeat the experiment for different numbers of prominent latent dimensions (1, 8, 16 . . . , 40) non-zero  $z_i \in \mathcal{Z}$ . In Figure 4(a) and Figure 4(b) the FID score and the mean-squared-error for the different number of prominent dimensions are plotted respectively. From Figure 4, it is evident that the CMF (solid lines) shows better results as compared with RNF (dotted lines) consistently, when using fewer prominent latent dimensions. This suggests that the CMF is able to use the latent space more efficiently. Furthermore, the RNF lines exhibit some oscillation whereas the CMF lines are smoother, indicating that there is also some additional inherent importance ordering in the latent dimensions as expected from sparse learning. Samples generated with different numbers of latent dimensions can be found in the appendix.
(a) Fashion-MNIST trained with RNF. MACS=0.03.
(b) Fashion-MNIST trained with CMF. MACS=0.02.
(c) Omniglot trained with RNF. MACS=0.04.
Figure 3: Sparsity and orthogonalization encouraged by CMF. Average diagonal elements of the metric tensor (top) and average cosine similarity between the basis vectors,  $\{\partial x / \partial z^i\}_{i = 1}^d$  (bottom), with  $\mathbb{R}^{d = 10}$ . Trained for different datasets with RNF and CMF methods. The sparse metric tensor indicates specialization, and cosine similarity closer to zero indicates orthogonalization. MACS stands for mean absolute cosine similarity. The top plots show the CMF method yielding a sparser transformation, while the bottom ones demonstrate its ability to learn more orthogonal transformations, evidenced further by the reduced MACS values.
(d) Omniglot trained with CMF. MACS=0.03.
5.3 Tabular data
The standard normalizing flow benchmarks of tabular datasets from [3] are used to compare the RNF, MFlow and CMF methods. We fit all models with latent dimension  $d = 2$ , 4, 10 and 21 to POWER, GAS, HEMPMASS and MINIBOONE datasets, respectively. The experiments here are similar to those in [7], however, we keep the latent dimension the same for all the methods. Further details can be found in the appendix. We then computed FID-like scores, as was done in compared works [7] to compare the three models. The results are shown in Table 2. CMF appears to have lower FID-like metric scores for POWER, HEMPMASS and MINIBOONE datasets. The GAS dataset
Table 1: Image datasets FID scores, lower is better.
Method	MNIST	FMNIST	OMNIGLOT	SVHN	CIFAR10	CELEBA
Rd=10	Rd=30
MFlow	77.2	663.6	174.9	102.3	541.2	N/A
RNF	68.8	545.3	150.1	95.6	544.0	9064
CMF	63.4	529.8	147.2	88.5	532.6	9060
Rd=20	Rd=40
Mflow	76.211	421.401	141.0	110.7	535.7	N/A
RNF	49.476	302.678	139.3	94.3	481.3	9028
CMF	49.023	297.646	135.6	72.8	444.6	8883
(a) Frechet Inception Distance.
(b) Mean-squared-error.
Figure 4: FID test (a) score and mean-squared-error (b) for fully trained CMF (solid line) and RNF (dotted line) methods on Fashion-MNIST and MNIST with  $\mathbb{R}^{d = 40}$ . The  $x$ -axis corresponds to sampling from the models using a different number of latent dimensions, while setting others to zero. The chosen latent dimensions are the most prominent ones, defined as those with the largest  $|G_{kk}|$  values.
showed mixed results in the second run set, as seen in Table 2. The performance ordering has shifted, suggesting challenges in learning this dataset, possibly leading to random outcomes. However, the other tabular datasets showed consistent performance rankings across run sets. See the appendix for further investigation into the tabular dataset training pathologies. Training on tabular datasets with doubled latent dimensions yielded no significant differences, Table 2.
Our method optimizes the learned manifold  $\mathcal{M} \subset \mathbb{R}^d$  of manifold learning flows by constraining the learning to a canonical manifold  $\mathfrak{M}$  as in Definition 4.1 at no significant additional computational cost as compared to RNFs [7]. A table of indicative training times can be found in the appendix. Nevertheless, the nature of likelihood-based training, where the full  $J^T J$  needs to be calculated, dictates that such a class of methods is computationally expensive for higher dimensions, despite the efficient approximate methods implemented here and in RNFs. Therefore, higher dimensional datasets can be prohibitively computationally expensive to train with this class of methods. Further computational efficiency advancements are expected to be achieved in future works. A notable fact is that the original Mflow method [6] relies on the fundamental approximation such that the full  $J^T J$  calculation can be avoided, resulting in inferior log-likelihood calculation and therefore generations, albeit it being noticeably more computationally efficient.
Additionally, CMF can only model a manifold that is a homeomorphism to  $R^d$ , and when such mapping preserves the topology. Therefore, if, for example, there is a topological mismatch between
Table 2: Tabular data FID-like metric scores, lower is better.
METHOD	POWER	GAS (1ST RUN)	GAS (2ND RUN)	HEMPMASS	MINIBOONE
MFLOW	0.258 ± 0.045	0.219 ± 0.016	0.433 ± 0.071	0.741 ± 0.052	1.650 ± 0.105
RNF	0.074 ± 0.012	0.283 ± 0.031	0.470 ± 0.101	0.628 ± 0.046	1.622 ± 0.121
CMF	0.053 ± 0.005	0.305 ± 0.059	0.373 ± 0.065	0.574 ± 0.044	1.508 ± 0.082
DOUBLE LATENT DIMENSIONS
RNF	0.432 ± 0.022	0.386 ± 0.044	0.608 ± 0.023
CMF	0.205 ± 0.098	0.367 ± 0.007	0.519 ± 0.063
the components of the image space  $R^D$  described directly by  $\mathcal{M} \subset R^d$ , it could lead to pathological training. This drawback, however, is shared across all manifold learning flow methods. We also acknowledge the possibility that an orthogonal basis may not always be an optimal representation and difficult to learn: enforcing strict complete orthogonality can restrict the expressivity of the transformation.
Furthermore, other families of state-of-the-art generative modeling methods working in full dimensional representations, for example denoising diffusion models [34], can obtain superior FID scores overall as compared to the class of methods described in this work. However, by default they use a latent space as large as the image space. Notably, for any future use that entails isolating prominent components, it would be beneficial to perform a proper intrinsic dimensionality estimation.
7 Summary and Conclusions
In this work, we attempt to improve on the learning objective of manifold learning normalizing flows. Having observed that MLFs do not assign always an optimal latent space representation and motivated from compact latent space structuring [15] and sparse learning [14], while taking advantage of the Riemannian manifold nature of homeomorphic transformations of NFs, we propose an effective yet simple to implement enhancement to MLFs that does not introduce additional computational cost.
Specifically, our method attempts to learn a canonical manifold, as defined in Section 4.1, i.e., a compact and non-degenerate intrinsic manifold basis. Canonical manifold learning is made possible by adding to the optimization objective the minimization of the off-diagonal elements of the metric tensor defining this manifold. We document this learned basis for low-dimensional data in Figure 2 and for image data in Figure 3. Furthermore, we showcase for image data that higher quality in generated samples stems from the preferable, compact latent representation of the data, Table 1.
Distinctly, the off-diagonal manifold metric elements when minimized by an  $\ell_1$  loss facilitates sparse learning and/or non-strict local orthogonality. This means that the chart's transformation can be significantly non-linear while retaining some of the benefits of orthogonal representations. This idea can be used in other optimization schemes, which mitigates many of the drawbacks of previous methods in a simple yet theoretically grounded way. Indeed, standard approaches often focus solely on the diagonal elements or strictly enforcing the off-diagonals to be zero, as in [32], limiting expressivity. For instance, in [8], an isometric embedding, essentially a predefined constrained transformation, is presented. Although it allows for direct density estimation and is a form of PCA, it is a strict constraint. Furthermore, as indicated by the synthetic data examples, our present method can be implemented in principle as a full dimensional flow, naturally assigning dimensions accordingly.
The learned manifold representations can be useful as a feature extraction procedure for a downstream task. For example, improved out-of-distribution detection could be a consequence of the feature extraction. Additionally, data that require orthogonal basis vectors, like solutions of a many-body-physics quantum Hamiltonian, can show improved learning performance with the current method.
Manifold learning flows are noteworthy due to their mathematical soundness, particularly in the case of methods such as RNF and CMF, which do not implement any fundamental approximations in the density calculations using lower dimensional latent space. We anticipate that the method will stimulate further research into the nature and importance of the latent space representation of MLFs, enhancing our understanding of generative modelling and paving the way to general unsupervised learning.
This project was supported by grants #2018-222, #2022-274 and #2022-643 of the Strategic Focus Area "Personalized Health and Related Technologies (PHRT)" of the ETH Domain (Swiss Federal Institutes of Technology)
[1] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning, pages 1530-1538. PMLR, 2015.
[2] George Papamakarios, Eric T Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1-64, 2021.
[3] George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density estimation. Advances in neural information processing systems, 30, 2017.
[4] Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arXiv preprint arXiv:1410.8516, 2014.
[5] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv:1605.08803, 2016.
[6] Johann Brehmer and Kyle Cranmer. Flows for simultaneous manifold learning and density estimation. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 442-453. Curran Associates, Inc., 2020.
[7] Anthony L Caterini, Gabriel Loaiza-Ganem, Geoff Pleiss, and John P Cunningham. Rectangular flows for manifold learning. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 30228-30241. Curran Associates, Inc., 2021.
[8] Eike Cramer, Felix Rauh, Alexander Mitsos, Raul Tempone, and Manuel Dahmen. Nonlinear isometric manifold learning for injective normalizing flows. arXiv preprint arXiv:2203.03934, 2022.
[9] Christian Horvat and Jean-Pascal Pfister. Denoising Normalizing Flow. In Advances in Neural Information Processing Systems, volume 34, pages 9099–9111. Curran Associates, Inc., 2021.
[10] Danilo Jimenez Rezende, George Papamakarios, Sebastien Racaniere, Michael Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres. In Hal Daume III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 8083-8092. PMLR, July 2020.
[11] Brendan Ross and Jesse Cresswell. Tractable density estimation on learned manifolds with conformal embedding flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 26635-26648. Curran Associates, Inc., 2021.
[12] Nutan Chen, Alexej Klushyn, Francesco Ferroni, Justin Bayer, and Patrick Van Der Smagt. Learning flat latent manifolds with VAEs. In Proceedings of the 37th International Conference on Machine Learning, ICML'20. JMLR.org, 2020.
[13] Brendan Leigh Ross, Gabriel Loaiza-Ganem, Anthony L. Caterini, and Jesse C. Cresswell. Neural implicit manifold learning for topology-aware generative modelling, 2022.
[14] Michael E Tipping. Sparse Bayesian learning and the relevance vector machine. Journal of machine learning research, 1(Jun):211-244, 2001.
[34] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 6840-6851. Curran Associates, Inc., 2020.
[35] Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008.
[36] M.F. Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics - Simulation and Computation, 19(2):433-450, 1990.
[37] Stephen Wright, Jorge Nocedal, et al. Numerical optimization. Springer Science, 35(67-68):7, 1999.
A Riemannian geometry
The Latin indices run over the spatial dimensions and Einstein summation convection is used for repeated indices.
A  $D$  dimensional curved space is represented by a Riemannian manifold  $\mathbf{M}$ , which is locally described by a smooth diffeomorphism  $\mathbf{h}$ , called the chart. The set of tangential vectors attached to each point  $\mathbf{y}$  on the manifold is called the tangent space  $T_{\mathbf{y}}M$ . In the fluid model, all the vector quantities are represented as elements of  $T_{\mathbf{y}}M$ . The derivatives of the chart  $\mathbf{h}$  are used to define the standard basis  $(\mathbf{e}_1,\dots,\mathbf{e}_D) = \frac{\partial\mathbf{h}}{\partial x^1},\dots,\frac{\partial\mathbf{h}}{\partial x^D}$ .
$$
g _ {i j} (x) = \mathbf {e} _ {i} (x) \cdot \mathbf {e} _ {j} (x) = \frac {\partial \mathbf {h}}{\partial x ^ {i}} \cdot \frac {\partial \mathbf {h}}{\partial x ^ {j}}, \tag {9}
$$
For a given metric tensor, the vector  $v = v^{i}\mathbf{e}_{i} \in T_{\mathbf{y}}M$  has a norm  $||v||_g = \sqrt{v^ig_{ij}v^j}$  and a corresponding dual vector  $v^{*} = v^{i}\mathbf{e}_{i} \in T_{\mathbf{y}}^{*}M$  in the cotangent space, which is spanned by the differential 1-forms  $dx^{i} = g(\mathbf{e}_{i},\cdot)$ . The coefficients  $v_{i}$  of the dual vector are typically denoted by a lower index and are related to the upper-index coefficients  $v^{i}$  by contraction with the metric tensor  $v_{i} = g_{ij}v^{j}$  or equivalently,  $v^{i} = g^{ij}v_{j}$ , where  $g^{ij}$  denotes the inverse of the metric tensor. The upper-index coefficients  $v^{i}$  of a vector  $v$  are typically called contravariant components, whereas the lower-index coefficients  $v_{i}$  of the dual vectors  $v^{*}$  are known as the covariant components.