ViT Classification Forward

Implement vit_classify: turn an image batch into patch tokens, apply the supplied transformer blocks, then classify using the final CLS token. The weights are supplied; this is a forward pass, not a training loop.

Inputs and output

• x: (N, C, H, W), a batch of images.
• w_proj: (CPP, D), projecting each square patch of side P into width D.
• cls_token: (D,), the shared classification token.
• pos_embed: (M + 1, D), where M = (H/P)*(W/P) is the patch count.
• blocks_weights: (L, 6, D, D), containing L encoder blocks.
• w_head: (D, num_classes), the classifier matrix.
• num_heads: a positive integer dividing D.

N, C, H, W, P, D and num_classes are positive. P divides H and W. All shapes are consistent, and all tensors share a floating-point dtype and device. L may be zero.

Return raw logits of shape (N, num_classes), preserving dtype and device. Do not apply softmax or choose a class index. Leave inputs unchanged.

Patch tokens and positions

P is not a separate argument. Infer it from w_proj.shape[0] = CPP (or equivalently from the image area and patch count).

Extract non-overlapping patches with no padding. Visit patch rows from top to bottom and patch columns from left to right. Within each patch, flatten in this order: channel, row within patch, column within patch.

For a single-channel image with rows [1, 2, 3, 4] and [5, 6, 7, 8], P=2 produces patches [1, 2, 5, 6] and [3, 4, 7, 8]. With multiple channels, append each next channel's patch values after the previous one's.

Multiply patch vectors by w_proj. Prepend the same cls_token to every image's sequence, then add pos_embed to every sequence. Position 0 belongs to CLS; positions 1 through M belong to the patches in the order above. The resulting sequence has shape (N, M + 1, D).

Encoder blocks

Apply blocks in order from 0 to L-1, passing each result into the next. The six matrices of block i are packed as:

blocks_weights[i] = [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2]

The feed-forward width is D in this exercise. Each block performs:

h = LN(seq)
q, k, v = h @ w_q, h @ w_k, h @ w_v
attention_output = multi_head_attention(q, k, v) @ w_o
y = seq + attention_output
seq = y + GELU(LN(y) @ w_mlp1) @ w_mlp2

LN normalizes each token over its D features, using population variance and (token - mean) / sqrt(variance + 1e-5), with no learned scale or bias. For attention, split projected features into consecutive heads of width D/num_heads. Scale scores by sqrt(D/num_heads), normalize over key positions, then concatenate head outputs in their original order. Every token, including CLS, attends to all tokens in its own image. Keep batch examples separate and use no causal mask.

Use the tanh GELU approximation elementwise:

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

Build normalization, attention, stable softmax and GELU from basic tensor operations in either framework, without packaged layer-normalization or attention helpers. For softmax, subtract each score row's maximum before exponentiating and dividing by the row sum.

Classifier

After the last block, return seq[:, 0, :] @ w_head. Do not average patch tokens. This simplified model has no final normalization, biases or dropout. When L=0, the result is (cls_token + pos_embed[0]) @ w_head for every image.

See An Image Is Worth 16x16 Words. The visibility illustration below corresponds to bidirectional attention within each image's token sequence.