Self-Supervised Learning: Exercises¶
Test your understanding of self-supervised learning principles, InfoNCE loss, masked autoencoders, and paradigm tradeoffs.
Prerequisites. Read theory.md and work through first_principles.ipynb before attempting these.
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, PatchMasking
Exercise 1: Hand-Computing InfoNCE¶
Given a batch size $N=2$ (so $2N=4$ augmented views total). Let the embeddings for the views be $z_1, z_2$ for sample A and $z_3, z_4$ for sample B. Suppose the cosine similarity matrix $S$ (shape 4x4) is:
$$S = \begin{bmatrix} 1.0 & 0.8 & -0.1 & -0.2 \\ 0.8 & 1.0 & -0.2 & 0.1 \\ -0.1 & -0.2 & 1.0 & 0.9 \\ -0.2 & 0.1 & 0.9 & 1.0 \end{bmatrix}$$
For $\tau = 0.5$, compute the InfoNCE loss for the anchor $z_1$. Ignore self-similarity.
Hint: Use the formula $\ell = -\log \frac{\exp(sim(z_1, z_2)/\tau)}{\exp(sim(z_1, z_2)/\tau) + \exp(sim(z_1, z_3)/\tau) + \exp(sim(z_1, z_4)/\tau)}$.
# Exercise 1 inputs — derive the anchor-z1 loss by hand before opening the solution below.
S = np.array([
[1.0, 0.8, -0.1, -0.2],
[0.8, 1.0, -0.2, 0.1],
[-0.1, -0.2, 1.0, 0.9],
[-0.2, 0.1, 0.9, 1.0],
])
tau = 0.5
Solution 1¶
For anchor $z_1$ the positive is $z_2$ and the negatives are $z_3, z_4$ (self-similarity $S_{11}$ is excluded).
- Scale the relevant similarities by $\tau = 0.5$ (division rule): $S_{12}/\tau = 0.8/0.5 = 1.6$, $\quad S_{13}/\tau = -0.1/0.5 = -0.2$, $\quad S_{14}/\tau = -0.2/0.5 = -0.4$.
- Exponentiate each term: $e^{1.6} \approx 4.9530$, $\quad e^{-0.2} \approx 0.8187$, $\quad e^{-0.4} \approx 0.6703$.
- Softmax probability of the positive (sum the denominator first): $p = \dfrac{4.9530}{4.9530 + 0.8187 + 0.6703} = \dfrac{4.9530}{6.4421} \approx 0.7689$.
- Negative log-likelihood: $\ell_1 = -\log p \approx -\log(0.7689) \approx 0.2629$.
Result: $\ell_1 = -\log \dfrac{e^{1.6}}{e^{1.6} + e^{-0.2} + e^{-0.4}} \approx 0.2629$
# Deterministic verification of the hand derivation
exp_pos = np.exp(S[0, 1] / tau)
exp_negs = np.exp(S[0, 2] / tau) + np.exp(S[0, 3] / tau)
loss_z1 = -np.log(exp_pos / (exp_pos + exp_negs))
print(f"Computed loss for anchor z1: {loss_z1:.4f}")
assert np.isclose(loss_z1, 0.2629, atol=1e-4), "Hand-derived value does not match"
Computed loss for anchor z1: 0.2629
Exercise 2: Implementing NT-Xent Symmetry¶
The InfoNCE loss in SimCLR is computed symmetrically for both views (i.e., using view 1 as anchor, then view 2 as anchor).
Implement a function symmetric_nt_xent that takes the similarity matrix directly and computes the total average loss over all $2N$ anchors. Verify that the result matches expectations.
def symmetric_nt_xent(sim_matrix, tau=1.0):
"""
sim_matrix: (2N, 2N) pairwise cosine similarity matrix
tau: temperature
Assumes the block layout: the first N rows are view 1, the next N rows
are view 2, so the positive partner of anchor i is (i + N) mod 2N.
"""
_2N = sim_matrix.shape[0]
N = _2N // 2
# Scale by tau
sim = sim_matrix / tau
# Mask self-similarity
np.fill_diagonal(sim, -np.inf)
# Construct labels (positive pairs)
labels = np.zeros(_2N, dtype=int)
labels[:N] = np.arange(N, _2N)
labels[N:] = np.arange(N)
# Softmax
exp_sim = np.exp(sim)
probs = exp_sim / np.sum(exp_sim, axis=1, keepdims=True)
# Gather positive probabilities
pos_probs = probs[np.arange(_2N), labels]
return -np.mean(np.log(pos_probs))
# Exercise 1 lists the embeddings as (z1, z2, z3, z4) = (A-view1, A-view2, B-view1, B-view2),
# while symmetric_nt_xent expects the block layout (A-view1, B-view1 | A-view2, B-view2).
# Permute rows and columns of S into that layout with the order (z1, z3, z2, z4).
order = np.array([0, 2, 1, 3])
S_blocks = S[np.ix_(order, order)]
loss = symmetric_nt_xent(S_blocks, tau=tau)
print(f"Symmetric NT-Xent loss: {loss:.4f}")
# Deterministic check: the mean of the four per-anchor losses (anchor z1 alone gave 0.2629)
assert np.isclose(loss, 0.2696, atol=1e-4), "Symmetric loss does not match the hand value"
# Cross-check against the unit-tested package reference on seeded embeddings,
# where the setup matches: the package takes the two views directly and computes
# the same symmetric NT-Xent (L2-normalizing internally).
v1 = rng.standard_normal((6, 3))
v2 = rng.standard_normal((6, 3))
v1_unit = v1 / np.linalg.norm(v1, axis=1, keepdims=True)
v2_unit = v2 / np.linalg.norm(v2, axis=1, keepdims=True)
Z = np.vstack([v1_unit, v2_unit])
loss_matrix = symmetric_nt_xent(Z @ Z.T, tau=0.5)
loss_ref = InfoNCELoss(temperature=0.5).forward(v1, v2)
print(f"symmetric_nt_xent: {loss_matrix:.6f} | package InfoNCELoss: {loss_ref:.6f}")
assert np.isclose(loss_matrix, loss_ref, atol=1e-8), "Does not match the package reference"
Symmetric NT-Xent loss: 0.2696 symmetric_nt_xent: 2.987693 | package InfoNCELoss: 2.987693
Exercise 3: Temperature and Hard-Negative Weighting (Hand Calculation)¶
theory.md (Section 3.3) shows that the InfoNCE gradient with respect to a negative similarity is $\frac{\partial \ell_i}{\partial (z_i^T z_k)} = \frac{1}{\tau} P(k \mid i)$, so the softmax mass $P(k \mid i)$ decides which negatives drive learning.
An anchor has positive similarity $s^{+} = 0.8$ and negative similarities $s^{-} = (0.6, \; 0.1, \; -0.2)$. For $\tau = 1.0$ and $\tau = 0.1$, compute by hand:
- The softmax distribution $P(\cdot \mid i)$ over the four candidates (one positive, three negatives).
- The hardest-negative share: the fraction of the total negative probability mass carried by the hardest negative ($s = 0.6$).
- The per-anchor loss $\ell = -\log P(\text{pos} \mid i)$.
What do the two temperatures imply about which negatives dominate the gradient?
Expected results (4 dp):
| $\tau$ | $P(\text{pos} \mid i)$ | hardest-negative share | $\ell$ |
|---|---|---|---|
| 1.0 | 0.3727 | 0.4864 | 0.9870 |
| 0.1 | 0.8801 | 0.9930 | 0.1278 |
Solution 3¶
Case $\tau = 1.0$. Scale (division rule), then exponentiate each candidate:
$e^{0.8} \approx 2.2255$, $\quad e^{0.6} \approx 1.8221$, $\quad e^{0.1} \approx 1.1052$, $\quad e^{-0.2} \approx 0.8187$.
Denominator (sum of all four): $2.2255 + 1.8221 + 1.1052 + 0.8187 = 5.9715$. Softmax rule:
$$P(\text{pos} \mid i) = \frac{2.2255}{5.9715} \approx 0.3727, \qquad P(\text{neg} \mid i) \approx (0.3051, \; 0.1851, \; 0.1371).$$
Hardest-negative share $= \dfrac{0.3051}{0.3051 + 0.1851 + 0.1371} \approx 0.4864$; loss $\ell = -\log 0.3727 \approx 0.9870$ (negative log-likelihood).
Case $\tau = 0.1$. Scaled logits are $(8, 6, 1, -2)$:
$e^{8} \approx 2980.96$, $\quad e^{6} \approx 403.43$, $\quad e^{1} \approx 2.7183$, $\quad e^{-2} \approx 0.1353$; denominator $\approx 3387.24$.
$$P(\text{pos} \mid i) = \frac{2980.96}{3387.24} \approx 0.8801, \qquad \text{hardest share} = \frac{403.43}{403.43 + 2.7183 + 0.1353} \approx 0.9930,$$
and $\ell = -\log 0.8801 \approx 0.1278$.
Interpretation. At $\tau = 1.0$ the three negatives share the gradient roughly evenly
(49% / 30% / 22%). At $\tau = 0.1$ the hardest negative absorbs 99.3% of the negative mass:
small temperatures turn the softmax into a near-max, focusing the update on the single
closest negative — powerful for tight separation, but brittle when that "negative" is a
false negative from the same class.
Result: hardest-negative share $\approx 0.4864$ at $\tau = 1.0$ vs $\approx 0.9930$ at $\tau = 0.1$.
# Deterministic verification of the hand table (stable log-space softmax)
s_pos, s_negs = 0.8, np.array([0.6, 0.1, -0.2])
hand = {1.0: (0.3727, 0.4864, 0.9870), 0.1: (0.8801, 0.9930, 0.1278)}
for tau_x, (p_pos_hand, share_hand, loss_hand) in hand.items():
logits = np.concatenate([[s_pos], s_negs]) / tau_x
logits -= logits.max() # max-subtraction: exp cannot overflow
probs = np.exp(logits) / np.sum(np.exp(logits))
p_pos, p_negs = probs[0], probs[1:]
share = p_negs[0] / p_negs.sum()
loss = -np.log(p_pos)
print(f"tau={tau_x:>3}: P(pos)={p_pos:.4f} hardest share={share:.4f} loss={loss:.4f}")
assert np.allclose([p_pos, share, loss], [p_pos_hand, share_hand, loss_hand], atol=1e-4), (
"Hand-derived temperature table does not match"
)
tau=1.0: P(pos)=0.3727 hardest share=0.4864 loss=0.9870 tau=0.1: P(pos)=0.8801 hardest share=0.9930 loss=0.1278
Exercise 4: Implementing the MAE Un-Shuffle¶
PatchMasking.mask_patches (the package MAE masker, also built in
first_principles.ipynb) returns
sequence_kept— the visible patches, in shuffled order, shape(batch, len_keep, dim);mask—1at masked positions,0at visible ones, in original order;ids_restore— for each original position, its index in the shuffled order.
An MAE decoder must invert this bookkeeping: rebuild the full-length sequence with a learned mask token at masked positions before adding positional embeddings.
Task. Implement
def restore_sequence(sequence_kept, mask, ids_restore, mask_token):
... # -> (batch, seq_len, dim)
such that visible positions carry their original patch values and masked positions carry
mask_token. Verify on a (2, 8, 3) batch with mask_ratio=0.75 that:
mask.mean()is exactly0.75(withseq_len = 8,len_keep = 2);- the restored sequence equals the input at every visible position (
np.allclose,atol=1e-12); - every masked position equals
mask_token.
Hint. The shuffled order is full_shuffled = concat([sequence_kept, mask_tokens]);
np.take_along_axis(full_shuffled, ids_restore[:, :, None], axis=1) maps it back.
Solution 4¶
mask_patches shuffles positions with ids_shuffle = argsort(noise) and keeps the first
len_keep entries, so shuffled slot $j$ holds original position ids_shuffle[j]. Because
ids_restore = argsort(ids_shuffle) is the inverse permutation, original position $i$ lives at
shuffled slot ids_restore[i]. Appending mask tokens recreates the full shuffled sequence
(slots len_keep, ..., seq_len - 1 are exactly the masked ones), and gathering along axis 1
with ids_restore applies the inverse permutation — no loops needed.
def restore_sequence(sequence_kept, mask, ids_restore, mask_token):
"""Invert PatchMasking: full (batch, seq_len, dim) with mask_token at masked slots."""
batch_size, seq_len = ids_restore.shape
dim = sequence_kept.shape[2]
n_masked = seq_len - sequence_kept.shape[1]
mask_tokens = np.broadcast_to(mask_token, (batch_size, n_masked, dim))
full_shuffled = np.concatenate([sequence_kept, mask_tokens], axis=1)
return np.take_along_axis(full_shuffled, ids_restore[:, :, None], axis=1)
patches = rng.standard_normal((2, 8, 3))
masker = PatchMasking(mask_ratio=0.75, random_state=SEED)
kept, mask, ids_restore = masker.mask_patches(patches)
mask_token = np.full(3, -1.0) # sentinel clearly outside the data range
restored = restore_sequence(kept, mask, ids_restore, mask_token)
visible = mask == 0
print(f"kept shape: {kept.shape} | mask ratio: {mask.mean():.2f}")
print(f"visible patches recovered exactly: {np.allclose(restored[visible], patches[visible], atol=1e-12)}")
assert restored.shape == patches.shape, "Restored sequence has the wrong shape"
assert np.isclose(mask.mean(), 0.75, atol=1e-12), "Mask ratio should be exactly 0.75"
assert np.allclose(restored[visible], patches[visible], atol=1e-12), "Visible patches corrupted"
assert np.allclose(restored[~visible], mask_token, atol=1e-12), "Masked slots must hold mask_token"
kept shape: (2, 2, 3) | mask ratio: 0.75 visible patches recovered exactly: True
Exercise 5: Representation Collapse (Failure Analysis)¶
Suppose the encoder collapses: it maps every input to the same unit vector, $f(x) = c$ with $\Vert c \Vert = 1$, for all $x$.
- Show that the InfoNCE loss over a batch of $N$ positive pairs ($2N$ views) equals $\log(2N - 1)$, independent of the temperature $\tau$.
- Now delete the negatives, keeping only the positive alignment term $\ell = -s^{+}/\tau$ with cosine similarity $s^{+}$. Show that collapse attains the global minimum of this loss.
- Why does gradient descent not escape the collapsed state under InfoNCE?
Solution 5¶
1. Loss at collapse. With every embedding equal to $c$, every pairwise cosine similarity is $c^T c = 1$. For any anchor, all $2N - 1$ non-self logits equal $1/\tau$, so the softmax is uniform over the candidates:
$$P(\text{pos} \mid i) = \frac{e^{1/\tau}}{(2N - 1)\, e^{1/\tau}} = \frac{1}{2N - 1} \qquad \Rightarrow \qquad \ell = -\log \frac{1}{2N - 1} = \log(2N - 1).$$
The $e^{1/\tau}$ factors cancel between numerator and denominator, so $\tau$ drops out. For $N = 6$: $\ell = \log 11 \approx 2.3979$ — the chance-level loss of a $(2N-1)$-way classifier.
2. Positive-only loss. Cosine similarity satisfies $s^{+} \le 1$, so $\ell = -s^{+}/\tau \ge -1/\tau$, with equality iff each pair of views coincides exactly — which collapse achieves for every pair simultaneously. Collapse is therefore a global minimizer: the denominator's negatives are the only part of InfoNCE that makes collapse costly (it charges $\log(2N-1)$ instead of rewarding it). This is why BYOL, which removes negatives, needs the extra asymmetry machinery of Exercise 6 to avoid the same trivial solution.
3. Plateau. The collapsed configuration is perfectly symmetric: every anchor sees identical positive and negative similarities, so the attraction toward the positive and the repulsion from the negatives cancel exactly and the gradient through the similarities vanishes. Training sits on a plateau with the loss pinned at $\log(2N - 1)$ — the flat-loss signature demonstrated empirically in the collapse experiment of first_principles.ipynb.
Result: $\ell_{\text{collapse}} = \log(2N - 1)$ for any $\tau$; without negatives, collapse is the global minimum $-1/\tau$.
# Numerical check: collapsed embeddings hit exactly log(2N - 1), for any temperature
N, dim = 6, 4
c = rng.standard_normal(dim)
c /= np.linalg.norm(c)
z_collapsed = np.tile(c, (N, 1)) # every view is the same unit vector
for tau_x in (0.07, 0.5, 2.0):
loss_c = InfoNCELoss(temperature=tau_x).forward(z_collapsed, z_collapsed)
print(f"tau={tau_x:>4}: collapsed InfoNCE = {loss_c:.6f} (log(2N-1) = {np.log(2 * N - 1):.6f})")
assert np.isclose(loss_c, np.log(2 * N - 1), atol=1e-8), "Collapse loss must equal log(2N-1)"
# Without negatives, collapse is the global optimum of the alignment-only loss
tau_x = 0.5
pos_only_collapsed = -np.sum(z_collapsed * z_collapsed, axis=1).mean() / tau_x
print(f"positive-only loss at collapse: {pos_only_collapsed:.4f} (lower bound -1/tau = {-1 / tau_x:.4f})")
assert np.isclose(pos_only_collapsed, -1 / tau_x, atol=1e-12), "Collapse must attain -1/tau"
tau=0.07: collapsed InfoNCE = 2.397895 (log(2N-1) = 2.397895) tau= 0.5: collapsed InfoNCE = 2.397895 (log(2N-1) = 2.397895) tau= 2.0: collapsed InfoNCE = 2.397895 (log(2N-1) = 2.397895) positive-only loss at collapse: -2.0000 (lower bound -1/tau = -2.0000)
Exercise 6: Conceptual Analysis of SSL Paradigms¶
Question: Compare Contrastive Learning (e.g., SimCLR), Non-Contrastive Learning (e.g., BYOL), and Masked Image Modeling (e.g., MAE) across the following dimensions:
- Dependence on Data Augmentation
- Need for Negative Samples
- Semantic Level of Representations
Answer:
- Dependence on Data Augmentation:
- SimCLR/BYOL: Highly dependent. They rely on augmentations to define semantic invariance. Poor augmentations (e.g., missing color jitter) lead to degenerate representations.
- MAE: Low dependence. Masking serves as the primary data corruption; complex structural augmentations are less critical.
- Need for Negative Samples:
- SimCLR: Yes. Requires massive batch sizes or momentum queues to provide sufficient negatives to prevent collapse.
- BYOL/MAE: No. BYOL uses asymmetric architectures (EMA) to prevent collapse, while MAE prevents collapse via the generative pixel-reconstruction objective.
- Semantic Level of Representations:
- SimCLR/BYOL: Captures highly abstracted, global semantic features (great for classification) but often loses fine-grained spatial information.
- MAE: Captures dense, localized features useful for pixel-level tasks (e.g., segmentation) as well as global structures, bridging the gap between localized reconstruction and global understanding.