Bag of Image Patch Embedding Behind the Success of Self-Supervised Learning

TL;DR

Addressing why joint-embedding SSL works, the paper shows it effectively models patch co-occurrence through the spectral contrastive loss $L_S$ and proves its equivalence to a co-occurrence loss $L_C$ with $\lambda=\frac{w}{2}$. It introduces BagSSL, where fixed-scale patches are embedded as $h=f(x;\theta)$ and projected to $z=g(h;\psi)$, with the image representation given by $R_{img} = \text{mean}_{HW}(h_{11},\dots,h_{HW})$, enabling linear aggregation to recover the whole-image embedding. Empirically, BagSSL matches or surpasses baselines on CIFAR-10/100 and ImageNet-100, and 32×32 patches with patch-averaging reach 62% top-1 accuracy on ImageNet-1K, illustrating the practicality of patch-based SSL. The work provides an interpretable view of SSL by highlighting patch locality at fine scales and showing that patch-based evaluation can further enhance representation quality.

Abstract

Self-supervised learning (SSL) has recently achieved tremendous empirical advancements in learning image representation. However, our understanding of the principle behind learning such a representation is still limited. This work shows that joint-embedding SSL approaches primarily learn a representation of image patches, which reflects their co-occurrence. Such a connection to co-occurrence modeling can be established formally, and it supplements the prevailing invariance perspective. We empirically show that learning a representation for fixed-scale patches and aggregating local patch representations as the image representation achieves similar or even better results than the baseline methods. We denote this process as BagSSL. Even with 32x32 patch representation, BagSSL achieves 62% top-1 linear probing accuracy on ImageNet. On the other hand, with a multi-scale pretrained model, we show that the whole image embedding is approximately the average of local patch embeddings. While the SSL representation is relatively invariant at the global scale, we show that locality is preserved when we zoom into local patch-level representation. Further, we show that patch representation aggregation can improve various SOTA baseline methods by a large margin. The patch representation is considerably easier to understand, and this work makes a step to demystify self-supervised representation learning.

Figures (12)

• Figure 1: BagSSL pipeline. From each image, fixed-size image patches are extracted, color-augmented, encoded to embedding and projection space. During training, patch projections of co-occurring patches (from the same image) are pulled together while an anti-collapse regularization is applied to push non co-occurring patch projections away. After training, patch embeddings $\{\vec{h}\}$ from the same image are averaged to form the image representation $\vec{R}_{img}$.
• Figure 2: Linear probing and kNN evaluation on CIFAR-10 are consistent. We evaluate the performance of a linear classifier (a) and a k-NN classifier (b) for pretraining with various patch sizes and various evaluation setups. During pretraining, BagSSL uses fixed-scale patches, and baseline methods use multi-scale augmentation RandomResizedCrop(min_scale, max_scale). For example, RandomResizedCrop(0.08, 1.0) corresponds to randomly and uniformly select crops ranging from $9\times 9$ to $32\times 32$. The patch size and augmentation crop size range are marked in the figure. The "Central" evaluation is the standard evaluation protocol where the classifier is trained and evaluated on single fixed central patches of the image, which is the entire image for CIFAR-10. For the $n$ patch evaluation, the classifier is trained and evaluated on the linearly-aggregated embedding of $n$ patches, sampled with the same scale factor as during pretraining. Please note that the "central" evaluation is expected to perform poorly on fix-scale pretraining as the model has never seen the entire image during pretraining.
• Figure 3: (a) Patch embedding convergence to the instance embedding. For a baseline multi-scale pretrained VICReg model, we show that the patch embedding aggregation converges to the whole-image embedding as the number of aggregated patches increases. We evaluate the cosine similarity between the aggregation of $N$ patch embeddings and the whole-image embedding. $N$ is selected from $1, 2, 4, 8, 16$ and all possible patches in the image. (b) Linear evaluation on ImageNet for various RandomResizedCrop scales. We show the performance of a linear classifier for various pretraining and evaluation settings. BagSSL uses fixed-scale patches, and the baseline method, VICReg, uses multi-scale augmentation RandomResizedCrop(0.08, 1.0), which corresponds to randomly and uniformly select crop scale ranging from $64\times 64$ to $224\times 224$. For each pretraining setting, we provide three different evaluations: 1) "Central": the standard evaluation, which takes a $224\times 224$ crop; 2) "$1$ patch": this evaluation takes $1$ patch or crop following the corresponding pretraining setting; 3) "$16$ patches": this evaluation takes $16$ randomly selected patches or crops following the corresponding pretraining setting.
• Figure 4: Visualization of cosine similarity in the projection space and the embedding space. Query patch is indicated by red dash. Projection and Embedding cosine-similarity heatmaps use the same color scaling. The projection vectors are significantly more invariant compared to the embedding ones, and the embedding space contains localized information that is shared among similar patches, when the size of the patches is small enough. We can see that the embedding space tends to preserve more locality compared to the projection space.
• Figure 5: Visualization of kNN in the projection space and the embedding space for CIFAR10. Distance is calculated by cosine similarity. Query patch is in the top left corner encircled by red-dash box, green box indicates patches from other image of the same class. Patches without surrounding box is from the same image as the query. While the nearest neighbors are both from same-category instances, we can see that the embedding space tends to preserve the local part information, whereas the projection space may collapse different parts of the same category.

Theorems & Definitions (2)

• Proposition 3.1
• proof