SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Babak Esmaeili*
Generative AI Group  
Eindhoven University of Technology  
b.esmaeili@tue.nl
Heiko Zimmermann
Amsterdam Machine Learning Lab  
University of Amsterdam  
h.zimmermann@uva.nl
Robin Walters*
Khoury College of Computer Sciences Northeastern University r.walters@northeastern.edu
Jan-Willem van de Meent
Amsterdam Machine Learning Lab  
University of Amsterdam  
j.w.vandemeent@uva.nl
Incorporating geometric inductive biases into models can aid interpretability and generalization, but encoding to a specific geometric structure can be challenging due to the imposed topological constraints. In this paper, we theoretically and empirically characterize obstructions to training encoders with geometric latent spaces. We show that local optima can arise due to singularities (e.g. self-intersection) or due to an incorrect degree or winding number. We then discuss how normalizing flows can potentially circumvent these obstructions by defining multimodal variational distributions. Inspired by this observation, we propose a new flow-based model that maps data points to multimodal distributions over geometric spaces and empirically evaluate our model on 2 domains. We observe improved stability during training and a higher chance of converging to a homeomorphic encoder.
A well-established idea in machine learning research is that geometric inductive biases can help us learn representations that reflect the underlying structure of a dataset [Bronstein et al., 2021, Higgins et al., 2022]. A key intuition behind this line of research is that such representations make it easier to reason about the similarities of different instances in the dataset, for example by relating inputs using a rotation or translation, which in turn aids interpretability and data-efficiency. Geometric inductive biases have been explored in a wide variety of forms, including models that are defined on hyperbolic or spherical Riemannian manifolds [Lezcano-Casado and Martinez-Rubio, 2019, Ganea et al., 2018], models that are equivariant or invariant with respect to particular symmetry groups [Kondor and Trivedi, 2018, Cohen and Welling, 2016], and work that leverages symmetries to learn disentangled representations that factorize into distinct axes of variation [Higgins et al., 2018].
In this paper, we consider symmetry-based approaches to learning representations in an unsupervised manner by imposing geometric inductive biases on the representation space. In this context, our notion of a representation that "reflects the underlying structure of the data" is a representation that is homeomorphic to the true generative factors. We are specifically interested in the optimization challenges that arise when encoding to geometric spaces such as Lie groups. While a common intuition is that an inductive bias that matches the underlying topology of the data will guide a model towards a homeomorphic representation, there are also indications that certain inductive biases can make a model more difficult to train in practice [Park et al., 2022, Falorsi et al., 2018, Batson et al., 2021], which limits their practical utility.
Figure 1: A GF-VAE consists of an encoder network  $h_{\phi}$  that maps data  $x$  to an intermediate parameter space  $\mathcal{V}$ . The encoded vector  $y$  is split into  $K$  parts, where each sub-vector  $y_{k}$  corresponds to the parameters of a normalizing flow at layer  $k$ . We can sample from the variational distribution  $q_{\phi}$  by first sampling from a uniform prior defined on the Lie group (SO(2) in this case), followed by a sequence of  $K$  bijective transformations  $r(\cdot; y_k)$ . The output of the normalizing flows, denoted as  $z_{K}$ , is then passed to a decoder  $f_{\theta}^{*}$  that maps from the Lie group  $\mathcal{Z}$  to the data space  $\mathcal{X}$ .
To understand why encoding to geometric structures can give rise to optimization challenges, we formalize topological defects that can occur in a randomly initialized encoder, such as discrepancies in the winding number or crossing number relative to those in a homeomorphic encoder. We show that these topological defects will be preserved under continuous optimization, which suggests that escaping these local optima relies on the discrete jumps that are employed during optimization.
To circumvent these obstacles to optimization, we propose Group-Flow Variational Autoencoders (GF-VAEs), which leverage normalizing flows to model complex multimodal distributions on Riemannian manifolds (Figure 1). We show that if we define the mode of the variational distribution to be the representation, normalizing flows can circumvent some of the challenges associated with local optima due to their multimodal nature. Experiments demonstrate that GF-VAEs can escape local optima during the early stages of training, resulting in more reliable convergence to a homeomorphic mapping and a greater degree of continuity after training.
We summarize the main contributions of this paper as follows:
- We characterize topological defects that can arise in encoders that map onto spaces with geometric structure. We show that some obstructions that arise from topological defects cannot be resolved using continuous optimization.
- We define evaluation criteria based on the winding number, crossing number, and continuity to measure topological defects and homeomorphism violations in the encoder.
- We propose GF-VAEs, a new VAE-based model that that employs normalizing flows to define complex distribution on Lie groups. We empirically show that GF-VAEs are able to aid in circumventing identified optimization obstructions.
Homeomorphic Encoders. Our goal is to learn representations in domains where we have prior knowledge of the geometric structure, specifically structure in the form of a symmetry group that can be associated with the input data. We assume that data lies on a low-dimensional manifold  $\mathcal{M}$  that is embedded in a higher-dimensional space  $\mathcal{X} \coloneqq \mathbb{R}^n$  via a mapping  $g_{\mathcal{X}}: \mathcal{M} \to \mathcal{X}$ . This is commonly known as the manifold hypothesis [Bengio et al., 2013]. We denote the image of the mapping by  $\mathcal{M}_x \coloneqq g_{\mathcal{X}}(\mathcal{M}) \subseteq \mathcal{X}$ . Then  $g_{\mathcal{X}}$  is a homeomorphism, or topological isomorphism, onto its image  $g \colon \mathcal{M} \stackrel{\sim}{\longrightarrow} \mathcal{M}_x$ . That is,  $g$  is continuous, bijective, and has a continuous inverse.
We wish to learn an encoder  $f_{\phi} \colon \mathcal{X} \to \mathcal{Z}$  such that its restriction  $f_{\phi}|_{\mathcal{M}_x} \colon \mathcal{M}_x \to \mathcal{Z}$  is a homeomorphism. Following Falorsi et al. [2018], we will define this mapping in terms of a network  $h_{\phi} \colon \mathcal{X} \to \mathcal{Y}$
that maps to an intermediate space  $\mathcal{Y} \coloneqq \mathbb{R}^d$ , followed by a known projection  $\pi: \mathcal{Y} \to \mathcal{Z}$
$$
f _ {\phi}: \mathcal {X} \xrightarrow {h _ {\phi}} \mathcal {Y} \xrightarrow {\pi} \mathcal {Z}. \tag {1}
$$
Lie Groups. Our work focuses the specific case where the manifolds  $\mathcal{M}$  and  $\mathcal{Z}$  are Lie groups. Any set of symmetries may be formally described by a group  $G$ , which is a set of invertible transformations which may be composed using a binary operation  $\cdot : G \times G \to G$ . A Lie group is a group that is also a differentiable manifold. Lie groups describe continuous symmetries such as rotations and translations, and are therefore a natural mathematical setting for any system with spatio-temporal symmetries.
Reasoning about Lie groups is difficult due to their non-flat structure. Lie algebras provide a way to study Lie groups by considering the tangent space  $\mathfrak{g}$  of the manifold at the identity element. The exponential map  $\exp : \mathfrak{g} \to G$  maps points the Lie algebra to points on the group manifold as such. A vector  $v \in \mathfrak{g}$  defines a vector field on  $G$  using the group action to transport  $v$  around  $G$ . Then  $\exp(v)$  is defined to be the point reached by flowing along this vector field for unit time. For more details on Lie groups, we refer the readers to [Hall and Hall, 2013].
Variational Autoencoders on Lie Groups. In this paper, we primarily focus on variational autoencoder (VAE) [Kingma and Welling, 2014, Rezende et al., 2014] as a means to learn a homeomorphic embedding. In this setting, we define a generative model by composing a uniform prior  $p(z) = \mathcal{U}(z)$  on  $\mathcal{Z}$  with a likelihood model  $p_{\theta}(x|z) = \mathcal{N}(x; f_{\theta}^{*}(z), \sigma_x^2 I_n)$  that is defined in terms of a decoder network  $f_{\theta}^{*}: \mathcal{Z} \to \mathcal{X}$ . We use the encoder  $f_{\phi}$  to define a variational distribution  $q_{\phi}(z|x)$ , whose design we discuss below and train the encoder and decoder jointly by optimizing the variational lower bound,
$$
\mathcal {L} _ {\phi , \theta} ^ {\mathrm {V A E}} (x) = \mathbb {E} _ {q _ {\phi} (x | z)} [ \log p _ {\theta} (x | z) ] - D _ {\mathrm {K L}} [ q _ {\phi} (z | x) \| p (z) ] \leq \log p _ {\theta} (x). \tag {2}
$$
Defining a variational distribution on a manifold is generally not straightforward as it requires finding an expression of the density on the manifold or keeping track of the change of volume when projecting to the manifold from the tangent space. Falorsi et al. [2019] define a reparameterized construction for sampling from a Gaussian distribution on  $\mathcal{Z}$  by first sampling  $\epsilon \sim \mathcal{N}(0, I_p)$  from the  $p$ -dimensional Lie algebra associated with  $\mathcal{Z}$ , then rescaling  $\epsilon$  by way of element-wise multiplication  $\sigma_{\phi}(x) \odot \epsilon$  using a network  $\sigma_{\phi}: \mathcal{X} \to \mathbb{R}^p$ , and computing the corresponding element  $z_{\epsilon}$  on the group manifold using the exponential map. By left multiplying this randomly sampled  $z_{\epsilon}$  by the group element  $f_{\phi}(x) \in \mathcal{M}$  to move the mode of the final distribution to its intended location, we obtain the construction
$$
\epsilon \sim \mathcal {N} (\cdot ; 0, I _ {p}), \quad z _ {\epsilon} = \exp \left(\sigma_ {\phi} (x) \odot \epsilon\right), \quad z = f _ {\phi} (x) \cdot z _ {\epsilon}, \quad p \left(z _ {\epsilon}\right) = p (\epsilon) \left| \det  \frac {\partial \exp (\epsilon)}{\partial \epsilon} \right| ^ {- 1}. \tag {3}
$$
We account for the possible change in volume using the change of variable formula. Note that composing group elements  $f_{\phi}(x) \cdot z_{\epsilon}$  does not change the density of the resulting point on the manifold because the Haar measure, a standard choice for Lie groups, is left invariant.
Running Example. As a concrete running example, which we will use throughout this paper, we will consider data on the circle  $\mathcal{M} = \mathrm{SO}(2)$  that is embedded into a space of images  $\mathcal{X}$ . The subset  $\mathcal{M}_x$  of images generated by a function  $g\colon \mathrm{SO}(2)\to \mathcal{M}_x$  corresponds to images of an object that is subject to a one-dimensional rotation. We will consider the case of in-plane rotations of images, as well images of a 3-dimensional object that is rotated around a single axis.
The special orthogonal Lie group  $\mathrm{SO}(2)$  is defined as
$$
\mathrm {S O} (2) := \{z \mid z \in \mathrm {G L} (2), z ^ {T} z = I, \det  (z) = 1 \} = \left\{A (y) := \left[ \begin{array}{c c} y _ {1} & - y _ {2} \\ y _ {2} & y _ {1} \end{array} \right] \mid y \in \mathbb {R} ^ {2}, \| y \| = 1 \right\},
$$
where  $\mathrm{GL}(2)$  is the general linear group, which is the group of invertible  $2\times 2$  matrices under matrix multiplication. The set of images  $\mathcal{M}_x$  lives on a manifold that is homeomorphic to the rotation group  $\mathrm{SO}(2)$ , embedded into the space of images by  $g$ . Because we assume we know the underlying manifold  $\mathcal{M}$ , we can design  $\mathcal{Z}$  to have the same structure. When  $\mathcal{M} = \mathrm{SO}(2)$ , we can use  $\mathcal{V} := \mathbb{R}^2$ . The function  $h_\phi$  in Equation 1 is then an ordinary neural network, and the projection to  $\mathcal{Z}$  is simply the projection onto the unit circle:  $\mathbb{R}^2\to \mathrm{SO}(2)\coloneqq y\mapsto A(y / \| y\|)$ .
Figure 2: Example of topological defects in learned encoders for a VAE with data on  $\mathcal{X} = \mathbb{R}^{32x32}$  in the form of rotations L-shaped tetrominoes and  $\mathcal{Z} = \mathrm{SO}(2)$  (the unit circle). Our goal is to learn an encoder  $f_{\phi}$  that defines a homeomorphism (a continuous bijection with continuous inverse) between the manifold of images of L-shaped tetrominoes  $\mathcal{M}_x \subset \mathcal{X}$  and that of the latent space  $\mathcal{Z} = \mathrm{SO}(2)$ . The encoder  $f_{\phi} = \pi \circ h_{\phi}$  is defined by composing a network  $h_{\phi}: \mathcal{X} \to \mathcal{Y}$  with a projection  $\pi: \mathcal{Y} \to \mathcal{Z}$ . On the right, we show the intermediate space  $\mathcal{Y}_x = h_{\phi}(\mathcal{M}_x)$  and its projection  $\mathcal{Z}_x = \pi(\mathcal{Y}_x)$  for 3 random seeds after convergence. Colours indicate the angle associated with each data point on the manifold  $\mathcal{M}$ . Optimization obstructions can arise when the network  $h_{\phi}$  maps data onto a trajectory  $\mathcal{Y}_x$  that exhibits topological defects, such as the crossing in a "figure 8" shape, which gives rise to discontinuities in the projection  $\mathcal{Z}_x$  onto the latent space.
3 Optimization Obstructions
In practice, training homeomorphic encoders can give rise to optimization challenges. To develop intuition for these challenges, we will consider the running example  $\mathcal{M} = \mathrm{SO}(2)$  with data in the form of in-plane rotated images on  $\mathcal{X} = \mathbb{R}^{32\times 32}$  (Figure 2). We train a VAE in which the network  $h_\phi$  is a standard convolutional network which is paired with a deconvolutional decoder  $f^*:\mathrm{SO}(2)\to \mathbb{R}^{32\times 32}$  (see Appendix A). Figure 2 shows encodings in the intermediate space  $\mathcal{V}$  after training with ELBO with 3 random seeds. For the first two seeds, we see that  $h_\phi (\mathcal{M}_x)$  crosses over itself, resulting in an figure "8" shape. When projected onto  $\mathcal{Z}\coloneqq \mathrm{SO}(2)$ , this results in discontinuities at the crossover points. The second seed in addition also exhibits a sparse region in the intermediate space, leading to a gap in the SO(2) projection. Only the third initialization has converged to the correct topology. These local optima are not unique to this example; obstructions have also been encountered in Falorsi et al. [2018] when trying to learn 3D orientations of a rotating multi-color cube from 2D images using a homeomorphic VAE. Park et al. [2022] show that the homeomorphic VAE cannot generalize well to other shapes and tends to learn a degenerate embedding to a small part of SO(3).
The main observation that we make in this paper is that imposing a geometric structure on the latent space can introduce topological obstructions during optimization. This insight is distinct from the homological obstructions identified by de Haan and Falorsi [2018], who describe the obstructions that emerge from the choice of parameterization on the Lie group. We will refer to the topological defects that we identify in this paper as "optimization obstructions". The problem that we identify here is that randomly-initialized layers have a high probability of exhibiting topological defects (degree, crossing, etc.) that cannot be resolved under continuous optimization using gradient flow. Removal of these topological defects is thus only possible by relying on the jumps coming from performing SGD with a large enough learning rate. This implies that while escaping such local minima is possible, it is difficult and may require many epochs to do so, dramatically slowing training and undercutting the advantages of learning a homeomorphic representation.
We now discuss several specific optimization obstructions. We focus on the case where  $\mathcal{M}$  is the Lie group  $\mathrm{SO}(2)$ , with the same example setting described in Section 2. All the optimization obstructions we consider occur in this case and in the case of higher-dimensional manifolds  $\mathcal{M}$  as well, but are simpler to describe for  $S^1$ . See Appendix B for proofs.
3.1 Figure Eight Local Minima
To more precisely describe obstructions that might arise during optimization, we consider continuous-time training along a gradient flow. We denote the weights of the initialized encoder as  $\phi(0)$  and the trained weights as  $\phi(1)$ . We consider the idealized setting in which  $\phi(t)$  is a continuous function of  $t$ .
Empirically, either at initialization or after some training, we often observe a "figure 8" pattern in  $\mathcal{V}$  of roughly the form  $(h^{\infty} \circ g)(\theta) = (\cos \theta \sin \theta, \sin \theta)^{\top}$ . In such a case, the embedding into  $\mathcal{Z}$  contains two disjoint pieces  $[-3\pi/4, -\pi/4] \cup [\pi/4, 3\pi/4]$ . One half of the circle is mapped to one piece at
the top of the circle and the other half of the circle is mapped to the other disjoint piece at the bottom of circle. Such a mapping takes advantage of the singularity at  $(0,0)^2$  into order to embed  $S^1$  in two continuous pieces. The resulting mapping is nearly bijective, failing to reconstruct on only a small region near the two discontinuities. It is also mostly continuous, having only two discontinuities.
Once this local minimum is obtained, it is very difficult to move out of it using gradient descent. It is unlikely for the two pieces to join together and become a homeomorphic embedding since this would require passing one disjoint segment through the other and reversing its orientation which would violate bijectivity and increase the reconstruction loss.
We make this observation precise by noting that continuous optimization preserves the ordering of points on the circle. Let  $(z_{1}, z_{2}, z_{3}, z_{4})$  denote the four end points of the two disjoint intervals of  $f_{\phi(0)}$ . The ordering of these points is only defined up to cyclic permutations mod  $C_{4}$ . Proposition 3.1 states that continuous optimization must preserve ordering mod  $C_{4}$ . Figure 3 illustrates the proof.
Proposition 3.1 Assume that  $f_{\phi(t)}$  undergoes continuous optimization. Assume that  $f_{\phi(t)} \circ g$  is injective for all  $t$ . The cyclic ordering induced on  $k$  points by  $f_{\phi(0)}$  is equal to  $f_{\phi(1)}$ . Thus a figure 8 embedding, which corresponds to cyclic order  $(z_1, z_2, z_4, z_3) \mod C_4$ , cannot be transformed to a homeomorphic embedding, which has cyclic order  $(z_1, z_2, z_3, z_4) \mod C_4$  or  $(z_4, z_3, z_2, z_1) \mod C_4$ .
In other words, transition from a figure 8 embedding to a homeomorphic embedding is impossible without violating continuity during optimization. This indicates that escaping a figure 8 local optimum during training would need to rely on discrete jumps and likely the stochasticity of the gradient estimate.
Figure 3: The figure 8 pattern in  $\mathcal{V}$  (grey) maps to two disconnected components in  $\mathcal{Z}$ . The cyclic order of these 4 endpoints is preserved by homotopy. Following the parameterization of the data manifold the cyclic order is  $(z_{1},z_{2},z_{4},z_{3})$ , which is distinct from the cyclic order of a homeomorphic embedding, either  $(z_{1},z_{2},z_{3},z_{4})$  or  $(z_{4},z_{3},z_{2},z_{1})$ .
3.2 Degree Obstructions
A second class of topological defects that can arise are encodings with a discrepancy in the winding number. We can compute the degree, or winding number, of a map  $\psi : S^1 \to S^1$  around the origin by summing up the differentials along its path on the sphere,
$$
\omega (\psi) = \frac {1}{2 \pi} \int_ {S ^ {1}} d \psi (\theta).
$$
This concept can be extended to arbitrary connected oriented manifolds, where it is usually referred to as the degree of a mapping. Intuitively, it describes the number of times that the domain manifold wraps around the co-domain manifold.
If the embedded image  $h_{\phi(t)}(\mathcal{M}_x)$  does not contain the origin, then the mapping factors through  $\mathbb{R}^2 \setminus \{(0,0)\}$  and is consequently continuous. Therefore, a continuous path in  $\phi$  yields a continuous path in  $\mathcal{Z}$ . If  $f_{\phi}|_{\mathcal{M}_x}$  is continuous, then the function  $\psi_{\phi} \coloneqq \pi \circ h_{\phi} \circ g : S^1 \to S^1$  has a well-defined degree, also known as winding number  $w(\psi_{\phi}) \in \mathbb{Z}$ . In order for  $f_{\phi}|_{\mathcal{M}_x}$  to be a homeomorphism, the winding number must be  $w(\psi_{\phi}) \in \{-1,1\}$ . Under random initialization, however, the initial network may have winding number equal to any integer. Assuming continuous optimization and continuous embeddings, then  $h_{\phi(t)}(x)$  is a continuous function of both  $x, t$ . We assume that  $h_{\phi(t)}(x) \neq (0,0)$  for any  $t, x$  and thus winding number is defined for any time. The following proposition thus holds.
Proposition 3.2 Winding numbers of the initialized and final model are equal  $w(\psi_{\phi(0)}) = w(\psi_{\phi(1)})$ .
In practice, neither the continuous optimization assumption nor the avoidance of the origin holds. Rather  $h_\phi$  is updated by SGD in discrete jumps and  $h_\phi \circ g$  may map to the origin. Thus, empirically, we do see that the winding number may change during training. However, if the initialization avoids
In practice, we never exactly cross  $(0,0)$  but only pass by it with a small distance.
Figure 4: The initialized encoder  $(t = 0)$  may contain defects as described in Section 3, such as a figure 8 pattern. With a standard VAE, where the learned representation (red point) is the mean of a conditional gaussian, these defects are unlikely to be resolved during optimization, as it would require passing through intermediate parameters (and hence representations) with higher reconstruction loss. Using a multimodal variational distribution, the parameters and corresponding representations are less tightly coupled, in the sense that a continuous change in the parameter space can result in a discontinues change in representation (i.e. the mode of the distribution) without passing through high-loss areas of the parameters space.
the origin, then due to the tendency of the magnitude of the unnormalized embeddings  $h_{\phi}(x)$  to grow during optimization (see Section 3.3), the winding number changing becomes more unlikely. This means that defects in the winding number pose significant obstruction to learning homeomorphic embeddings. The winding number is also the primary optimization obstruction which makes it impractical to remove the hard projection  $\pi$ . If instead we decode directly from  $y \in \mathcal{V}$  but push embeddings to the unit circle using the loss  $||y|| - 1|$ , then it is far more likely we converge to discontinuous embeddings with the incorrect winding number (Appendix D).
3.3 Magnitude Growth in  $\mathcal{V}$
Empirically, we observe the values of the embeddings in  $\mathcal{Y} \coloneqq \mathbb{R}^2$  continually grow during training. This phenomenon makes it more difficult for the embedded data manifold  $h_{\phi}(\mathcal{M}_x)$  to cross the origin and for the winding number to change. We give a theoretical explanation for this behavior.
Consider what would happen if the embedding  $y$  were updated directly based on the gradient of the loss  $\nabla_y\mathcal{L}$  with respect to  $y$ . We assume the loss depends only on  $z = y / \| y\|$ , and so has level sets which are unions of radial rays from the origin. The gradient  $\nabla_y\mathcal{L}$  must then be tangent to a circle about the origin. That is, for  $y = (a,b)$ , the gradient  $\nabla_y\mathcal{L} = (\pm b,\mp a)$ . Under gradient flow, the evolution of  $y$  in time  $y_{t}$  would thus flow along circles of fixed radius and so  $\| y_0\| = \| y_t\|$ . Under gradient descent, however, due to the convexity of the flow lines, which are circular, the embeddings  $y$  will tend to grow in magnitude. For  $\eta \in \mathbb{R}_{>0}$ , we compute