SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Definition 3.5. (Pullback metric) Suppose  $\mu : \mathcal{Z} \to \mathbb{R}^d$  is an immersion onto a Riemannian submanifold  $\mathcal{X} \subset \mathbb{R}^d$  equipped with a Riemannian metric  $\mathbf{I}_d$ . Then, for tangent vectors  $\mathbf{u}_{\mathbf{x}}, \mathbf{v}_{\mathbf{x}} \in \mathcal{T}_{\mathbf{x}}\mathcal{X}$  at each  $\mathbf{x} = \mu(\mathbf{z}) \in \mathcal{X} = \mu(\mathcal{Z})$ , the following holds:
$$
\langle \mathbf {u} _ {\mathbf {x}}, \mathbf {v} _ {\mathbf {x}} \rangle_ {\mathbf {H} _ {\mathcal {X}} = \mathbf {I} _ {d}} = \mathbf {u} _ {\mathbf {x}} ^ {\mathrm {T}} \mathbf {v} _ {\mathbf {x}} = (d \mu_ {\mathbf {z}} \mathbf {u} _ {\mathbf {z}}) ^ {\mathrm {T}} d \mu_ {\mathbf {z}} \mathbf {v} _ {\mathbf {z}} = \mathbf {u} _ {\mathbf {z}} ^ {\mathrm {T}} [ (d \mu_ {\mathbf {z}}) ^ {\mathrm {T}} d \mu_ {\mathbf {z}} ] \mathbf {v} _ {\mathbf {z}} \triangleq \mathbf {u} _ {\mathbf {z}} ^ {\mathrm {T}} \mathbf {H} _ {\mu} (\mathbf {z}) \mathbf {v} _ {\mathbf {z}} = \langle \mathbf {u} _ {\mathbf {z}}, \mathbf {v} _ {\mathbf {z}} \rangle_ {\mathbf {H} _ {\mu}}.
$$
Therefore,  $\mathcal{Z}$  can be viewed as a Riemannian manifold equipped with  $\mathbf{H}_{\mu}(\mathbf{z})$  called a pullback metric. In addition, because  $\mathcal{Z} \subset \mathbb{R}^n$  and  $\mathcal{X} \subset \mathbb{R}^d$ , the pushforward  $d\mu_{\mathbf{z}}$  is equal to the standard Jacobian  $\mathbf{J}_{\mu}(\mathbf{z})$  of  $\mu$ , i.e.,  $\mathbf{H}_{\mu}(\mathbf{z}) = \mathbf{J}_{\mu}^{\mathrm{T}}(\mathbf{z})\mathbf{J}_{\mu}(\mathbf{z})$ .
Definition 3.6. (Isometric immersion) Suppose  $\mathcal{Z}$  is a Riemannian manifold equipped with a Riemannian metric  $\mathbf{H}_{\mathcal{Z}}(\mathbf{z})$  at each  $\mathbf{z} \in \mathcal{Z}$ . A smooth immersion  $\mu : \mathcal{Z} \to \mathbb{R}^d$  onto  $\mathcal{X} \subset \mathbb{R}^d$  is a Riemannian isometry if the metric  $\mathbf{H}_{\mathcal{Z}}(\mathbf{z})$  is equal to the pullback metric  $\mathbf{H}_{\mu}(\mathbf{z})$ , i.e.,  $\mathbf{H}_{\mathcal{Z}}(\mathbf{z}) = \mathbf{J}_{\mu}^{\mathrm{T}}(\mathbf{z})\mathbf{J}_{\mu}(\mathbf{z})$  for every  $\mathbf{z} \in \mathcal{Z}$ . In this case, the length of curves on  $\mathcal{X}$  is equivalent with that on  $\mathcal{Z}$ :
$$
\begin{array}{l} L [ \gamma_ {\mathcal {X}} ] = \int_ {0} ^ {1} d t \sqrt {\langle \dot {\gamma} _ {\mathcal {X}} (t) , \dot {\gamma} _ {\mathcal {X}} (t) \rangle} = \int_ {0} ^ {1} d t \sqrt {\dot {\gamma} _ {\mathcal {X}} ^ {\mathrm {T}} (t) \dot {\gamma} _ {\mathcal {X}} (t)} = \int_ {0} ^ {1} d t \sqrt {\dot {\gamma} _ {\mathcal {Z}} ^ {\mathrm {T}} (t) \mathbf {J} _ {\mu} ^ {\mathrm {T}} (\gamma_ {\mathcal {Z}} (t)) \mathbf {J} _ {\mu} (\gamma_ {\mathcal {Z}} (t)) \dot {\gamma} _ {\mathcal {Z}} (t)} \\ = \int_ {0} ^ {1} d t \sqrt {\dot {\gamma} _ {\mathcal {Z}} ^ {\mathrm {T}} (t) \mathbf {H} _ {\mathcal {Z}} (\gamma_ {\mathcal {Z}} (t)) \dot {\gamma} _ {\mathcal {Z}} (t)} = \int_ {0} ^ {1} d t \sqrt {\langle \dot {\gamma} _ {\mathcal {Z}} (t) , \dot {\gamma} _ {\mathcal {Z}} (t) \rangle_ {\mathbf {H} _ {\mathcal {Z}}}} = L [ \gamma_ {\mathcal {Z}} ], \\ \end{array}
$$
where  $\gamma_{\mathcal{Z}}:[0,1]\to \mathcal{Z}$  and  $\gamma_{\mathcal{X}}:[0,1]\to \mathcal{X}$  are respectively curves from  $\mathbf{z}_0 = \gamma_{\mathcal{Z}}(0)$  to  $\mathbf{z}_1 = \gamma_{\mathcal{Z}}(1)$  and from  $\mathbf{x}_0 = \gamma_{\mathcal{X}}(0)$  to  $\mathbf{x}_1 = \gamma_{\mathcal{X}}(1)$ , when  $\mathbf{x}_0 = \mu (\mathbf{z}_0)$ ,  $\mathbf{x}_1 = \mu (\mathbf{z}_1)$ ,  $\dot{\gamma}_{\mathcal{Z}}(t)\in T_{\gamma_{\mathcal{Z}}(t)}\mathcal{Z}$ , and accordingly  $\dot{\gamma}_{\mathcal{X}}(t)\in T_{\mu (\gamma_{\mathcal{Z}}(t))}\mu (\mathcal{Z})$
With some mild assumptions [36], a sufficiently well-trained decoder  $\mu_{\phi}:\hat{\mathcal{Z}}\to \mathbb{R}^{d}$  of the QVAE can be viewed as a smooth immersion onto the quotient of  $\hat{\mathcal{X}} = \mathcal{X} / G$ . Then, from Definition 3.6,
the following regularization should be added to (5) to make the decoder be a Riemannian isometry between the latent representation  $\hat{\mathcal{Z}}$  and quotient  $\hat{\mathcal{X}}$ :
$$
\mathcal {L} _ {\mathrm {I S O}} (\theta , \phi ; \mathbf {x}, \mathbf {H} _ {\hat {\mathcal {Z}}}, \lambda) = \lambda \mathbb {E} _ {\hat {\mathbf {z}} \sim q _ {\theta} (\hat {\mathbf {z}} | \mathbf {x})} \| \mathbf {J} _ {\mu_ {\phi}} ^ {\mathrm {T}} (\hat {\mathbf {z}}) \mathbf {J} _ {\mu_ {\phi}} (\hat {\mathbf {z}}) - \mathbf {H} _ {\hat {\mathcal {Z}}} (\hat {\mathbf {z}}) \| _ {F}, \tag {6}
$$
where  $\mathbf{J}_{\mu_{\phi}}(\hat{\mathbf{z}})$  is Jacobian of  $\mu_{\phi}$  at  $\hat{\mathbf{z}}$ ,  $\lambda$  is a hyper-parameter, and  $\mathbf{H}_{\hat{\mathcal{Z}}}(\hat{\mathbf{z}})$  is a metric tensor of  $\hat{\mathcal{Z}}$ .
There are two things that should be clarified for practical computation of (6). The first aspect to consider is the selection of the Riemannian metric  $\mathbf{H}_{\hat{\mathcal{Z}}}(\hat{\mathbf{z}})$ . A commonly used and particularly useful choice of  $\mathbf{H}_{\hat{\mathcal{Z}}}(\hat{\mathbf{z}})$  is the scaled Euclidean metric in  $\mathbb{R}^n$ , i.e.,  $\mathbf{H}_{\hat{\mathcal{Z}}}(\hat{\mathbf{z}}) = \mathbf{C}_n = c^2\mathbf{I}_n$  where  $c$  is a constant. This is because it is the case that using the Euclidean distance on the latent representation  $d_{\hat{\mathcal{Z}}}(\hat{\mathbf{z}}_0,\hat{\mathbf{z}}_1) = \sqrt{(\hat{\mathbf{z}}_1 - \hat{\mathbf{z}}_0)^{\mathrm{T}}\mathbf{I}_n(\hat{\mathbf{z}}_1 - \hat{\mathbf{z}}_0)}$  can preserve the geodesic distance on the quotient manifold  $d_{\hat{\mathcal{X}}}(\hat{\mathbf{x}}_0,\hat{\mathbf{x}}_1)$  up to  $c$ . The constant  $c$  can be regarded as either a pre-defined hyper-parameter or a learnable one. More generally, one can use any parameterized sym
Algorithm 1 IQVAEs
Input: data  $\{\mathbf{x}_i\}_{i = 1}^N$  , hyper-parameters  $(\beta ,\lambda)$  group  $G$ $G$  invariant encoders  $(\mu_{\theta}^{G},\sigma_{\theta}^{G})$  , decoder  $\mu_{\phi}$
Initialize  $\theta ,\phi ,\mathbf{C}_n$
while training do
Sample  $\{\alpha_{i}\sim [0,1]\}_{i}^{N},\{\epsilon_{i}\sim \mathcal{N}(\mathbf{0},\mathbf{I})\}_{i}^{N}$
Compute  $\{\mu_{\theta}^{i},\sigma_{\theta}^{i}\}_{i}^{N} = \{\mu_{\theta}^{G}(\mathbf{x}_{i}),\sigma_{\theta}^{G}(\mathbf{x}_{i})\}_{i}^{N}$
Sample  $\{\mathbf{z}_i\}_{i = 1}^N = \{\mu_\theta^i +\sigma_\theta^i\odot \epsilon_i\}_{i = 1}^N$
Shuffle  $\{\mathbf{z}_j\}_{j = 1}^N$  is shuffle(  $\{\mathbf{z}_i\}_{i = 1}^N$
Augment  $\{\tilde{\mathbf{z}}_i\}_{i = 1}^N = \{(1 - \alpha_i)\mathbf{z}_i + \alpha_i\mathbf{z}_j\}_{i = j = 0}^{i = j = N}$
Compute  $\mathcal{L}_{\mathrm{QAE}} = \sum_{i = 1}^{N}\min_{g\in G}\| g*\mathbf{x}_i - \mu_\phi (\mathbf{z}_i)\| _2^2$
Compute  $\mathcal{L}_{\mathrm{KL}} = \sum_{i=1}^{N} D_{\mathrm{KL}}(\mathcal{N}(\mu_{\theta}^{i}, \mathrm{diag}[\sigma_{\theta}^{i}]^{2}) \| \mathcal{N}(\mathbf{0}, \mathbf{I}_{n}))$
Compute  $\{\mathbf{J}_{\mu_{\phi}}^{i}\}_{i = 1}^{N} = \{\mathbf{J}_{\mu_{\phi}}(\tilde{\mathbf{z}}_i)\}_{i}^{N}$
Compute  $\mathcal{L}_{\mathrm{ISO}} = \sum_{i = 1}^{N}\| (\mathbf{J}_{\mu_{\phi}}^{i})^{\mathrm{T}}\mathbf{J}_{\mu_{\phi}}^{i} - \mathbf{C}_{n}\|_{F}$
Optimize  $(\mathcal{L}_{\mathrm{QAE}} + \beta \mathcal{L}_{\mathrm{KL}} + \lambda \mathcal{L}_{\mathrm{ISO}}) / N$  w.r.t  $\theta ,\phi$
metric positive-definite matrix  $\mathbf{C}_n$  for  $\mathbf{H}_{\hat{\mathcal{Z}}}(\hat{\mathbf{z}})$ . Note that such matrices can be achieved from an arbitrary learnable parameterized  $n\times n$  matrix  $\mathbf{M}_n$  by using the Cholesky decomposition as follows:
$$
\mathbf {C} _ {n} = \mathbf {L} _ {n} \left(\mathbf {L} _ {n}\right) ^ {\mathrm {T}}, \quad \mathbf {L} _ {n} = \operatorname {l o w e r} \left[ \mathbf {M} _ {n} \right] - \operatorname {d i a g} \left[ \mathbf {M} _ {n} \right] + \operatorname {d i a g} \left[ \mathbf {M} _ {n} \right] ^ {2} + \varepsilon^ {2} \mathbf {I} _ {n}, \tag {7}
$$
where  $\mathbf{L}_n$  is a real lower triangular matrix with all positive diagonal entries and  $\varepsilon$  is a small nonzero constant. In this case, using the Mahalanobis distance on the latent space  $d_{\hat{\mathcal{Z}}}(\hat{\mathbf{z}}_0,\hat{\mathbf{z}}_1) = \sqrt{(\hat{\mathbf{z}}_1 - \hat{\mathbf{z}}_0)^{\mathrm{T}}\mathbf{C}_n(\hat{\mathbf{z}}_1 - \hat{\mathbf{z}}_0)}$  preserves the geodesic distance of the quotient space of the observation manifold  $d_{\hat{\mathcal{X}}}(\hat{\mathbf{x}}_0,\hat{\mathbf{x}}_1)$ .
The second aspect to consider is the sampling from  $q_{\theta}(\hat{\mathbf{z}}|\mathbf{x})$ . It can be considered as Gaussian vicinal distribution for finite observation samples that smoothly fills the latent space where data is missing. In addition to that, following [8], we use the mix-up vicinal distribution [39] to effectively regularize the entire space of interest to be isometric. As a result, the tractable form of (6) is equal to:
$$
\lambda \mathbb {E} _ {\mathbf {x} _ {i, j} \sim p _ {\mathcal {X}}, \epsilon_ {i, j} \sim \mathcal {N}, \alpha \sim [ 0, 1 ]} \| \mathbf {J} _ {\mu_ {\phi}} ^ {\mathrm {T}} \left(f _ {\theta} ^ {\alpha} \left(\mathbf {x} _ {i, j}, \epsilon_ {i, j}\right)\right) \mathbf {J} _ {\mu_ {\phi}} \left(f _ {\theta} ^ {\alpha} \left(\mathbf {x} _ {i, j}, \epsilon_ {i, j}\right)\right) - \mathbf {C} _ {n} \| _ {F}, \tag {8}
$$
where  $f_{\theta}^{\alpha}(\mathbf{x}_{i,j},\epsilon_{i,j}) = (1 - \alpha)\hat{\mathbf{z}}_i + \alpha \hat{\mathbf{z}}_j$  is the latent mix-up for  $\hat{\mathbf{z}}_{i,j} = \mu_{\theta}^{G}(\mathbf{x}_{i,j}) + \sigma_{\theta}^{G}(\mathbf{x}_{i,j})\odot \epsilon_{i,j}$ . Practically, it is computed by sampling a mini-batch of latent variables, i.e.,  $\{\hat{\mathbf{z}}_i\}_{i = 1}^N = \{\mu_\theta^G (\mathbf{x}_i) + \sigma_\theta^G (\mathbf{x}_i)\odot \epsilon_i\}_{i = 1}^N$ , shuffling it for  $\{\hat{\mathbf{z}}_j\}_{j = 1}^N = \mathrm{shuffle}(\{\hat{\mathbf{z}}_j\}_{i = 1}^N)$ , and mixing-up them.
IQVAE. We define the IQVAE as a class of QVAEs whose objective function is given by the sum of the variational quotient objective  $\mathcal{L}_{\mathrm{QVAE}}$  (5) and Riemannian isometry  $\mathcal{L}_{\mathrm{ISO}}$  (8) as follows:
$$
\mathcal {L} _ {\mathrm {I Q V A E}} (\theta , \phi ; \mathbf {x}, \beta , \lambda , G) \triangleq \mathcal {L} _ {\mathrm {Q V A E}} (\theta , \phi ; \mathbf {x}, \beta , G) + \mathcal {L} _ {\mathrm {I S O}} (\theta , \phi , \mathbf {C} _ {n}; \mathbf {x}, \lambda). \tag {9}
$$
The proposed IQVAE optimization procedure is summarized in Algorithm 1 (see Section A of SM for additional tips). In the IQVAE, we typically set  $\mathbf{C}_n = c^2\mathbf{I}_n$ . For clarity, we will refer to the IQVAE with (7) as IQVAE-M, to distinguish it from the basic version.
We compared our proposed QVAE and IQVAE with six different competitors: the auto-encoder (AE), VAE,  $\beta$ -VAE [19], consistency regularized VAE (CRVAE) [37], QAE [29], and FMVAE [8]. We used the same convolutional architecture<sup>3</sup>, hyper-parameters, optimizer [21], and training scheme for
<sup>1</sup>We assume the straight path between  $\hat{\mathbf{z}}_0$  and  $\hat{\mathbf{z}}_1$  is an on-manifold path for every pair  $\hat{\mathbf{z}}_0, \hat{\mathbf{z}}_1 \in \hat{\mathcal{Z}}$ .
$^{2}$ In (8), we temporarily retrieve the expectation on observations to make things more explicit.
<sup>3</sup>We used  $G$ -orbit pooling (see Section B of SM for details), which can achieve  $G$ -invariance in combination with arbitrary convolutional models, for encoders of QAE, QVAE, and IQVAE.
Table 1: The averaged clustering ARIs and classification test accuracies for  $(2\pi /3)$ -rotated (left) and  $2\pi$ -rotated (right) MNIST across various models, each repeated five times.
METHOD	k-MEANS	GMM	SVM	RF	METHOD	k-MEANS	GMM	SVM	RF
AE	21.8±2.3	29.5±2.2	86.4±0.4	86.2±0.5	AE	10.0±1.0	12.5±2.0	68.3±2.6	71.2±2.1
VAE	23.1±0.5	28.4±0.7	87.4±0.4	87.4±0.4	VAE	11.4±1.2	14.7±1.6	72.8±1.3	73.3±0.9
β-VAE	21.5±1.0	26.5±0.8	88.5±0.5	88.5±0.5	β-VAE	10.1±0.8	11.9±1.1	72.5±0.5	73.5±0.5
CRVAE	34.9±1.3	37.4±2.4	93.0±0.2	93.0±0.3	CRVAE	26.1±2.7	33.4±3.6	78.3±2.4	80.3±2.0
QAE	45.2±4.2	65.5±3.6	93.1±0.3	92.7±0.3	QAE	36.9±5.2	52.9±3.1	91.8±0.6	91.7±0.3
FMVAE	19.0±1.2	25.7±2.0	85.6±0.5	85.1±0.6	FMVAE	10.1±0.9	11.9±1.0	67.6±0.7	69.2±0.4
QVAE	59.3±6.5	72.6±5.2	93.0±0.4	92.5±0.5	QVAE	53.9±7.2	66.3±6.5	93.0±0.1	92.5±0.1
IQVAE	72.8±2.9	74.4±1.6	93.2±0.2	92.9±0.1	IQVAE	70.9±2.9	72.9±1.5	92.9±0.3	92.7±0.4
all models. Implementation details can be found in Section D of SM. After training, we encoded all samples and extracted their latent representations across models. Then, we conducted latent-based downstream clustering tasks by using the simple  $k$ -means and Gaussian mixture model (GMM) [33]. Similarly, we conducted downstream classification tasks by using the support vector machine (SVM) [17] and random forest (RF) classifiers. We used the same hyper-parameters for downstream models regardless of the benchmarked latent representations. Then, we computed the adjusted Rand index (ARI) [38] and test classification accuracy for clustering and classification models, respectively.
4.1 Rotated MNIST
Rotated MNIST consists of randomly rotated images of handwritten digits. It has the rotational group symmetry of  $G = \mathrm{SO}(2)$ . We padded the original MNIST images to size  $32 \times 32$  for improved usability. Then, we generated  $\varphi$ -rotated MNIST by randomly rotating each sample in the MNIST with a uniform distribution of  $[-\varphi / 2, \varphi / 2]$ . We considered both  $\varphi = 2\pi / 3$  and  $\varphi = 2\pi$  cases (see Section C of SM for examples). We used 60,000 samples for training and 10,000 samples for testing.
Basic evaluation. For quantitative analysis, we summarize downstream task performances across the models in Table 1. It clearly shows the proposed QVAE and IQVAE outperform other competitors by a large margin because they consider the group symmetry explicitly. Moreover, the IQVAE shows greatly improved clustering performances even over the QVAE when  $\varphi = 2\pi$ . We deduce that the Riemannian isometry helps the IQVAE recognize a meaning-
ful distance relationship between controversial samples, e.g., digits 6 and 9, as depicted in Figure 2. For qualitative analysis, we also visualize the learned latent representations for the test dataset of the rotated MNIST by using the uniform manifold approximation and projection (UMAP) [28] (see Figure 3, and Section E of SM for all models). It further supports that the IQVAE succeeds in finding
meaningful representations of the dataset.
Aligned reconstructions. Figure 4 shows six images of digit 7 and their reconstructions based on the proposed IQVAE demonstrating that the IQVAE can reconstruct rotated digits in the aligned frame thanks to learning the quotient up to rotations.
Figure 2: Clustering confusion matrices of (left) QVAE and (right) IQVAE for  $\varphi = 2\pi$
Figure 3: UMAP visualizations of (left) VAE and (right) IQVAE for  $\varphi = 2\pi$
ies, e.g., digits 6 and 9, as depicted in Figure 2.1 and latent representations for the test dataset of aproximation and projection (UMAP) [28] (see below) supports that the IQVAE succeeds in finding
Figure 4: Six images of digit 7 and their reconstructions based on the proposed IQVAE.
Better interpolation. Figure 5 compares linear interpolations between digits 1 and 7 for the VAE, QVAE, and IQVAE. It is conducted by parameterizing a linear path as  $\mathbf{z}(t) = \mathbf{z}(0) + (\mathbf{z}(10) - \mathbf{z}(0))\cdot (t / 10)$  and then reconstructing  $\mathbf{x}(t) = \mu_{\phi}(\mathbf{z}(t))$  for  $t\in \{0,1,\dots,10\}$ . As shown in Figure 5 (a), the VAE shows a meaningless interpolation from  $t = 5$  to  $t = 9$  (see Section F of SM for other examples).
Figure 5: (a) Linear interpolations between digits 1 and 7 for the VAE, QVAE, and IQVAE. (b) The geometric volume measures along each linear path of the QVAE and IQVAE.
While the QVAE shows a more convincing interpolation compared to the VAE, it also shows abrupt changes at  $t = 7$  and  $t = 8$ . The IQVAE shows the smoothest interpolation between digits 1 and 7 compared to the other models owing to the isometry regularization. The smooth interpolation of the IQVAE is also quantitatively confirmed in Figure 5 (b) which compares the geometric volume measures along each linear path of the QVAE and IQVAE. The volume measure is computed via  $\sqrt{\operatorname*{det}\mathbf{J}_{\mu_{\phi}}^{\mathrm{T}}(\mathbf{z}(t))\mathbf{J}_{\mu_{\phi}}(\mathbf{z}(t))}$  and can be viewed as a local infinitesimal volume ratio between the latent and observation spaces. The volume measure of the QVAE shows a clear peak near  $t = 8$ , while that of the IQVAE is nearly constant.
Evidence for isometry. To demonstrate that the learned decoding function of IQVAEs upholds the desired Riemannian isometry, we evaluated the condition numbers of the pullback metric tensors and the Riemannian volume measures for all test samples across models. The condition number is defined as the ratio  $\lambda_{\mathrm{max}} / \lambda_{\mathrm{min}}$  where  $\lambda_{\mathrm{max}}$  and  $\lambda_{\mathrm{min}}$  respectively represent the maximum and minimum eigenvalues of the pull-back metric, denoted as  $\mathbf{J}_{\mu_{\phi}}^{\mathrm{T}}(\mathbf{z})\mathbf{J}_{\mu_{\phi}}(\mathbf{z})$ . All computed volume measures
Figure 6: (left) Condition numbers and (right) volume measures for the  $2\pi$  -rotated MNIST for the VAE, QVAE, and IQVAE.
are normalized by their average values. Therefore, condition numbers and normalized volume measures approaching 1.0, accompanied by minimal variances, indicate that the learned latent representation is isometrically isomorphic to Euclidean space. In other words, the learned pull-back metric aligns with the Euclidean metric. As shown in Figure 6, the IQVAE presents condition numbers and volume measures close to 1.0, exhibiting trivial variances compared to other models.
Sample efficiency comparison. Table 2 compares the sample efficiency of IQVAE and the most comparable baseline, QAE, in terms of classification accuracy. The results demonstrate the effectiveness of the proposed method as IQVAE shows more robust performance with respect to variations in the training sample size.
Table 2: Classification test accuracies for  ${2\pi }$  -rotated MNIST with varying training sample sizes.
SAMPLE SIZE	QAE + SVM	IQVAE + SVM
5,000	83.9±0.7	87.6±0.6
10,000	85.5±0.8	89.8±1.0
60,000	91.8±0.6	92.9±0.3
Test-time orbit augmentation. The current QVAE and IQVAE require more computational cost than vanilla VAE due to the  $G$ -orbit pooling encoders, which augment the number of observations  $|G|$  times larger by expanding a data point  $\mathbf{x}$  as a  $G$ -orbit  $G*\mathbf{x}$ . It might be a bottleneck, especially when training the proposed model with larger-scale datasets. To resolve this issue, we suggest the test-time orbit augmentation strategy, which involves using a coarse step size when discretizing the group parameters (e.g., rotation angle) during training, and
Table 3: Training wall-clock time per epoch and clustering ARIs of VAEs, QVAEs, and IQVAEs. We used a single V100 32GB GPU.
METHOD	TIME [s]	k-MEANS	GMM
VAE	1.12	11.4±1.2	14.7±1.6
QVAE (36-36)	10.17	53.9±7.2	66.3±6.5
QVAE (12-12)	4.06	50.8±2.6	59.2±1.5
QVAE (12-36)	4.06	54.6±2.8	63.8±3.9
IQVAE (36-36)	13.22	70.9±2.9	72.9±1.5
IQVAE (12-12)	6.11	59.6±2.7	60.9±1.8
IQVAE (12-36)	6.11	64.3±2.1	65.0±2.1
then switching to a finer step size during the inference phase. Table 3 compares wall-clock training
times and downstream task performances for  $2\pi$ -rotated MNIST datasets for the VAE, QVAE, and IQVAE. The former and latter numbers in parentheses for QVAEs and IQVAEs in the table indicate discretized group sizes used in the training and inference phases respectively. The results show that the proposed technique can scale the training time without severe performance degradation.
Effect of mix-up. We employed the mix-up technique (8) to regularize the entire latent space. By augmenting the latent space with the mix-up, it is possible to smoothly fill gaps in the latent space regions where data might be missing. This becomes particularly crucial when the number of available training samples is limited. To validate this, we trained IQVAEs both with and without the mix-up technique in two scenarios: one with a training sample size of 60,000 and the other with 10,000. Subsequently, we computed the condition numbers of the pull-back metric tensors and mean-normalized volume measures
Figure 7: (left) Condition numbers of the pullback metrics and (right) volume measures of  $2\pi$ -rotated MNIST test samples, with and without the mix-up for IQVAEs. The numbers in parentheses indicate the training sample size.
for test samples. Figure 7 shows that using the mix-up reduces the variances of both condition numbers and volume measures for pull-back metric tensors, demonstrating its efficacy in regularizing the entire latent space. Furthermore, the impact of the mix-up approach becomes more noticeable when the training sample size is smaller.
4.2 MixedWM38
MixedWM38 is a dataset containing 38,015 52  $\times 52$  semiconductor wafer maps, separated into 38 defect categories including 1 normal type, 8 single-defect types, and 29 mixed-defect types [42] (see Section C of SM for examples). We assume rotation and reflection (i.e.,  $G = \mathrm{O}(2)$ ) symmetries are present in the MixedWM38. We used 30,720 samples for training and 7,295 samples for testing.
Basic evaluation. We performed an experiment on the MixedWM38 dataset, similar to the
Table 4: Clustering ARIs and classification test F-scores for MixedWM38 (repeated five times).
METHOD	k-MEANS	GMM	SVM	RF
AE	18.7±0.8	22.0±0.6	81.7±0.6	83.8±0.2
VAE	21.3±0.7	24.9±0.6	82.5±0.3	84.1±0.5
β-VAE	21.3±0.4	27.1±1.1	82.5±0.6	84.5±0.5
CRVAE	21.4±0.6	24.6±0.6	82.3±0.9	84.9±1.0
QAE	22.7±0.9	30.7±1.5	84.4±0.3	85.1±0.2
FMVAE	17.6±0.4	25.8±0.7	81.4±0.2	83.1±0.4
QVAE	23.3±0.5	31.1±0.5	85.0±0.4	85.1±0.1
IQVAE	27.5±0.2	35.3±1.4	84.4±0.2	85.6±0.3
one previously conducted on rotated MNIST. Unlike the rotated MNIST, the MixedWM38 dataset requires a multi-label classification task, which is a key difference between the two experiments. Thus, we use the F-score as a classification accuracy for this problem. Table 4 summarizes the downstream task performances of eight models on the MixedWM38 dataset (see Section E of SM for UMAP visualizations). As shown in the table, the proposed methods have a higher performance on the downstream task than the other models.
OoD evaluation. We further demonstrate the effectiveness of the proposed IQVAE in detecting OoD samples, an important task in the semiconductor engineering field. In order to mimic a real-world setting, we divided the 7,015 single-defect type data (8 classes: C, D, EL, ER, L, NF, S, R; see Section C of SM) into two sets: C-EL-L-S as the in-distribution set, and D-ER-NF-R as the OoD set. Then, we randomly rotated all samples in the C-EL-L-S subset, and considered the rotated samples an in-distribution set, aligning with the nature of the semiconductor manufacturing process [43]. Using auto-encoding models trained on the original C-EL-L-S data, we formulated an OoD detection task, the goal of which is to differentiate between samples from the rotated C-EL-L-S set (in-distribution) and
Table 5: AUCROCs and AUPRCs for MixedWM38 across the models (repeated five times).
METHOD	AUCROC	AUPRC
AE	47.3±3.6	27.2±3.3
VAE	62.4±3.1	34.3±2.4
β-VAE	66.1±4.1	41.4±6.8
CRVAE	80.5±6.1	60.4±10.2
QAE	90.2±3.7	81.8±3.4
FMVAE	69.8±5.6	42.5±7.2
QVAE	90.3±1.1	79.5±3.9
IQVAE	93.8±0.4	88.3±1.7
samples from the D-ER-NF-R set (OoD). We assessed the standard distance-based OoD detection: it labels an instance as OoD if its latent vector is farther than a pre-defined threshold from any
class-conditional mean latent vector of learned in-distribution sets, in terms of the Mahalanobis distance [24] (see Section G of SM for details). Table 5 shows the comparison of competing models using standard metrics (AUCROC and AUPRC) for threshold-based OoD detection, demonstrating that the proposed IQVAE excels in the OoD tasks as it accurately learns the data structure.
4.3 SIPaKMeD