"""Original reduced dual-encoder CLIP objective lesson. One width-eight, two-head pre-LN block per encoder. Images are 8x8 RGB, patches 4x4. Text has four tokens; the final position stands for EOS. Vocabulary 16. Both encoders project into four normalized embedding features. """ import math import torch import torch.nn as nn import torch.nn.functional as F class EncoderBlock(nn.Module): def __init__(self, causal=False): super().__init__() self.causal = causal self.norm = nn.LayerNorm(8) self.query = nn.Linear(8, 8) self.key = nn.Linear(8, 8) self.value = nn.Linear(8, 8) self.softmax = nn.Softmax(dim=-1) self.project = nn.Linear(8, 8) self.ff_norm = nn.LayerNorm(8) self.up = nn.Linear(8, 32) self.gelu = nn.GELU() self.down = nn.Linear(32, 8) def forward(self, x): b, t, _ = x.shape h = self.norm(x) q = self.query(h).reshape(b, t, 2, 4).transpose(1, 2) k = self.key(h).reshape(b, t, 2, 4).transpose(1, 2) v = self.value(h).reshape(b, t, 2, 4).transpose(1, 2) scores = q @ k.transpose(-2, -1) / 2 if self.causal: scores = scores.masked_fill( torch.ones(t, t, device=x.device).triu(1).bool(), float("-inf") ) mixed = (self.softmax(scores) @ v).transpose(1, 2).reshape(b, t, 8) x = x + self.project(mixed) return x + self.down(self.gelu(self.up(self.ff_norm(x)))) class ImageEncoder(nn.Module): def __init__(self): super().__init__() self.patches = nn.Conv2d(3, 8, 4, stride=4, bias=False) self.class_token = nn.Parameter(torch.randn(1, 1, 8) * 0.02) self.positions = nn.Parameter(torch.randn(1, 5, 8) * 0.02) self.pre_norm = nn.LayerNorm(8) self.block = EncoderBlock() self.post_norm = nn.LayerNorm(8) self.project = nn.Linear(8, 4, bias=False) def forward(self, images): patches = self.patches(images).flatten(2).transpose(1, 2) x = torch.cat([self.class_token.expand(images.shape[0], -1, -1), patches], 1) x = self.block(self.pre_norm(x + self.positions)) return F.normalize(self.project(self.post_norm(x[:, 0])), dim=-1) class TextEncoder(nn.Module): def __init__(self): super().__init__() self.tokens = nn.Embedding(16, 8) self.positions = nn.Parameter(torch.randn(1, 4, 8) * 0.02) self.block = EncoderBlock(causal=True) self.norm = nn.LayerNorm(8) self.project = nn.Linear(8, 4, bias=False) def forward(self, tokens): x = self.block(self.tokens(tokens) + self.positions) return F.normalize(self.project(self.norm(x[:, -1])), dim=-1) class PairSimilarity(nn.Module): def __init__(self): super().__init__() self.log_scale = nn.Parameter(torch.tensor(math.log(1 / 0.07))) def forward(self, image, text): return self.log_scale.exp() * image @ text.T class TinyCLIP(nn.Module): def __init__(self): super().__init__() self.image_encoder = ImageEncoder() self.text_encoder = TextEncoder() self.similarity = PairSimilarity() def forward(self, images, tokens): image = self.image_encoder(images) text = self.text_encoder(tokens) return self.similarity(image, text) def objective(logits): targets = torch.arange(logits.shape[0], device=logits.device) return (F.cross_entropy(logits, targets) + F.cross_entropy(logits.T, targets)) / 2 def setup(): torch.manual_seed(31) model = TinyCLIP().double() images = torch.randn(3, 3, 8, 8, dtype=torch.float64) tokens = torch.tensor([[1, 3, 4, 15], [1, 5, 6, 15], [1, 7, 8, 15]]) return model, images, tokens def verify(): model, images, tokens = setup() image, text = model.image_encoder(images), model.text_encoder(tokens) torch.testing.assert_close(image.norm(dim=-1), torch.ones(3, dtype=torch.float64)) torch.testing.assert_close(text.norm(dim=-1), torch.ones(3, dtype=torch.float64)) logits = model(images, tokens) perm = torch.tensor([2, 0, 1]) torch.testing.assert_close( objective(model(images[perm], tokens[perm])), objective(logits) ) torch.testing.assert_close(model(images, tokens[perm]), logits[:, perm]) # The first causal text state cannot consult a changed final token. x = model.text_encoder.tokens(tokens) + model.text_encoder.positions changed = x.clone() changed[:, -1] += torch.arange(8, dtype=x.dtype) torch.testing.assert_close( model.text_encoder.block(x)[:, 0], model.text_encoder.block(changed)[:, 0] ) manual = ( -(logits.log_softmax(-1).diag().mean() + logits.log_softmax(0).diag().mean()) / 2 ) torch.testing.assert_close(objective(logits), manual) optimizer = torch.optim.SGD(model.parameters(), lr=0.0001) before = objective(logits).item() optimizer.zero_grad() objective(logits).backward() assert model.image_encoder.patches.weight.grad.norm() > 0 assert model.text_encoder.tokens.weight.grad.norm() > 0 assert model.similarity.log_scale.grad.abs() > 0 optimizer.step() assert objective(model(images, tokens)).item() < before return [ "Both projected modalities have unit norm; logits are scaled pairwise cosines", "Joint pair permutations preserve the objective; text-only permutations move similarity columns", "The text attention is causal and the symmetric loss equals both diagonal log-probabilities", "Both encoders and the learned log-temperature receive gradients; a small SGD step lowers this batch loss", ] def experiment(): model, images, tokens = setup() cases = [] with torch.no_grad(): cosines = model.image_encoder(images) @ model.text_encoder(tokens).T for identity, label, scale, perm in [ ("paired", "Original three pairs", 1 / 0.07, [0, 1, 2]), ("shuffled", "Reverse only the text batch", 1 / 0.07, [2, 1, 0]), ("warm", "Higher temperature · 0.5", 2, [0, 1, 2]), ("sharp", "Lower temperature · 0.03", 1 / 0.03, [0, 1, 2]), ]: cosine = cosines[:, perm] logits = cosine * scale pi, pt = logits.softmax(-1), logits.T.softmax(-1) cases.append( dict( id=identity, label=label, target="similarity", note=( "Rows are images; columns are texts in the displayed order. Diagonal targets encode the supplied pairing, not model knowledge. " + ( "Only text order changes; diagonal targets now ask for different pairings. " if identity == "shuffled" else "" ) + "Temperature rescales the same cosines before both softmax directions. These random encoders have no learned image–text alignment." ), matrices=[ dict( label="Cosine similarity · images × texts", values=cosine.tolist(), ), dict(label="Image → text probabilities", values=pi.tolist()), dict(label="Text → image probabilities", values=pt.tolist()), ], vectors=[ dict(label="Text column order · original index", values=perm), dict(label="Target column per image", values=[0, 1, 2]), ], metrics=[ dict(label="Temperature", value=1 / scale), dict( label="Symmetric contrastive loss", value=objective(logits).item(), ), dict( label="Mean image-row entropy · nats", value=(-(pi * pi.log()).sum(-1).mean()).item(), ), ], ) ) return dict( kind="matrices", title="Pair the batch in both directions.", description="Recorded similarities for three synthetic image/token pairs. Compare pairing order and temperature using the same untrained encoder weights.", controlLabel="Contrastive example", cases=cases, )