SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\| y - \eta \nabla_ {y} \mathcal {L} \| ^ {2} = (a \mp \eta b) ^ {2} + (b \pm \eta a) ^ {2} = (a ^ {2} + b ^ {2}) (1 + \eta^ {2}) > \| y \| ^ {2}.
$$
In practice, however, we do not update  $y$  based on  $\nabla_y\mathcal{L}$  but rather based on the gradient with respect to model parameters  $\phi$ . Let  $F\colon \Phi \to \mathcal{V}$  be the map from model parameters  $\phi$  to  $y$  given fixed input data  $x$ . Then the actual update to  $y$  is  $\tilde{\nabla}_y\mathcal{L} = dF^T\circ \nabla_\phi \mathcal{L}$  where  $dF$  is the total derivative or Jacobian of the map  $F$ . Since  $\nabla_{\phi}\mathcal{L} = (dF)\nabla_y\mathcal{L}$  we have  $\tilde{\nabla}_y\mathcal{L} = dF^T dF\nabla_y\mathcal{L}$ . The angle between  $\tilde{\nabla}_y\mathcal{L}$  and  $\nabla_y\mathcal{L}$  is bounded by some  $\theta$  a quantity depending on the eigenvalues of the operator  $dF$ . Given that  $\nabla_y\mathcal{L}$  is tangential to the circle, assuming for simplicity  $\tilde{\nabla}_y\mathcal{L}$  has constant length  $L$  and uniform distribution  $[- \theta ,\theta ]$  in angle to  $\nabla_y\mathcal{L}$ , the norm of  $y$  still grows in expectation.
Proposition 3.3 Assume a circle of radius  $R$ . Let  $v$  be a random vector at  $y$  on the circle of length  $L$  with angle to the tangent uniform in  $[-\theta, \theta]$ . Then
$$
\mathbb {E} \left[ \| y + v \| ^ {2} \right] = L ^ {2} + R ^ {2} > R ^ {2} = \| y \| ^ {2}.
$$
4 GroupFlow-VAE
The topological obstructions to optimization that we identified in Section 3 cannot be resolved under continuous optimization. Moreover, even if we allow for a degree of discontinuity, resolving obstructions like a figure 8 will require violating bijectivity, which implies that the reconstruction loss
must increase to escape this defect. This suggests that a VAE with a reparameterized construction as described in Equation 3 will be susceptible to local optima, which aligns with the empirical observation that homeomorphic VAEs can be difficult to train.
This raises the question of whether we can make VAEs less susceptible to topological obstructions by employing a different parameterization from that of Equation 3. More concretely, we will consider a parameterization that admits multiple modes in the variational distribution  $q_{\phi}(z|x)$ , rather than the unimodal construction in Equation 3, and define  $f_{\phi}(x)$  in terms of the mode,
$$
f _ {\phi} (x) := \underset {z} {\arg \max } q _ {\phi} (z | x). \tag {4}
$$
The intuition behind this approach is illustrated in Figure 4. If the variational distribution contains multiple modes, rather than a single peak centered at  $\pi(h_{\phi}(x))$ , then it may become possible for small changes to result in non-local changes, such as the reordering of points that is needed to untwist a figure 8, by switching between modes in the variational distribution. Moreover, increasing the number of parameters of  $q_{\phi}(z|x)$  is likely beneficial for escaping some of the topological obstructions as some defects such as self-intersection are less likely to occur in high-dimensional spaces.
Motivated by this intuition, we consider a parameterization of  $q_{\phi}(z|x)$  that employs normalizing flows [Rezende and Mohamed, 2015]. In a normalizing flow, the sample  $z$  is defined as a push forward of sequence of smooth bijective transformations, which makes it possible to reshape a simple unimodal distribution into a more complex multimodal distribution. The probability of a sample from the final density can be computed by repeatedly applying the rule for change of variables. Concretely, given a base distribution  $p(z_0)$  and a sequences of bijective transformations  $r_k: \mathcal{Z} \to \mathcal{Z}$ , we obtain a sample  $z = z_K$  and probability by first sampling  $z_0 \sim p(z_0)$  and defining a sequence of transformations
$$
z _ {k} = r _ {k} \left(z _ {k - 1}\right), \quad \log p \left(z _ {k}\right) = \log p \left(z _ {k - 1}\right) - \log \left| \det  \frac {\partial r _ {k} \left(z _ {k - 1}\right)}{\partial z _ {k - 1}} \right| \quad \text {f o r} k = 1 \dots K. \tag {5}
$$
Normalizing flows have been used to define distributions on geometric structures such as Lie groups by either defining a flow on the Lie algebra and computing the push-forward density of the exponential map [Falorsi et al., 2019], or designing structure-specific transformations [Rezende et al., 2020].
In the GF-VAE, we will define a construction in which the encoder network returns the parameters of the flow. We show an overview of the architecture in Figure 1. We define a network  $h_{\phi}:\mathcal{X}\to \mathcal{Y}:= \mathbb{R}^{K\times l}$  where  $K$  and  $l$  are the number of flow layers and parameters respectively, a sequence of bijective transformations  $\{r(\cdot ;y_k)\}_{k = 1}^K$  parameterized by  $\{y_{k}\}_{k = 1}^{K}$ , and a base distribution distribution  $p(z_0)$  which we define as a distribution on the group. The sequence  $\{r(\cdot ;y_k)\}_{k = 1}^K$  is then used to define a new distribution on  $\mathcal{Z}$  given an  $x$  by transforming the base distribution. This defines a conditional flow that can be trained using a stand lower bound (Eq. 2),
$$
z _ {0} \sim p (z _ {0}), \qquad z _ {k} = r (z _ {k - 1}; y _ {k}), \qquad \log q (z _ {k} | x) = \log q (z _ {k - 1} | x) - \log \left| \det \frac {\partial r (z _ {k - 1} ; y _ {k})}{\partial z _ {k - 1}} \right|.
$$
The choice of the flow  $r$  is very important here as normalizing flows are typically defined on flat spaces. This means that, for a specific manifold  $\mathcal{M}$ , additional care must be taken when designing them to ensure that they are a diffeomorphism from  $\mathcal{Z}$  to itself [Rezende et al., 2020, Durkan et al., 2019, Mathieu and Nickel, 2020]. In this paper, we employ a similar method as Falorsi et al. where we define an affine layer followed by a single layer of spline flow, and a Tanh layer multiplied by  $\pi$  as the last layer to push all the probability density between  $-\pi$  and  $\pi$ .
We perform a series of experiments to evaluate the difficulty of learning a homeomorphic encoder. Concretely, we investigate how often we fall into one of the failure cases described in Section 3 in standard VAEs with geometric latent spaces in practice. Subsequently, we examine how well GF-VAE can circumvent these topological obstructions during training.
Baselines. Throughout our experiments, we compare against (1) a standard VAE, and (2) a supervised VAE where in addition to maximizing the ELBO, the encoder is trained to predict the ground truth representations. These two scenarios serve as extremes on the spectrum of guiding the model to the right representation. We also compare against a deterministic autoencoder (AE) which does not
Figure 5: Top: Latent traversals in the decoder for a VAE and GF-VAE after training. Bottom: The representations in latent space  $\mathcal{Z}$  for VAEs and GF-VAEs initialized with 15 different random seeds. For VAEs, the gray line inside the circle shows the (scaled-down)  $y$ -traversals. For GF-VAEs, we it is difficult to visualize  $\mathcal{Y}$  space as it is high-dimensional space. Therefore, to inspect for obstructions, we show the traversal of the  $z$  vectors instead.
regularize the embedding space. We also tried regularizing the y-space to be close  $S_{1}$  (by penalizing  $(\|y\|_2 - 1)^2$ ) in order to mitigate the optimization problem discussed in Subsection 3.3, which we refer to as "reg-y" loss. Finally, we evaluate the  $\beta$ -VAE objective [Higgins et al., 2017] which increases the regularization on the latent space by upweighting the KL term in Eq 2. All models employ a 4-layer CNN architecture for the encoder and decoder with LeakyReLU activations. For the decoder, we also experiment with action-decoder proposed by Falorsi et al., which we found helpful for learning a homeomorphic mapping. The action-decoder uses a special first layer where the group action is applied to a set of learned Fourier coefficients rather than directly being passed as input to the architecture. For further details regarding our experiments, please refer to Appendix A.
Evaluation. We evaluate all models based on two criteria: (1) Has the encoder learned a homeomorphic mapping? and (2) Has the decoder learned a good model of the data? Assessing whether a learned mapping is homeomorphic is challenging. To verify homeomorphism, we follow the evaluation proposed in Falorsi et al. [2018] by examining whether the encoder yields a continuous path when interpolating in the data manifold from  $-\pi$  to  $\pi$ . Details on evaluating continuity are provided in Appendix C. To determine how often the models encounter the topological obstructions described in Section 3, we also report crossing and winding numbers in Appendix E. A crossing number greater than 0 implies a "figure 8" obstruction, and a winding number that is not equal to 1 or  $-1$  implies winding number obstruction. Lastly, we measure the log-likelihood to assess how well each model approximates the data manifold. If the model has diverged during training due to posterior collapse, we report it as a non-homeomorphic mapping and ignore its continuity score in the average.
5.1 Images: SO(2)
In our first experiment, we train on images of an L-shaped tetromino [Bozkurt et al., 2021], a teapot, and an airplane [Shilane et al., 2004]. The SO(2) manifold corresponding to each object is made by rotating the image of the object around the center. We report our findings in Tables 1 and 3.
Table 1: Number of learned homeomorphic encoders and their continuity score for different objectives trained with 15 different random seeds.
L-shaped Tetromino	Teapot	Airplane
# H.	Continuity	# H.	Continuity	# H.	Continuity
AE	2/15	137.28 ± 54.15	0/15	173.37 ± 21.99	0/15	132.95 ± 40.61
VAE (β = 1)	0/15	117.02 ± 23.17	6/15	14.21 ± 7.30	0/15	86.72 ± 65.12
VAE (β = 4)	2/15	132.45 ± 58.79	15/15	1.22 ± 0.05	5/15	67.88 ± 59.20
VAE + y-reg	1/15	152.71 ± 74.20	1/15	14.51 ± 81.	3/15	124.10 ± 93.53
GF-VAE (β = 1)	5/15	50.51 ± 58.18	7/15	21.07 ± 24.42	0/15	63.94 ± 34.43
GF-VAE (β = 4)	9/15	24.04 ± 29.14	13/15	8.16 ± 17.89	7/15	27.70 ± 29.23
Action-GF-VAE (β = 4)	10/15	19.35 ± 23.37	13/15	3.14 ± 1.74	9/15	33.26 ± 40.03
Sup-VAE	15/15	5.74±0.49	13/15	5.67±0.34	0/15	96.78 ± 38.
Figure 6: Latent traversals in the decoder of a GF-VAE trained on Tetrominoes with a torus latent spaces. In each traversal, we keep the value for one Lie group fixed and do a geodesic interpolation across the other group. The traversals in the latent space highlighted are by red and orange.
We observe that even though the types of obstructions vary across images, both VAE and AE in general fail to learn a homeomorphic encoder. The GF-VAE objective improves performance noticeably across both metrics. We also did not find  $y$  regularization to be very helpful as it mainly stabilized training towards whatever mapping that was achieved at the early stages of training. We observed the main failure case for most of the non-homeomorphic encoders was due to discontinuity emerging from figure-eight obstruction. Moreover, looking at the learning curves in Figure 9, we observe that the winding number is susceptible to change during training which suggests that the continuity assumption is of importance in proposition 3.2. GF-VAE resolves both issues, which allows us to interpolate nicely in the  $\mathcal{Z}$ -space (Figure 5). What is very surprising is that in the case of airplanes, we see that even supervised objective fails to overcome these optimization obstructions, while a GF-VAE is able to achieve this at a much better rate.
We also observe that increasing the  $\beta$  value for the KL term generally helps. This is perhaps unsurprising given that a high  $\beta$  value encourages the latent space to cover the prior and therefore discourages the winding numbers from being 0 and improves the performance on the continuity metric. In the case of the teapots, we in fact observe that increasing  $\beta$  is sufficient to learn a homeomorphic encoder. However, GF-VAE still generally scores better in terms of continuity and winding numbers.
5.2 Images: Torus
In this experiment, we consider a ring torus as the latent space, which is homeomorphic to the Lie group  $\mathrm{SO}(2) \times \mathrm{SO}(2)$ . We create a dataset homeomorphic to this group by independently rotating the L-shaped tetromino in both orientation and color. All models can be extended to Tori by simply duplicating the latent space and learning multiple encodings to  $\mathrm{SO}(2)$  in parallel. This is done by defining a factorized density  $q_{\phi}(z^{(1)}, z^{(2)}|x) = q_{\phi}(z^{(1)}|x)q_{\phi}(z^{(2)}|x)$ , where each distribution  $q_{\phi}(z^{(i)}|x)$  defines a distribution on a  $S^1$ . Ideally, we want one subgroup to correspond to colour and the other to orientation. In this setting, we evaluate homeomorphism by picking 10 different values in either colour or orientation and measuring the continuity of the encoded path when interpolating in the data manifold on the other attribute. We identify the encoder as homeomorphic if the average continuity score of all 10 paths is below a certain threshold (see Appendix C). Unsurprisingly, we found that learning a homeomorphic mapping in tori is more challenging compared to circles. Out of 15 runs, the VAE models learn a homeomorphic mapping 1 time, while a GF-VAE manages to learn a homeomorphic mapping 7 times. To be able to align each Lie group with the corresponding attribute, we also used the weakly supervised proposed by Locatello et al., where the model receives a sequence of images in which only the angle or the colour is changing. We show the latent traversal in the decoder model for one of the successful runs of GF-VAE in Figure 6. As we can see, the model has successfully managed to disentangle colour and orientation in the latent space.
Learning Geometric Representations. There has been a large amount of work concerned with learning representations from data with geometric structure in the unsupervised or weakly supervised setting. One line of work has focused on the use of geometric spaces such as lie groups in latent space [Davidson et al., 2018, Falorsi et al., 2018, Perez Rey et al., 2020, Vadgama et al., 2022]. In [Miao et al.], it was argued that that geometric inductive biases in VAEs should be Incorporated via a deterministic mapping rather than the prior which is consistent with our work. As we show in this
work, naively incorporating a geometric bias leads to topological obstructions. Another line of work has focused on inferring the latent structure by exploring the local connectivity information [Moor et al., 2020, Chen et al., 2021, Lee et al., 2021, Pfau et al., 2020].
Learning Disentangled Representations. Topological group structure can be used to define a notion of disentanglement that is based on equivariant properties under group transformations [Higgins et al., 2018]. There exists a body of work that aims to learn both disentangled and equivariant group representations. One set of methods relies on agent actions to predict the group element [Caselles-Dupré et al., 2019, Quessard et al., 2020]. An adjacent line of work is to regularize encoders to be equivariant with respect to the group action, using triplets of the form  $(x_{t},m_{t},x_{t + 1})$  or longer sequences [Guo et al., 2019, Dupont et al., 2020, Tonnaer et al., 2022]. Some approaches focus on learning symmetry-based disentangled representations in fully unsupervised settings or by enforcing commutativity in the latent Lie group [Yang et al., 2022, Zhu et al., 2021]. Another related area disentangles class and pose in the latent space [Marchetti et al., 2022, Winter et al., 2022]. Our work does not focus on disentanglement, but the topological obstructions that we describe are relevant to this domain.
Topological Obstructions in Learning. The topological obstructions that we consider in this work are distinct from the homological obstructions that have been characterized in prior work [de Haan and Falorsi, 2018, Batson et al., 2021, Falorsi et al., 2018, Perez Rey et al., 2020]. Theorem 1 by de Haan and Falorsi [2018] defines homological obstructions as follows: For any latent space  $\mathcal{Z}$  with non-trivial topology, it is possible to learn an encoder  $f_{\phi}$  that is continuous when restricted to  $\mathcal{M}_x \subset \mathcal{X}$ , but this encoder must be discontinuous on the full space  $\mathcal{X}$ . For this reason, Falorsi et al. [2018] and others [Xu and Durrett, 2018, Meng et al., 2019] use the two-part encoder from equation 1, inserting discontinuous layers  $\pi$  when mapping to circles, spheres,  $\mathrm{SO}(n)$ , or other manifolds. This explicit discontinuity circumvents the homological obstruction without forcing the linear layers of the network to approximate discontinuities using large weights, which we and others find leads to instability during training and inferior reconstructions (Section 5).
Normalizing Flows on Manifolds. In recent years, there has been a surge of interest in extending normalizing flows, originally formulated for Euclidean spaces, to Riemannian manifolds [Rezende et al., 2020, Mathieu and Nickel, 2020, Kohler et al., 2021, Durkan et al., 2019]. One approach involves leveraging the Lie group structure of the manifold to define a parametrization of the flow [Rezende et al., 2020], which is also the strategy we adopt in this work. These recent advancements have paved the way for applying normalizing flows to the group  $\mathrm{SO}(3)$  in order to learn pose estimation in molecular structures [Kohler et al., 2023] and images [Liu et al., 2023, Murphy et al., 2021]. Our work differs from these efforts in that the primary objective of our method is not to target flexible distributions on Riemannian manifolds, but rather to demonstrate that utilizing a flow as a variational distribution aids optimization in learning a homeomorphic embedding.
In this paper, we investigate obstructions to optimization that can arise when learning encoders for topological spaces. We classify different types of obstructions, provide evidence these are encountered in practice, and give mathematical explanations for how they occur under certain assumptions. We propose GF-VAE, a novel model that employs normalizing flows as variational distributions to help circumvent these issues and show its effectiveness across several datasets when encoding to circles and tori. This work contains several limitations. Firstly, our theoretical analysis is limited by the idealized assumptions necessary to analyze the method using topological tools which do not exactly match those encountered in practice. Secondly, the metrics we define such as winding and crossing numbers are harder to define and compute for higher dimensional manifolds. Future work includes expanding our analysis and techniques to a wider array of Lie groups and non-group manifolds.
This work was supported by NSF grants #2107256 and #2134178. We also would like to thank Patrick Forre, Sharvaree Vadgama, Marco Federici, and Erik Bekker for helpful discussions.
We train all our models for 150 epochs with a batch-size of 600. For optimization, we use the RAdam optimizer [Liu et al., 2019] with a learning rate of 5e-4. For all image datasets, we use a 4-layer CNN with kernel, stride, and padding of size 4, 2 and 1 respectively followed by a leakyReLU activation (Table 2) for the encoder and a Sigmoid activation for the decoder. The network  $\sigma_{\epsilon}^{2}$ , shares the same architecture with the only difference being the last layer, which is a fully-connected network followed by a Softplus activation which is common in standard VAEs. In all our experiments, we used  $K = 1$  for the GF-VAE models as it was sufficient to avoid the optimization obstructions mentioned in the paper. All models were initialized and trained with  $15^{3}$  different random seeds.
Table 2: Architecture of the encoders and decoders employed for all image datasets.
Encoder
Input 32 × 32 images
4 × 4 conv. 32 stride 2, LeakyReLU.
4 × 4 conv. 32 stride 2, LeakyReLU.
4 × 4 conv. 64 stride 2, LeakyReLU.
4 × 4 conv. 64 stride 2, LeakyReLU.
F.C. 2, π(y) := y/||y|.
Decoder
Input z ∈ R2 s.t. ||z||2 = 1
4 × 4 deconv. 64, stride 2, ELU.
4 × 4 deconv. 64, stride 2, ELU.
4 × 4 deconv. 32, stride 2, ELU.
4 × 4 deconv. 3, stride 2, Sigmoid.
We include the proofs for the propositions in the main text.
B.1 Figure Eight Local Minimum
Proposition B.1 Assume that  $f_{\phi(t)}$  undergoes continuous optimization and is thus continuous in  $t$ . Assume that  $\pi \circ f_{\phi(t)} \circ g$  is injective for all  $t$ . The cyclic ordering induced on  $k$  points by  $f_{\phi(0)}$  is equal to that induced by  $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$ .
Proof: Since we assume  $\pi \circ f_{\phi(t)} \circ g$  is injective for all  $t$ , the path  $\mathbf{z}(t) = (\pi \circ f_{\phi(0)} \circ g(\theta_i))_{i=1}^4$  is inside the  $k$ -fold configuration space on  $S^1$  defined  $\operatorname{Conf}_k(S^1) = \{(z_1, \ldots, z_k) \in (S^1)^k : z_i \neq z_j \text{ for } i \neq j\}$ . In order to prove the claim, we will show that the path-connected components of  $\operatorname{Conf}_k(S^1)$  correspond to cyclic orderings of  $(z_1, \ldots, z_k)$  and thus the start and end point of every path share a cyclic ordering.
Mapping  $(z_{1},\ldots ,z_{k})\mapsto (z_{k},(z_{k}^{-1}z_{1},\ldots ,z_{k}^{-1}z_{k - 1}))$  gives a homeomorphism  $\operatorname{Conf}_k(S^1)\cong$  SO(2)  $\times$  Conf  $k - 1(S^{1}\setminus \{1\})\cong$  SO(2)  $\times$  Conf  $k - 1(\mathbb{R})$  . Let  $\tilde{z}_i = z_k^{-1}z_i$  . Consider  $D =$ $\{(\tilde{z}_1,\dots ,\tilde{z}_{k - 1}):\tilde{z}_1 <   \dots <  \tilde{z}_{k - 1}\} \subset \mathrm{Conf}_{k - 1}(\mathbb{R}).$
We can identify the connected components of  $\mathrm{Conf}_{k-1}(\mathbb{R})$ . The set  $D$  is a fundamental domain for the action of the symmetric group  $S_{k-1}$  on  $\mathrm{Conf}_{k-1}(\mathbb{R})$ . Thus  $\mathrm{Conf}_{k-1}(\mathbb{R}) = \coprod_{\sigma \in S_k} \sigma(D)$  is a disjoint union. Linear interpolation shows  $D$  is connected. The sets  $D$  and  $\sigma(D)$  for  $\sigma \in S_k$  are not connected. Consider a path from  $\mathbf{z} = (z_1, \ldots, z_k) \in D$  to  $\sigma(\mathbf{z}) \in \sigma(D)$ . The element  $\sigma$  must reverse the order of at least two elements  $z_j < z_i$ . Thus the function  $f(\mathbf{z}) = z_i - z_j$  must take the value 0 over the path by intermediate value theorem. Hence the path cannot be in  $\mathrm{Conf}_{k-1}(\mathbb{R})$ . Thus the connected components of  $\pi_0(\mathrm{Conf}_{k-1}(\mathbb{R})) \cong S_{k-1}$ .
Since  $\mathrm{SO}(2)$  is connected,  $\pi_0(\operatorname{Conf}_k(S^1)) \cong S_{k-1}$ . That is each connected component of  $\operatorname{Conf}_k(S^1)$  is labeled by an element of  $S_{k-1}$  describing the ordering of  $\tilde{z}_1, \ldots, \tilde{z}_{k-1}$  in  $\mathbb{R}$ . Each ordering of  $(\tilde{z}_1, \ldots, \tilde{z}_{k-1})$  in turn corresponds to a different cyclic ordering of  $z_1, \ldots, z_k$  in  $S^1$ , that is, a different element of  $S_k / C_k$ . Thus two  $k$ -point configurations are homotopic if and only if they have the same cyclic ordering.
<sup>3</sup>We used a higher number of random seeds than normal to account for the training instability.
B.2 Degree Obstruction
Proposition B.2 The winding number of the initialized model and final model are equal  $w(\pi' \circ h_{\phi(1)} \circ g) = w(\pi' \circ h_{\phi(1)} \circ g)$ .
Proof: The winding number of a map is a continuous function  $t \mapsto w(\pi' \circ h_{\phi(t)} \circ g)$ . Since the output space  $\mathbb{Z}$  is discreet, the winding number must be constant in  $t$ .
B.3 Magnitude Growth in  $\mathcal{V}$