SwAV (Swapping Assignments between Views)

SwAV computes cluster assignments online using the Sinkhorn-Knopp algorithm, an iterative normalization procedure that enforces an equipartition constraint. This ensures that samples within a batch are evenly distributed across clusters, preventing the trivial solution where all representations collapse into a single cluster. The algorithm operates directly on the transport polytope, producing soft assignments that respect both row and column normalization constraints. This approach replaces the need for a large memory bank or queue of negative samples, significantly reducing memory footprint while maintaining representation quality.

The core innovation of SwAV is the swapped prediction loss: the model predicts the cluster assignment of one augmented view using the feature representation of another view, and vice versa. Given two views of the same image, the encoder produces features that are mapped to a set of trainable prototype vectors. The loss enforces that the assignment computed from view A matches the features of view B, creating a cross-prediction task. This swapped mechanism implicitly compares all views against all prototypes, providing a rich training signal without requiring pairwise feature comparisons or negative samples.

SwAV leverages a multi-crop strategy that generates multiple views at different resolutions from a single image. Standard crops (e.g., 2x224²) capture global context, while smaller crops (e.g., Vx96²) capture local details at reduced computational cost. The loss is applied only between standard-resolution crops, while smaller crops are used as additional targets. This enforces local-to-global consistency: the model learns that a small patch should map to the same cluster as the full image, dramatically improving representation quality without a proportional increase in compute.

SwAV maintains a set of learnable prototype vectors that serve as cluster centroids in the embedding space. These prototypes are updated via backpropagation alongside the encoder weights, allowing the clustering structure to adapt during training. The number of prototypes (typically 3,000) is a hyperparameter that controls the granularity of the learned representation. Each prototype captures a distinct visual concept, and the soft assignment of features to these prototypes creates a rich, structured representation that transfers effectively to downstream medical imaging tasks such as organ segmentation and lesion classification.

SwAV prevents representation collapse through the equipartition constraint enforced by the Sinkhorn-Knopp algorithm. Unlike contrastive methods that rely on negative pairs to push representations apart, SwAV ensures diversity by requiring that each prototype is assigned an equal number of samples on average across the batch. This constraint operates at the batch level and is enforced through an entropy regularization term. The result is a self-organizing system where the model cannot cheat by assigning all samples to a single prototype, maintaining a balanced and informative clustering structure throughout training.

SwAV is particularly well-suited for medical image pre-training due to its ability to learn from unlabeled scans without requiring large batch sizes or negative pair curation. In radiology and pathology, where annotated data is scarce but unlabeled DICOM archives are abundant, SwAV can pre-train on chest X-rays, CT volumes, or whole slide images to learn clinically relevant features. The prototype vectors often converge to represent anatomical structures and tissue textures, providing a strong initialization for downstream tasks such as tumor detection and disease staging with limited labeled data.

The core paradigm behind SwAV, where cluster assignments are computed on-the-fly for each batch rather than over the entire dataset. This contrasts with offline clustering methods like DeepCluster. The system maintains a set of prototype vectors that act as cluster centroids, and the Sinkhorn-Knopp algorithm ensures balanced assignment across these prototypes to prevent the trivial solution where all representations collapse to a single cluster.

A fast, GPU-friendly iterative normalization procedure that transforms a raw similarity matrix into a doubly stochastic matrix. In SwAV, it enforces an equipartition constraint, ensuring that each prototype receives an equal share of assignments across the batch. This prevents representation collapse by forcing the model to utilize all clusters uniformly, acting as a regularizer without requiring explicit negative pairs.

A data augmentation strategy integral to SwAV's efficiency. Instead of using only two full-resolution crops, it generates multiple views per image:

• 2 standard large crops (e.g., 224x224) capturing global context
• V additional small crops (e.g., 96x96) capturing local details The loss enforces consistency between cluster assignments of large crops and predictions from small crops, dramatically increasing training signal density without proportional compute cost.

SwAV fundamentally differs from contrastive methods like SimCLR and MoCo by eliminating the need for explicit negative pairs. Instead of comparing features directly in the embedding space, SwAV compares their cluster assignment codes. This 'swapped prediction' mechanism—predicting the code of one view from the features of another—achieves comparable or superior performance while being computationally more efficient, as it avoids the quadratic cost of pairwise comparisons.

A learnable matrix of K trainable vectors of dimension D, where K is typically set to 3000. These prototypes represent the cluster centroids in the projection space. During training, the model computes the similarity between encoded features and all prototypes, then applies Sinkhorn-Knopp to generate soft assignment codes. The prototypes are optimized via backpropagation alongside the encoder, learning to represent the most discriminative visual concepts in the dataset.

The namesake operation of SwAV. Given two views of an image, the model computes a cluster assignment code from one view and predicts it using the features of the other view. This 'swapping' forces the model to learn representations that are invariant to augmentations while being discriminative enough to map to consistent clusters. The loss is a cross-entropy between the predicted softmax probabilities and the Sinkhorn-generated target code.