"""Original reduced Vision Transformer: patches, class token and learned positions. 16x16 RGB image, 4x4 patches, width 16, two attention heads, one pre-LN encoder. Untrained teaching architecture with a linear three-class head. """ import math import torch import torch.nn as nn import torch.nn.functional as F class PatchEmbedding(nn.Module): def __init__(self): super().__init__() self.project = nn.Conv2d(3, 16, kernel_size=4, stride=4) def forward(self, image): return self.project(image).flatten(2).transpose(1, 2) class Attention(nn.Module): def __init__(self): super().__init__() self.query = nn.Linear(16, 16) self.key = nn.Linear(16, 16) self.value = nn.Linear(16, 16) self.softmax = nn.Softmax(dim=-1) self.project = nn.Linear(16, 16) def forward(self, x): b, t, _ = x.shape q = self.query(x).reshape(b, t, 2, 8).transpose(1, 2) k = self.key(x).reshape(b, t, 2, 8).transpose(1, 2) v = self.value(x).reshape(b, t, 2, 8).transpose(1, 2) scores = q @ k.transpose(-2, -1) / math.sqrt(8) weights = self.softmax(scores) return self.project((weights @ v).transpose(1, 2).reshape(b, t, 16)) class Encoder(nn.Module): def __init__(self): super().__init__() self.norm1 = nn.LayerNorm(16) self.attention = Attention() self.norm2 = nn.LayerNorm(16) self.up = nn.Linear(16, 32) self.gelu = nn.GELU() self.down = nn.Linear(32, 16) def forward(self, x): hidden = x + self.attention(self.norm1(x)) return hidden + self.down(self.gelu(self.up(self.norm2(hidden)))) # tensorviz: input=1,3,16,16 class TinyViT(nn.Module): def __init__(self): super().__init__() self.patches = PatchEmbedding() self.class_token = nn.Parameter(torch.randn(1, 1, 16) * 0.02) self.positions = nn.Parameter(torch.randn(1, 17, 16) * 0.02) self.encoder = Encoder() self.norm = nn.LayerNorm(16) self.head = nn.Linear(16, 3) def forward(self, image): patches = self.patches(image) cls = self.class_token.expand(image.shape[0], -1, -1) sequence = torch.cat([cls, patches], dim=1) + self.positions encoded = self.encoder(sequence) return self.head(self.norm(encoded)[:, 0]) def classify(model, sequence): return model.head(model.norm(model.encoder(sequence))[:, 0]) def verify(): torch.manual_seed(13) model = TinyViT().double().eval() x = torch.randn(1, 3, 16, 16, dtype=torch.float64) with torch.no_grad(): patches = model.patches(x) flattened = F.unfold(x, kernel_size=4, stride=4).transpose(1, 2) independent = ( flattened @ model.patches.project.weight.flatten(1).T + model.patches.project.bias ) torch.testing.assert_close(patches, independent) assert patches.shape == (1, 16, 16) and model(x).shape == (1, 3) sequence = torch.cat([model.class_token, patches], dim=1) permutation = torch.tensor([0] + list(range(16, 0, -1))) torch.testing.assert_close( classify(model, sequence), classify(model, sequence[:, permutation]), rtol=1e-10, atol=1e-10, ) original = classify(model, sequence + model.positions) permuted = classify(model, sequence[:, permutation] + model.positions) assert not torch.allclose(original, permuted) changed = x.clone() changed[..., 0:4, 0:4] += 3 changed_patches = model.patches(changed) torch.testing.assert_close(patches[:, 1:], changed_patches[:, 1:]) assert not torch.allclose(model(x), model(changed)) return [ "Strided patch convolution equals independent flatten-and-linear projection", "Sixteen patch embeddings plus one class token form a seventeen-token sequence", "Without positions, patch permutation preserves the class output; fixed learned positions break this symmetry", "Changing one patch changes only its embedding, but global attention can change the class output", ] def experiment(): torch.manual_seed(13) model = TinyViT().double().eval() cases = [] with torch.no_grad(): for identity, label, row, column in [ ("top-left", "Bright patch at top left", 0, 0), ("middle", "Bright patch near the middle", 1, 1), ("bottom-right", "Bright patch at bottom right", 3, 3), ]: image = torch.zeros(1, 3, 16, 16, dtype=torch.float64) image[..., row * 4 : (row + 1) * 4, column * 4 : (column + 1) * 4] = 1 patches = model.patches(image) sequence = torch.cat([model.class_token, patches], dim=1) + model.positions normalized = model.encoder.norm1(sequence) attention = model.encoder.attention q = attention.query(normalized).reshape(1, 17, 2, 8).transpose(1, 2) k = attention.key(normalized).reshape(1, 17, 2, 8).transpose(1, 2) weights = (q @ k.transpose(-2, -1) / math.sqrt(8)).softmax(-1) cases.append( { "id": identity, "label": label, "target": "patches", "note": "Move a bright 4×4 patch in a synthetic 16×16 image. The attention map shows head 0's class-token weights to the sixteen patch tokens, excluding its self-weight. These are untrained weights and are not a saliency explanation or object detector.", "matrices": [ { "label": "Input image · red channel, exact pixels", "values": image[0, 0].tolist(), }, { "label": "Class-to-patch attention · head 0", "values": weights[0, 0, 0, 1:].reshape(4, 4).tolist(), }, ], "vectors": [ { "label": "Bright patch embedding · 16 features", "values": patches[0, row * 4 + column].tolist(), }, { "label": "Three untrained class logits", "values": model(image)[0].tolist(), }, ], "metrics": [ {"label": "Patch tokens", "value": 16}, {"label": "Tokens including class token", "value": 17}, { "label": "Class-token self-attention weight", "value": weights[0, 0, 0, 0].item(), }, { "label": "Total class attention row weight", "value": weights[0, 0, 0].sum().item(), }, ], } ) return { "kind": "matrices", "title": "An image becomes a sequence of patches.", "description": "Recorded image-to-token conversion and class-token attention. A stride-four patch projection is equivalent to flattening each RGB patch and applying the same linear map. Learned positions preserve where each patch belongs.", "controlLabel": "Synthetic image", "cases": cases, }