Vision Transformer, Part 3 — The Full Encoder Block and the Prediction

Abstract. With attention understood, this part assembles the complete transformer block and reaches the first prediction. It covers multi-head attention (running several attentions in parallel, for free), the LayerNorm and residual connections that wrap each sub-layer, and the token-wise MLP that refines each token on its own — then shows how the final block’s [CLS] token becomes 10 class logits. Includes the real attention patterns the four heads learned and the actual learned γ and β of LayerNorm.

In this 6-part series of Vision Transformer:

1. Part 1 — Embeddings: from pixels to a token sequence
2. Part 2 — Attention: how tokens relate to each other (single head)
3. Part 3 — The full block: multi-head, LayerNorm, residuals, MLP, and the prediction — you are here
4. Part 4 — Building the ViT in PyTorch, trained on CIFAR-10
5. Part 5 — Do overlapping patches help? A controlled experiment
6. Part 6 — How it learns: loss, backprop, the chain rule, the weight update

Part 3 of the series. Part 1 built the token sequence; Part 2 built single-head attention (figures 1–8). Here we complete the forward pass: what wraps around attention inside a block, and how the last block becomes 10 class logits. Figures use real weights from the trained model.

Multi-head attention

The limitation of a single head

A single head produces, for each token, one softmax distribution over the other tokens. One notion of “relevant.” But a patch usually relates to others in several ways at once — which are near me, which share my colour, which are part of the same object. With one head those must be squeezed into a single weighting, and because the output is a weighted average, competing signals blur together.

The fix is simple: run several attentions at once and let each specialise.

1 · The split — same space, partitioned

Figure 1: Attention head splitting to get Multi-head attention

The surprise: going multi-head adds no parameters and does not widen the model. The Q/K/V projections still map 128 → 128; we just reshape each into 4 slices of 32 and run attention inside each slice independently:

Q : [65, 128]  ──reshape──▶  [65, 4, 32]      4 heads × 32 dims = 128

What you buy is not capacity — it’s independence. Four separate softmaxes mean four attention patterns can coexist instead of being averaged into one.

One detail changes: the √d scaling from Part 2’s figure 4 now uses the head dimension — √32 ≈ 5.66, not √128 ≈ 11.3 — because each head's dot product sums only 32 terms.

2 · The pipeline, end to end

Figure 2: Multi-head attention and each QKV matrix sizes

Each head runs the complete single-head procedure (Part 2, figures 1–6) on its own 32-dim slice → its own 65×65 matrix → its own [65,32] output. Then two steps stitch them back: concatenate the four outputs → [65,128], and project through a learned W_o. That final projection lets the model mix what the heads found before the result re-enters the block. Shape in = shape out = [65,128].

3 · What our heads actually learned

Figure 3: Different spatial perspective or visual association being learned

From the trained model (last layer, what CLS attends to): head 1 locks onto the horse’s body/back, head 2 onto the neck and head, heads 3–4 stay diffuse. That division of labour is exactly what the split makes possible.

Honest caveat: one image, one layer, one model. Head specialisation is real and widely observed, but not guaranteed to be this clean — some heads genuinely are near-redundant, which is why head-pruning research exists.

LayerNorm

Before attention (and before the MLP), each token vector passes through a LayerNorm. It exists because activations drift in scale as they flow through layers, and unstable scales make training brittle. LayerNorm resets each token to a standard shape.

4 · Normalize, then rescale with learned γ and β

Figure 4: LayerNorm — a) Zero mean & unity std-deviation, b) γ, β are paramters being learned

Three steps, applied to one token’s own features:

1. compute this token's mean μ and std σ (over its 128 values)
2. x̂ = (x − μ) / √(σ² + ε)          → mean 0, std 1
3. y  = γ ⊙ x̂ + β                    → rescaled by learned γ, β

The key point: γ and β are vectors — one value per dimension (128 each), applied element-wise. Follow any single dimension across the figure:

y = γ · x̂ + β  =  (+0.66) × (−0.62) + (−0.00)  =  −0.41

Two things are easy to conflate, so keep them separate:

