E2E: Tiny Vision Transformer - End-to-End Implementations

End-to-End Tiny Vision Transformer (ViT)

What you'll build

TinyViT: a Vision Transformer classifier with one pre-norm encoder block, including patch extraction, the [CLS] token, position embeddings, the loss and a training step.

Background

A Vision Transformer (Dosovitskiy et al., 2021) treats an image as a sequence of patches: cut it into P × P tiles, embed each tile linearly, prepend a learnable [CLS] token, add position embeddings, run a Transformer encoder, and classify from the final CLS token. With no convolutions, everything the model knows about 2-D layout comes from the patch order and the position embeddings.

Architecture

x (B, C, H, W)
  │ patchify                  (B, N, P*P*C)       N = (H/P) * (W/P)
  │ @ W_patch + b_patch       (B, N, D)
  │ prepend CLS token         (B, N+1, D)
  │ + pos_emb (N+1, D)        (only if use_pos_emb)
  ▼
  t = t + MHA(LN1(t))         pre-norm attention with residual
  t = t + MLP(LN2(t))         MLP: Linear(D, mlp_dim) -> GELU -> Linear(mlp_dim, D)
  ▼
  cls = LN_f(t[:, 0])         (B, D)      <- cls_features(x)
  logits = cls @ W_head + b_head  (B, num_classes)

Details

• Patch order: patches in row-major order over the (H/P, W/P) grid; inside a patch, flatten as (p1, p2, c), i.e. einops 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)'. This matches a stride-P convolution.
• MHA: qkv = h @ W_qkv (no bias) split into Q, K, V of width D; reshape to num_heads heads of size D / num_heads; softmax(Q K^T / sqrt(d_head)) V; merge heads; output projection @ W_o (no bias). No mask: every token attends to every token.
• LayerNorm: over the last dimension, eps = 1e-5, with learnable gain (init 1) and bias (init 0).
• GELU: tanh approximation, 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))).
• Initialization: weights ~ N(0, 1/fan_in), biases 0, CLS token and position embeddings ~ N(0, 0.02^2).

Why position embeddings matter

Self-attention is permutation-equivariant, and the CLS token attends to all patches with the same weights regardless of where they are. So without position embeddings, shuffling the patches of an image leaves the CLS output (and the prediction) unchanged. Position embeddings are the only source of spatial order.

Method contract

Method	Returns
TinyViT(img_size=8, patch_size=4, in_chans=1, dim=16, num_heads=2, mlp_dim=32, num_classes=2, use_pos_emb=True)	creates all parameters (initialization above)
patchify(x)	(B, C, H, W) → (B, N, P·P·C)
cls_features(x)	final-LayerNormed CLS token, (B, D)
forward(x)	logits (B, num_classes)
loss(x, y)	float, mean cross-entropy with integer labels y (B,)
train_step(x, y, lr=0.1)	float, the loss before one plain SGD step

train_step differs by track:

• NumPy track: update only the classification head (W_head, b_head); the encoder stays frozen (a linear probe). Gradient: dlogits = (softmax(logits) − onehot(y)) / B, dW_head = clsᵀ @ dlogits, db_head = dlogits.sum(0).
• PyTorch track: autograd updates every parameter.

In the NumPy track, parameter names and storage are up to you (the tests only use the methods above); the PyTorch starter names the patch embedding patch_embed.

Constraints

NumPy only, no loops over batch, heads or tokens.

Follow-up

• Attention over N patches costs O(N²): what happens to the cost if you halve the patch size?
• Why do ViTs need more data (or strong augmentation) than CNNs? (They lack the convolutional inductive bias of locality and translation equivariance.)

Input: TinyViT(img_size=8, patch_size=4, in_chans=1, dim=16, num_heads=2, num_classes=3).forward(np.random.randn(4, 1, 8, 8))
Output: logits of shape (4, 3)

Input: model.patchify(np.arange(32.).reshape(1, 2, 4, 4)) with patch_size=2
Output: shape (1, 4, 8); first patch [0, 16, 1, 17, 4, 20, 5, 21]

import numpy as np

class TinyViT:
    def __init__(self, img_size=8, patch_size=4, in_chans=1, dim=16, num_heads=2,
                 mlp_dim=32, num_classes=2, use_pos_emb=True):
        """
        Create all parameters (weights ~ N(0, 1/fan_in), biases 0,
        CLS token and position embeddings ~ N(0, 0.02^2), LayerNorm gains 1).
        """
        self.patch_size = patch_size
        self.num_heads = num_heads
        self.use_pos_emb = use_pos_emb
        # Your code here

    def patchify(self, x):
        """(B, C, H, W) -> (B, N, P*P*C), patches row-major, flattened as (p1, p2, c)."""
        # Your code here
        pass

    def cls_features(self, x):
        """Embed patches, prepend CLS, add pos-emb, one pre-norm block, final LN on CLS -> (B, D)."""
        # Your code here
        pass

    def forward(self, x):
        """Logits (B, num_classes)."""
        # Your code here
        pass

    def loss(self, x, y):
        """Mean cross-entropy (float)."""
        # Your code here
        pass

    def train_step(self, x, y, lr=0.1):
        """One SGD step on the head (W_head, b_head); returns the loss before the step."""
        # Your code here
        pass