Topic 22: Self-Supervised Learning — First Principles¶

Goal: Implement core self-supervised learning algorithms from scratch in pure NumPy, including InfoNCE (NT-Xent) contrastive loss, an encoder/projector MLP, SimCLR training loop, and patch masking utilities for Masked Autoencoders (MAE).

Prerequisites: 13 Neural Networks, 17 Autoencoder

Theory Link: theory.md

In [1]:

import random

import matplotlib.pyplot as plt
import numpy as np

%load_ext autoreload
%autoreload 2

SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)

plt.rcParams["figure.dpi"] = 90

from ml_first_principles.ssl_models import InfoNCELoss
from sklearn.datasets import make_blobs

import random import matplotlib.pyplot as plt import numpy as np %load_ext autoreload %autoreload 2 SEED = 42 random.seed(SEED) rng = np.random.default_rng(SEED) plt.rcParams["figure.dpi"] = 90 from ml_first_principles.ssl_models import InfoNCELoss from sklearn.datasets import make_blobs

1. Problem Setup — WHY¶

In supervised learning, models rely on human labels $y$ to group inputs $x$. In contrastive Self-Supervised Learning (SSL), we build artificial supervision by generating multiple stochastic views of the same sample ($\tilde{x}_i, \tilde{x}_j$). The network must learn to pull representations of positive views together while pushing representations of different samples apart.

sklearn.datasets.make_blobs below is used for synthetic data generation only — all the SSL math in this notebook stays pure NumPy.

In [2]:

# Generate synthetic clusters representing underlying semantic concepts
# (300 samples keeps the 2N x 2N contrastive matrices fast enough for a < 60 s run)
X, y = make_blobs(n_samples=300, centers=4, cluster_std=0.5, random_state=SEED)

def augment(X, noise_scale=0.3):
    # Stochastic data augmentation: add random Gaussian noise
    return X + rng.normal(0, noise_scale, size=X.shape)

plt.figure(figsize=(6, 4))
plt.scatter(X[:, 0], X[:, 1], c=y, cmap='tab10', alpha=0.6, s=20)
plt.title("Synthetic Unlabeled Data Clusters")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
plt.grid(True, linestyle='--', alpha=0.5)
plt.show()

# Generate synthetic clusters representing underlying semantic concepts # (300 samples keeps the 2N x 2N contrastive matrices fast enough for a < 60 s run) X, y = make_blobs(n_samples=300, centers=4, cluster_std=0.5, random_state=SEED) def augment(X, noise_scale=0.3): # Stochastic data augmentation: add random Gaussian noise return X + rng.normal(0, noise_scale, size=X.shape) plt.figure(figsize=(6, 4)) plt.scatter(X[:, 0], X[:, 1], c=y, cmap='tab10', alpha=0.6, s=20) plt.title("Synthetic Unlabeled Data Clusters") plt.xlabel("Feature 1") plt.ylabel("Feature 2") plt.grid(True, linestyle='--', alpha=0.5) plt.show()

2. Mathematical Core — WHAT¶

For a batch of $N$ samples, we generate 2 augmented views per sample ($2N$ points total). For positive pair $(i, j)$:

$$\ell_{i,j} = -\log \frac{\exp(\text{sim}(z_i, z_j) / \tau)}{\sum_{k=1}^{2N} \mathbb{1}_{[k \neq i]} \exp(\text{sim}(z_i, z_k) / \tau)}$$

where $z_i = g(f(\tilde{x}_i)) / \|g(f(\tilde{x}_i))\|_2$ is the normalized projection and $\text{sim}(u, v) = u^T v$.

3. Solution Method — HOW¶

1. InfoNCE Loss: Compute the full $2N \times 2N$ cosine similarity matrix, mask self-similarities with $-\infty$, apply softmax, and calculate cross-entropy over positive pairs.
2. Encoder + Projector: Build a 2-layer MLP with L2 normalization on the output projection space.
3. SimCLR Loop: Iterate over batches, generate positive views, compute InfoNCE loss, and backpropagate gradients.
4. Patch Masking Utility: Implement random patch masking (75%) for Vision Transformer inputs.

