SlowGuess/NeurIPS_2023_OCR huggingface dataset:
A necessary feature for the description of objects moving on the manifold is parallel transport of vectors along the manifold. The tangent space is equipped with a covariant derivative  $\nabla$  (Levi-Civita connection), which connects the tangent spaces at different points on the manifold and thus allows for the transportation of a tangent vector from one tangent space to the other along a given curve  $\gamma(t)$ . The covariant derivative can be viewed as the orthogonal projection of the Euclidean derivative  $\partial$  onto the tangent space, such that the tangency of the vectors is preserved during the transport. In local coordinates, the covariant derivative is fully characterized by its connection coefficients  $\Gamma_{jk}^{i}$  (Christoffel symbols), which are defined by the action of the covariant derivative on the basis vector,  $\nabla_{j}\mathbf{e}_{k} = \Gamma_{jk}^{i}$ . In the standard basis,  $\mathbf{e}_{i} = \frac{\partial\mathbf{h}}{\partial x^{i}}$ , the Christoffel symbols are related to the metric by
$$
\Gamma_ {j k} ^ {i} = \frac {1}{2} g ^ {i j} \left(\partial_ {j} g _ {k l} + \partial_ {k} g _ {j l} - \partial_ {l} g _ {j k}\right). \tag {10}
$$
Acting on a general vector  $v = v^{i}\mathbf{e}_{i}$ , the covariant derivative becomes:
$$
\nabla_ {k} v = \left(\partial_ {k} v ^ {i} + \Gamma_ {k j} ^ {i} v ^ {j}\right) \mathbf {e} _ {i}, \tag {11}
$$
where the product rule has been applied, using that the covariant derivative acts as a normal derivative on the scalar functions  $v^i$ . Extended to tensors of higher rank, for example the second order tensors  $T = T^{ij}$ , it becomes
$$
\nabla_ {k} T = \left(\partial_ {k} T ^ {i j} + \Gamma_ {k l} ^ {i} T ^ {l j} + \Gamma_ {k l} ^ {j} T ^ {i l}\right) \mathbf {e} _ {i} \otimes \mathbf {e} _ {j} \tag {12}
$$
; in this work, the basis vectors  $\mathbf{e}_i$  are generally dropped. Compatibility of the covariant derivative with the metric tensor implies that  $\nabla_k g^{ij} = \nabla_k g_{ij} = 0$ . This property allows us to commute the covariant derivative with the metric tensor for the raising or lowering of tensor indices in derivative expressions.
The motion of the particle can be described by the curve  $\gamma (t)$ , which parameterizes the position of the particle at time  $t$ . The geodesic equation,  $\nabla_{\dot{\gamma}}\dot{\gamma} = 0$ , in local coordinates  $\gamma (t) = \gamma^i (t)\mathbf{e}_i$  is defined by
$$
\ddot {\gamma} ^ {i} + \Gamma_ {j k} ^ {i} \dot {\gamma} ^ {j} \dot {\gamma} ^ {k} = 0. \tag {13}
$$
The geodesic equation can be interpreted as the generalization of Newton's law of inertia to curved space. The solutions of Eq. (13) represent lines of constant kinetic energy on the manifold, i.e. the geodesics. The Riemann curvature tensor  $R$  can be used to measure curvature, or more precisely, it measures curvature-induced change of a tangent vector  $\boldsymbol{v}$  when transported along a closed loop.
$$
R \left(\mathbf {e} _ {i}, \mathbf {e} _ {j}\right) v = \nabla_ {i} \nabla_ {j} v - \nabla_ {j} \nabla_ {i} v. \tag {14}
$$
In a local coordinate basis  $\mathbf{e}_i$ , the coefficients of the Riemann curvature tensor are given by
$$
\begin{array}{l} R _ {i j k} ^ {l} = g \left(R \left(\mathbf {e} _ {i}, \mathbf {e} _ {j}\right) \mathbf {e} _ {k}, \mathbf {e} _ {l}\right) (15) \\ = \partial_ {j} \Gamma_ {i k} ^ {l} - \partial_ {k} \Gamma_ {i j} ^ {l} + \Gamma_ {j m} ^ {l} \Gamma_ {i k} ^ {m} - \Gamma_ {k m} ^ {l} \Gamma_ {i j} ^ {m}. (16) \\ \end{array}
$$
Contraction of  $R_{jkl}^{i}$  to a rank 2 and 1 tensor yields the Ricci-tensor  $R_{ij} = R_{ikj}^{k}$  and the Ricci-scalar  $R = g^{ij}R_{ij}$  respectively, which can also be used to quantify curvature.
The gradient is defined as  $\nabla^i f = g^{ij}\partial_jf$ , the divergence as  $\nabla_{i}v^{i} = \frac{1}{\sqrt{g}}\partial_{i}(\sqrt{g} v^{i})$ , and the integration over curved volume as  $V = \int_{V}QdV$ , where  $dV = \sqrt{g} dx^{1}\ldots dx^{D} =: \sqrt{g} d^{D}x$  denotes the volume element.  $\sqrt{g}$  denotes the square root of the determinant of the metric tensor.
B Extension to Related Work and further discussions
Normalizing flows have direct access to the latent distribution due to homeomorphism defining the transformation [4]. Nevertheless, the very nature of the homeomorphism can result in pathological behaviour when manifolds with complex topology are learned. For instance, a two-dimensional data space  $p_{\mathcal{X}}$  can have different numbers and types of connected components, 'holes' and 'knots', and images have complicated shapes and topology due to the sheer number of dimensions. Therefore, as usually  $p_{\mathcal{Z}}$  is modelled by a simple distribution, like a Gaussian, it cannot trivially match the topology of  $p_{\mathcal{X}}$ . [33] demonstrates a stable solution to this topological pathology of NFs, albeit with  $\mathbb{R}^D$  support.
Effectively learning manifolds with complicated topologies remains a challenge, even with the full-dimensional latent space. However, the RNF is based on CIF [33] which is an NF that already tries to solve the learning pathologies of complicated manifolds: "CIFs are not subject to the same topological limitations as normalizing flows" [33]. Furthermore, in that regard, we estimate there is no practical limitation in setting, as also seen empirically. In principle, CMF can be a solution for higher-dimensional data. However, in practice it is very computationally expensive to train a full latent dimension (e.g.  $J^T J$  calculation) and the solution will take long to converge. Additionally, empirical knowledge suggests that a lower-dimensional representation can enhance expressivity, as discussed in the context of Mflows. Consequently, in practice using the current method, one can start from a predefined lower dimension and let the network optimize at will.
PCAflow [32] relies on a lower-bound estimation of the probability density to bypass the  $J^t J$  calculation. This bound is tight when complete principal component flow, as defined by them, is achieved. The fundamental distinction here is that CMF does not confine itself to a pure PCAflow scenario, which has the potential to restrict expressivity. Since our method only loosely enforces orthogonality, or mutual information  $\rightarrow 0$ , we anticipate the contour lines to lie somewhere between NF and PCAflow. Furthermore, the introduction of sparsity adds an extra layer of complexity, making direct comparison challenging. As an outlook, we acknowledge that DNF [9] and PCAflow, can undergo quantitative comparison with our method. However, we estimate that RNF is precisely in alignment with CMF, providing a more representative comparison.
Generally, the meaning of "canonical manifold" can vary depending on the context in which it is used, therefore we use Definition 4.1 to precisely define the term "canonical manifold". Definition 4.1 aligns closely with the assumption of PCAflow. However, it is vital to note that our method does not enforce this strictly. To put it in perspective, our approach seeks a 'partly canonical manifold,' if you will. Additionally, our method encompasses sparsity—meaning the diagonal elements can also approach zero.
In principle, full-dimensional flow  $d = D$  can be also solution for higher-dimensional datasets, as the CMF loss will ensure that the number of prominent latent dimensions should correspond to the true intrinsic dimensionality of the data. However, in practice it is very computationally expensive to train a full latent dimension (e.g.  $J^t J$  calculation) and the solution will take long to converge. Additionally, empirical knowledge suggests that a lower-dimensional representation can enhance expressivity, as discussed in the context of Mflows. Consequently, in practice, one can start from a predefined lower-dim and let the network optimize at will. Nevertheless, the efficient approximate methods indicate that CMF could be implemented for a full-dimensional flow. Further exploration of this potential could be undertaken in future studies.
C FID score
For a dataset  $\{x_{1},\ldots x_{n}\} \subset \mathbb{R}^{D}$  and sampled generated data  $\{\tilde{x}_1,\dots \tilde{x}_n\} \subset \mathbb{R}^D$  for a statistic  $F$  the mean and variance are:
$$
\mu = \frac {1}{n} \sum_ {1} ^ {n} F \left(x _ {i}\right), \Sigma = \frac {1}{n - 1} \sum_ {1} ^ {n} \left(F \left(x _ {i}\right) - \mu\right) \left(F \left(x _ {i}\right) - \mu\right) ^ {T}, \tag {17}
$$
$$
\tilde {\mu} = \frac {1}{n} \sum_ {1} ^ {n} F \left(\tilde {x} _ {i}\right), \tilde {\Sigma} = \frac {1}{n - 1} \sum_ {1} ^ {n} \left(F \left(\tilde {x} _ {i}\right) - \tilde {\mu}\right) \left(F \left(\tilde {x} _ {i}\right) - \tilde {\mu}\right) ^ {T}. \tag {18}
$$
The FID score takes  $F$  as the last hidden layer of a pre-trained network and evaluates generated sample quality by comparing generated moments against data moments. This is achieved by calculating the squared Wasserstein-2 distance between Gaussian and corresponding moments by:
$$
\left\| \mu - \tilde {\mu} \right\| _ {2} ^ {2} + t r (\Sigma + \tilde {\Sigma} - 2 (\Sigma \tilde {\Sigma}) ^ {1 / 2}. \tag {19}
$$
This would reach zero if they match. For tabular data,  $F$  is taken to be the identity  $F(x) = x$ , as in [7].
D Optimizing the objective and calculating the metric tensor
The third term of canonical manifold learning optimization objective, (Equation (8)):
$$
\begin{array}{l} \phi^ {*} = \operatorname {a r g m a x} _ {\phi} \sum_ {i = 1} ^ {n} \left[ \log p _ {\mathcal {Z}} \left(g _ {\phi} ^ {- 1} \left(x _ {i}\right)\right) - \log | \det  \mathbf {J} _ {h _ {\eta}} ^ {T} \left(g _ {\phi} ^ {- 1} \left(x _ {i}\right)\right) | \right. \\ - \frac {1}{2} \log | \det {\bf J} _ {f _ {\theta}} ^ {T} (f _ {\theta} ^ {- 1} (x _ {i})) {\bf J} _ {f _ {\theta}} (f _ {\theta} ^ {- 1} (x _ {i})) | - \beta \left\| x _ {i} - g _ {\phi} (g _ {\phi} ^ {- 1} (x _ {i})) \right\| _ {2} ^ {2} \biggr ] - \gamma \| G _ {j \neq k} \| _ {1} ^ {1}, \\ \end{array}
$$
is not so trivial to backpropagate through. Here we use the method suggested by [7]. First,  $J_{\theta}^{T}J_{\theta}$  is calculated efficiently with the Forward-Backward automatic differentiation trick as explained in Section 4.3 of [7]. For the simulated low-dimensional datasets the exact method is used, and for images the approximate stochastic gradient method with  $K = 1$  is used. Second, to back-propagate for the approximate method, a stochastic unbiased estimator is used. Namely, from matrix calculus [35] it is known that:
$$
\frac {\partial}{\partial \theta_ {j}} \operatorname {l o g} \det J _ {\theta} ^ {T} J _ {\theta} = t r \left[ \left(J _ {\theta} ^ {T} J _ {\theta}\right) ^ {- 1} \frac {\partial}{\partial \theta_ {j}} J _ {\theta} ^ {T} J _ {\theta} \right], \tag {20}
$$
where  $tr$  is the trace and  $\theta_{j}$  is the  $j$ -th component. Then the Hutchinson trace estimator is used [36]:
$$
\frac {\partial}{\partial \theta_ {j}} \log \det J _ {\theta} ^ {T} J _ {\theta} \approx \frac {1}{K} \sum_ {k = 1} ^ {K} \epsilon_ {k} ^ {T} \left(J _ {\theta} ^ {T} J _ {\theta}\right) ^ {- 1} \frac {\partial}{\partial \theta_ {j}} J _ {\theta} ^ {T} J _ {\theta} \epsilon_ {k}, \tag {21}
$$
where  $\varepsilon_1, \ldots, \varepsilon_k$  are sampled from a Gaussian distribution. Access to  $J_{\theta}^{T} J_{\theta}$  as explained above allows for the direct computation of Equation (21). Nevertheless, the first term  $\epsilon_{k}^{T} (J_{\theta}^{T} J_{\theta})^{-1} = [(J_{\theta}^{T} J_{\theta})^{-1\varepsilon}]^{T}$  is calculated by an iterative method solver of the matrix inverse problem, the conjugate gradient method CG [37].
$$
\frac {\partial}{\partial \theta_ {j}} \log \det J _ {\theta} ^ {T} J _ {\theta} \approx \frac {1}{K} \sum_ {k = 1} ^ {K} C G \left(\left(J _ {\theta} ^ {T} J _ {\theta}; \epsilon_ {k}\right) ^ {T} \frac {\partial}{\partial \theta_ {j}} J _ {\theta} ^ {T} J _ {\theta} \epsilon_ {k}. \right. \tag {22}
$$
Lastly, the metric tensor is obtained from  $J_{\theta}^{T}J_{\theta}$  as stated in Section 2.2. Back-propagation of  $\| G_{j\neq k}\| _1^1$  is straightforward by taking the derivative w.r.t.  $\theta$  by an (E.g. Adams) optimizer, similar to the reconstruction loss in Equation (8).
(a) HEMPMASS trained with RNF. MACS=0.33.
(b) HEMPMASS trained with CMF. MACS=0.30.
(c) Miniboone trained with RNF. MACS=0.43.
(d) Miniboone trained with CMF. MACS=0.25.
Figure 5: 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 tabular 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.
E Tabular datasets
In Figure 5, we visualize the metric tensor for the tabular datasets trained with sufficiently high latent dimensions and present the mean absolute cosine similarity (MACS) in the supplementary PDF. For the specific case of GAS with  $d = 2$ , the plot might not offer significant insights. We observe that the mean absolute cosine similarity is lower for CMF, although, considering the error, there remains some overlap, which is within expectations.
In order to investigate the cause for the performance gap in GAS, we visualize the validation FID-like scores during training, Figure 6. Upon comparing the results of HEMPMASS and GAS, it is apparent that learning converges until a certain point, after which instability may arise for all three methods. Despite this instability, it is important to note that our implementation selects the best validation step checkpoint for testing, mitigating the impact of training instability. However, it is plausible that these instabilities contribute to the variations between different runs. Furthermore, the training curves follow similar behaviors, converging up to a certain point and then displaying a degree of instability. This trend holds true for all three methods when trained on the GAS dataset.
(a) GAS trained with CMF and RNF.
Figure 6: FID-like validation score training curves. Different colors correspond to different repetitions with both methods, CMF and RNF.
(b) HEMPMASS trained with CMF and RNF.
F Image generations
F.1 Samples from prominent dimensions
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 MNIST and Fashion-MNIST datasets. We choose a varying number of prominent latent dimensions from  $\mathcal{Z} \in \mathbb{R}^{40}$  as those with the largest  $|G_{kk}|$ , this is mathematical intuition that the largest eigenvalues (if an eigen-decomposition is possible) correspond to the dimensions with the highest weights.
We generate samples from the models using only the prominent dimensions, by sampling from the prominent dimensions and setting all others to zero, subsequently FID scores are calculated for the
generated samples. Additionally, we reconstruct real samples using only the prominent dimensions and compute reconstruction loss using mean squared error, as shown in the main text. We repeat the experiment for different number of prominent latent dimensions  $(1,8,16\ldots ,40)$ , or similarly, non-zero  $z_{i}\in \mathcal{Z}$ . Here, in the appendix, we show samples when the number of prominent dimensions sampled from increases.
We would like to note that potentially for any future use that entails isolating prominent components, a proper ID estimation should be carried out.
For example, in Figure 7, different samples are generated along the horizontal. This is repeated for increasing number of prominent dimensions used for sampling, top to bottom along the vertical (4,8,12,16...,40). In Figure 7, FashionMNIST models were trained with  $\mathbb{R}^{d = 40}$  and for the CMF and RNF methods. Similarly to the FID plots in the main text, the CMF shows a bit less variation when going along the columns from top to bottom; this is an indication that fewer prominent dimensions are needed for a faithful generation.
(b)
Figure 7: Generated samples for increasing prominent dimensions used for sampling from top to bottom along the columns. The model is trained on Fashion-MNIST with  $\mathbb{R}^{d = 40}$ , (a) RNF, (b) CMF. Rows represent different samples. The CMF captures the generations almost fully, even with the first prominent dimensions.
To visualize the sparsity of learned manifold, we generate samples from individual latent space dimensions for fully trained models on FashionMNIST with the RNF and CMF. In Figure 8, the trained RNF and CMF models exhibit FIDs of  $\sim 302$  and 297 respectively, see main text. Across the rows, different samples can be found as sampled from each of 10 prominent component subgroups, i.e.  $(\{1,2\}, \{3,4\}, \{5,6\}, \ldots, \{19,20\})$  out of the 20 latent dimensions. Here, only these subgroups are used for sampling at once and the other 18 dimensions are set to zero. The dimensions sampled from at the bottom are the most significant ones; the rest follow in order towards the top. It can be clearly seen that the CMF method manages to capture most of the information of the manifold in the first two dimension, as the generations in the bottom row are good and the rest of the columns seem to have less relevance, indicating strong sparse learning. For the RNF, again there is an importance distinction, as expected, but even dimensions of less prominence seem to be generating variation in the samples. Combining these observations with the higher FID of the CMF, efficient sparse learning is evident.
In Figure 9, we plot a sampled manifold between -3 and 3 for the two most prominent latent dimensions from  $\mathcal{Z} \in \mathbb{R}^{20}$  as those have the largest  $|G_{kk}|$ . This is done for an RNF and CMF trained MNIST model on the left and right of Figure 9. The CMF method appears to split the latent space of  $z_{1}$  without much variation in the  $z_{2}$  direction, indicating perpendicularity.
Figure 8: Generated samples for 10 different hierarchical subgroups of prominent dimensions  $\{1,2\}, \{3,4\}, \{5,6\} \ldots, \{19,20\}$  along the columns from bottom to top, while setting all other dimensions to zero. Models were trained on Fashion-MNIST with  $\mathbb{R}^{d=20}$  with (a) RNF, (b) CMF. Rows represent different samples. Sparse learning is evident for the CMF method.
(a)
(b)
Figure 9: Sampled manifold for the two most prominent dimensions as calculated from  $G_{kk}$ , (a) RNF, (b) CMF.
F.2 Out of distribution detection
We evaluate the performance of CMF for out-of-distribution detection (OoD). Table 3 quantifies the results, where accuracy here is the log-likelihood. The results are similar for the previous methods as expected.
The training was performed on a GPU cluster with various GPU nodes including Nvidia GTX 1080, Nvidia GTX 1080 Ti, Nvidia Tesla V100, Nvidia RTX 2080, Nvidia Titan RTX, Nvidia Quadro RTX 6000, Nvidia RTX 3090. Multiple GPUs were used in parallel across the batch.
Table 3: Decision stump for OoD accuracy (higher is better)
METHOD	MNIST → FMNIST	FMNIST → MNIST
M-FLOW	94 %	84 %
RNF	96 %	78 %
CMF	97 %	77 %
Figure 10: OoD dection with CMF, trained on Fashion-MNIST.
There exists no substantial computational cost distinction between the RNF and CMF methods, as both encounter the calculation bottleneck of the Jacobian-transpose-Jacobian. Any other option was set to the default as found in [7]. The approximate training times can be found in Table 4. Currently, manifold learning flows are computationally demanding; nevertheless, their potential is 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.
Table 4: Approximate training times for  $\mathbb{R}^{d = 40|20}$ . D and H indicate days and hours respectively. MNIST and FMNIST required around 120 epochs, SVHN, around 350, CelebA around 220, CIFAR10 around 500 epochs of training.
METHOD	SVHN	CIFAR10	CELEBA	MNIST	FMNIST
M-FLOW	N/A	N/A	N/A	1D	20H
RNF	12D	20D	28D	3D	2D 18H
CMF	14D	21D	27D	3D 4H	2D 17H
G.1 Images
To train for images, 4 or 5 nodes of less than 10GB each were used in parallel across the batch size. The hyperparameters were kept constant for the three different methods compared, unless a distinction was made for the manifold learning flow as explained in [7].
Parameter	Setting
Learning rate	1 × 10-4
Optimizer	Adams
Likelihood annealing	between epochs 25 and 50
Early stopping	MNIST, FashionMNIST, Omniglot: min 121 epochsSVHN, CelebA, CIFAR10: min 201 epochs
Batch size	120
Latent dimensions	{5,10,20,30,40}
Hyperparameters	Equation (8): β = 5, γ = 0.1, 0.01
Calculation method	Jacobian-transpose-Jacobian: Hutchinson, Gaussian, K=1
CG tolerance	0.001
D-dim flow coupler	8x64
d-dim flow layers	10
G.2 Simulated data
The parameter choice is also kept consistent for training for simulated data. Training was performed on a single GPU or CPU and training times were observed to be 1-4 hours depending on epochs and if early stopping was used. Again, similar training times between RNF and CMF were observed.
Parameter	Setting
Learning rate	1 × 10-4
Optimizer	Adams
Likelihood annealing	False
Epochs	5000 (without early stopping)
Batch size	1000
Latent dimensions	{2,3}
Hyperparameters	Equation (8): β = 1, γ = 1
Calculation method	Jacobian-transpose-Jacobian: ’exact’
D-dim flow coupler	RealNVP
d-dim flow layers	4
G.3 Tabular data
We train for the tabular datasets as found in Papamakrios et al. [3]. Parameters and settings were chosen as follows, as per RNF.
Parameter	Setting
Learning rate	1 × 10-4
Optimizer	Adams
Likelihood annealing	between epochs 25 and 50
Early stopping	Yes
Batch size	POWER: 5000, GAS: 2500, HEMPMASS: 750 ,MINIBOONE: 400
Hyperparameters	Equation (8): β = 5, γ = 0.01
Calculation method	Jacobian-transpose-Jacobian: Hutchinson, Gaussian, K=1
CG tolerance	0.001
D-dim flow coupler	4x128 with 10 layers
d-dim flow	RealNVP 2x32 with 5 layers
Latent dimension	d=D/2
For each run set, we performed 5 repetitions, of which we report an average FID score and a standard error in Table 2. We train models with latent dimension  $d = 2$ , 4, 10 and 21 to POWER, GAS, HEMPMASS and MINIBOONE datasets, respectively.