| name | Self-Supervised Learning |
| category | ai |
| description | Learning representations from unlabeled data using pretext tasks and contrastive methods |
Self-Supervised Learning
What I do
I enable machines to learn useful representations from large amounts of unlabeled data by creating supervised signals from the data itself. I eliminate the dependency on human-annotated labels by designing pretext tasks that encourage models to learn meaningful features. The learned representations can then be transferred to downstream tasks with minimal fine-tuning.
When to use me
- Pre-training on large unlabeled datasets when labeled data is scarce
- Building representation learning systems for transfer learning
- Learning from web-scale data without manual annotation
- Reducing labeling costs for domain-specific applications
- Pre-training foundation models for various downstream tasks
- Learning from images, text, audio, and other data types
- Building models that generalize across domains without supervision
Core Concepts
-
Pretext Tasks: Artificially created supervised tasks that don't require labels but teach useful representations.
-
Contrastive Learning: Learning by pulling similar examples together and pushing dissimilar examples apart in embedding space.
-
Momentum Contrast (MoCo): Using a momentum queue to maintain consistent negative samples during contrastive learning.
-
SimCLR Framework: Simple contrastive learning framework with data augmentation and large batch sizes.
-
Masked Prediction: Learning by predicting masked portions of input, used successfully in BERT and MAE.
-
BYOL and Bootstrap Your Own Latent: Methods that avoid negative samples through online and target networks.
-
Cluster Assignment: Methods like SwAV that assign cluster predictions as supervision signal.
-
Information Maximization: Learning representations that maximize mutual information between views.
Code Examples
import torch
import torch.nn as nn
import torch.nn.functional as F
class SimCLRProjectionHead(nn.Module):
def __init__(self, input_dim=2048, hidden_dim=2048, output_dim=128):
super().__init__()
self.head = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim)
)
def forward(self, features):
return self.head(features)
def nt_xent_loss(z_i, z_j, temperature=0.5):
batch_size = z_i.size(0)
z = torch.cat([z_i, z_j], dim=0)
similarity = torch.matmul(z, z.t()) / temperature
mask = torch.eye(2 * batch_size, device=z.device)
similarity = similarity.masked_fill(mask.bool(), -float('inf'))
pos_mask = torch.zeros(2 * batch_size, 2 * batch_size, device=z.device)
pos_mask[:batch_size, batch_size:] = torch.eye(batch_size, device=z.device)
pos_mask[batch_size:, :batch_size] = torch.eye(batch_size, device=z.device)
loss_i = -torch.log(
torch.exp(similarity[:batch_size, batch_size:2*batch_size].diag()) /
torch.exp(similarity[:batch_size]).sum(dim=1)
)
loss_j = -torch.log(
torch.exp(similarity[batch_size:, :batch_size].diag()) /
torch.exp(similarity[batch_size:]).sum(dim=1)
)
return (loss_i + loss_j).mean()
import torch
import torch.nn as nn
import torch.nn.functional as F
class MoCoQueue:
def __init__(self, queue_size=65536, embed_dim=128, momentum=0.999):
self.queue_size = queue_size
self.embed_dim = embed_dim
self.momentum = momentum
self.register_buffer("queue", torch.randn(queue_size, embed_dim))
self.queue = F.normalize(self.queue, p=2, dim=1)
self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
@torch.no_grad()
def update(self, keys):
batch_size = keys.size(0)
ptr = int(self.queue_ptr)
self.queue[ptr:ptr+batch_size] = keys
ptr = (ptr + batch_size) % self.queue_size
self.queue_ptr[0] = ptr
def forward(self, query, key, momentum_encoder):
query = F.normalize(query, p=2, dim=1)
with torch.no_grad():
key = momentum_encoder(key)
key = F.normalize(key, p=2, dim=1)
self.update(key)
positive = torch.sum(query * key, dim=-1, keepdim=True)
negative = torch.matmul(query, self.queue.t())
logits = torch.cat([positive, negative], dim=1) / 0.07
labels = torch.zeros(logits.size(0), dtype=torch.long, device=query.device)
return logits, labels
class MoCoModel(nn.Module):
def __init__(self, base_encoder, embed_dim=128, m=0.999):
super().__init__()
self.encoder_q = base_encoder
self.encoder_k = base_encoder
self.projector_q = SimCLRProjectionHead(output_dim=embed_dim)
self.projector_k = SimCLRProjectionHead(output_dim=embed_dim)
for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
param_k.data.copy_(param_q.data)
param_k.requires_grad = False
self.moco_queue = MoCoQueue(embed_dim=embed_dim, momentum=m)
def forward(self, im_q, im_k):
q = self.encoder_q(im_q)
q = self.projector_q(q)
logits, labels = self.moco_queue.forward(q, im_k, self.encoder_k)
return logits, labels
@torch.no_grad()
def _momentum_update_key_encoder(self):
for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)
import torch
import torch.nn as nn
import torch.nn.functional as F
class BYOLPredictor(nn.Module):
def __init__(self, input_dim=2048, hidden_dim=512, output_dim=128):
super().__init__()
self.predictor = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim)
)
def forward(self, z):
return self.predictor(z)
class BYOLModel(nn.Module):
def __init__(self, base_encoder, input_dim=2048, hidden_dim=512, output_dim=128, m=0.996):
super().__init__()
self.encoder = base_encoder
self.projector = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim)
)
self.predictor = BYOLPredictor(input_dim, hidden_dim, output_dim)
self.online_encoder = nn.Sequential(self.encoder, self.projector)
self.target_encoder = nn.Sequential(
type(self.encoder)(**self.encoder_kwargs),
type(self.projector)(**self.projector_kwargs)
)
for param in self.target_encoder.parameters():
param.requires_grad = False
self.m = m
def forward(self, v1, v2):
h1 = self.online_encoder(v1)
h2 = self.online_encoder(v2)
z1 = F.normalize(h1, dim=1)
z2 = F.normalize(h2, dim=1)
p1 = self.predictor(z1)
p2 = self.predictor(z2)
with torch.no_grad():
t1 = self.target_encoder(v1)
t2 = self.target_encoder(v2)
z1_t = F.normalize(t1, dim=1)
z2_t = F.normalize(t2, dim=1)
loss = 2 - F.cosine_similarity(p1, z2_t, dim=-1).mean() - F.cosine_similarity(p2, z1_t, dim=-1).mean()
return loss
@torch.no_grad()
def update_target(self):
for param_o, param_t in zip(self.online_encoder.parameters(), self.target_encoder.parameters()):
param_t.data = param_t.data * self.m + param_o.data * (1 - self.m)
import torch
import random
from typing import Tuple, List
class SimCLRTransforms:
def __init__(self, size=224):
self.color_jitter = torch.nn.ColorJitter(
brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1
)
self.gaussian_blur = GaussianBlur(kernel_size=23)
self.size = size
def __call__(self, img):
t1 = self._transform(img)
t2 = self._transform(img)
return t1, t2
def _transform(self, img):
img = self._random_resize_crop(img)
img = self.color_jitter(img)
img = self._random_grayscale(img)
if random.random() < 0.5:
img = self.gaussian_blur(img)
img = self._normalize(img)
return img
def _random_resize_crop(self, img):
i, j, h, w = transforms.RandomResizedCrop.get_params(img, (0.08, 1.0), (3/4, 4/3))
return transforms.functional.resized_crop(img, i, j, h, w, self.size)
def _random_grayscale(self, img):
if random.random() < 0.2:
return transforms.functional.to_grayscale(img)
return img
def _normalize(self, img):
return transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])(img)
class GaussianBlur:
def __init__(self, kernel_size=23):
self.kernel_size = kernel_size
def __call__(self, img):
sigma = random.uniform(0.1, 2.0)
return transforms.functional.gaussian_blur(img, self.kernel_size, sigma=[sigma])
import torch
import torch.nn as nn
class MAEEncoder(nn.Module):
def __init__(self, patch_size=16, embed_dim=768, depth=12, num_heads=12):
super().__init__()
self.patch_embed = nn.Conv2d(3, embed_dim, kernel_size=patch_size, stride=patch_size)
self.pos_embed = nn.Parameter(torch.zeros(1, 196, embed_dim))
self.blocks = nn.ModuleList([
nn.TransformerEncoderLayer(
d_model=embed_dim,
nhead=num_heads,
dim_feedforward=4*embed_dim,
dropout=0.0
) for _ in range(depth)
])
self.norm = nn.LayerNorm(embed_dim)
self._init_weights()
def _init_weights(self):
nn.init.trunc_normal_(self.pos_embed, std=0.02)
nn.init.xavier_uniform_(self.patch_embed.weight)
def forward(self, x):
x = self.patch_embed(x)
x = x.flatten(2).transpose(1, 2)
x = x + self.pos_embed
for block in self.blocks:
x = block(x)
return self.norm(x)
class MAEDecoder(nn.Module):
def __init__(self, embed_dim=768, depth=8, num_heads=16, patch_size=16, num_patches=196):
super().__init__()
self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
self.blocks = nn.ModuleList([
nn.TransformerEncoderLayer(
d_model=embed_dim,
nhead=num_heads,
dim_feedforward=4*embed_dim
) for _ in range(depth)
])
self.head = nn.Linear(embed_dim, 3 * patch_size ** 2)
self._init_weights()
def _init_weights(self):
nn.init.normal_(self.mask_token, std=0.02)
nn.init.trunc_normal_(self.pos_embed, std=0.02)
def forward(self, x, ids_restore, ids_keep):
x = self._unpatchify(x)
return x
def _unpatchify(self, x):
C = 3
p = 16
h = w = int(x.shape[1] ** 0.5)
x = x.reshape(shape=(x.shape[0], h, w, p, p, C))
x = torch.einsum('nhwpqc->nchpwq', x)
return x.reshape(x.shape[0], C, h * p, w * p)
def mae_loss(reconstructed, original, mask):
loss = (reconstructed - original) ** 2
loss = loss.mean(dim=-1)
loss = (loss * mask).sum() / mask.sum()
return loss
Best Practices
-
Use stronger augmentations for contrastive learning to create more informative positive pairs.
-
Increase batch size for SimCLR as it directly affects the number of negative samples.
-
Use learning rate warmup and cosine scheduling for stable contrastive training.
-
Monitor representation similarity metrics to track learning progress.
-
Use projection heads to map features to contrastive space, not raw backbone outputs.
-
Apply symmetric loss terms when using two views to ensure balanced optimization.
-
Use momentum encoders to maintain consistency in MoCo and BYOL.
-
Choose appropriate temperature values (typically 0.1) for contrastive losses.
-
Use sufficient epochs for pre-training as contrastive learning converges slower than supervised.
-
Evaluate using linear probing on frozen features to assess representation quality.