4. Implementation — BUILD¶

In [3]:

def info_nce_loss(z1, z2, temperature=0.5):
    """
    Computes NT-Xent (InfoNCE) loss and gradients for a batch of positive pairs.
    """
    N = z1.shape[0]
    Z = np.vstack([z1, z2])  # (2N, D)
    
    # Pairwise similarity matrix
    sim_matrix = (Z @ Z.T) / temperature  # (2N, 2N)
    np.fill_diagonal(sim_matrix, -np.inf)
    
    # Targets: positive pair for i is i+N (for i < N) and i-N (for i >= N)
    labels = np.zeros(2 * N, dtype=int)
    labels[:N] = np.arange(N, 2 * N)
    labels[N:] = np.arange(N)
    
    # Softmax probabilities
    exp_sim = np.exp(sim_matrix - np.max(sim_matrix, axis=1, keepdims=True))
    probs = exp_sim / np.sum(exp_sim, axis=1, keepdims=True)
    
    loss = -np.mean(np.log(probs[np.arange(2 * N), labels] + 1e-8))
    
    # Gradients
    d_sim = probs.copy()
    d_sim[np.arange(2 * N), labels] -= 1.0
    d_sim = d_sim / (2 * N * temperature)
    dZ = d_sim @ Z + d_sim.T @ Z
    
    return loss, dZ

def info_nce_loss(z1, z2, temperature=0.5): """ Computes NT-Xent (InfoNCE) loss and gradients for a batch of positive pairs. """ N = z1.shape[0] Z = np.vstack([z1, z2]) # (2N, D) # Pairwise similarity matrix sim_matrix = (Z @ Z.T) / temperature # (2N, 2N) np.fill_diagonal(sim_matrix, -np.inf) # Targets: positive pair for i is i+N (for i < N) and i-N (for i >= N) labels = np.zeros(2 * N, dtype=int) labels[:N] = np.arange(N, 2 * N) labels[N:] = np.arange(N) # Softmax probabilities exp_sim = np.exp(sim_matrix - np.max(sim_matrix, axis=1, keepdims=True)) probs = exp_sim / np.sum(exp_sim, axis=1, keepdims=True) loss = -np.mean(np.log(probs[np.arange(2 * N), labels] + 1e-8)) # Gradients d_sim = probs.copy() d_sim[np.arange(2 * N), labels] -= 1.0 d_sim = d_sim / (2 * N * temperature) dZ = d_sim @ Z + d_sim.T @ Z return loss, dZ

In [4]:

class SimpleMLP:
    def __init__(self, input_dim, hidden_dim, output_dim):
        self.W1 = rng.standard_normal((input_dim, hidden_dim)) * np.sqrt(2.0 / input_dim)
        self.b1 = np.zeros(hidden_dim)
        self.W2 = rng.standard_normal((hidden_dim, output_dim)) * np.sqrt(2.0 / hidden_dim)
        self.b2 = np.zeros(output_dim)
        
    def forward(self, X):
        self.X = X
        self.z1 = X @ self.W1 + self.b1
        self.a1 = np.maximum(0, self.z1)
        self.z2 = self.a1 @ self.W2 + self.b2
        self.norms = np.linalg.norm(self.z2, axis=1, keepdims=True)
        self.out = self.z2 / (self.norms + 1e-8)
        return self.out
        
    def backward(self, d_out, lr=0.01):
        d_z2 = (d_out - self.out * np.sum(d_out * self.out, axis=1, keepdims=True)) / (self.norms + 1e-8)
        dW2 = self.a1.T @ d_z2
        db2 = np.sum(d_z2, axis=0)
        d_a1 = d_z2 @ self.W2.T
        d_z1 = d_a1 * (self.z1 > 0)
        dW1 = self.X.T @ d_z1
        db1 = np.sum(d_z1, axis=0)
        self.W1 -= lr * dW1
        self.b1 -= lr * db1
        self.W2 -= lr * dW2
        self.b2 -= lr * db2