So each token normalises by its own μ/σ, but all 65 are rescaled by the same learned γ/β. And note: the normalization step gives mean-0/std-1, but the output does not — γ and β deliberately let the model pick a different final scale (this model learned γ ≈ 0.64, so it shrinks the vector to ~64%).

Get Chandrashekhar Padole’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

The model has 13 LayerNorms — two per block (before attention, before the MLP) × 6, plus one final norm before the head — each with its own γ/β.

Residual connections and the block

Now assemble a block. Each of the two sub-layers (attention, MLP) is wrapped the same way — pre-norm with a residual:

── one encoder block ─────────────────────────────────────
  x                                       block input [65,128]
  ├─ LayerNorm₁ → Multi-Head Attention → ⊕ x     → x'   (sub-layer 1)
  └─ LayerNorm₂ → MLP                   → ⊕ x'    → x''  (sub-layer 2)
  x''                                     block output [65,128]
───────────────────────────────────────────────────────────

Read it carefully — three things people get wrong:

1. The residual carries the un-normalized input. LayerNorm’s output goes into the sub-layer; what gets added back is the original x. That's what "pre-norm" means — the norm is on the branch, not the skip path.
2. Each residual adds that block’s own input, not the encoder’s original input. Block 2 adds block-1’s output, and so on.
3. Two LayerNorms and two residual adds per block — with different learned γ/β for each LayerNorm.

Why add the input back at all? The skip path lets each block learn a small adjustment to x rather than replace it, and it gives gradients a direct route back — which is what makes a deep stack trainable.

The block is shape-preserving ([65,128] in and out), so blocks stack cleanly. Our model stacks 6, each with its own weights (they share a structure, not parameters):

x₀ → block 1 → block 2 → block 3 → block 4 → block 5 → block 6 → x₆

The MLP

The second sub-layer of every block is a small feed-forward network.

5 · Expand → GELU → contract, per token

Figure 5: MLP layer — Expand & Contract Network

It’s two linear layers with a non-linearity between: 128 → 512 → GELU → 128. Expand each token 4× into a wider working space, apply GELU (a smooth gate that suppresses negatives), project back. 66,048 + 65,664 = 131,712 params per block.

The one point that matters most: the MLP is applied independently to each token, with shared weights. Unlike attention, it never mixes tokens — token 7’s output depends only on token 7’s input. That’s the clean division of labour in a block:

• Attention = mix information between tokens
• MLP = refine each token on its own

A block alternates the two: gather from neighbours, then think privately.

From the last block to the prediction

After block 6 we have x₆ = [65, 128]. Two small steps turn it into logits.

6 · CLS → final LayerNorm → Linear → logits

Figure 6: CLS only prediction — a) LayerNorm, b) Linear Network & c) Logits

x₆ [65,128]  → take token 0 (CLS)  → final LayerNorm  → Linear(128→10)  → 10 logits
              the 64 patch tokens are discarded

The head is a single Linear layer — 128 → 10, no hidden layer, no activation, no dropout — with 1,280 weights + 10 biases = 1,290 params. Its 10 output nodes match the 10 classes (the length of the one-hot label). Each node is a dot product of the CLS vector with that class's 128 weights, plus a bias.

The outputs are logits — raw, unbounded scores, not probabilities. On our horse: horse 3.80, everything else below 0.5.

• The prediction = argmax(logits) → horse. Just pick the biggest.
• The confidence = softmax(logits) → 88%.

The elegant closing detail: the whole ~1.2M-parameter model funnels into a 1,290-parameter head that reads only the CLS token. Everything the 6 blocks did to the 64 patch tokens mattered only insofar as it shaped that one CLS vector — exactly the job the CLS token was added to do.

The forward pass, complete

pixels → patch embed → +CLS +positions          (Part 1)
       → attention                               (Part 2)
       → 6 × [ LN → multi-head attn → +residual ,  LN → MLP → +residual ]   (Part 3)
       → take CLS → final LN → Linear → 10 logits → argmax / softmax        (Part 3)

That’s the entire prediction path. Part 4 builds it in code and trains it; Part 6 runs the same path backwards to show how the weights are learned.