Medium End-to-End

ViT Encoder Block

Implement vit_block: one Vision Transformer encoder block with multi-head self-attention followed by a feed-forward network. Both stages normalize their input before computing an update, then add that update to the unnormalized input. This order is called pre-layer normalization.

Inputs and output

x: (N, T, D), containing N independent sequences of T token vectors.
w_q, w_k, w_v, w_o: (D, D), the attention projection matrices.
w_mlp1: (D, F) and w_mlp2: (F, D), the feed-forward matrices.
num_heads: a positive integer H that divides D. Each head has width D/H.

All dimensions are positive. Tensors share a floating-point dtype and device. Return shape (N, T, D), preserving dtype and device and leaving inputs unchanged. The tokens are already embedded; do not extract patches, add a CLS token or add position embeddings.

Normalize before each stage

Normalize each token over its D features, independently of other tokens and batch examples. Use population variance, dividing by D rather than D-1:

LN(x) = (x - mean) / sqrt(variance + 1e-5)

There are no learned normalization scales or biases.

Multi-head attention

Compute Q = LN(x) @ w_q, K = LN(x) @ w_k and V = LN(x) @ w_v. Split each projected vector into H consecutive chunks of width D/H. For each batch example and head independently:

scores = Q_head @ K_head.T / sqrt(D/H)   # (T, T)
weights = softmax(scores, axis=-1)
head_output = weights @ V_head         # (T, D/H)

Use stable softmax across key positions. Every token can attend to every token in its own sequence, including itself; there is no causal mask. Concatenate the head outputs in their original order to get (N, T, D), then compute attention_output = concatenated @ w_o.

Add the first residual: y = x + attention_output.

Feed-forward stage

Normalize y, not the original x, then compute:

output = y + GELU(LN(y) @ w_mlp1) @ w_mlp2

Use this tanh approximation to GELU elementwise, where ** means exponentiation:

GELU(z) = 0.5 * z * (1 + tanh(sqrt(2/pi) * (z + 0.044715 * z**3)))

There are no projection biases, dropout or final layer normalization. If both stage outputs are zero, the residual paths return x unchanged.

Build normalization, attention and GELU from basic tensor operations in either framework. Do not call packaged layer-normalization, multi-head attention or scaled-dot-product-attention helpers.

This exercise implements a simplified block, not a complete image classifier. See An Image Is Worth 16x16 Words. The visualization below illustrates visibility; use bidirectional attention for this encoder.

Loading visualization…

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Turn image patches into tokens

A Vision Transformer divides an image into patches, flattens each patch, and projects it into a feature vector called a token. For image height H, width W, and non-overlapping square patches of side P, the number of patches is (H/P) * (W/P) when both dimensions are divisible by P.

For an 8-by-8 image with 4-by-4 patches, there are four patch tokens. With three color channels, each patch contains 4 * 4 * 3 = 48 values before projection. Preserve the required pixel, channel, and patch order: obtaining the right final shape does not prove that each token contains the right pixels.

Position embeddings identify where the patches came from. Some models also prepend a learned class token, making the example above a five-token sequence. Attention allows this token to gather information from the patch tokens. A classifier can then use its final representation; other architectures pool the patch tokens instead.

Transformer blocks combine attention with feed-forward layers. Follow the task's normalization placement and classification head, since these choices change the model.

Mix images and labels with the same coefficient

Mixup forms a training image as

$$\tilde x=\lambda x_a+(1-\lambda)x_b.$$

It combines the two label distributions with the same weights. At lambda = 0.25, the image and target each contain one quarter of example a and three quarters of example b. For this one model prediction, ordinary cross-entropy with the mixed target equals the weighted sum of the losses against the two original targets.

During manual backpropagation, shared class-token and position-embedding parameters receive contributions from every example. Add those contributions, including any scaling already required by the batch loss. Remember the extra class-token position when checking sequence lengths.

Detecting runtime…