class SimpleMLP: def __init__(self, input_dim, hidden_dim, output_dim): self.W1 = rng.standard_normal((input_dim, hidden_dim)) * np.sqrt(2.0 / input_dim) self.b1 = np.zeros(hidden_dim) self.W2 = rng.standard_normal((hidden_dim, output_dim)) * np.sqrt(2.0 / hidden_dim) self.b2 = np.zeros(output_dim) def forward(self, X): self.X = X self.z1 = X @ self.W1 + self.b1 self.a1 = np.maximum(0, self.z1) self.z2 = self.a1 @ self.W2 + self.b2 self.norms = np.linalg.norm(self.z2, axis=1, keepdims=True) self.out = self.z2 / (self.norms + 1e-8) return self.out def backward(self, d_out, lr=0.01): d_z2 = (d_out - self.out * np.sum(d_out * self.out, axis=1, keepdims=True)) / (self.norms + 1e-8) dW2 = self.a1.T @ d_z2 db2 = np.sum(d_z2, axis=0) d_a1 = d_z2 @ self.W2.T d_z1 = d_a1 * (self.z1 > 0) dW1 = self.X.T @ d_z1 db1 = np.sum(d_z1, axis=0) self.W1 -= lr * dW1 self.b1 -= lr * db1 self.W2 -= lr * dW2 self.b2 -= lr * db2

In [5]:

def mask_image_patches(image, patch_size=4, mask_ratio=0.75):
    """Simulates 75% random patch masking for Masked Autoencoder (MAE)."""
    H, W = image.shape
    num_patches = (H // patch_size) * (W // patch_size)
    num_keep = int(num_patches * (1 - mask_ratio))
    indices = rng.permutation(num_patches)
    keep_indices = indices[:num_keep]
    mask = np.ones(num_patches, dtype=bool)
    mask[keep_indices] = False
    mask_2d = mask.reshape(H // patch_size, W // patch_size)
    mask_pixel = np.kron(mask_2d, np.ones((patch_size, patch_size)))
    masked_img = image.copy()
    masked_img[mask_pixel == 1] = 0
    return masked_img

def mask_image_patches(image, patch_size=4, mask_ratio=0.75): """Simulates 75% random patch masking for Masked Autoencoder (MAE).""" H, W = image.shape num_patches = (H // patch_size) * (W // patch_size) num_keep = int(num_patches * (1 - mask_ratio)) indices = rng.permutation(num_patches) keep_indices = indices[:num_keep] mask = np.ones(num_patches, dtype=bool) mask[keep_indices] = False mask_2d = mask.reshape(H // patch_size, W // patch_size) mask_pixel = np.kron(mask_2d, np.ones((patch_size, patch_size))) masked_img = image.copy() masked_img[mask_pixel == 1] = 0 return masked_img

5. Library Comparison — VERIFY¶

We verify our vectorized info_nce_loss against an explicit step-by-step 4-anchor symmetric evaluation.

No external library pins NT-Xent exactly (implementations differ in normalization and anchor averaging), so the unit-tested package implementation ml_first_principles.ssl_models.InfoNCELoss serves as the pinned reference.

In [6]:

z1 = np.array([[1.0, 0.0], [0.0, 1.0]])
z2 = np.array([[0.9, 0.435], [0.435, 0.9]])
z1 = z1 / np.linalg.norm(z1, axis=1, keepdims=True)
z2 = z2 / np.linalg.norm(z2, axis=1, keepdims=True)
tau = 1.0

loss, _ = info_nce_loss(z1, z2, temperature=tau)

# Step-by-step 4-anchor calculation
Z = np.vstack([z1, z2])
p0 = np.exp(np.dot(Z[0], Z[2])/tau) / (np.exp(np.dot(Z[0], Z[2])/tau) + np.exp(np.dot(Z[0], Z[1])/tau) + np.exp(np.dot(Z[0], Z[3])/tau))
p1 = np.exp(np.dot(Z[1], Z[3])/tau) / (np.exp(np.dot(Z[1], Z[3])/tau) + np.exp(np.dot(Z[1], Z[0])/tau) + np.exp(np.dot(Z[1], Z[2])/tau))
p2 = np.exp(np.dot(Z[2], Z[0])/tau) / (np.exp(np.dot(Z[2], Z[0])/tau) + np.exp(np.dot(Z[2], Z[3])/tau) + np.exp(np.dot(Z[2], Z[1])/tau))
p3 = np.exp(np.dot(Z[3], Z[1])/tau) / (np.exp(np.dot(Z[3], Z[1])/tau) + np.exp(np.dot(Z[3], Z[2])/tau) + np.exp(np.dot(Z[3], Z[0])/tau))
manual_loss = -0.25 * (np.log(p0) + np.log(p1) + np.log(p2) + np.log(p3))

print(f"Vectorized InfoNCE Loss: {loss:.6f}")
print(f"Manual 4-Anchor Loss:    {manual_loss:.6f}")
assert np.isclose(loss, manual_loss, atol=1e-5), "InfoNCE calculation mismatch!"
print("InfoNCE loss computation verified!")

z1 = np.array([[1.0, 0.0], [0.0, 1.0]]) z2 = np.array([[0.9, 0.435], [0.435, 0.9]]) z1 = z1 / np.linalg.norm(z1, axis=1, keepdims=True) z2 = z2 / np.linalg.norm(z2, axis=1, keepdims=True) tau = 1.0 loss, _ = info_nce_loss(z1, z2, temperature=tau) # Step-by-step 4-anchor calculation Z = np.vstack([z1, z2]) p0 = np.exp(np.dot(Z[0], Z[2])/tau) / (np.exp(np.dot(Z[0], Z[2])/tau) + np.exp(np.dot(Z[0], Z[1])/tau) + np.exp(np.dot(Z[0], Z[3])/tau)) p1 = np.exp(np.dot(Z[1], Z[3])/tau) / (np.exp(np.dot(Z[1], Z[3])/tau) + np.exp(np.dot(Z[1], Z[0])/tau) + np.exp(np.dot(Z[1], Z[2])/tau)) p2 = np.exp(np.dot(Z[2], Z[0])/tau) / (np.exp(np.dot(Z[2], Z[0])/tau) + np.exp(np.dot(Z[2], Z[3])/tau) + np.exp(np.dot(Z[2], Z[1])/tau)) p3 = np.exp(np.dot(Z[3], Z[1])/tau) / (np.exp(np.dot(Z[3], Z[1])/tau) + np.exp(np.dot(Z[3], Z[2])/tau) + np.exp(np.dot(Z[3], Z[0])/tau)) manual_loss = -0.25 * (np.log(p0) + np.log(p1) + np.log(p2) + np.log(p3)) print(f"Vectorized InfoNCE Loss: {loss:.6f}") print(f"Manual 4-Anchor Loss: {manual_loss:.6f}") assert np.isclose(loss, manual_loss, atol=1e-5), "InfoNCE calculation mismatch!" print("InfoNCE loss computation verified!")

Vectorized InfoNCE Loss: 0.816813
Manual 4-Anchor Loss:    0.816813
InfoNCE loss computation verified!

The check below pins our implementation to the package reference on seeded data, including a tiny temperature. At $\tau = 0.001$ the scaled similarities reach $|s|/\tau \approx 10^3$ and naive np.exp overflows float64 to inf (anything beyond $\approx 709$ does), so the reference evaluates the softmax denominator in log-space via logsumexp — subtract the row maximum, then exponentiate — and stays finite.

In [7]:

# Pinned-reference check on seeded embeddings: our info_nce_loss vs the package InfoNCELoss
z_a = rng.standard_normal((16, 8))
z_b = z_a + 0.1 * rng.standard_normal((16, 8))
z_a_unit = z_a / np.linalg.norm(z_a, axis=1, keepdims=True)
z_b_unit = z_b / np.linalg.norm(z_b, axis=1, keepdims=True)

for temp in [0.07, 0.5, 1.0]:
    loss_ours, _ = info_nce_loss(z_a_unit, z_b_unit, temperature=temp)
    loss_ref = InfoNCELoss(temperature=temp).forward(z_a, z_b)  # L2-normalizes internally
    print(f"tau = {temp:<5} ours = {loss_ours:.6f}  package = {loss_ref:.6f}")
    assert np.isclose(loss_ours, loss_ref, atol=1e-5), f"Mismatch against package reference at tau={temp}"

# Tiny temperature: log-space evaluation stays finite where naive exp would overflow
loss_tiny = InfoNCELoss(temperature=0.001).forward(z_a, z_b)
print(f"tau = 0.001 package = {loss_tiny:.4f}")
assert np.isfinite(loss_tiny), "Log-space InfoNCE must stay finite at tau = 0.001"
print("Notebook InfoNCE matches the pinned package reference.")

# Pinned-reference check on seeded embeddings: our info_nce_loss vs the package InfoNCELoss z_a = rng.standard_normal((16, 8)) z_b = z_a + 0.1 * rng.standard_normal((16, 8)) z_a_unit = z_a / np.linalg.norm(z_a, axis=1, keepdims=True) z_b_unit = z_b / np.linalg.norm(z_b, axis=1, keepdims=True) for temp in [0.07, 0.5, 1.0]: loss_ours, _ = info_nce_loss(z_a_unit, z_b_unit, temperature=temp) loss_ref = InfoNCELoss(temperature=temp).forward(z_a, z_b) # L2-normalizes internally print(f"tau = {temp:<5} ours = {loss_ours:.6f} package = {loss_ref:.6f}") assert np.isclose(loss_ours, loss_ref, atol=1e-5), f"Mismatch against package reference at tau={temp}" # Tiny temperature: log-space evaluation stays finite where naive exp would overflow loss_tiny = InfoNCELoss(temperature=0.001).forward(z_a, z_b) print(f"tau = 0.001 package = {loss_tiny:.4f}") assert np.isfinite(loss_tiny), "Log-space InfoNCE must stay finite at tau = 0.001" print("Notebook InfoNCE matches the pinned package reference.")

tau = 0.07  ours = 0.046231  package = 0.046231
tau = 0.5   ours = 1.832620  package = 1.832620
tau = 1.0   ours = 2.554461  package = 2.554461
tau = 0.001 package = -0.0000
Notebook InfoNCE matches the pinned package reference.

6. Experiments and Failures — VERIFY¶

In [8]:

def train_contrastive(temperature=0.1, epochs=150, lr=0.3):
    model = SimpleMLP(input_dim=2, hidden_dim=16, output_dim=2)
    losses = []
    for epoch in range(epochs):
        idx = rng.permutation(len(X))
        v1, v2 = augment(X[idx]), augment(X[idx])
        z1, z2 = model.forward(v1), model.forward(v2)
        loss, dZ = info_nce_loss(z1, z2, temperature)
        losses.append(loss)
        N = len(v1)
        dZ1, dZ2 = dZ[:N], dZ[N:]
        model.forward(v1); model.backward(dZ1, lr=lr)
        model.forward(v2); model.backward(dZ2, lr=lr)
    return model, losses

model, losses = train_contrastive(temperature=0.1, epochs=150, lr=0.3)

plt.figure(figsize=(7, 4))
plt.plot(losses, color='#1f77b4', lw=2)
plt.title("SimCLR InfoNCE Training Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.grid(True, linestyle='--', alpha=0.5)
plt.show()

def train_contrastive(temperature=0.1, epochs=150, lr=0.3): model = SimpleMLP(input_dim=2, hidden_dim=16, output_dim=2) losses = [] for epoch in range(epochs): idx = rng.permutation(len(X)) v1, v2 = augment(X[idx]), augment(X[idx]) z1, z2 = model.forward(v1), model.forward(v2) loss, dZ = info_nce_loss(z1, z2, temperature) losses.append(loss) N = len(v1) dZ1, dZ2 = dZ[:N], dZ[N:] model.forward(v1); model.backward(dZ1, lr=lr) model.forward(v2); model.backward(dZ2, lr=lr) return model, losses model, losses = train_contrastive(temperature=0.1, epochs=150, lr=0.3) plt.figure(figsize=(7, 4)) plt.plot(losses, color='#1f77b4', lw=2) plt.title("SimCLR InfoNCE Training Loss") plt.xlabel("Epoch") plt.ylabel("Loss") plt.grid(True, linestyle='--', alpha=0.5) plt.show()

In [9]:

# Visualizing learned embeddings on the unit circle
embeddings = model.forward(X)
plt.figure(figsize=(6, 6))
plt.scatter(embeddings[:, 0], embeddings[:, 1], c=y, cmap='tab10', alpha=0.7)
circle = plt.Circle((0, 0), 1, color='gray', fill=False, linestyle='--', lw=1.5)
plt.gca().add_patch(circle)
plt.title("Learned Unsupervised Representations (Unit Hypersphere)")
plt.xlabel("Embedding dimension 1")
plt.ylabel("Embedding dimension 2")
plt.xlim(-1.2, 1.2)
plt.ylim(-1.2, 1.2)
plt.grid(True, linestyle='--', alpha=0.3)
plt.show()

# Visualizing learned embeddings on the unit circle embeddings = model.forward(X) plt.figure(figsize=(6, 6)) plt.scatter(embeddings[:, 0], embeddings[:, 1], c=y, cmap='tab10', alpha=0.7) circle = plt.Circle((0, 0), 1, color='gray', fill=False, linestyle='--', lw=1.5) plt.gca().add_patch(circle) plt.title("Learned Unsupervised Representations (Unit Hypersphere)") plt.xlabel("Embedding dimension 1") plt.ylabel("Embedding dimension 2") plt.xlim(-1.2, 1.2) plt.ylim(-1.2, 1.2) plt.grid(True, linestyle='--', alpha=0.3) plt.show()

In [10]:

# Temperature Ablation Study
fig, axes = plt.subplots(1, 3, figsize=(15, 4.5))
taus = [0.01, 0.1, 1.0]
for ax, tau in zip(axes, taus):
    m, _ = train_contrastive(temperature=tau, epochs=100, lr=0.3)
    emb = m.forward(X)
    ax.scatter(emb[:, 0], emb[:, 1], c=y, cmap='tab10', alpha=0.7)
    circle = plt.Circle((0, 0), 1, color='gray', fill=False, linestyle='--')
    ax.add_patch(circle)
    ax.set_title(rf"Temperature $\tau = {tau}$")
    ax.set_xlabel("Embedding dimension 1")
    ax.set_ylabel("Embedding dimension 2")
    ax.set_xlim(-1.2, 1.2)
    ax.set_ylim(-1.2, 1.2)
    ax.grid(True, linestyle='--', alpha=0.3)
plt.suptitle("Temperature Hyperparameter Sensitivity", y=1.03, fontsize=13)
plt.tight_layout()
plt.show()

# Temperature Ablation Study fig, axes = plt.subplots(1, 3, figsize=(15, 4.5)) taus = [0.01, 0.1, 1.0] for ax, tau in zip(axes, taus): m, _ = train_contrastive(temperature=tau, epochs=100, lr=0.3) emb = m.forward(X) ax.scatter(emb[:, 0], emb[:, 1], c=y, cmap='tab10', alpha=0.7) circle = plt.Circle((0, 0), 1, color='gray', fill=False, linestyle='--') ax.add_patch(circle) ax.set_title(rf"Temperature $\tau = {tau}$") ax.set_xlabel("Embedding dimension 1") ax.set_ylabel("Embedding dimension 2") ax.set_xlim(-1.2, 1.2) ax.set_ylim(-1.2, 1.2) ax.grid(True, linestyle='--', alpha=0.3) plt.suptitle("Temperature Hyperparameter Sensitivity", y=1.03, fontsize=13) plt.tight_layout() plt.show()

In [11]:

# MAE Masking Visualization
img = np.indices((32, 32)).sum(axis=0) % 2 * 255.0
fig, axes = plt.subplots(1, 3, figsize=(12, 4))
axes[0].imshow(img, cmap='gray'); axes[0].set_title("Original Image")
axes[1].imshow(mask_image_patches(img, mask_ratio=0.50), cmap='gray'); axes[1].set_title("Masked (50%)")
axes[2].imshow(mask_image_patches(img, mask_ratio=0.75), cmap='gray'); axes[2].set_title("Masked (75% - MAE default)")
for ax in axes:
    ax.set_xlabel("Pixel column")
    ax.set_ylabel("Pixel row")
plt.tight_layout()
plt.show()

# MAE Masking Visualization img = np.indices((32, 32)).sum(axis=0) % 2 * 255.0 fig, axes = plt.subplots(1, 3, figsize=(12, 4)) axes[0].imshow(img, cmap='gray'); axes[0].set_title("Original Image") axes[1].imshow(mask_image_patches(img, mask_ratio=0.50), cmap='gray'); axes[1].set_title("Masked (50%)") axes[2].imshow(mask_image_patches(img, mask_ratio=0.75), cmap='gray'); axes[2].set_title("Masked (75% - MAE default)") for ax in axes: ax.set_xlabel("Pixel column") ax.set_ylabel("Pixel row") plt.tight_layout() plt.show()

Failure Case: Representation Collapse¶

The canonical SSL failure: the encoder maps every input to the same embedding vector. All $2N \times 2N$ similarities then become identical, each softmax row is uniform over its $2N - 1$ candidates, and the loss sits at exactly $\log(2N - 1)$ for any temperature — positives are indistinguishable from negatives.

This plateau is why contrastive methods need negatives: the InfoNCE denominator turns the collapsed configuration into a high-loss point rather than a minimum, but only because other samples are there to contrast against. Negative-free methods (BYOL, SimSiam) must block the same degenerate optimum through architectural asymmetry instead — predictor head, stop-gradient, EMA target. See theory.md §7 for the collapse analysis.

In [12]:

# Representation collapse: every input maps to the SAME unit embedding vector
N = 256
z_collapsed = np.tile(np.array([[1.0, 0.0]]), (N, 1))

loss_collapsed, _ = info_nce_loss(z_collapsed, z_collapsed, temperature=0.1)
expected = np.log(2 * N - 1)

print(f"InfoNCE at collapse: {loss_collapsed:.6f}")
print(f"log(2N - 1):         {expected:.6f}")
assert np.isclose(loss_collapsed, np.log(2 * N - 1), atol=1e-4), "Collapse must pin the loss at log(2N - 1)"
print("Collapsed representations sit exactly on the uniform-similarity plateau log(2N - 1).")

# Representation collapse: every input maps to the SAME unit embedding vector N = 256 z_collapsed = np.tile(np.array([[1.0, 0.0]]), (N, 1)) loss_collapsed, _ = info_nce_loss(z_collapsed, z_collapsed, temperature=0.1) expected = np.log(2 * N - 1) print(f"InfoNCE at collapse: {loss_collapsed:.6f}") print(f"log(2N - 1): {expected:.6f}") assert np.isclose(loss_collapsed, np.log(2 * N - 1), atol=1e-4), "Collapse must pin the loss at log(2N - 1)" print("Collapsed representations sit exactly on the uniform-similarity plateau log(2N - 1).")

InfoNCE at collapse: 6.236364
log(2N - 1):         6.236370
Collapsed representations sit exactly on the uniform-similarity plateau log(2N - 1).

7. Connections & Takeaways¶

1. Contrastive Learning: InfoNCE maximizes a lower bound on mutual information $I(X; Y) \ge \log(K) - \mathcal{L}_{InfoNCE}$.
2. Projector Head: Keeps the representation $h$ rich while allowing the loss on $z$ to discard augmentation-variant features.
3. Masked Image Modeling: MAE forces ViTs to learn semantic visual priors by reconstructing 75% masked pixels.