Advanced Contrastive Learning
Contrastive learning learns representations by pulling positive pairs together and pushing negative pairs apart. This module covers the theoretical foundations, negative sampling strategies, augmentation design, and advanced variants with mathematical depth.
1. Information-Theoretic Foundations
1.1 Mutual Information Maximization
Contrastive learning maximizes a lower bound on mutual information :
1.2 InfoNCE Bound
The InfoNCE loss (van den Oord et al., 2018) provides a lower bound:
where is the number of negative samples.
Proof: The InfoNCE loss is:
By the data processing inequality, where .
1.3 Tightness of the Bound
The bound is tight up to . As :
The gap is , which is negligible for large .
2. Alignment and Uniformity
2.1 Wang & Isola's Decomposition
The contrastive loss decomposes into alignment and uniformity:
Theorem: Minimizing contrastive loss simultaneously:
- Alignment: Brings positive pairs close in embedding space
- Uniformity: Distributes embeddings uniformly on the hypersphere
2.2 Uniformity on Hypersphere
A distribution on is uniform if:
Power-law uniformity: The uniformity loss corresponds to:
which is minimized when is uniform on .
2.3 Alignment Loss
The alignment loss requires:
for some small .
3. Negative Sampling
3.1 Uniform Negative Sampling
Standard approach: sample negatives uniformly from the dataset:
3.2 Hard Negative Mining
Sample negatives that are semantically similar to the anchor:
Benefits: More informative gradients, faster convergence Risks: May learn spurious features, amplifies label noise
3.3 Memory Bank Negatives
Store all dataset embeddings in a memory bank:
Update with momentum:
3.4 Cluster-Based Negatives
Use cluster assignments to weight negatives:
where is the prototype of the cluster containing .
4. Augmentation Design
4.1 Augmentation Space
Define augmentation distribution :
where each is a composition of transformations.
4.2 Augmentation Strength
The augmentation strength affects representation quality:
Too weak: Insufficient invariance, trivial solutions Too strong: Semantic content destroyed, poor generalization
4.3 Adaptive Augmentation
Learn augmentation policy jointly with representation:
AutoAugment: Search for optimal augmentation policy per dataset.
4.4 View Distribution
Global-local views (DINO): Multiple small crops + 2 large global crops.
Multi-scale views: Different scales reveal different semantic levels.
5. Instance Discrimination
5.1 Theoretical Foundation
Instance discrimination treats each instance as its own class:
5.2 Connection to Kernel Methods
The contrastive kernel is:
which is a learned kernel function. The embedding is the feature map of this kernel.
5.3 NCE (Noise Contrastive Estimation)
Gutmann & Hyvärinen (2010) estimate the density ratio:
The NCE loss:
6. Clustering-Based Methods
6.1 SwAV
Sinkhorn-Knopp assignment with online prototypes:
Sinkhorn normalization: Iterate until doubly stochastic:
6.2 DeepCluster
Alternating between k-means clustering and representation learning:
- Assign clusters:
- Learn representation:
6.3 Prototypical Contrastive Learning
PCL (Li et al., 2021): Cluster prototypes as negative samples:
7. Implementation
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class ContrastiveLoss(nn.Module):
def __init__(self, temperature=0.07, negative_sampling='uniform'):
super().__init__()
self.temperature = temperature
self.negative_sampling = negative_sampling
def info_nce(self, z1, z2, negatives=None):
batch_size = z1.shape[0]
z1 = F.normalize(z1, dim=1)
z2 = F.normalize(z2, dim=1)
if negatives is None:
# Uniform negative sampling
z_all = torch.cat([z1, z2], dim=0)
sim_matrix = z_all @ z_all.T / self.temperature
mask = ~torch.eye(2 * batch_size, dtype=bool, device=z1.device)
sim_matrix = sim_matrix.masked_select(mask).view(2 * batch_size, -1)
else:
# With explicit negatives
neg = F.normalize(negatives, dim=1)
pos_sim = (z1 * z2).sum(dim=1) / self.temperature
neg_sim = z1 @ neg.T / self.temperature
sim_matrix = torch.cat([pos_sim.unsqueeze(1), neg_sim], dim=1)
labels = torch.zeros(z1.shape[0], dtype=torch.long, device=z1.device)
loss = F.cross_entropy(sim_matrix, labels)
return loss
def alignment_loss(self, z1, z2):
return (z1 - z2).norm(dim=1).pow(2).mean()
def uniformity_loss(self, z, t=2):
z = F.normalize(z, dim=1)
sq_pdist = torch.pdist(z, p=2).pow(2)
return sq_pdist.mul(-t).exp().mean().log()
def forward(self, z1, z2, negatives=None):
align = self.alignment_loss(z1, z2)
unif = (self.uniformity_loss(z1) + self.uniformity_loss(z2)) / 2
return align + unif
class AugmentationPipeline:
def __init__(self):
self.transforms = [
RandomResizedCrop(224, scale=(0.2, 1.0)),
RandomHorizontalFlip(),
ColorJitter(0.4, 0.4, 0.4, 0.1),
RandomGrayscale(p=0.2),
GaussianBlur(kernel_size=23, sigma=(0.1, 2.0)),
Solarization(p=0.2),
]
def __call__(self, x):
for t in self.transforms:
if np.random.random() < t.p:
x = t(x)
return x
class MemoryBank:
def __init__(self, size, dim, momentum=0.5):
self.size = size
self.momentum = momentum
self.memory = F.normalize(torch.randn(size, dim), dim=1)
def update(self, indices, features):
features = F.normalize(features, dim=1)
with torch.no_grad():
self.memory[indices] = (
self.momentum * self.memory[indices] +
(1 - self.momentum) * features
)
def get_negatives(self, indices, exclude=None):
mask = torch.ones(self.size, dtype=bool, device=indices.device)
mask[indices] = False
if exclude is not None:
mask[exclude] = False
return self.memory[mask]
8. SVG: Contrastive Pairs Visualization
9. SVG: Augmentation Pipeline
10. Advanced Topics
10.1 Adaptive Temperature
Learn temperature per sample:
where is a small network predicting the optimal temperature.
10.2 Multi-Scale Contrastive Learning
where each scale uses different augmentation strengths.
10.3 Contrastive Learning with Structured Negatives
Use class hierarchy or semantic relationships:
References
- van den Oord, A., Li, Y., & Vinyals, O. (2018). Representation Learning with Contrastive Predictive Coding. arXiv.
- Wang, T., & Isola, P. (2020). Understanding Contrastive Representation Learning through Alignment and Uniformity. ICML.
- Gutmann, M., & Hyvärinen, A. (2010). Noise-Contrastive Estimation. JMLR.
- Caron, M., et al. (2020). Unsupervised Learning of Visual Features by Contrasting Cluster Assignments. NeurIPS.
- Li, J., et al. (2021). Prototypical Contrastive Learning of Unsupervised Representations. ICML.