Vision Transformer Encoder

A Vision Transformer encoder is a stack of transformer blocks that turns a sequence of patch tokens into a sequence of contextual tokens. The patch embedding hands it 197 tokens, each a vector with no awareness of any other patch. To recognize a cat, every patch has to learn which patches hold whiskers, which hold background, and which hold the eye. Attention builds that awareness one layer at a time. Without it, the patch front end is a clever tokenizer and nothing more.

The standard recipe at ViT-Base scale is twelve blocks deep and twelve heads wide, with pre-LayerNorm placement, GELU activation, and a 4x feed-forward expansion. That block shape (scaled in width and depth) is the spine of CLIP ViT-L, SigLIP, DINOv2, the Qwen-VL family, InternVL, and the other open-weight vision encoders of 2025-2026. It is stable enough that you can read any of those papers and assume this block shape unless they say otherwise.

Pre-LN vs post-LN

Pre-LN means LayerNorm runs before each sub-layer, inside the residual branch. Post-LN, the original Transformer placement, runs it after the residual add. Modern vision-language models use pre-LN because it trains stably without learning-rate warm-up tricks. The difference is one line in the forward pass, and the gradient flow at depth 12+ is night and day.

Multi-head self-attention

Multi-head self-attention runs several attention heads in parallel, and each head projects the token vector to its own (query, key, value) triple of dimension head_dim = hidden / num_heads. With hidden = 768 and heads = 12, each head has dim = 64. The 12 heads attend in parallel, their outputs concatenate back to dimension 768, and an output projection mixes them. Splitting the work this way lets one head learn "attend to the cat eye" while another learns "attend to the background gradient" without interference.

Why the 4x feed-forward expansion

The feed-forward sub-layer (FFN) goes hidden -> 4 * hidden -> hidden with GELU in the middle. The factor 4 is empirical and has held across language and vision transformers since 2017. Smaller (2x) underfits, and larger (8x) overfits at a fixed data budget. The MLP is where the model stores most of its learned facts, and the wider middle is where they sit.

The table counts weight matrices only. Biases and the front end's patch projection and position table add a small remainder, which is why the demo's printed total lands near, not exactly on, 86M.

ViT-Base is an 86M-parameter encoder. That is small by 2026 standards (SigLIP-So400M is 400M, the Qwen-VL ViT is 675M), but the architecture is the same up to width and depth.

Causal mask or not?

Vision Transformers are encoder-only and bidirectional: token i may attend to token j for any pair, so there is no mask. The decoder-side cross-attention in lesson 61 uses a causal mask, but inside the vision encoder attention is fully connected.

What the CLS token learns

The CLS token is a learned parameter with no patch content of its own. It picks up information through attention across every block, and by the final layer its row is a vector summary of the whole image. Downstream heads project that single vector into class logits, contrastive embeddings, or cross-attention keys for a text decoder.

code/main.py implements:

• MultiHeadSelfAttention, with qkv and output projections, the scaled-dot-product attention math, and a (B, N, D) input check that raises ValueError on any other rank.
• FeedForward, the 4x-expansion GELU MLP.
• Block, a pre-LN block that composes the attention and feed-forward sub-layers with residuals.
• ViT, a stack of 12 blocks with a final LayerNorm.
• VisionEncoder, which wires VisionFrontEnd from lesson 58 to the ViT stack. Its forward() returns the contextual sequence and the pooled CLS vector.
• A demo that runs a synthesized 224x224 fixture image through the full encoder and prints the config, front-end, ViT and total parameter counts, the input and output shapes, and the CLS norm at every other layer. It also prints a softmax row-sum check on block 0 and the gradient norms reaching the patch projection and the CLS token.

Run it:

python3 code/main.py

The fixture encodes to a (1, 197, 768) tensor. The CLS norm drifts upward as the layers compose, then stabilizes at the final LayerNorm. Total parameters come to about 86M.

Up to width and depth, the encoder defined here is the same block stack that ships inside every open-weight VLM in 2025-2026. The differences live in:

• Width and depth. ViT-Large is hidden=1024, depth=24, heads=16; SigLIP So400M is hidden=1152, depth=27, heads=16. Same block.
• Pooling head. CLS pooling (this lesson) vs average pooling (SigLIP) vs attention pooling (later VLMs).
• Position handling. Fixed sinusoidal (lesson 58) vs learned 1D vs ALiBi vs 2D RoPE. The block math is unchanged.
• Register tokens. DINOv2 prepends 4 extra learned tokens. One line of code.

This block stack is the substrate, and lessons 60-63 build on top of it.

code/test_main.py has eight tests:

• attention preserves the input shape, and its attention rows sum to one along the key axis (softmax sanity)
• the feed-forward layer widens to hidden * mlp_ratio and contracts back to hidden
• a block's residual path is wired: the output keeps the input shape and differs from the input
• a 4-layer VisionEncoder forward pass returns the right shapes for the token sequence and the CLS vector
• gradients flow to the patch projection from a loss on the CLS output
• head_dim divides hidden, and a hidden that heads do not divide is rejected with ValueError

The gradient test backprops (cls * cls).sum(), not cls.sum(). The encoder ends in a LayerNorm with weight 1 and bias 0, so each output vector sums to zero for any input, and the gradient of a plain sum is exactly zero even when the wiring is fine. Squaring first makes the loss depend on the values.

Nothing tests invariance to batch size, a zero input, or the ValueError on a non-(B, N, D) input. Treat those as unchecked.

Run them:

python3 -m unittest code/test_main.py

1. Add register tokens (4 learned vectors prepended after CLS) and rerun. Compare attention map smoothness via the entropy of the softmax distribution on the last layer.

2. Swap pre-LN for post-LN and train for one epoch on a synthetic shape classifier. Observe which one trains stably without LR warm-up.

3. Implement causal masking as an attn_mask argument so the same block can be reused as a decoder block. The mask shape is (seq, seq), lower-triangular.

4. Profile a forward pass at batch sizes 1, 8, 64 with torch.profiler. Check which of the MLP and attention dominates wall time. The FFN holds about two thirds of each block's parameters, so it is the natural suspect, but let the profiler say.

5. Replace one attention head's q-k-v projections with a low-rank LoRA adapter, freeze the rest, and verify the gradient only flows where you expect.

• An Image is Worth 16x16 Words (ViT, 2021) for the encoder recipe.
• DINOv2 (2023) for register tokens and the self-supervised pretraining objective.
• SigLIP (2023) for the average-pooling variant and the sigmoid contrastive loss used in lesson 62.