SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Qing Su $^{1*}$ , Anton Netchaev $^{2}$ , Hai Li $^{3}$ , and Shihao Ji $^{1}$
<sup>1</sup>Georgia State University, <sup>2</sup>U.S. Army ERDC, <sup>3</sup>Duke University
Current self-supervised learning (SSL) methods (e.g., SimCLR, DINO, VICReg, MOCOV3) target primarily on representations at instance level and do not generalize well to dense prediction tasks, such as object detection and segmentation. Towards aligning SSL with dense predictions, this paper demonstrates for the first time the underlying mean-shift clustering process of Vision Transformers (ViT), which aligns well with natural image semantics (e.g., a world of objects and stuffs). By employing transformer for joint embedding and clustering, we propose a bi-level feature clustering SSL method, coined Feature-Level Self-supervised Learning (FLSL). We present the formal definition of the FLSL problem and construct the objectives from the mean-shift and  $k$ -means perspectives. We show that FLSL promotes remarkable semantic cluster representations and learns an encoding scheme amenable to intra-view and inter-view feature clustering. Experiments show that FLSL yields significant improvements in dense prediction tasks, achieving 44.9 (+2.8)% AP and 46.5% AP in object detection, as well as 40.8 (+2.3)% AP and 42.1% AP in instance segmentation on MS-COCO, using Mask R-CNN with ViT-S/16 and ViT-S/8 as backbone, respectively. FLSL consistently outperforms existing SSL methods across additional benchmarks, including UAV object detection on UAVDT, and video instance segmentation on DAVIS 2017. We conclude by presenting visualization and various ablation studies to better understand the success of FLSL. The source code is available at https://github.com/ISL-CV/FLSL.
Following its success in natural language processing (NLP) [47, 5, 20], self-supervised learning (SSL) with transformer [58, 22] has emerged as a highly effective strategy and a popular model choice over the CNN-based counterparts in vision tasks. The remarkable performance achieved by SSL has been demonstrated by SimCLR [14], MOCOV3 [16], DINO [10], VICReg [3], SwAV [9], BYOL [27], and among others. Without relying on manual supervision, a successful paradigm of SSL promotes semantic representations conducive to the downstream tasks, e.g., classification, detection and segmentation. However, most existing SSL methods operate at the instance-level, where an encoder is trained to maximize the agreement of the representations of multiple augmented views of an image. Though demonstrating strong performance on the classification tasks [14, 29], the instance-level SSL is inherently misaligned with the dense prediction tasks, such as object detection, where the lower level semantic information plays a bigger role than the instance-level semantic information. This leads to inferior transferability to those dense prediction tasks.
Recent attempts to bridge the semantic gap are mainly based on region [50], patch [69, 21], or pixel (i.e., dense feature) matching tasks [63, 73, 38] with optional instance-level objectives. However, learning of distinct representation for each image patch or region still mismatches the natural semantics within an image (referred to as local semantics), where features of the same semantics should be highly correlated other than being distinct. Semantics can range from features of high similarity, features of the same object, to more complex semantic structures. In light of this, methods such as SoCo [65], ORL [70] and DetCon [32] leverage the off-the-shelf algorithms, e.g., selective
*To whom correspondence should be addressed: qsu3@gsu.edu
Figure 1: The bi-level clustering of FLSL. An object or stuff in an image is essentially a cluster of features. Hence, their representations can be extracted as cluster representatives, e.g., modes. In FLSL, we aim to make these representations both locally and globally semantic via a bi-level clustering process. On the first level, the locally semantic representations are fostered by driving features of various concepts (book, person, plant, etc.) closer to their cluster modes  $\hat{z}_{\mathrm{cls}}$ s and far away from features of other concepts within an image(intra-view clustering). On the second level, cluster modes serving as representations  $\hat{z}_{\mathrm{cls}}$ s are pushed closer to their positive samples  $\hat{z}_{\mathrm{cls}}^{+}$ s in  $X^{+}$ , which is augmented via a random transformation  $t \sim \mathcal{T}$  (inter-view clustering). In such a way, those representations encode the same category information and become globally semantic.
search [55] and Felzenszwalb-Huttenlocher algorithm [25] to impose the semantic constraint to the contrastive learning pipeline. Nonetheless, the inclusion of a non-trainable region proposal module in those methods restricts the model's ability to learn the distinct representations for those RoIs from the rest of an image. This ability is vital in representation learning for object detection.
Existing SSL methods targeting dense prediction primarily focus on the learning of globally semantic representations of image sub-regions, such as RoIs, patches, or pixels. However, these methods fall short with limited consideration for the alignment of those representations with local semantics. This observation leads us to ask the following question: Can we learn a representation that is both locally and globally semantic for a group of features (e.g., representing an object or stuff) in an end-to-end trainable SSL approach? To this end, we propose the Feature Level Self-supervised Learning (FLSL). It leverages the mean-shift clustering process inherent in transformer to extract modes as representations and incorporates  $k$ -means-based SSL approach to ensure that the extracted representations are semantically coherent both locally and globally. Figure. ① illustrates an overview of FLSL with details to be discussed in Sec. ④
Contributions This paper takes a step forward to bridge the gap between the current SSL methods and downstream dense prediction tasks. Our contributions are summarized as follows:
1. We demonstrate for the first time the connection between the attention mechanism and mean-shift clustering, and reinterpret vision transformer from the perspective of mean-shift.
2. By employing transformer for joint embedding and feature clustering, we propose FLSL, an end-to-end trainable SSL method that promotes the representations of feature clusters to be semantic at two levels: (i) intra-view: within an image, and (ii) inter-view: over an entire dataset.
3. The derivation and construction of the FLSL objectives is rooted in mean-shift and the non-empty  $k$ -means clustering. Semantic representations on the first level are encouraged by optimizing the intra-cluster affinity with a self-attention layer, while the second-level semantic representations are fostered via non-empty  $k$ -means clustering with positive samples retrieved through a cross-attention layer.
4. We validate the synergy between FLSL and ViT, and show significant improvement in transferability of the learned features to dense prediction tasks, including object detection and semantic segmentation. FLSL-pretrained ViT on ImageNet-1k (IN1k) demonstrates superior performance compared to the state-of-the-art ADCLR-IN1k [76] and MAE [40] pretrained counterparts. Moreover, it consistently outperforms existing SSL methods across additional benchmarks, including UAV object detection on UAVDT, and video instance segmentation on DAVIS 2017.
SSL for dense prediction Recent attempts to bridge the gap between common SSL and dense prediction tasks focus primarily on sub-region matching tricks. For example, DenseCL [63] applies contrastive learning on pairs of patches with highest similarity. However, the patch-matching trick leads to distinct representations with low correlation among patches, which is not well-suited for the semantics of a natural image. On top of the instance-level objective, PixPro [73] and LC-loss [35] factor in agreement between positive pixel pairs which are assigned through thresholded-distance in
PixPro and position projection in LC-loss. ReSim [68] maximizes the agreement between sliding-window-pooled representations in the overlapped region of two augmented views. DetCo [69] further incorporates instance-patch level contrastive losses along with instance level and patch level losses. To learn representations at object level, SoCo [65] and ORL [70] employ selective search to crop out RoIs. ORL further enables inter-object representation learning via BYOL [27] using top-ranked RoI pair. In contrast, SCRL [50] relaxes the semantic constraint using random crops within the intersection area of augmented views as RoIs. As discussed in Sec. 1, all of these methods focus on learning globally semantic representations for image sub-regions, and do not touch on local semantics that are necessary for dense prediction.
Self-supervised vision transformer In pioneering works, self-supervised training of transformer for vision tasks generally follow the paradigm of masked autoencoder in NLP [47, 20]. For instance, iGPT [13] features reconstruction of masked pixels as one of its objectives. In general, SSL for ViT can be classified into two categories: the joint-embedding strategy epitomized by DINO [10] and MoCov3 [16], and the generative approaches represented by MAE [28]. The crossover of the two strategies is demonstrated by iBOT [78]. Regarding dense prediction, EsViT [38], designed for Swin Transformer [43], follows the region-matching strategy and applies the DINO loss to the probabilities of positive pairs determined by highest similarity. Instead of finding the best-matching patch, SelfPatch [75] considers the direct neighbors as its positive patches. However, with limited semantics contained in a fixed small area (e.g., 8-connected neighbors), the method still suffers from semantic misalignment. To address the sub-region mismatch issue of DINO, ADCLR [76] constructs query tokens from random sub-regions and treats them as extra class tokens in the DINO objective. This promotes region-aware semantic representations that better aligned with the local semantics, and leads to substantial improvement in dense prediction.
3 Intuition: the connection between mean-shift and attention
As discussed in Sec. 1, the misalignment between the current SSL methods and dense prediction tasks lies in the clustering bias at the semantic level. Instead of setting a fixed granularity, such as instance-level or fix-sized patch-level, a desired semantic representation scheme should be able to represent from a single patch to a cluster of patches or even an entire image. The representation space of an image can be considered as an empirical probability density function of features, and the modes (local maxima) therefore can be regarded as the representatives of clusters [11, 17, 18]. These modes can then be readily retrieved via clustering algorithms, particularly, non-parametric kernel density estimation (KDE) methods [62] when the image composition (e.g., number of objects and stuffs) is unknown. One typical KDE-based method is the mean-shift clustering [33]. In the following, we first give an overview of self-attention (SA) mechanism of transformer and the mean-shift algorithm. We then show that the mean-shift update rule conforms to the SA mechanism of transformer.
Attention mechanism First introduced to recurrent neural networks as a context extractor for machine translation [2], attention has premised major breakthroughs in NLP with the emergence of transformer that relies solely on the scaled dot-product attention mechanism [58] given by
$$
\text {a t t e n t i o n} (Q, K, V) = V \operatorname {s o f t m a x} \left(Q ^ {\top} K / \sqrt {D _ {q k}}\right), \tag {1}
$$
where  $Q$ ,  $K$  and  $V$  denote query, key and value matrices which pack together sets of query, key and value vectors, respectively.  $D_{qk}$  denotes the dimension of query and key vectors, and  $\mathrm{softmax}(Z)_{ij} = \exp (Z_{ij}) / \sum_k\exp (Z_{ik})$ . As a special case of attention, SA matches a sequence  $Z$  with itself to extract the semantic dependencies among its components, i.e.,  $Q = W_{Q}Z$ ,  $K = W_{K}Z$ ,  $V = W_{V}Z$ , where the projections  $\mathbf{W}_{-}$ 's are the parameter matrices.
Mean-shift clustering and attention Given  $N$  data points  $\{\pmb{z}_i\}_{i=1}^N \subset \mathbb{R}^D$ , the kernel density estimate of  $p(\pmb{z})$  with kernel  $K(t)$  can be defined as
$$
p (\boldsymbol {z}) = \sum_ {i = 1} ^ {N} p \left(\boldsymbol {z} _ {i}\right) p \left(\boldsymbol {z} \mid \boldsymbol {z} _ {i}\right) = \sum_ {i = 1} ^ {N} \pi_ {i} \frac {1}{T _ {i}} K \left(d \left(\boldsymbol {z}, \boldsymbol {z} _ {i}; \boldsymbol {\Sigma} _ {i}\right)\right), \tag {2}
$$
where  $p(\boldsymbol{z}_i) = \pi_i$  is the mixing proportion of point  $\boldsymbol{z}_i$ , s.t.  $\sum_{i=1}^{N} \pi_i = 1$ ,  $T_i$  denotes the normalization term dependent only on the covariance matrix  $\boldsymbol{\Sigma}_i$ , e.g., for a Gaussian kernel  $T_i = |2\pi \boldsymbol{\Sigma}_i|^{1/2}$  and  $d(\boldsymbol{z}, \boldsymbol{z}_i; \boldsymbol{\Sigma}_i) = (\boldsymbol{z} - \boldsymbol{z}_i)^T \boldsymbol{\Sigma}_i^{-1} (\boldsymbol{z} - \boldsymbol{z}_i)$  is the Mahalanobis distance. Finding the modes of  $p(\boldsymbol{z})$  is
to seek stationary points by equating the gradient of  $p(z)$  to zero,  $\frac{\partial p(z)}{\partial z} = 0$ , which arrives at
$$
\hat {\boldsymbol {z}} = \boldsymbol {f} (\boldsymbol {z}) = \sum_ {i = 1} ^ {N} p \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z}\right) \boldsymbol {z} _ {i}, \quad \text {w i t h} p \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z}\right) = \frac {\pi_ {i} \frac {1}{T _ {i}} K ^ {\prime} \left(d \left(\boldsymbol {z} , \boldsymbol {z} _ {i} ; \boldsymbol {\Sigma} _ {i}\right)\right) \boldsymbol {\Sigma} _ {i} ^ {- 1}}{\sum_ {j = 1} ^ {N} \pi_ {j} \frac {1}{T _ {j}} K ^ {\prime} \left(d \left(\boldsymbol {z} , \boldsymbol {z} _ {j} ; \boldsymbol {\Sigma} _ {j}\right)\right) \boldsymbol {\Sigma} _ {j} ^ {- 1}}, \tag {3}
$$
where  $K' = dK / dt$ . The above fixed-point iterative scheme is the mean-shift algorithm. Practically, on  $\ell_2$ -normalized vectors, for a homoscedastic Gaussian kernel with constant mixing proportion and isotropic covariances (e.g.,  $\pi_i = 1 / N$ ,  $1 / \sigma^2 = \tau$ ), Eq. [3] further simplifies to
$$
\hat {\boldsymbol {z}} = \text {m e a n s h i f t} (\boldsymbol {z}, \tau) = \sum_ {i = 1} ^ {N} \frac {\exp \left(\tau \boldsymbol {z} ^ {\top} \boldsymbol {z} _ {i}\right)}{\sum_ {j = 1} ^ {N} \exp \left(\tau \boldsymbol {z} ^ {\top} \boldsymbol {z} _ {j}\right)} \boldsymbol {z} _ {i} \Longrightarrow \hat {\boldsymbol {Z}} = \boldsymbol {Z} \text {s o f t m a x} \left(\tau \boldsymbol {Z} ^ {\top} \boldsymbol {Z}\right), \tag {4}
$$
which conforms to the attention function (Eq. [1]) with identity projection matrices, i.e.,  $W_{Q} = W_{K} = W_{V} = \mathbf{I}$ , and  $\tau = 1 / \sqrt{D_{qk}}$ . Conversely, the conventional SA mechanism can be viewed as a generalized mean-shift:
$$
\hat {\boldsymbol {Z}} = \operatorname {S A} (\boldsymbol {Z}) = \mathbf {W} _ {V} \boldsymbol {Z} \text {s o f t m a x} \left(1 / \sqrt {D _ {q k}} \boldsymbol {Z} ^ {\top} \left(\mathbf {W} _ {Q} ^ {\top} \mathbf {W} _ {K}\right) \boldsymbol {Z}\right), \tag {5}
$$
with learnable distance measure  $Z^{\top}(\mathbf{W}_Q^\top \mathbf{W}_K)Z$  and projection  $\mathbf{W}_V$ . Unlike GMM and  $k$ -means, mean-shift is capable of modeling clusters of complex non-convex shape with cluster number automatically determined by local scale (proscribed by covariance) [33]. Hence, it is well-aligned with the semantics of natural images.
ViT from the perspective of mean-shift In ViT [22], images are initially tokenized and then processed through a sequence of transformer layers. Each transformer layer is comprised of a skip-connected multi-head SA (MHSA) and a skip-connected MLP. MHSA can be constructed from Eq. 5 with  $m$  projections in parallel, i.e.,  $[\mathbf{W}_Q^h, \mathbf{W}_K^h, \mathbf{W}_V^h]$ ,  $h = 1, \dots, m$ . The  $m$  returned modes are then concatenated along channel dimension and reprojected to a single return through
$$
\hat {\boldsymbol {Z}} = \operatorname {M H S A} (\boldsymbol {Z}) = \mathbf {W} _ {O} \operatorname {c o n c a t} \left(\left[ \left[ \hat {\boldsymbol {Z}} ^ {1} \right], \dots , \left[ \hat {\boldsymbol {Z}} ^ {m} \right] \right]\right) + \mathbf {b} _ {O}. \tag {6}
$$
Note that the  $\ell_2$  normalization assumed in Eq. 4 is moderately relaxed via layer normalization (LN) to incorporate the extra degree of freedom in the vector magnitude. With skip connection and the one-step mean-shift update described in Eqs. 5, 6, a transformer layer essentially finds the local centroid for each query  $z$  and drives them closer to the re-projected centroids through  $z = z + \hat{z}$ , followed by an MLP processing step with skip connection. ViT iterates the process multiple times (e.g., 12 or 24 layers) to capture the contextual and semantic information of an image.
The clustering process above concords with one inductive bias of the attention mechanism represented by the sparse variable creation [24], i.e., an SA head learns a sparse function that only depends on a small subset of input coordinates. In the context of clustering, the subset of input corresponds to the modes of density  $p(z)$ . As the high-level semantic information is typically spatially sparse (e.g., the representation for a RoI in object detection, a single label for a region in segmentation, or a scene-graph, etc.), it is natural to leverage transformer for joint embedding and clustering to learn semantically meaningful representations.
FLSL features a bi-level clustering process (Figure 1), which is formally described as follows.
Given a dataset  $\mathcal{X}$  (e.g., a set of images), FLSL learns an encoding scheme  $f_{\theta}:\mathcal{X}\to \mathcal{Z},\forall \mathbf{X}\in \mathcal{X},\mathbf{Z} = f_{\theta}(\mathbf{X})$  .  $\mathbf{Z}$  can be formulated as  $\mathbf{Z} = \bigcup_{c}^{N_c}\tilde{\mathbf{z}}^c$  , where  $\tilde{z}^c$  is a subset of  $\mathbf{Z}$  forming a cluster,  $N_{c}$  is the number of clusters determined by a clustering scheme, e.g., mean-shift, and  $N_{c}\leq |\mathbf{Z}|$  . FLSL aims to encourage the following properties:
(i) Intra-view: encodings corresponding to a semantic concept (as a cluster),  $z \in \tilde{z}^c$ , are close to the cluster representative (e.g., mode)  $\hat{z}^c$  and far away from the encodings of other clusters;
(ii) Inter-view: the cluster representatives  $\hat{z}$ s of the positive regions in  $X$ s over  $\mathcal{X}$  are pushed closer to each other.
The FLSL-extracted features should be well-aligned with dense prediction tasks, such as object detection, where the representation of an object or stuff (i.e., cluster of features) are desired to be (i) well-separated from others in an image (locally semantic), and (ii) close to its positive samples in the dataset (globally semantic). In this section, we present the objectives for both levels of clustering, which are then combined to form the final objective.
Figure 2: Overview of the FLSL framework. Similar to DINO [10], FLSL is comprised of a teacher network and a student network, which have the same architecture - a ViT encoder  $f$  and a projection head  $g$  - but with different parameters. Two mean-shift operations: a non-parametric self-attention (SA) and a non-parametric cross-attention (CA) are applied to the last layer of  $f_{t}$ ,  $f_{s}$  before  $g_{t}$ ,  $g_{s}$ , respectively, and the CA takes output of  $f_{s}$  as queries. The two networks are trained to maximize the agreement between the probability distributions  $p_{i}$  s and  $p_{i}^{+}$ s and the agreement between features  $z_{i}$  s and their cluster representatives  $\hat{z}_{i}$  s.
4.1 Intra-view clustering with mean-shift
As discussed in Sec. 3, local semantics of an image can be captured by non-parametric clustering such as mean-shift. Hence, with mean-shift update rule Eq. 4, it can be proved that the probability of  $z_{j}$  given point  $z_{i}$ ,  $p(z_{j}|z_{i}) = [\mathrm{softmax}(\tau z_{i}^{\top}\mathbf{Z})]_{j}$ , should satisfy:
$$
p \left(\boldsymbol {z} _ {j} \mid \boldsymbol {z} _ {i}\right) \geq 1 / \left(\left(\sum_ {k \in c _ {i}} e ^ {\left(\boldsymbol {z} _ {i} ^ {\top} \boldsymbol {z} _ {k} - \boldsymbol {z} _ {i} ^ {\top} \boldsymbol {z} _ {j}\right) \tau}\right) + (N - | c _ {i} |) e ^ {- \Delta_ {i j} \tau}\right), \forall j \in c _ {i} \tag {7}
$$
where  $N = |\mathbf{Z}|$ ,  $c_{i}$  is the set of indices of points in the same cluster including  $\mathbf{z}_{i}$ , and  $\Delta_{ij}$  is the degree of separability defined as  $\Delta_{ij} = \mathbf{z}_i^\top \mathbf{z}_j - \max_{k\in [N]\setminus c_i}\mathbf{z}_i^\top \mathbf{z}_k$ , such that larger  $\Delta_{c_i} = \sum_{j\in c_i}\Delta_{ij}$  indicates better separation. For locally semantic encodings, we desire the in-cluster points to be close to each other, or equivalently, to be close to its cluster representative, and stay far away from the out-cluster points, which indicates a large  $\Delta$  value. As  $\Delta$  becomes sufficiently large, the RHS of Eq. 7 can be approximated as  $1 / \sum_{k\in c_i}\exp ((\mathbf{z}_i^\top \mathbf{z}_k - \mathbf{z}_i^\top \mathbf{z}_j)\tau)$ , and for out-cluster points, the probability  $p(\mathbf{z}_{j\notin c_i}|\mathbf{z}_i)$  approaches to 0. This results in a semantics-aligned cluster representative via mean-shift - a weighted sum of only in-cluster points. This can be realized by contrasting among points using attention map as soft cluster mask to drive the query point  $\mathbf{z}_{i}$  closer to the returned mode  $\hat{\mathbf{z}}_i$ . It leads to the intra-view clustering objective:
$$
\min  _ {f _ {\theta}} \sum_ {i = 1} ^ {N} \| \boldsymbol {z} _ {i} - \hat {\boldsymbol {z}} _ {i} \| _ {2} ^ {2}. \tag {8}
$$
Proof of Eq.7 and detailed explanation is provided in Appendix A
4.2 Interview clustering with k-means
To learn globally semantic representations, similar to the existing SSL methods, we formulate the problem as a variant of  $k$ -means clustering. For  $\hat{z}$ s extracted from an entire dataset, the  $k$ -means objective with generalized non-empty cluster constraint  $\square$  can be expressed as
$$
\min  _ {\mathcal {M}} \frac {1}{N ^ {\prime}} \sum_ {\hat {\boldsymbol {z}} \in \hat {\mathcal {Z}}} \sum_ {k = 1} ^ {K} \delta_ {k k (\hat {\boldsymbol {z}})} \| \hat {\boldsymbol {z}} - \boldsymbol {\mu} _ {k (\hat {\boldsymbol {z}})} \| _ {2} ^ {2} + D _ {\mathrm {K L}} (\bar {\boldsymbol {p}} \| \boldsymbol {\pi}), \tag {9}
$$
where  $\mathcal{M}$  is a set of  $K$  centroids  $\{\pmb{\mu}_1,\dots ,\pmb{\mu}_K\}$ ,  $\hat{\mathcal{Z}}$  is a set of cluster representatives over the entire dataset,  $N^{\prime} = |\hat{\mathcal{Z}} |$ ,  $k(\hat{z}) = \arg \min_{k}\| \pmb{\mu}_{k} - \hat{\pmb{z}}\|_{2}$ ,  $\delta_{ij}$  is the Kronecker delta, with  $\delta_{ij} = 1$  iff  $i = j$ , and 0 otherwise,  $[\bar{\pmb{p}} ]_{[i]} = 1 / N'\sum_{\hat{z}}\delta_{ik(\hat{z})}$ , and  $\pi$  is the prior, e.g., a vector of the preset proportion for each cluster. With positive pairs  $(\hat{z}^{+},\hat{z})$  created via data augmentation, the objective can then be constructed as  $k$ -means clustering with extra separation margin for  $\hat{z}^{+}$ :
$$
\min  _ {\mathcal {M}} \frac {1}{N ^ {\prime}} \sum_ {\hat {z} \in \mathcal {Z}} \left(\sum_ {k = 1} ^ {K} \delta_ {k k (\hat {z})} \| \hat {z} - \boldsymbol {\mu} _ {k (\hat {z})} \| _ {2} ^ {2} + \left(1 - \delta_ {k (\hat {z} ^ {+}) k (\hat {z})}\right) \| \hat {z} ^ {+} - \boldsymbol {\mu} _ {k (\hat {z})} \| _ {2} ^ {2}\right) + D _ {\mathrm {K L}} (\bar {\boldsymbol {p}} \| \boldsymbol {\pi}). \tag {10}
$$
A common approach to tackle the optimization problem above is to relax the hard cluster assignment constraint  $\delta_{ij} \in \{0,1\}$  to  $[0,1]$  via a classification head to  $\hat{\mathbf{z}}$  with a small temperature ( $\ll 1$ ). This relaxes Eq. [9] to a more general Gaussian Mixture Model (GMM) formulation (cf. Appendix B).