Building the Encoder: Many Heads, Positions, and Depth

One head isn't enough: multi-head attention

In guide 2 we built a single self-attention operation: every patch produced a query, a key, and a value, and we let each patch decide how much to listen to every other patch. That one operation is powerful, but it forces all of the comparison to happen in a single shared 'space'. Imagine asking one person to simultaneously track colour similarity, spatial closeness, AND texture across the whole image. They can do it, but they will end up blending those very different notions of 'related' into one muddy averaged judgement.

Multi-head attention fixes this by running several attention operations in parallel, called heads. Each head gets its own learned Q/K/V projection matrices, so each head looks at the patches through its own lens. One head might learn to attend to patches of similar colour, another to patches that are spatially adjacent, a third to repeated texture. The analogy: instead of one overworked generalist, you convene a panel of specialists, each examining the same scene from their own angle, and then you pool their notes at the end.

Multi-head attention as parallel attention heads, concatenated and projected.

Let us unpack every symbol. There are h heads (the number of specialists). For head i, the matrices W_Q^i, W_K^i, W_V^i are learned projections that map the full D-dimensional token down to a smaller dimension d_k = D/h. Inside each head, \mathrm{Attention}(\cdot) is exactly the scaled dot-product formula from guide 2: \mathrm{Attention}(Q,K,V)=\mathrm{softmax}\!\big(\tfrac{QK^\top}{\sqrt{d_k}}\big)V — the only change is that Q,K,V are now the head's own smaller projections. Each head outputs an N \times d_k result. We then `Concat` the h heads side by side, which stitches the h pieces of width d_k back into width h \cdot d_k = D. Finally W_O (a D \times D matrix) mixes the heads' findings into one coherent output per token.

Concrete numbers, using ViT-Base. The model dimension is D = 768 and we use h = 12 heads. Then each head works in dimension d_k = 768 / 12 = 64. So the giant 768-wide vector is sliced into 12 strips of width 64; each strip runs its own attention; the 12 outputs of width 64 are concatenated back to 12 \times 64 = 768; and W_O (768×768) blends them. Notice the cost is roughly the same as one big 768-dimensional attention — we did not pay 12× more compute, we just reorganized the same budget into 12 specialized sub-spaces.

Where am I? Positional encoding

Here is a subtle but crucial problem with attention as we have built it so far: it is permutation-invariant. Attention computes how much each patch relates to every other patch purely from the content of their vectors — it never looks at where a patch sat in the image. If you shuffled the patches into a random order, attention would produce the exact same set of relationships, just reordered. To attention, the patches are an unordered bag, not a grid.

For vision this is a disaster. The patch in the top-left corner and the patch in the bottom-right corner are very different things even if they happen to contain similar pixels; a nose above a mouth is a face, the same parts shuffled are not. The analogy: someone hands you the panels of a comic strip in a shuffled pile. The drawings are all there, but without panel numbers you cannot reassemble the story. Positional encoding adds those numbers back — it injects a position-dependent vector into each patch embedding so the model knows where each patch came from.

The input sequence to the encoder: patch embeddings plus a class token, with positional encodings added.

Reading this left to right: x_p^j is the j-th flattened patch (recall guide 1, where we cut the image into N patches), and multiplying by the patch embedding matrix E projects it into a D-dimensional token. We prepend one extra learnable vector x_{\text{class}} — the class token — which carries no pixels of its own but will later gather a summary of the whole image (we use it in section 5). The square brackets just mean 'stack these N+1 tokens into a sequence'. Then we add E_{\text{pos}}, a learnable table of shape (N{+}1) \times D: row k of this table is the position vector for slot k. The key detail is that we add the position vector (element-wise) rather than concatenate it. Adding keeps the dimension at D, so every downstream layer still sees width-D tokens; concatenating would have grown the width and forced every other matrix to change shape.

There are two common ways to fill E_{\text{pos}}. ViT typically uses learned 1D positions: E_{\text{pos}} is simply trained from scratch like any other parameter, and the model figures out useful position vectors on its own ('1D' means we number the patches 0, 1, 2, … in raster order rather than encoding their 2D row/column separately). The alternative, from the original Transformer, is a fixed sinusoidal encoding given below.

The fixed sinusoidal alternative: each dimension is a sine or cosine wave of a different frequency.

Here pos is the integer position of the token in the sequence (0, 1, 2, …) and i indexes the dimension within the d-dimensional vector (so 2i and 2i{+}1 are a paired sine/cosine slot). The trick is that low dimensions use a long wavelength (they change slowly across positions) while high dimensions use a short wavelength (they change quickly), so the combination of all the sines and cosines gives every position a unique 'fingerprint' — much like binary digits combine to count. Sinusoidal encodings need no training and can extrapolate to longer sequences, but for the fixed-resolution images ViT works on, the learned table usually performs at least as well and is simpler, which is why ViT defaults to it.

Thinking it over: the feedforward MLP block

Attention is great at mixing information across tokens — it lets each patch gather context from the others. But after that gathering, each token holds a freshly updated vector that deserves some deeper, private processing. That is the job of the feedforward MLP block. It is applied to each token's vector independently — the same little neural network runs on every token separately, with no further talking between tokens. The analogy: attention is the group discussion where everyone shares what they heard; the MLP is each person going home afterwards to quietly think it all through on their own.

Two linear layers with a GELU nonlinearity between them.

