Train ViT Classifier End-to-End

Implement train_vit_step: classify an image batch, compute gradients by hand, and take one SGD step. Train every supplied parameter, including the CLS token and position embeddings. Return the updated parameters as one flat tensor; do not return the logits or loss.

Inputs

• x: (N, C, H, W), a batch of images.
• y: (N,), class IDs in [0, K). They arrive as floats; cast them to integers for indexing.
• w_proj: (CPP, D), the patch projection.
• cls_token: (D,), the shared classification token.
• pos_embed: (M + 1, D), where M = (H/P)*(W/P) is the number of patches.
• blocks_weights: (L, 6, D, D), containing L encoder blocks.
• w_head: (D, K), the classifier matrix; K is the number of classes.
• num_heads: a positive integer dividing D.
• lr: a finite, nonnegative learning rate.

N, C, H, W, P, D and K are positive; L may be zero. P divides H and W. Infer P from w_proj.shape[0] = CPP, or from the image area and patch count. All shapes are consistent. The floating-point tensors are finite and share a dtype and device. Leave all inputs unchanged.

Forward pass

Extract non-overlapping square patches, visiting patch rows then columns. Within a patch, flatten by channel, row within patch, then column within patch. This gives shape (N, M, CPP). Project with w_proj, prepend cls_token to each image's sequence, and add pos_embed. Position 0 belongs to CLS; positions 1 through M follow the patch order.

Each block contains [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2], in that order. All six matrices have shape (D, D), so the feed-forward width is D here. Apply the L blocks in order:

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

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

Use the tanh GELU approximation:

GELU(u) = 0.5 * u * (1 + tanh(sqrt(2/pi) * (u + 0.044715 * u**3)))
cls_repr = seq[:, 0, :]
logits = cls_repr @ w_head

This simplified classifier has no biases, dropout or final normalization. Build normalization, attention, softmax and GELU from basic tensor operations, without packaged layer helpers. Subtract each score row's maximum before exponentiating to compute stable softmax.

Loss and backward pass

Use cross-entropy averaged over the N examples. Do not divide by the number of classes or tokens. Start with:

loss = -sum(log_softmax(logits)[i, y[i]] for i in range(N)) / N
dlogits = (softmax(logits) - one_hot(y, K)) / N
dw_head = cls_repr.T @ dlogits
dcls_repr = dlogits @ w_head.T

Create a zero gradient for the final sequence and put dcls_repr at position 0. Then visit blocks in reverse order. Within each block, differentiate the feed-forward branch first, followed by attention. A residual addition sends the incoming gradient along both paths; add the two contributions when they meet again at the branch input.

For a shared matrix product B = A @ W, flatten A and its incoming gradient G = dB into rows before accumulating the weight gradient:

A_rows = A.reshape(-1, A.shape[-1])
G_rows = G.reshape(-1, G.shape[-1])
dW = A_rows.T @ G_rows
dA = G @ W.T

For example, dw_mlp2 uses the cached GELU output as A and the incoming feed-forward gradient as G. Multiplication by w_mlp2.T belongs to the activation gradient dA, not to dw_mlp2.

After all blocks, dseq is the gradient at the initial token sequence. Sum over the batch for position embeddings and the shared CLS token:

dpos_embed = sum(dseq, axis=0)
dcls_token = sum(dseq[:, 0, :], axis=0)
dpatches = dseq[:, 1:, :]
dw_proj = patches_raw.reshape(-1, C*P*P).T @ dpatches.reshape(-1, D)

Cache the forward values needed for these derivatives. Use manual chain-rule calculations throughout; do not use backward(), automatic differentiation, jax.grad, or optimizer helpers in either framework.

Update and return

Compute every gradient using the original parameters, then apply new_parameter = parameter - lr * gradient. Flatten each updated parameter in row-major order and concatenate in this order:

[w_proj_flat, cls_token_flat, pos_embed_flat, blocks_weights_flat, w_head_flat]

Preserve dtype and device. When lr=0, the result is the original parameters concatenated in this order. With no encoder blocks, logits depend only on cls_token, pos_embed[0] and w_head, so patch-projection weights and the remaining position embeddings stay unchanged.