Online Prototype Clustering in SwAV Using Sinkhorn-Knopp

This article explores how Swapping Assignments between Views (SwAV) achieves efficient self-supervised learning by combining online clustering with contrastive methods. It details the mechanics of extracting features from augmented image views, mapping these features to learnable prototype vectors, and utilizing the Sinkhorn-Knopp algorithm to normalize cluster assignments in real-time without suffering from mode collapse. Finally, it explains how the swapped prediction mechanism forces the neural network to learn consistent, high-level semantic representations across different views of the same image.

Generating Views and Feature Representations

SwAV begins by taking a single input image and applying different random augmentations to generate multiple distinct views. These views are passed through a convolutional neural network encoder, such as a ResNet, to extract high-dimensional feature vectors. These vectors are then projected onto the unit sphere using L2 normalization. Instead of directly comparing the features of the different views—which requires massive memory banks or large batch sizes—SwAV maps the normalized features to a set of \(K\) learnable prototype vectors. These prototypes act as the centers of clusters in the embedding space and are updated via backpropagation alongside the encoder weights.

Computing Codes and the Collapse Problem

When a batch of \(B\) images is processed, the network computes the dot product between the normalized feature vectors of the images and the \(K\) learnable prototypes. This produces a matrix of similarity scores (logits). To convert these similarity scores into cluster assignments—referred to as "codes"—the system must decide which prototype each image view belongs to.

In an online setting, simply assigning each image to its nearest prototype leads to a trivial solution known as mode collapse. The network quickly learns to assign every single image in the batch to the exact same prototype, maximizing similarity without learning any meaningful visual features.

Sinkhorn-Knopp Normalization for Online Assignment

To prevent mode collapse and process batches online, SwAV treats the assignment of features to prototypes as an Optimal Transport problem. The goal is to maximize the similarity between the features and the prototypes while enforcing an equipartition constraint: on average, each prototype in the batch must be selected an equal number of times.

SwAV achieves this using the Sinkhorn-Knopp algorithm. The algorithm takes the similarity matrix of the current batch and iteratively normalizes the rows and columns.

1. First, the similarity scores are exponentiated to make them positive, scaled by a temperature parameter.
2. The algorithm then alternates between normalizing the rows (ensuring each image view's assignment weights sum to 1) and normalizing the columns (ensuring the total weight assigned to each prototype across the batch equals \(B/K\)).

After just a few iterations (typically three), the Sinkhorn-Knopp algorithm converges. It outputs a doubly stochastic matrix of "soft" codes, meaning images are assigned fractional memberships to multiple prototypes rather than a hard, single-cluster assignment. Because every prototype is forced to receive an equal amount of assignment weight across the batch, mode collapse is mathematically impossible. If the batch size is too small relative to the number of prototypes, SwAV caches features from previous batches in a small queue to provide enough data for the Sinkhorn-Knopp normalization to remain stable.

The Swapped Prediction Loss

Once the Sinkhorn-Knopp algorithm generates the soft assignment codes for the different views of an image, SwAV executes its core learning mechanism: swapped prediction.

If an image is transformed into View 1 and View 2, they should contain the same semantic information and therefore belong to the same prototype clusters. SwAV sets up a cross-entropy loss where the network uses the feature representation of View 1 to predict the Sinkhorn-Knopp code generated from View 2, and simultaneously uses the features of View 2 to predict the code of View 1.

During backpropagation, gradients flow through the predicted probabilities to update both the encoder network and the learnable prototype vectors. Note that gradients do not flow through the Sinkhorn-Knopp algorithm itself; the codes are treated as fixed targets for the opposite view. By continuously forcing different views of the same image to map to the same uniformly distributed prototypes, the model learns rich, invariant visual representations entirely online.