Unpacking it: z is one token's vector of width D. The first linear layer W_1 (with bias b_1) expands it from D up to a wider hidden dimension, conventionally 4D. Then \mathrm{GELU}, a smooth nonlinearity, is applied element-wise — it behaves like a softened version of 'keep positive values, suppress negative ones', and the smoothness (versus a hard ReLU) tends to help training. Finally the second linear layer W_2 (with bias b_2) contracts the 4D hidden vector back down to D. The two biases b_1, b_2 are just learnable offsets. Because both linear layers share the same weights across every token, this block has no notion of position or order — all the cross-token mixing already happened in attention.

Concrete dims for ViT-Base: 768 \to 3072 \to 768. So W_1 is 768 \times 3072 and W_2 is 3072 \times 768. Why blow the width up 4× before shrinking it back? The wide hidden layer is scratch space: it gives the network room to compute many intermediate features in parallel and recombine them, much like rough working-out on a large sheet of paper before you write down a one-line answer. A direct 768 \to 768 layer would be far more limited in the transformations it could express. In fact, the MLP blocks hold the majority of a Transformer's parameters, and most of the model's raw 'thinking capacity' lives here.

Putting it together: the encoder block

We now have the two engines — multi-head attention (MHA) and the feedforward MLP block. A transformer encoder block wires them together with two more ingredients: LayerNorm (LN) before each engine, and a residual (skip) connection around each engine. ViT uses the pre-norm arrangement, meaning the LayerNorm comes before the sub-layer, not after. The block is: LN → MHA → add the input back; then LN → MLP → add again.

The two sub-layers of a pre-norm encoder block, each wrapped in a residual connection.

Let us walk one full pass through block number \ell. The input is z_{\ell-1}, the sequence of tokens coming out of the previous block (shape N{\times}D). First sub-layer: we normalize it with \mathrm{LN}, feed that into multi-head attention, and then add the original z_{\ell-1} back — that `+ z_{\ell-1}` is the residual connection. The result z'_\ell is the input with an attention-based update layered on top. Second sub-layer: we normalize z'_\ell, run it through the MLP, and add z'_\ell back, giving z_\ell. Here \ell is just the block index (block 1, 2, 3, …), \mathrm{LN} is layer normalization, and each + is a residual add. Crucially, both inputs and outputs are shape N{\times}D — the block takes in N tokens of width D and hands back N tokens of width D, only with richer contents.

And why LayerNorm? As signals flow through many layers, the scale of the activations can drift — some grow huge, some shrink to nothing — which destabilizes training. \mathrm{LN} rescales each token's vector to have a controlled mean and variance (across its D features) before it enters a sub-layer, keeping the numbers in a healthy range. Because LN, the residual, attention, and the MLP all preserve the N{\times}D shape, the output of one block is a drop-in input for the next — which is exactly what lets us stack identical blocks in section 5.

def encoder_block(z, params):
    # z has shape (N, D) and stays (N, D) throughout
    # --- sub-layer 1: multi-head attention with residual ---
    a = multi_head_attention(layer_norm(z))   # LN, then MHA
    z = z + a                                 # residual add
    # --- sub-layer 2: feedforward MLP with residual ---
    m = mlp(layer_norm(z))                     # LN, then MLP (GELU inside)
    z = z + m                                 # residual add
    return z                                   # still (N, D)

Stacking depth into a full ViT

Because every encoder block maps N{\times}D to N{\times}D, we can stack L identical blocks one on top of another, feeding each block's output straight into the next. Each block refines the token representations a little more: early blocks tend to capture local, low-level relationships, later blocks compose them into global, semantic ones. ViT-Base uses L = 12 blocks; larger variants like ViT-Large use 24. This is what 'depth' means here — repeated rounds of gather-then-think.

After the final block, how do we turn N{+}1 token vectors into a single prediction? We pull out just the class token — the special slot we prepended back in section 2. Throughout all L blocks, attention let this token gather information from every patch, so by the end it holds a summary of the whole image. We pass that one vector through a small classification head — a single linear layer followed by softmax — to produce probabilities over the label classes. This completes the vision transformer.

The classification head: take the final class token, normalize, project to class logits, softmax.

Symbol by symbol: z_L^0 is the class token's vector after the final (the L-th) block — the superscript 0 picks slot 0, the class token, and the subscript L says 'after all L blocks'. We apply a final \mathrm{LN} to it, then multiply by W_{\text{head}}, the classification weight matrix of shape D \times C where C is the number of classes; this maps the width-D summary into C raw scores called logits. Finally \mathrm{softmax} exponentiates and normalizes those C logits so they become a probability distribution that sums to 1 — the model's confidence in each label. For example, with D{=}768 and C{=}1000 ImageNet classes, W_{\text{head}} is 768 \times 1000, and y is a length-1000 probability vector whose largest entry is the predicted class.

1. Patchify (guide 1): cut the image into N fixed-size patches, flatten each, and project with the patch embedding matrix E to get N tokens of width D.
2. Prepend the class token and add the learned positional encodings E_pos, giving the input sequence z_0 of shape (N+1)×D.
3. Pass z_0 through L identical encoder blocks; each does LN → multi-head attention → residual, then LN → MLP → residual, keeping the shape (N+1)×D.
4. Take the class token from the final block, z_L^0, apply LayerNorm, the classification head W_head, and softmax to read off the predicted label probabilities y.

That is the entire forward pass, from raw pixels to a predicted label, built across guides 1 to 3: patches and embeddings, multi-head self-attention, positional encoding, the MLP, LayerNorm and residuals, L stacked encoder blocks, and the class-token head. The architecture is now complete and, reassuringly, fairly uniform — almost everything is the same block repeated. What we have not yet discussed is how to actually train this thing well. ViTs are famously data-hungry, and that is exactly the problem guide 4 (DeiT, distillation, and hybrids) will tackle, before guide 5 reworks the architecture itself with the hierarchical Swin Transformer to scale to large, real-world images.