SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ido Ben-Shaul
Department of Applied Mathematics Tel-Aviv University & eBay Research ido.benshaul@gmail.com
Ravid Shwartz-Ziv*
New York University ravid.shwartz.ziv@nyu.edu
Tomer Galanti*
Massachusetts Institute of Technology galanti@mit.edu
Shai Dekel
Department of Applied Mathematics Tel-Aviv University shaidekel6@gmail.com
New York University & Meta AI, FAIR
yann@cs.nyu.edu
Self-supervised learning (SSL) is a powerful tool in machine learning, but understanding the learned representations and their underlying mechanisms remains a challenge. This paper presents an in-depth empirical analysis of SSL-trained representations, encompassing diverse models, architectures, and hyperparameters. Our study reveals an intriguing aspect of the SSL training process: it inherently facilitates the clustering of samples with respect to semantic labels, which is surprisingly driven by the SSL objective's regularization term. This clustering process not only enhances downstream classification but also compresses the data information. Furthermore, we establish that SSL-trained representations align more closely with semantic classes rather than random classes. Remarkably, we show that learned representations align with semantic classes across various hierarchical levels, and this alignment increases during training and when moving deeper into the network. Our findings provide valuable insights into SSL's representation learning mechanisms and their impact on performance across different sets of classes.
Self-supervised learning (SSL) [12, 5] has made significant progress over the last few years, almost reaching the performance of supervised baselines on many downstream tasks [42, 4, 31, 30, 47, 34, 15, 37, 66, 16]. However, understanding the learned representations and their underlying mechanisms remains a persistent challenge due to the complexity of the models and the lack of labeled training data. Moreover, the pretext tasks used in self-supervision frequently lack direct relevance to the specific downstream tasks, further complicating the interpretation of the learned representations.
In supervised classification, however, the structure of the learned representations is often very simple. For instance, a recent line of work [50, 36] has identified a universal phenomenon called neural collapse, which occurs at the terminal stages of training. Neural collapse reveals several properties of the top-layer embeddings, such as mapping samples from the same class to the same point, maximizing the margin between the classes, and certain dualities between activations and the top-layer weight matrix. Later studies [10, 28, 26, 56, 62] characterized related properties at intermediate layers of the network.
Figure 1: SSL training induced semantic clustering. UMAP visualization of the SSL representations before and after training in different hierarchies. (top) Augmentations of five different samples, each sample colored distinctly. (middle) Samples from five different classes within the standard CIFAR-100 dataset. (bottom) Samples from five different superclasses within the dataset.
Compared to traditional classification tasks that aim to accurately categorize samples into specific classes, modern SSL algorithms typically minimize loss functions that combine two essential components: one for clustering augmented samples (invariance constraint) and another for preventing representation collapse (regularization constraint). For example, contrastive learning methods [14, 48, 63, 46] train representations to be indistinguishable for different augmentations of the same sample and at the same time to distinguish augmentations of different samples. On the other hand, non-contrastive approaches [6, 66, 34] use regularizers to avoid representation collapse.
Contributions. In this paper, we provide an in-depth analysis of representation learning with SSL through a series of carefully designed experiments. Our investigation sheds light on the clustering process that occurs during training. Specifically, our findings reveal that augmented samples exhibit highly clustered behavior, forming centroids around the mean embedding of augmented samples sharing the same original image. More surprisingly, we observe that clustering also occurs with respect to semantic labels, despite the absence of explicit information about the target task. This demonstrates the capability of SSL to capture and group samples based on semantic similarities.
- Clustering: We investigate the clustering properties of SSL-trained representations. In Figure 2, we show that akin to supervised classification [50], SSL-trained representation functions exhibit a centroid-like geometric structure, where feature embeddings of augmented samples belonging to the same image tend to cluster around their respective means. Interestingly, at later stages of training, a similar trend appears with respect to semantic classes.
- The role of regularization: In Figure 3 (left) we show that the accuracy of extracting (i.e., linear probing) semantic classes from SSL-trained representations continuously improves even after the model accurately clusters the augmented training samples based on their sample identity. As can be seen, the regularization constraint plays a key role in inducing the clustering of data into semantic attributes at the embedding space.
- Impact of randomness: We argue that SSL-trained representation functions are highly correlated with semantic classes. To support this claim, we study how the degree of randomness in the targets affects the ability to learn them from a pretrained representation. As we show in Figure 4, the alignment between the targets and the learned representations substantially improves as the targets become less random and more semantically meaningful.
- Learning hierarchic representations: We study how SSL algorithms learn representations that exhibit hierarchic classes. We demonstrate that throughout the training process, the ability to cluster and distinguish between samples with respect to different levels of semantic targets continually improves at all layers of the model. We consider several levels of hierarchy, such as the sample level (the sample identity), semantic classes (such as different types of animals and vehicles), and high-level classes (e.g., animals, and vehicles).
Furthermore, we analyze the ability of each layer to capture distinct semantic classes. As shown in Figure 5, we observe that the clustering and separation ability of each layer improves as we move deeper into the network. This phenomenon closely resembles the behavior observed in supervised learning [10, 26, 56, 62], despite the fact that SSL-trained models are trained without direct access to the semantic targets.
In Section 4, we investigate various clustering properties in the top layer of SSL-trained networks. In Section 5, we conduct several experiments to understand what kinds of functions are encoded in SSL-trained neural networks and can be extracted from their features. In Section 6, we study how SSL algorithms learn different hierarchies and how this occurs at intermediate layers.
2.1 Self-Supervised Learning
SSL is a family of techniques that leverages unlabeled data to learn representation functions that can be easily adapted to new downstream tasks. Unlike supervised learning, which relies on labeled data, SSL employs self-defined signals to establish a proxy objective. The model is pre-trained using this proxy objective and then fine-tuned on the supervised downstream task. However, a major challenge of this approach is to prevent 'representation collapse', where the model maps all inputs to the same output [39].
Several methods have been proposed to address this challenge. Contrastive learning methods such as SimCLR [14] and its InfoNCE criterion [64] learn representations by maximizing the agreement between positive pairs (augmentations of the same sample) and minimizing the agreement between negative pairs (augmentations of different samples). In contrast, non-contrastive learning methods [6, 66, 34] replace the dependence on contrasting positive and negative samples, by applying certain regularization techniques to prevent representation collapse. For instance, some methods use stop-gradients and extra predictors to avoid collapse [16, 34], while others use an additional clustering constraints [13].
Variance-Invariance-Covariance Regularization (VICReg). A widely used method for SSL training [6], which generates representations for both an input and its augmented counterpart. It aims to optimize two key aspects:
- Invariance loss - The mean-squared Euclidean distance between pairs of embedding, serving to ensure consistency between the representation of the original and augmented inputs.
- Regularization - Comprising two elements, the variance loss, which promotes increased variance across the batch dimension through a hinge loss restraining the standard deviation, and the covariance loss, which penalizes off-diagonal coefficients of the covariance matrix of the embeddings to foster decorrelation among features.
This leads to the VICReg loss function:
$$
L (f) = \underbrace {\lambda s \left(Z , Z ^ {\prime}\right)} _ {\text {I n v a r i a n c e}} + \underbrace {\mu \left[ v (Z) + v \left(Z ^ {\prime}\right) \right] + \nu \left[ c (Z) + c \left(Z ^ {\prime}\right) \right]} _ {\text {R e g u l a r i z a t i o n}}, \tag {1}
$$
where  $L(f)$  is minimized over batches of samples  $I$ .  $s(Z, Z')$ ,  $v(Z)$ , and  $c(Z)$  denote the invariance, variance, and covariance losses respectively, and  $\lambda, \mu, \nu > 0$  are hyperparameters controlling the balance between these loss components.
A Simple Framework for Contrastive Learning of Visual Representations (SimCLR). The SimCLR method [14], another popular approach, minimizes the contrastive loss function [35]. Given two batches of embeddings  $Z$ ,  $Z'$ , we can also decompose the objective into an invariance term and a
regularization term:
$$
L(f) = -\frac{1}{B}\sum_{i = 1}^{B}\underbrace{(\operatorname{sim}(Z_{i},Z_{i}^{\prime}) / \tau)}_{\text{Invariance}} - \underbrace{\log\sum_{j\neq i}\exp(\operatorname{sim}(Z_{i},Z_{j}^{\prime}) / \tau)}_{\text{Regularization}},
$$
where  $\operatorname{sim}(a, b) := \frac{a^\top b}{\|a\| \cdot \|b\|}$  is the cosine similarity between vectors  $a$  and  $b$ , and  $\tau > 0$  is a 'temperature' parameter that controls the sharpness of the distribution. Intuitively, minimizing the loss encourages the representations  $Z_i$  and  $Z_j'$  of the same input sample to be similar (i.e., have a high cosine similarity) while pushing away the representations  $Z_i$  and  $Z_j'$  of different samples (i.e., have a low cosine similarity) for all  $i, j \in [B]$ .
2.2 Neural Collapse and Clustering
In a recent paper [50], it was demonstrated that deep networks trained for classification tasks exhibit a notable behavior: the top-layer feature embeddings of training samples of the same class tend to concentrate around their respective class means, which are maximally distant from each other. This behavior is generally regarded as desirable since max-margin classifiers tend to exhibit better generalization guarantees (e.g., [2, 7, 33]), and clustered embedding spaces are useful for few-shot transfer learning (e.g., [32, 28, 27, 29]).
In this paper, we aim to explore whether similar clustering tendencies occur in SSL-trained representation functions at both the sample level (with respect to augmentations) and at the semantic level. To address this question, we first introduce some key aspects of neural collapse. Specifically, we will focus on class-feature variance collapse and nearest class-center (NCC) separability.
To measure these properties, we use the following metrics. Suppose we have a representation function  $f: \mathbb{R}^d \to \mathbb{R}^p$  and a collection of unlabeled datasets  $S_1, \ldots, S_C \subset \mathbb{R}^d$  (each corresponding to a different class). The feature-variability collapse measures to what extent the samples from the same class are mapped to the same vector. To quantify this property we use the averaged class-distance normalized variance (CDNV) [28]  $\mathrm{Avg}_{i \neq j}[V_f(S_i, S_j)] = \mathrm{Avg}_{i \neq j}\left[\frac{\mathrm{Var}_f(S_i) + \mathrm{Var}_f(S_j)}{2\|\mu_f(S_i) - \mu_f(S_j)\|^2}\right]$ , where  $\mu_f(S)$  and  $\mathrm{Var}_f(S)$  are the mean and variance of the uniform distribution  $U[\{f(x) \mid x \in S\}]$ .
A weaker notion of clustering is the nearest class-center (NCC) separability [26], which measures to what extent the feature embeddings of training samples form centroid-like geometry. The NCC separability asserts that, during training, the penultimate-layer's feature embeddings of training samples can be accurately classified with the 'nearest class-center (NCC) classifier',  $h(x) = \arg \min_{c\in [C]}\| f(x) - \mu_f(S_c)\|$ . To measure this property, we compute the NCC train and test accuracies, which are simply the accuracy rates of the NCC classifier.
A recent paper [22] argued that idealized SSL-trained representations simultaneously cluster data into multiple equivalence classes. This implies that there are various ways to categorize the data that result in high NCC accuracy. However, achieving a very small CDNV for multiple different categorizations is impossible since a zero CDNV would mean that feature embeddings of two samples from the same category are identical. Consequently, an intriguing question arises: with respect to which sets of classes do the learned feature embeddings cluster?
2.3 Reverse Engineering Neural Networks
Reverse engineering neural networks have recently garnered attention as an approach to explain how neural networks make predictions and decisions. While substantial work has been dedicated to understanding the functionalities of intermediate layers of neural network classifiers [23, 9, 17, 50, 51, 60, 56, 3], characterizing the functionalities of representation functions generated by SSL algorithms remains an open problem.
A major source of complexity is the reliance on a pretext task and image augmentations. For instance, to fully understand the success of SSL it is essential to understand the relationship between the pretext task and the downstream task [58, 65]. It is yet unclear how to properly train representations with SSL, and therefore, SSL algorithms are quite complicated and include many different types of losses [14, 66, 6] and optimization tricks [14, 16, 37]. Of note, it is not obvious what are the differences between representations that are learned with contrastive learning and non-contrastive
Figure 2: SSL algorithms cluster the data with respect to semantic targets. (left) The normalized NCC train accuracy, computed by dividing the accuracy values by their value at initialization. (middle) The ratio between the NCC test accuracy and the linear test accuracy for per-sample and semantic classes normalized by its value at initialization. (right) The linear test accuracy rates, normalized by their values at initialization. All experiments are conducted on CIFAR-100 with VICReg training.
Figure 3: (left) The regularization and invariance losses together with the original target linear test accuracy of an SSL-trained model over the course of training. (middle) Compression of the mutual information between the input and the representation during the course of training. (right) SSL training learns clustered representations. The NCC train and test accuracy in supervised (with/without augmentation), and SSL. All experiments are conducted on CIFAR-100 with VICReg training.
learning methods [52]. In this work, we make new strides to understand SSL-trained representations, by investigating their clustering properties with respect to various types of classes.
Self-supervised learning (SSL) is frequently used to pretrain models, preparing them for adaptation to new downstream tasks. This raises an essential question: How does SSL training influence the learned representation? Specifically, what underlying mechanisms are at work during this training, and which classes can be learned from these representation functions? To investigate this, we train SSL networks across various settings and analyze their behavior using different techniques. All the training details can be found in Appendix A.
Data and augmentations. Throughout all of the experiments (in the main text) we used the CIFAR-100 [41] image classification dataset, which contains 100 classes of semantic objects (original classes), which are also divided into 20 superclasses, with each superclass containing 5 classes (for example, the 'fish' superclass contains the 'aquarium fish', 'flatfish', 'ray', 'shark' and 'trout' categories). In order to train our models we used the image augmentation protocol introduced in SimCLR [14]. Each SSL training session is carried out for 1000 epochs, using the SGD optimizer with momentum. Usually, SSL-trained models are tested on a variety of downstream datasets to assess their performance. To simplify our analysis and minimize the influence of potential distribution shifts, we specifically evaluate the models' performance on the CIFAR-100 dataset. Such evaluations are typically done in SSL research [6, 34, 66], usually training on and evaluating on ILSVRC (ImageNet). Additional evaluations on various datasets (CIFAR-10 [41], FOOD101 [11], Aircrafts [45], Tiny Magnet [43], and Imagenet [19]) can be found in Appendix B.2.
Figure 4: SSL continuously learns semantic targets over random ones. (left) The linear test accuracy for targets with varying levels of randomness from the last layers at different epochs. (middle) The linear test accuracy for targets with varying levels of randomness for the trained model. (right) The ratios between non-random and random targets for various clustering metrics. All experiments are conducted on CIFAR-100 with VICReg training.
Figure 5: SSL efficiently learns semantic classes throughout intermediate layers. The linear test accuracy of different layers of the model at various epochs (left) With respect to the 100 original classes. (middle) With respect to the 20 superclasses. (right) The ratio between the superclass and the original classes. All experiments are conducted on CIFAR-100 with VICReg training.
Backbone architecture. In all our experiments, we utilized the RES-  $L - H$  architecture as the backbone, coupled with a two-layer multi-layer perceptron (MLP) projection head. The RES-  $L - H$  architecture is a variant of the ResNet architecture [38] where the width  $H$  remains constant across all  $L$  residual blocks as introduced in [26].
Linear probing. To evaluate the effectiveness of extracting a given discrete function (e.g., categories) from a representation function, we employ linear probing. In this approach, we train a linear classifier, also known as a linear probe, on top of the representation, using a set of training samples. To train the linear classifier, we used the Linear Support Vector Classification (LSVC) algorithm from the scikit-learn library [54]. The "linear accuracy" is determined by measuring the performance of the trained linear classifier on a validation set.
Sample level classification. We aim to evaluate the capacity of a given representation function to correctly classify different augmentations originating from the same image and distinguish them from augmentations of different images. To facilitate this, we constructed a new dataset specifically designed to evaluate sample-level separability.
The training dataset comprises 500 random images sourced from the CIFAR-100 training set. Each image represents a particular class and undergoes 100 different augmentations. As such, the training dataset contains 500 classes with a total of 50000 samples. For the test set, we utilize the same 500 images but apply 20 different augmentations drawn from the same distribution. This results in a test set consisting of 10000 samples. To measure the linear or NCC accuracy rates of a given representation function at the sample level, we first compute a relevant classifier using the training data and subsequently evaluate its accuracy rates on the corresponding test set.
4 Unraveling the Clustering Process in Self-Supervised Learning
The clustering process has long played a significant role in aiding the analysis of deep learning models [1, 17, 50]. To gain intuition behind the SSL training, in Figure 1 we present a UMAP [57]
visualization of the network's embedding space for the training samples, both pre- and post-training, across different hierarchies.
As anticipated, the training process successfully clusters together samples at the per-sample level, mapping different augmentations of the same image (as depicted in the first row). This outcome is unsurprising given that the objective inherently encourages this behavior (through the invariance loss term). Yet, more remarkably, the training process also clusters the original 'semantic classes' from the standard CIFAR-100 dataset, even though the labels were absent during the training process. Interestingly, the high-level hierarchy (the superclasses) is also effectively clustered. This example demonstrates that while our training procedure directly encourages clustering at the sample level, the SSL-trained representations of data are additionally clustered with respect to semantic classes across different hierarchies.
To quantify this clustering process further, we train a RES-10-250 network using VICReg. We measure the NCC train accuracy, which is a clustering measure (as described in Section 2.2), both at the sample level and with respect to the original classes. Notably, the SSL-trained representation exhibits neural collapse at a sample level (with an NCC train accuracy near 1.0), yet the clustering with respect to the semantic classes is also significant ( $\approx 0.41$  on the original targets).
As shown in Figure 2 (left), most of the clustering process with respect to the augmentations (on which the network was directly trained) occurs at the beginning of the training process and then plateaus, while the clustering with respect to semantic classes (which is not directly featured in the training objective) continues to improve throughout the training process.
A key property observed in [50] is that the top-layer embeddings of supervised training samples gradually converge towards a centroid-like structure. In our quest to better understand the clustering properties of SSL-trained representation functions, we investigated whether something similar occurs with SSL. The NCC classifier's performance, being a linear classifier, is bounded by the performance of the best-performing linear classifier. By evaluating the ratio between the accuracy of the NCC classifier and a linear classifier trained on the same data, we can investigate data clustering at various levels of granularity. In Figure 2 (middle), we monitor this ratio for the sample-level classes and for the original targets, normalized with respect to its value at initialization. As SSL training progresses, the gap between NCC accuracy and linear accuracy narrows, indicating that augmented samples progressively cluster more based on their sample identity and semantic attributes. Furthermore, the plot shows that the sample-level ratio is initially higher, suggesting that augmented samples cluster based on their identity until they converge to centroids (the ratio between the NCC accuracy and the linear accuracy is  $\geq 0.9$  at epoch 100). However, as training continues, the sample-level ratio saturates while the class-level ratio continues to grow and converges to  $\approx 0.75$ . This implies that the augmented samples are initially clustered with respect to their sample identities, and once achieved, they tend to cluster with respect to high-level semantic classes.
Implicit information compression in SSL training. Having examined the clustering process within SSL training, we next turn our attention to analyzing information compression during the learning process. As outlined in Section 2.2, effective compression often yields representations that are both beneficial and practical. Nevertheless, it is still largely uncharted territory whether such compression indeed occurs during the course of SSL training [59].