Training two encoders into one space
Each image and caption becomes a point on a circle. Training pulls matching pairs together and pushes the rest apart.
Image embedding (inner ring)Caption embedding (outer ring), labeled by initials
Training steps0
Loss (symmetric InfoNCE)3.182
Images matched to their caption2 of 9
Challenges
0 of 3 complete
- Train until every image is closest to its own caption. (not yet)
- Train at temperature 1 until all nine pairs match. How low does the loss go? (not yet)
- Get the loss below 0.1. (not yet)
Try this
- Press Train and watch the diagonal of the similarity matrix light up.
- Tap an image before and after training to see which caption it is closest to.
- Reset, raise the temperature to 1, and train again. Compare where the loss settles.
The contrastive loss in codePython with NumPy
import numpy as np
def unit(v):
return v / np.linalg.norm(v, axis=1, keepdims=True)
def infonce(sim, tau):
# Symmetric InfoNCE: each image must pick its own caption from its row,
# and each caption its own image from its column.
logits = sim / tau
rows = logits - np.log(np.exp(logits).sum(axis=1, keepdims=True))
cols = logits - np.log(np.exp(logits).sum(axis=0, keepdims=True))
return -(np.trace(rows) + np.trace(cols)) / (2 * len(sim))
# Three images and their three captions, embedded in 2 dimensions.
images = unit(np.array([[1.0, 0.2], [-0.3, 1.0], [-0.7, -0.8]]))
captions = unit(np.array([[0.9, -0.1], [0.1, 1.0], [-0.9, -0.5]]))
sim = images @ captions.T # cosine similarity: image i (row), caption j (column)
for row in sim:
print(" ".join(f"{v:5.2f}" for v in row))
for tau in (1.0, 0.3):
print(f"loss at temperature {tau}: {infonce(sim, tau):.3f}")
# Nine pairs, each image exactly on its caption, spread evenly around the
# circle. This is about where the demo's training settles.
angles = 2 * np.pi * np.arange(9) / 9
pairs = np.stack([np.cos(angles), np.sin(angles)], axis=1)
for tau in (1.0, 0.3, 0.05):
print(f"nine perfect pairs at temperature {tau}: loss {infonce(pairs @ pairs.T, tau):.3f}")0.95 0.29 -0.95
-0.39 0.92 -0.21
-0.57 -0.81 0.94
loss at temperature 1.0: 0.435
loss at temperature 0.3: 0.051
nine perfect pairs at temperature 1.0: loss 1.433
nine perfect pairs at temperature 0.3: loss 0.723
nine perfect pairs at temperature 0.05: loss 0.018The features and encoders are deliberately tiny; the loss and the training are the real thing. The lesson explains what the full-size model does differently.