SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kyriakos Flouris
and Electrical Engineering
kflouris@vision.ee.ethz.ch
Ender Konukoglu
kender@vision.ee.ethz.ch
Manifold learning flows are a class of generative modelling techniques that assume a low-dimensional manifold description of the data. The embedding of such a manifold into the high-dimensional space of the data is achieved via learnable invertible transformations. Therefore, once the manifold is properly aligned via a reconstruction loss, the probability density is tractable on the manifold and maximum likelihood can be used to optimize the network parameters. Naturally, the lower-dimensional representation of the data requires an injective-mapping. Recent approaches were able to enforce that the density aligns with the modelled manifold, while efficiently calculating the density volume-change term when embedding to the higher-dimensional space. However, unless the injective-mapping is analytically predefined, the learned manifold is not necessarily an efficient representation of the data. Namely, the latent dimensions of such models frequently learn an entangled intrinsic basis, with degenerate information being stored in each dimension. Alternatively, if a locally orthogonal and/or sparse basis is to be learned, here coined canonical intrinsic basis, it can serve in learning a more compact latent space representation. Toward this end, we propose a canonical manifold learning flow method, where a novel optimization objective enforces the transformation matrix to have few prominent and non-degenerate basis functions. We demonstrate that by minimizing the off-diagonal manifold metric elements  $\ell_1$ -norm, we can achieve such a basis, which is simultaneously sparse and/or orthogonal. Canonical manifold flow yields a more efficient use of the latent space, automatically generating fewer prominent and distinct dimensions to represent data, and consequently a better approximation of target distributions than other manifold flow methods in most experiments we conducted, resulting in lower FID scores. $^1$
Many emerging methods in generative modeling are based on the manifold hypothesis, i.e., higher-dimensional data are better described by a lower-dimensional sub-manifold embedded in an ambient space. These methods replace the bijective normalizing flow in the original construction [1-5] (NF) with an injective flow. NFs are constructed as a smooth bijective mapping, i.e. a homeomorphism, where the learned distribution density is supported on the set of dimension  $D$  equal to the data dimension, where  $\mathbb{R}^{\mathbb{D}}$  is the data space. In order to fulfill the manifold hypothesis, an emerging class of methods, namely manifold learning flow models (MLF) [6-9], model the latent distribution as a random variable in  $\mathbb{R}^d$  where  $d < D$ , i.e. realize injective flows. Consequently, the density lives on some  $d$ -dimensional manifold  $\mathcal{M}_d$  embedded in  $\mathbb{R}^D$  as required. If  $\mathcal{M}_d$  is known a-priori, [10, 11], the density transformation to  $\mathbb{R}^D$  in MLF is a relatively straight forward affair. If the manifold is to be
$^{1}$ Code is available at https://github.com/k-flouris/cmf/.
learned, the change-of-variable formula, i.e. the volume element, can be computationally inhibiting to calculate. Caterini et al [7] have recently solved this problem and established tractable likelihood maximization.
MLF methods have thus evolved to the point where the general manifold hypothesis is properly incorporated and the likelihood is efficiently calculated. However, it has been observed that the learned manifold can be suboptimal; for example, due to a nonuniform magnification factor [12] and the inability to learn complex topologies [13]. Furthermore, pushforward models such as MLFs, i.e.,  $d < D$ , are designed to avoid over-fitting, but  $d$  in  $\mathbb{R}^d$  is still arbitrary unless determined by some meta-analysis method. As well as, there is often degeneracy in the information stored in each of these latent dimensions  $d$ . For instance, as we show in Section 4.1, consider modeling a two-dimension line with considerable random noise, i.e. a 'fuzzy line'. If  $d = 2$ , both latent dimensions will attempt to encompass the entire manifold in a degenerative use of latent dimensions. Particularly, this behavior of storing degenerate representations is inefficient and can lead to over-fitting, topological, and general learning pathologies. Alternatively, we postulate that motivating sparsity [14] and local orthogonality during learning can enhance the manifold learning process of MLFs.
To this end, we propose an optimization objective that either minimizes or finds orthogonal gradient attributions in the learned transformation of the MLF. Consequently, a more compact [15] and efficient latent space representation is obtained without degenerate information stored in different latent dimensions, in the same spirit as automatic relevance determination methods [16] or network regularization [17]. The relevant dimensions correspond to locally distinct tangent vectors on the data manifold, which can yield a natural ordering among them in the generations. Borrowing the term 'canonical basis', i.e., the set of linearly independent generalized eigenvectors of a matrix, we introduce Canonical manifold learning flows (CMF). In CMF, the normalizing flow is encouraged to learn a sparse and canonical intrinsic basis for the manifold. By utilizing the Jacobians calculated for likelihood maximization, the objective is achieved without additional computational cost.
2 Theoretical background
2.1 Normalizing and rectangular normalizing flows
First, we define a data space as  $\mathcal{X}$  with samples  $\{x_{1},\dots,x_{n}\}$ , where  $x_{i}\in \mathbb{R}^{D}$ , and a latent space  $\mathcal{Z}\subset \mathbb{R}^d$ . If,  $d = D$ , then a normalizing flow is a  $\phi$  parameterized diffeomorphism  $q_{\phi}:\mathbb{R}^{D}\to \mathbb{R}^{D}$ , i.e., a differentiable bijection with a differentiable inverse, that transforms samples from the latent space to the data space,  $\mathcal{X}\coloneqq q_{\phi}(\mathcal{Z})$ . Optimization of  $\phi$  is achieved via likelihood maximization, where due to the invertibility of  $q_{\phi}$  the likelihood,  $p_{\mathcal{X}}(x)$ , can be calculated exactly from the latent probability distribution  $p_{\mathcal{Z}}(z)$  using the transformation of distributions formula [18],
$$
p _ {\mathcal {X}} (x) = p _ {\mathcal {Z}} \left(q _ {\phi} ^ {- 1} (x)\right) | \det  \mathbf {J} _ {q _ {\phi}} \left(q _ {\phi} ^ {- 1} (x)\right) | ^ {- 1}. \tag {1}
$$
$\mathbf{J}$  denotes the Jacobian functional such that  $\mathbf{J}_{q_{\phi}} = \partial x / \partial z$ , i.e., the Jacobian of  $q_{\phi}$  at  $x = q_{\phi}(z)$ . Here,  $\mathcal{Z}$  is modelled by a random variable  $z \sim p_{\mathcal{Z}}$  for a simple density  $p_{\mathcal{Z}}$ , e.g., a normal distribution. Therefore, for large  $n$ , a normalizing flow approximates well the distribution that gave rise to the samples in a dataset  $\{x_i\}_{i=1}^n$  by transforming the samples  $z \sim p_{\mathcal{Z}}$  with a trained  $q_{\phi^*}$  if  $\phi^* := \operatorname{argmax}_{\phi} \sum_{i=1}^{n} \log p_{\mathcal{X}}(x_i)$  [4].
In the context of manifold learning flows, or rectangular normalizing flows,  $d < D$  implies there is a lower dimensional manifold  $\mathcal{M} \subset \mathbb{R}^d$  that is embedded into the  $R^D$  with a smooth and injective transformation  $g_{\phi}: \mathbb{R}^d \to \mathbb{R}^D$ . In practice,  $g_{\phi} = h_{\eta} \circ \operatorname{pad} \circ f_{\theta}$ , where  $h_{\eta}: \mathbb{R}^d \to \mathbb{R}^d$ ,  $f_{\theta}: \mathbb{R}^D \to \mathbb{R}^D$  and  $\operatorname{pad}: \mathbb{R}^d \to \mathbb{R}^D$ . Padding can be used for trivial embedding, i.e.,  $\operatorname{pad}(\tilde{z}) = (\tilde{z}_0 \dots \tilde{z}_{d-1}, 0 \dots 0)$ . The rectangular Jacobian of  $\mathbf{J}_{g_{\phi}}$  is not invertible and thus,  $\mathbf{J}_{g_{\phi}}^T \mathbf{J}_{g_{\phi}} \neq \mathbf{J}_{g_{\phi}}^2$ , which implies that the volume element needs to be calculated fully, i.e., the square root of the determinant of  $\mathbf{J}_{g_{\phi}}^T \mathbf{J}_{g_{\phi}}$  which is an equivalent expression for the general metric tensor in this context - see the appendix. This leads to a more generic formulation
$$
p _ {\mathcal {X}} (x) = p _ {\mathcal {Z}} \left(g _ {\phi} ^ {- 1} (x)\right) | \det  \mathbf {J} _ {g _ {\phi}} ^ {T} \left(g _ {\phi} ^ {- 1} (x)\right) \mathbf {J} _ {g _ {\phi}} \left(g _ {\phi} ^ {- 1} (x)\right) | ^ {- 1 / 2}. \tag {2}
$$
However, note that maximizing  $p_{\mathcal{X}}(x)$  on its own would only maximize the likelihood of the projection of  $x$  on the latent space, not ensuring that the model can reconstruct  $x$  from  $z$ , i.e., the
modeled manifold may not be aligned with the observed data. To encourage  $x \in \mathcal{M}$ , namely to align the learned manifold to the data space, a reconstruction error needs to be minimized during optimization. The following loss is thus added to the maximum-likelihood loss function, as explained in [6]:
$$
\sum_ {i = 1} ^ {n} \left\| x _ {i} - g _ {\phi} \left(g _ {\phi} ^ {- 1} (x)\right) \right\| _ {2} ^ {2}. \tag {3}
$$
Combining the logarithm of Equation (2) and using the hyperparameter  $\beta > 0$  to adjust Equation (3) we arrive at the total Lagrangian to be maximized:
$$
\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. \tag {4} \\ \left. - \frac {1}{2} \log | \det  \mathbf {J} _ {f _ {\theta}} ^ {T} \left(f _ {\theta} ^ {- 1} \left(x _ {i}\right)\right) \mathbf {J} _ {f _ {\theta}} \left(f _ {\theta} ^ {- 1} \left(x _ {i}\right)\right) | - \beta \left\| x _ {i} - g _ {\phi} \left(g _ {\phi} ^ {- 1} \left(x _ {i}\right)\right) \right\| _ {2} ^ {2} \right]. \\ \end{array}
$$
2.2 Riemannian geometry and the metric tensor
A  $d$  dimensional curved space is represented by a Riemannian manifold  $\mathcal{M} \subset \mathbb{R}^d$ , 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}}\mathcal{M}$ . All the vector quantities are represented as elements of  $T_{\mathbf{y}}\mathcal{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 z_1}, \dots, \frac{\partial\mathbf{h}}{\partial z_d}$ .
The metric tensor  $g$  can be used to measure the length of a vector or the angle between two vectors. In local coordinates, the components of the metric tensor are given by
$$
G _ {i j} = \mathbf {e} _ {i} \cdot \mathbf {e} _ {j} = \frac {\partial \mathbf {h}}{\partial z _ {i}} \cdot \frac {\partial \mathbf {h}}{\partial z _ {j}}, \tag {5}
$$
where  $\cdot$  is the standard Euclidean scalar product.
The metric tensor  $G_{ij}$  describing the manifold  $\mathcal{M} \subset \mathbb{R}^D$  learned in manifold learning flows is equivalent to the section of the rectangular Jacobian-transpose-Jacobian  $J^T J$  that describes  $\mathcal{M}$ . As explained in [7] and the appendix, the  $J^T J$  can be very efficiently approximated in the context of MLFs. Thus, it can be seen that no additional computational cost is needed to calculate the metric tensor.
3 Related Work and motivation
Generative models with lower-dimensional latent space have long been established as favorable methods due to their efficiency and expressiveness. For instance, variational autoencoders [19] have been shown to learn data manifolds  $\mathcal{M}$  with complex topologies [20, 21]. However, as seen in [22], variational approaches exhibit limitations in learning the distribution  $p_{\mathcal{X}}$  on  $\mathcal{M}$ . Manifold learning flows [6] (Mflow) form a bridge between tractable density and expressiveness. Their support  $\in \mathbb{R}^d$  can inherently match the complex topology of  $\mathbb{R}^D$ , while maintaining a well-defined change of variables transformation. Furthermore, state-of-the-art methods can efficiently calculate the determinant of the Jacobian dot product, as in rectangular normalizing flow [7] (RNF). Denoising normalizing flow [9] (DNF) represents a progression beyond Mflow, where density denoising is used to improve on the density estimation. RNFs can be seen as a parallel approach to DNF but with a more direct methodology, circumventing potential ambiguities stemming from heuristic techniques like density denoising.
Despite the success of manifold learning flows in obtaining a tractable density whilst maintaining expressivity, the learned manifold is not optimized within its latent space representation [23], an outstanding problem for most generative models. For example, due to lack of constrains between latent variables it can be the case that some of them contain duplicate information, i.e., form a degenerate intrinsic basis to describe the manifold, also demonstrated in Section 4.1. Sparse learning [14] was an early attempt to limit this degeneracy and encourage learning of global parameters, introducing solutions such as the relevance vector machine [24, 25] (RVM). The RVM attempts to
limit learning to its minimum necessary configuration, while capturing the most relevant components that can describe major features in the data.
Even stricter approaches are principal [26, 27] and independent component analyses [28, 29], PCA and ICA, respectively. PCA and ICA can obtain a well-defined and optimized representation when convergent; however, PCA and ICA are linear in nature and their non-linear counterparts require kernel definitions. Here, the network-based approach is more general, reducing the need for handcrafting nonlinear mappings. Furthermore, pure PCA and ICA can be limited in modelling some data manifold topologies where strict global independent components are not necessarily a favorable basis [30]. Furthermore, [15] presents a method for efficient post-learning structuring of the latent space, showcasing the advantages of model compactness. We propose a canonical manifold learning flow where the model is motivated to learn a compact and non-degenerate intrinsic manifold basis.
Furthermore, it has been shown that tabular neural networks can show superior performance when regularized via gradient orthogonalization and specialization [17]. They focus on network regularization, which is not specific to one method but requires certain network structure and costly attribution calculations.
Recently, NF methods with flow manifested in the principal components have been proposed [31, 32]. Cranmer et al. [31] focus on datasets that have distinct PCAs that can be efficiency calculated and, therefore, utilize the computed PCAs for efficient density estimation. Cunningham et al. [32] (PCAflow) 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 our method does not confine itself to a pure PCAflow scenario, which has the potential to restrict expressivity. In contrast, our method only loosely enforces orthogonality, while encouraging sparsity, i.e. promoting learning of a canonical manifold basis.
4.1 Canonical intrinsic basis
In this section, in order to motivate the proposed development, we first demonstrate how rectangular manifold learning flows assign a latent space representation in comparison to canonical manifold learning flow. We generated synthetic data on a tilted line with noise in the perpendicular and parallel directions, i.e., a fuzzy line. We sample 1000  $x_{i}$ 's such that  $x_{1} \sim \mathcal{U}(-2.5, 2.5)$  and  $x_{2} = x_{1} + \epsilon$ , where  $\epsilon \sim \mathcal{U}(-0.5, 0.5)$ . In Figure 1(a), a two-dimensional manifold  $\mathcal{M}$  is learned with the vanilla-RNF method using the samples by minimizing the loss given in Equation 4. As expected and seen on the left of Figure 1(a), the probability density is learned satisfactorily, i.e., the model fit well the samples and can generate samples with accurate density. However, when samples on the manifold are generated from the two latent dimensions,  $z_{i}$ , individually, shown in Figure 1 (ii) and (iii), the sampling results are very similar. The implication being that the intrinsic basis is almost degenerate and information is duplicated amongst the latent variables. In contrast, when the same  $\mathcal{M}$  is learned with the proposed CMF, shown in Figure 2, the sampled latent dimensions are close to perpendicular to each other. Each basis is capturing non-degenerate information about  $\mathcal{M}$ , see Figure 1(b), i.e., the two major axes of the dataset.
4.2 Canonical manifold learning
As canonical manifold is not a standard term in literature, let us start by defining a canonical manifold for manifold learning flows.
Definition 4.1. Here, a canonical manifold,  $\mathfrak{M}$ , is a manifold that has an orthogonal and/or sparse basis  $\mathbf{e}_i$  such that  $\mathbf{e}_i \cdot \mathbf{e}_j = 0 \forall \mathbf{y} \in \mathfrak{M}$  and whenever  $i \neq j$ .
The name in Definition 4.1 is inspired from the "canonical basis", i.e., the set of linearly independent generalized eigenvectors of a matrix, as defined in linear algebra. We hypothesize that enforcing learning a canonical manifold as defined in Definition 4.1 during manifold flows will lead to a less degenerate use of the latent space, and consequently a better approximation of the data distribution. As in sparse learning [14], during training necessary dimensions in the latent space will be automatically determined, and those that are necessary will be used in such a way to model a manifold with an
(a) Density plot for a fuzzy line learned with RNF, via Equation 4.
(b) Density plot for a fuzzy line learned with CMF, via Equation 8.
Figure 1: Comparison of density plots for a fuzzy line learned with RNF (a) and CMF (b). Sampled  $1000x_{i}$ 's such that  $x_{1} \sim \mathcal{U}(-2.5, 2.5)$  and  $x_{2} = x_{1} + \epsilon$  with  $\epsilon \sim \mathcal{U}(-0.5, 0.5)$ . The black dots represent samples from the real data manifold. (i) Samples from the model using the full latent space where colors correspond to  $\log p(x)$ . (ii), (iii) Samples from the model using only the first and the second components  $z_{1/2} \in \mathcal{Z}$  while setting the other to zero, respectively. Colors correspond to  $\log p(x)$  once more. All densities have been normalized to  $[-1, 1]$  for visualization purposes. Note that compared to RNF, with CMF different latent variables could capture almost an orthogonal basis for this manifold.
orthogonal local basis to fit the data samples. This, we believe, will enhance the manifold learning process, avoiding overfitting and other learning pathologies associated with correlated attributions to dimensions of the latent space.
In order to obtain a canonical basis we first realize that, as described in Section 2 and the appendix, the metric tensor can be used as concise language to describe the manifold i.e. the transformation of the chart.
$$
G _ {i j} = \sum_ {k} \frac {\partial g _ {\phi} ^ {- 1 , k} (x)}{\partial z ^ {i}} \frac {\partial g _ {\phi} ^ {- 1 , k} (x)}{\partial z ^ {j}} = \sum_ {k} \frac {\partial x ^ {k}}{\partial z ^ {i}} \frac {\partial x ^ {k}}{\partial z ^ {j}}. \tag {6}
$$
This provides an opportunity to regularize the learned representation directly with a single optimization objective. The dependencies on  $G$  are dropped for brevity.
Leveraging the metric tensor representation, we propose to minimize the off-diagonal elements of the metric tensor, enforcing learning of a canonical manifold
$$
\left\| G _ {i \neq j} \right\| _ {1} ^ {1} \triangleq \sum_ {i} \sum_ {j \neq i} \left\| G _ {i j} \right\| _ {1} ^ {1}, \tag {7}
$$
where  $\| \cdot \| _1^1$  is the  $\ell_1$  norm. Using  $\| G_{i\neq j}\| _1^1$  as a cost to minimize serves both the sparsity and generating orthogonal basis. While the  $\ell_2$  norm could be used to enforce a specific embedding [8], it is not in the scope of this work. In order to minimize this cost, the network will have to minimize the dot products  $\partial x / \partial z^i\cdot \partial x / \partial z^j$ . There are two ways to accomplish this. The first is to reduce the magnitudes of the basis vectors  $\partial x / \partial z^i$ . Minimizing the magnitudes with the  $\ell_1$  norm would lead to sparser basis vectors. The second is to make the basis vectors as orthogonal as possible so the dot product is minimized, which serves towards learning a canonical manifold as well. As detailed in [7] and the appendix, the Jacobian product, the bottleneck for both RNF and our method, can be efficiently approximated. With complexity  $\mathcal{O}(id^2)$ , it is less than  $\mathcal{O}(d^3)$  when  $i <   <   d$ , where  $i$  is the conjugate gradients method iterations. Thus, encouraging a canonical manifold through  $G_{i\neq j}$  incurs minimal extra computational cost.
One can imagine an alternative optimization schemes. For instance, the diagonal elements of the metric tensor  $G_{kk}$  can be used to access the transformation of individual  $z_i \in \mathcal{Z}$ . Minimizing the  $\ell_1$  norm  $\sum_k \| G_{kk} \|_1^1$  will encourage sparse learning, akin to a relevance vector machine [14] or network specialization [17]. Nevertheless, minimizing the diagonal elements will clearly not motivate perpendicular components and notably there is no mixing constraint on these components. Consequently, only the magnitude of the transformation can be affected.
Furthermore, one can imagine minimizing the cosine similarity between the basis vectors to promote orthogonality. However, this would not serve towards automatically determining the necessary dimensions in the latent space. Combining the  $\ell_1$  of the diagonal elements and the cosine similarity may also be a solution; however, this would bring an additional weighing factor between the two
losses.  $\ell_1$  loss of the off-diagonal elements of the metric tensor brings together these two objectives elegantly.
Note that the injective nature of the flow, the arbitrary topology of the image space and the finite dimensionality of the chart imply that a canonical solution may not exist for the complete learned manifold. However, this is not prohibitive as the local nature of the proposed method should allow even for isolated, sparse and independent component realizations. Even such solutions can be more preferable when representing complex multidimensional image dataset manifolds. Additionally, the canonical solution is not absolutely but rather statistically enforced, i.e. it is not a strict constraint. This is similar to encouraging orthogonality and specialization of gradient attributions for network regularizations [17].
To this end, combining Equation (7) with the established manifold learning flow log likelihood and reconstruction loss from Section 2.1, i.e., Equation (4), we arrive at the following total optimization objective of the canonical manifold learning flow loss:
$$
\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} (x _ {i})\right) - \log | \det \mathbf {J} _ {h _ {\eta}} ^ {T} (g _ {\phi} ^ {- 1} (x _ {i})) | \right. \\ \left. - \frac {1}{2} \log \left| \det  \mathbf {J} _ {f _ {\theta}} ^ {T} \left(f _ {\theta} ^ {- 1} \left(x _ {i}\right)\right) \mathbf {J} _ {f _ {\theta}} \left(f _ {\theta} ^ {- 1} \left(x _ {i}\right)\right) \right| - \beta \left\| x _ {i} - g _ {\phi} \left(g _ {\phi} ^ {- 1} \left(x _ {i}\right)\right) \right\| _ {2} ^ {2} \right] - \gamma \| G _ {j \neq k} \| _ {1} ^ {1}, \\ \end{array}
$$