SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Thalles Silva
Institute of Computing
University of Campinas
thalles.silva@students.ic.unicamp.br
Adín Ramírez Rivera
adinr@uio.no
We present Consistent Assignment of Views over Random Partitions (CARP), a self-supervised clustering method for representation learning of visual features. CARP learns prototypes in an end-to-end online fashion using gradient descent without additional non-differentiable modules to solve the cluster assignment problem. CARP optimizes a new pretext task based on random partitions of prototypes that regularizes the model and enforces consistency between views' assignments. Additionally, our method improves training stability and prevents collapsed solutions in joint-embedding training. Through an extensive evaluation, we demonstrate that CARP's representations are suitable for learning downstream tasks. We evaluate CARP's representations capabilities in 17 datasets across many standard protocols, including linear evaluation, few-shot classification,  $k$ -NN,  $k$ -means, image retrieval, and copy detection. We compare CARP performance to 11 existing self-supervised methods. We extensively ablate our method and demonstrate that our proposed random partition pretext task improves the quality of the learned representations by devising multiple random classification tasks. In transfer learning tasks, CARP achieves the best performance on average against many SSL methods trained for a longer time.
Learning from unlabeled data has been one of the main challenges in computer vision. Recent approaches based on self-supervised learning (SSL) have significantly reduced the gap between supervised and unsupervised pre-trained representations. Nowadays, self-supervised pre-training on vast quantities of unlabeled data, prior to learning a downstream supervised task of interest, can be more effective than supervised pre-training for many tasks [7, 19, 22].
Current SSL methods can be divided into two classes: (1) self-supervised embedding prediction [9, 24, 39, 41] and (2) clustering [1, 5, 6, 29]. As the name suggests, embedding prediction methods work directly in the representation space. They are trained with either contrastive or non-contrastive loss functions. Instead of reconstructing the input signal, their loss function maximizes agreement between embeddings of the same view and optionally pushes representations from different views apart. On the other hand, clustering methods discretize the representation space by learning a finite set of prototypes. These prototypes aggregate representations from different images that are similar enough to be assigned together. Nevertheless, recent SSL methods build a joint-embedding architecture that can be pure siamese [4] or use different encoders.
Code at https://sthalles.github.io/carp/.
The most challenging and significant difference among these methods is how they avoid trivial solutions when training joint-embedding architectures. Contrastive methods avoid trivial solutions by explicitly pushing representations of negative samples away from the anchor representation, while non-contrastive methods avoid trivial solutions by regularization or architectural designs [10, 22].
Among self-supervised clustering methods, recent work proposed avoiding trivial solutions by using non-differentiable modules such as the Sinkhorn-Knopp algorithm [1, 7, 14] and classic machine learning methods [29, 42] such as  $k$ -Means clustering or  $k$ -Nearest Neighbor, to solve the cluster assignment problem.
Consistent assignments [37, 38] were recently proposed as a way to learn prototypes to improve the representations. Similarly, our learning objective imposes two constraints (1) consistent assignment of views over learnable prototypes and (2) a uniform distribution for the average predictions within a batch. However, we show that such a strategy does not scale well to enormous datasets containing millions of classes. In such situations, we need to model a large number of prototypes while enforcing consistent assignment between views and avoiding collapsed solutions. We show that a naive implementation of this strategy makes the learning problem challenging from a training stability perspective, where the model quickly settles for a trivial solution by assigning all views' embeddings to the same prototype.
To overcome these issues, we propose a novel self-supervised approach based on the consistent assignment of views over random partition sets (CARP). We train CARP to minimize a consistency loss, which encourages the model to assign different views of the same unlabeled example to the same prototype. We solve the dimensionality problem by enforcing smaller pseudo-classification problems through the introduction of random partitions that enforce consistency and regularize the model. The energy between the views' representations and the trainable prototypes (within random partitions) allows us to automatically bootstrap predictions and targets to our consistency loss. Our contributions are three-fold:
1. A novel and entirely online joint-embedding learning strategy based on self-supervised clustering, see Figure 1. We propose a divide-and-conquer pretext task based on randomly generated partitions of learnable prototypes. Our loss function allows stable training of joint-embedding architectures in a self-supervised context.
2. A framework that simplifies self-supervised training and does not require normalization techniques [7, 9] or the necessity of mining negatives for contrastive training [9, 31, 32].
3. A differentiable assigner module that generates soft pseudo-labels by comparing the representations of image views to prototypes within random partition subsets. To avoid trivial solutions, we enforce the average predictions over a batch to be non-informative over the set of prototypes within a random subset.
Self-supervised embedding prediction methods operate directly in the representation space by learning a metric such that embeddings from views of the same image are closer to one another while embeddings from views of different images are far away in the feature space. These methods can be trained using contrastive or non-contrastive loss functions. Methods that minimize a loss function with a contrastive term date back to 1990s [4, 13, 21]. They must explicitly find representations from non-correlated images to use as negatives. Recent contrastive methods include InstDisc [46], CPC [32], SimCLR [9] and MoCo [11, 24]. These methods learn unsupervised representations by minimizing nearly the same contrastive loss function, i.e., the InfoNCE [32]. CARP does not directly optimize the views' embeddings, nor is it a contrastive method. Instead, we learn a set of general prototypes using a random partition strategy that stabilizes the learning process and avoids trivial solutions commonly found when training joint-embedding SSL models.
On the other hand, non-contrastive methods work by approximating embeddings of views from the same image. The main advantage is not requiring explicit opposing representations in the loss formulation. To avoid trivial solutions, a common approach is to implement a "stop-gradient" operation that prevents the gradient signal from flowing to the two branches of the join-embedding architecture simultaneously. BYOL [22] and SimSiam [10] are examples of such methods. CARP takes advantage of non-contrastive training since it does not require mining negatives for optimization. Also, different from Grill et al.'s [22] work, CARP does not require a momentum encoder, though using it
Figure 1: CARP's training architecture. From two views, we train an encoder  $f(\cdot)$ , followed by a projection  $g_{(\cdot)}$  to produce representation vectors  $z^v$  and  $w^v$ , respective to the parameters  $\theta$  and  $\xi$  for each branch, for each view indexed by  $v$ . The representations are fed to an assigner function  $q_{(\cdot)}$  that produces normalized distributions of views w.r.t. the learnable prototypes. Note that  $\theta$  are the trainable weights, and  $\xi$  are an exponential moving average of  $\theta$ . We create partitions by randomly arranging the prototypes into a predefined number of blocks. E.g., from a total of  $K = 6$  prototypes, we create  $N_{\mathcal{P}} = 2$  blocks, each containing  $N_B = 3$  prototypes. Then, we enforce consistent assignment of views over prototypes within the blocks.
significantly improves the learned representations. CARP trains a joint-embedding architecture and uses the "stop-gradient" operation in conjunction with a regularized pretext task based on random partitions of prototypes to avoid mode collapse.
Self-supervised clustering methods do not work directly on the views' embeddings. Instead, they learn a set of prototypes that are used to solve subsequent pretext tasks in different ways. Caron et al. [5], for instance, used  $k$ -Means clustering at every epoch to cluster the representations from the entire dataset and produce pseudo-labels for the training images. Then, a classifier head is trained to predict the pseudo-labels. Caron et al. [6] proposed a method to combine the rotation prediction pretext task [18] with clustering. Li et al. [29] presented a method based on expectation-maximization (EM) that merges clustering with the contrastive learning framework from He et al. [24]. Recent work by Caron et al. [7] and Asano et al. [1] combine SSL with clustering. They utilize the non-differentiable Sinkhorn-Knopp algorithm to solve the cluster assignment problem without falling into collapsed solutions. Silva and Ramírez Rivera [38] proposed, CARL, an online clustering method that does not require non-differential algorithms to avoid trivial solutions.
Contrast to previous approaches. Instead of solving the cluster assignment problem using a non-differentiable algorithm such as the Sinkhorn-Knopp, CARP is trained end-to-end with gradient descent. Different from Caron et al.'s [8] work, our model does not require extra momentum encoders or data structures to store previous predictions as a way to avoid trivial solutions. Unlike CARL [37, 38], CARP avoids trivial solutions by posing the optimization problem at the level of random partitions of prototypes that solve the high-dimensionality nature of the task. Unlike Caron et al.'s [6] work, our method does not require clustering the entire dataset every epoch to generate pseudo-labels. Instead, CARP generates soft pseudo-labels in an online fashion from examples in a single minibatch.
Currently, the best self-supervised methods [8, 12] use vision transformers [15, 44] as their backbones. Recent methods [20] employ a masking pretext task where some patches of the views are manually hidden, and the network is tasked to infer the missing pieces. In this paper, we consider transformer-based methods to be out-of-scope of our compared backbones. Thus, we do not include them in our results to maintain a fair comparison.
3 Consistent Assignment of Views
From an image  $x_{i}$ , we create two views,  $\hat{x}_i^1 = T(x_i)$  and  $\hat{x}_i^2 = T(x_i)$ , using a stochastic function  $T$  that applies a set of random image transformations to  $x_{i}$  (cf. Appendix C.1). CARP is a joint-embedding architecture with two modules: a differentiable (student) and a non-differentiable (teacher)
Figure 2: Instead of posing a pseudo-classification problem overall prototypes (circles and squares), the views are assigned (colored dashed lines) to a subset of prototypes (blocks), devising multiple pseudo-classification problems. Then, we contrast their distributions of assignments. (For this example,  $K = 8$ ,  $N_B = 4$  with  $N_{\mathcal{P}} = 2$ .)
branch. Each module has its own set of weights and the same architectural design. Both contain an encoder  $f_{(\cdot)}$  and a projection head  $g_{(\cdot)}$ . The differentiable student receives a pair of views and produces embedding vectors  $z_{i}^{v} = g_{\theta}(f_{\theta}(\hat{x}_{i}^{v}))$ , for  $v \in \{1,2\}$ . Similarly, the non-differentiable teacher produces target embeddings  $w_{i}^{v} = g_{\xi}(f_{\xi}(\hat{x}_{i}^{v}))$ .
The objective is to learn a set of prototype vectors  $C$  to discretize the embedding space. These prototypes are not meant to represent the true classes of the data. Instead, they may be interpreted as anchors to attract views of a given image to a commonplace in the embedding space. The function  $q(\cdot, \cdot)$  receives the views' representations,  $z_i^v$  and  $w_i^v \in \mathbb{R}^{1 \times d}$ , as input and outputs normalized probability vectors relating the views' embeddings with the prototypes such that  $s_i^v = q(z_i^v, C)$  and  $t_i^v = q(w_i^v, C)$ , where  $s_i^v$  and  $t_i^v \in \mathbb{R}^{1 \times K}$  are the normalized probabilities of a view,  $\hat{x}_i^v$ , w.r.t. the prototypes  $C \in \mathbb{R}^{K \times d}$ . Note that  $d$  is the dimensionality of the embedding vector,  $K$  is the number of prototypes, and the assigner  $q(h, C) = \text{softmax}\left(h \cdot C^T\right)$ .
To avoid trivial solutions in the joint-embedding training, we need a loss function that prevents the assignment of all representation vectors  $z_{i}$  to a unique prototype. Unlike previous work, we seek a method that solves the cluster assignment problem in an online fashion using gradient descent.
We propose a loss function composed of two terms: consistency and entropy. The consistency term learns the relations between embedding vectors and prototypes. It enforces different views of the same image to be assigned to the same prototype with high confidence. For normalized probability vectors  $a$  and  $b$ , we define the consistency term as
$$
\mathcal {L} _ {c} (a, b) = - \log \langle a, b \rangle , \tag {1}
$$
where  $\langle \cdot ,\cdot \rangle$  is a dot product.
The consistency loss is optimized when the two views  $\hat{x}_i^1$  and  $\hat{x}_i^2$  are assigned to the same prototype with maximal confidence, i.e., when the probability distributions of the two views  $s_i^1$  and  $s_i^2$  resemble equal one-hot vectors.
If we optimize the consistency loss,  $\mathcal{L}_c$ , training collapses to a state where all views are assigned to the same prototype. A common approach to avoid such failure is to ensure that all the prototypes get roughly the same number of assignments. Let us define the function
$$
\operatorname {a v g} \left(\left\{\left(a _ {i}, b _ {i}\right) \right\} _ {i = 1} ^ {L}\right) = \frac {1}{L} \sum_ {i = 1} ^ {L} a _ {i} + b _ {i} \tag {2}
$$
as the average probability across the representations within a batch of size  $L$ . For our distributions, we define  $\bar{p} = \mathrm{avg}\left(\left\{(s_i^v,t_i^v)\right\}_{i = 1}^N\right)$ . If we maximize the entropy of the mean probabilities of a batch,  $H(\bar{p})$ , we will encourage the average predictions to be closer to a uniform distribution. Previous work [2, 26, 42] has used this entropy term in various scenarios, ranging from discriminative unsupervised clustering to semi-supervised learning. In our case, this term enforces a non-informative prior over the prototypes with a particular schedule learning. Thus, the final proposed objective to
minimize is
$$
\mathcal {L} = \frac {1}{N} \sum_ {i} ^ {N} \left(\mathcal {L} _ {c} \left(s _ {i} ^ {1}, t _ {i} ^ {2}\right) + \mathcal {L} _ {c} \left(s _ {i} ^ {2}, t _ {i} ^ {1}\right)\right) - \lambda_ {e} H (\bar {p}), \tag {3}
$$
where  $\lambda_{e} > 0$  trades off consistency at the view level with the average uniform assignment at the batch level. Note that the contribution of the entropy term decays as training progresses.
3.1 Limitations of Consistent Assignments
The formulation of the consistent assignments (3) is similar to CARL [37, 38], except for the additional teacher stream that stabilizes training and improves performance. Nevertheless, training an unsupervised system by minimizing the loss (3) is challenging. The main limitation is how to avoid trivial solutions in unsupervised training of joint-embedding architectures. For some cases, tuning the contribution of the entropy term might be enough to optimize the loss (3) stably. However, adjusting such a hyperparameter is difficult because one configuration does not hold for all training conditions. For instance, if the value of  $\lambda_{e}$  is too small, the consistency term wins the arms race, and the average distribution over the batch,  $H(\bar{p})$  becomes one-hot alike, i.e., all views end up assigned to the same prototype. If the value of  $\lambda_{e}$  is too large, the entropy term gets the upper hand, and collapse is avoided. However, the process of view assignment is neglected over the policy of distributing views uniformly, which results in poor performance of the learned representations, as shown by Table B.2.
For a small number of general prototypes, training is more stable, and the model avoids collapse with a simple tuning of the  $\lambda_{e}$  parameter. However, for a larger number of general prototypes, stability becomes an issue. The main problem lies with the entropy term. When the distribution is larger, i.e.,  $K\gg N$ , regular batch sizes, such as  $N = 64$  or  $N = 128$ , become too small to properly model the distribution, i.e., the signal is too weak for most prototypes. Consequently, to avoid collapse, we need to increase the strength of the entropy term or increase the batch size, which in turn decreases performance.
To address such limitation, we propose to decouple the loss function (3) into smaller sub-problems. Instead of enforcing both consistency and uniform assignments over all the general prototypes, we propose a pretext task over subsets or blocks of a random partition of the general prototype set  $C$ .
4 Assignment based on Random Partitions
Given the set of  $K$  trainable prototypes  $C = \{c_{1}, c_{2}, \dots, c_{K}\}$ , we define a partition of  $C$  as  $\mathcal{P} = \{B_{i} \subset C\}_{i=1}^{N_{\mathcal{P}}}$ , such that  $\emptyset \notin \mathcal{P}$ ,  $\bigcup_{i} B_{i} = \mathcal{P}$  where  $B_{i} \in \mathcal{P}$ , and  $B_{i} \cap B_{j} = \emptyset$  for all  $B_{i}, B_{j} \in \mathcal{P}$ , and  $i \neq j$ . We refer to  $B_{i}$  as a block or subset of the partition. We are interested in a partition set  $\mathcal{P} = \{B_{i}\}_{i=1}^{N_{\mathcal{P}}}$  of size  $N_{\mathcal{P}}$ , i.e.,  $|\mathcal{P}| = N_{\mathcal{P}}$ .
Using the concept of a partition of a set, we can define a framework of pretext tasks over partition blocks that satisfies the learning problem defined in Section 3. If the size of a partition block,  $B_{i}$ , equals the number of prototypes,  $N_{B} = K$  then the partition  $\mathcal{P}$  is trivial, i.e.,  $\mathcal{P} = \{B_1\} = \{C\}$ . If the size of the partition blocks equals  $N_{B} = 1$ , then we have  $K$  blocks in  $\mathcal{P}$ , and each block has a unique prototype. Here, the learning task is equivalent to multiple binary classification problems, where each output score, if normalized, expresses the likelihood of a data point  $x_{i}$  to independently belong to each prototype.
However, if the block size  $1 < N_B < K$ , and  $N_B$  divides  $K$ , then the partition  $\mathcal{P}$  will be composed of  $N_{\mathcal{P}} = \lfloor K / N_{\mathcal{P}} \rfloor$  blocks. We define  $\mathcal{P}$  by randomly assigning  $N_B$  prototypes  $c_j$ , for  $j = 0, 1, \ldots, N_B$ , to each block  $B_i = \{c_j\}_j$ , where  $i = 0, 1, \ldots, N_{\mathcal{P}}$ .
Instead of mapping a single representation  $z_{i}^{v}$  as a linear combination of all prototypes in  $C$ , we compare the view's representations  $z_{i}^{v}$  and  $w_{i}^{v}$  against all the prototypes in the  $j$ -th block. That is,  $s_{i,j}^{v} = q(z_{i}^{v}, B_{j})$  and  $t_{i,j}^{v} = q(w_{i}^{v}, B_{j})$ , for every block in the partition  $\mathcal{P}$ , to obtain the normalized probability distribution relating a view from image  $i$  with the prototypes of the  $j$ -th block of the random partition, where  $s_{i,j}^{v}$  and  $t_{i,j}^{v} \in \mathbb{R}^{1 \times 1 \times N_{P}}$ .
To ensure that views are consistent among the blocks, we optimize the views' distributions  $s_{i,j}^{v}$  and  $t_{i,j}^{v}$  over the prototypes of a block indexed by  $j$ , so that the two distributions are consistent with one another. Thus, the consistency term of our partition loss is  $\mathcal{L}_c(s_{i,j}^1, t_{i,j}^2)$ , where for each block
$B_{j}$ , the loss is minimized when the pair of views,  $\hat{x}_i^1$  and  $\hat{x}_i^2$ , gets assigned to the same prototypes across blocks. In other words, we look for the agreement between student and teacher assignments' probabilities across views of a given sample.
Following similar reasoning, the block-wise entropy term is defined as  $H(\bar{p}_j)$ , where  $\bar{p}_j = \mathrm{avg}(\{(s_{i,j}^v,t_{i,j}^v)\}_{i = 1}^N)$  is the average prediction over each block  $B_{j}$  for a batch of size  $N$ . Thus, the final objective for consistent assignment of random partition is,
$$
\mathcal {L} = \frac {1}{N N _ {\mathcal {P}}} \sum_ {i} ^ {N} \sum_ {j} ^ {N _ {\mathcal {P}}} \left(\mathcal {L} _ {c} \left(s _ {i, j} ^ {1}, t _ {i, j} ^ {2}\right) + \mathcal {L} _ {c} \left(s _ {i, j} ^ {2}, t _ {i, j} ^ {1}\right)\right) - H (\bar {p} _ {j}). \tag {4}
$$
Note that to fully use the pair of views at each iteration, we symmetrically use the  $\mathcal{L}_c$  consistency function. The probability vectors  $t_{i,j}^v$  come from the momentum teacher and are used as target distributions.
We can view the random partition pretext task as posing multiple pseudo-classification problems over subsets of prototypes. At each iteration, the pseudo-classification tasks change because the partitions are recreated with different prototypes, cf. Figure 2. One of the benefits of such a strategy is that we no longer require tuning the hyperparameter  $\lambda_{e}$  to avoid trivial solutions. The stochastic nature of the random partition pretext task, blended with the multiple prediction tasks over subsets of prototypes, provides a regularization effect that improves the learned representations and training stability.
Other clustering-based methods [7, 8] rely on sharpening the distributions to improve their self-supervised signals used as targets in a cross-entropy loss. On the contrary, our formulation does not require the temperature parameter for sharpening the predictions and guiding the learning of the student. We can think of the consistency loss as implicitly learning the temperature parameter to make the predictions sharper at each iteration. This is an important advantage of our consistency loss in contrast to previous methods.
5.1 Transfer learning evaluation
Table 1 shows that CARP's pre-trained representations are suitable for learning new downstream tasks over multiple datasets with distinct difficulty levels. We compare CARP's  $k$ -NN performance against nine SSL methods across eight datasets and report average results for  $k = \{10, 20, 100, 200\}$ , over all datasets. We advocate for  $k$ -NN instead of linear evaluation, where a linear classifier is trained on top of the frozen features, for the following reasons: (1)  $k$ -NN is faster, (2)  $k$ -NN demands fewer resources, and (3)  $k$ -NN requires less hyperparameter tuning. For individual references, Table 1 shows top-1 accuracy for  $k = 20$  on each dataset. For fixed  $k = 20$ , CARP outperforms competitors in 5 out of 8 datasets, and in the Flowers and Country datasets, it is a close second. Moreover, the average performances at  $k$  show that CARP performs comparably well on all datasets without significant performance differences when varying the number of labeled examples  $k$ . We show the detailed results in Appendix A.1.
5.2 Clustering evaluation
Table 2 reports clustering performance metrics of various clustering-based SSL methods on the ImageNet-1M [36], CIFAR-10/100 [27], and the GTSRB [40] datasets. For the ImageNet-1M, only  $1\%$  of the labeled data is used following the subset provided by Chen et al. [9]. See Appendix D.4 for the detailed  $k$ -means evaluation protocol.
5.3 Image retrieval and copy detection