SwAV: Unsupervised Learning of Visual Features by Contrasting Cluster Assignments

Contrastive learning methods like SimCLR learn representations by creating two augmented views of each image, passing them through an encoder, and then comparing their features. The model is trained to make the two views of the same image similar (positive pair) while making views of different images dissimilar (negative pairs).

The catch: you need a lot of negative pairs for this to work well. SimCLR uses batch sizes of 4096–8192, meaning each positive pair is compared against thousands of negatives. MoCo builds a momentum-maintained queue of 65,536 negatives. Both approaches work, but they introduce significant computational or architectural overhead.

What if you could get the benefits of contrastive learning — learning without labels by comparing views — without the cost of comparing every feature pair?

SwAV — Swapping Assignments between Views — replaces direct feature comparison with an elegant two-step dance. First, it maps every feature to a set of $K$ learnable prototype vectors, producing a soft cluster assignment (a "code"). Second, it enforces consistency by predicting the code of one view from the representation of another view.

Here is the pipeline: take an image $x_n$, create two augmented views $x_{n,t}$ and $x_{n,s}$ using different data augmentation transformations. Pass each through a shared encoder $f_\theta$ to get features $z_t$ and $z_s$. Each feature is then compared to all $K$ prototype vectors $\{c_1, \ldots, c_K\}$ to produce a code — a probability distribution over prototypes.

The key trick: the code for view $t$ is predicted using the feature of view $s$, and vice versa. If both views truly capture the same content, then knowing one view's representation should be enough to predict the other view's cluster assignment. This "swapped" prediction is what gives SwAV its name.

Before seeing the formula, let's understand what SwAV is optimizing. For each image, we have two views. Each view produces a feature vector and a code (soft cluster assignment). The loss asks: "Can I predict the code that view $t$ received, using only the feature vector of view $s$?" If the encoder and prototypes are good, the answer should be yes — because both views depict the same image.

The prediction works by computing the similarity between a feature $z_s$ and each prototype $c_k$ (via a dot product), applying a temperature-scaled softmax to get a probability $p_s^{(k)}$, and then measuring how well this probability matches the code $q_t$ of the other view using cross-entropy.

The prototype vectors $\{c_1, \ldots, c_K\}$ live in the same embedding space as the features. Think of them as $K$ "topic anchors" — fixed reference points that the features are compared against. Unlike k-means centroids that are recomputed from scratch, SwAV prototypes are learned end-to-end via backpropagation alongside the encoder.

In practice, $K = 3000$ prototypes are used. The prototypes form the columns of a matrix $C \in \mathbb{R}^{D \times K}$, where $D$ is the feature dimension. Computing the code for a feature $z$ is simply a matrix multiplication $C^\top z$ followed by the Sinkhorn-Knopp normalization.

Because the prototypes are shared across all images and all batches, they act as a persistent clustering structure. Each batch updates the prototypes slightly through gradient descent, so the clustering evolves smoothly over training.

A critical danger in clustering-based methods is collapse: the model assigns every image to the same cluster and achieves zero loss by ignoring all differences. This is the trivial solution — and it is catastrophic.

SwAV prevents collapse by formulating the code computation as an optimal transport problem. Instead of simply taking the argmax or softmax of $C^\top z$, SwAV requires that across all images in a batch, the cluster assignments are equipartitioned — every prototype gets roughly the same total assignment mass. This means no single prototype can hoard all the features.

The constraint is: the assignment matrix $Q$ must have rows that sum to $1/K$ (equal mass per prototype) and columns that sum to $1/B$ (each image gets a valid distribution). The Sinkhorn-Knopp algorithm finds this doubly-stochastic matrix by alternating row and column normalizations:

Standard contrastive methods create two views per image, both at full resolution (e.g. 224×224). Increasing the number of views quadratically increases compute in pairwise methods (more pairs to compare). But SwAV doesn't compare pairs — it compares features to prototypes. So adding more views only increases compute linearly.

The authors exploit this by proposing multi-crop: instead of 2 full-resolution views, use 2 full-resolution views (160×160) plus $V$ small views (96×96). The small views cover different parts of the image, capturing local details like textures, edges, and small objects. The full views capture global structure.

The key insight: the loss only computes codes from the two full-resolution views (which are high quality), but uses ALL views for the prediction term. This means the small views provide additional training signal without needing expensive code computation. In practice, $V=4$ extra views adds only ~20% compute but boosts performance by ~2% on ImageNet.

SwAV with ResNet-50 achieves 75.3% top-1 accuracy on ImageNet linear evaluation — a jump of +4.2% over SimCLR (71.1%) and +4.2% over MoCo-v2 (71.1%). This was the first self-supervised method to close the gap with supervised pretraining (76.5%) to just 1.2 percentage points.

On transfer tasks, the picture is even more striking. SwAV representations surpass supervised pretraining on all considered downstream benchmarks — object detection on VOC and COCO, instance segmentation on COCO, and image classification on Places205 and iNaturalist. This suggests that self-supervised features generalize better than supervised ones, which may overfit to ImageNet's label distribution.

The multi-crop strategy contributes significantly: without it, SwAV achieves 72.7%. Adding 4 small views ($V=4$) pushes it to 75.3%, a gain of +2.6% for minimal extra cost. The number of prototypes (from 3K to 100K) has little effect — performance varies by only ±0.3%, showing that SwAV is robust to this hyperparameter.

MoCo maintains a momentum-updated copy of the entire encoder, doubling memory. SimCLR needs batch sizes in the thousands. SwAV needs neither. The prototypes serve as a compact, persistent summary of the data distribution — replacing both the memory bank and the large batch requirement.

Because the prototypes are just a $D \times K$ matrix (e.g., $256 \times 3000$), they add negligible memory. The Sinkhorn-Knopp step operates on a $K \times B$ matrix, which is also tiny. The dominant cost is the forward pass through the encoder — the same as any other method.

This architectural simplicity also makes SwAV easier to scale. It works with batch sizes from 256 to 4096. It can use a small queue of recent features (2–3 batches) to stabilize the Sinkhorn step for small batches, but this queue stores only features, not gradients, so it is much cheaper than MoCo's momentum encoder.

SwAV stands at the crossroads of two traditions: offline clustering (DeepCluster, SeLa) and online contrastive learning (SimCLR, MoCo). By making clustering online through learnable prototypes and Sinkhorn-Knopp, and by replacing pairwise comparison with swapped prediction, SwAV combined the best of both worlds. Its multi-crop augmentation became a default tool, and the prototype-based architecture directly seeded DINO, one of the most influential self-supervised methods for vision.