Vision Transformer, Part 2 — Attention: From First Principles

Abstract. Attention is the operation that finally lets an image’s tokens share information — and it’s simpler than its reputation suggests. Building from nothing, this part shows what a dot product actually measures, how each token forms a query, key, and value, why the scores are divided by √d, and how softmax turns them into a weighted blend of the other tokens’ values. The payoff: every token walks away as a mixture of the tokens it found relevant. Single head throughout, with every number computed from a real trained model.

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) — you are here
3. Part 3 — The full block: multi-head, LayerNorm, residuals, MLP, and the prediction
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

The problem attention exists to solve

After patch embedding, an image is 65 independent vectors. Each one only knows about its own 4×4 corner of the picture. But “is this a cat?” can’t be answered from one 4×4 patch — you need the ear and the whisker and the body, which live in different patches.

So we need an operation that lets every token gather information from every other token, and — crucially — one that learns which tokens are worth gathering from. That is all attention is. Everything below builds it from nothing.

1 · The dot product measures agreement

Figure 1: Similarity between two Vectors using Dot Product (Each token is vector)

Take two vectors. Multiply them position by position, add up the results. That single number — the dot product — tells you how much the two vectors point the same way:

Table 1: Dot product significance

For unit-length vectors this is exactly cos(angle). That's the whole reason attention uses dot products: it's the cheapest possible way to ask "how similar are these two things?" — one multiply-and-add per dimension, and the answer falls out as a single number.

Hold onto that: dot product = a similarity meter.

2 · Query, Key, Value — and one relevance score

Before the arithmetic, the idea. From each token we compute three different vectors, using three learned weight matrices:

Table 2: QKV vector’s Significance

The split matters: a token’s advertisement (key) can differ from its content (value), and from what it wants (query). Those are three different jobs, so they get three different learned projections.

Now the arithmetic. To ask “how relevant is token B to token A?”, take A’s query and B’s key and dot them:

Figure 2: Dot product gives the correlation between q & k. This forms the attention matrix.

Multiply position by position, add it all up → one number. High score = “this key matches what I’m looking for.”

That’s the atom of attention. Everything else is bookkeeping around this.

3 · One query against every key → a row of scores

A token doesn’t compare itself to just one other token — it compares itself to all of them (including itself):

Figure 3: Example of Visualization of 8 tokens Attention Matrix

Repeat step 2 once per key and you get a row of scores — one number per token, saying how relevant each one is to this particular query.

Do that for every token as the query, and you’ve filled a full square table: 65 × 65 in our ViT. That table is the attention matrix. Row i = “what token i finds relevant.” Nothing mystical — it’s just every pairwise dot product, computed in one matrix multiply Q · Kᵀ.

4 · Why we divide by √d

Figure 4: Normalization to have unity deviation in Attention Matrix

Here’s a subtlety that most explanations skip entirely.

Get Chandrashekhar Padole’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

A dot product sums d products. With d = 128, you're adding 128 random-ish terms — and sums of many terms spread out. Concretely, if the components are roughly unit-variance, the dot product has standard deviation ≈ √d ≈ 11, as the left panel shows.

Why is that a problem? Because the next step (softmax) exponentiates. Feed it scores separated by 20 or 30 and e^score makes the biggest one utterly dominate: the weights collapse to essentially 1.0 on a single token, 0.0 everywhere else (orange bars, right panel). Two bad consequences:

• The token can only look at one place — you’ve thrown away the “weighted blend” that makes attention powerful.
• Softmax is flat in that saturated regime, so gradients vanish and the layer stops learning.

Dividing every score by √d rescales the spread back to σ ≈ 1 — scores in a range where softmax produces a genuinely soft, useful distribution (violet).

So 1/√d isn't a magic constant. It's variance control, and it follows directly from "we added up d terms."

5 · Softmax turns scores into weights

Figure 5: Softmax turns scores into weights, all positive & sum to 1.0

Raw scores are awkward: any magnitude, sometimes negative. We want proportions — “take 31% of this token, 12% of that one.” Softmax does exactly that in two moves:

1. Exponentiate each score → everything becomes positive (and bigger scores get disproportionately bigger, which is what “attending more” should mean).
2. Divide by the total → the numbers now sum to exactly 1.

softmax(z)_i  =  exp(z_i) / Σ_j exp(z_j)

The output is a set of mixing proportions. That’s it — softmax is a normalizer, converting arbitrary similarity scores into a recipe.

6 · The output is a weighted blend of the values

Figure 6: Combining weights with V vector

Now use the recipe. Multiply each token’s value vector by its weight, and add them all up:

output_i  =  Σ_j  weight_ij · v_j

Because the weights sum to 1, this is a weighted average of the value vectors — the output lands somewhere inside the cloud of values, pulled toward whichever tokens got the big weights.

This is the payoff, and it’s worth saying plainly:

Each token leaves the attention layer as a blend of the tokens it found relevant.

The whisker patch, having scored highly against the ear and fur patches, walks out carrying a mixture of their content. That’s how information moves between patches — and stacking 6 layers means blends of blends, which is how a model assembles “cat” out of 4×4 squares.

7 · Attention is order-blind — hence positional embeddings

Figure 7: Order-invariant Attention

Look back at every formula so far. q · k, softmax, weighted sum — not one of them mentions where a token sits. Attention treats its input as a bag of tokens, not a sequence.

Proof by picture: shuffle the input tokens and the outputs simply shuffle with them. Identical values. The mechanism genuinely cannot tell top-left from bottom-right.

For images that’s fatal — a face is its parts in the right arrangement. So before the encoder we add a learned positional embedding to each token, baking “where I am” into the vector itself. After that, q · k implicitly compares content and position, because position is now part of the content.

This is why positional embeddings aren’t an optional nicety — they’re repairing a structural blind spot in attention.

8 · The whole thing, with real numbers

Figure 8: QKV Matrix Computational Flow

Everything above, on one slide, with numbers you can check by hand. A 6-token slice (real Q and K from the trained model), run through the exact pipeline:

CLS row raw Q·Kᵀ :  -2.69  -1.27   3.96   8.35  -0.96   1.90
  ÷ √32           →  -0.48  -0.22   0.70   1.48  -0.17   0.34   (figure 4)
  softmax         →   0.06   0.08   0.20   0.44   0.08   0.14   (figure 5, Σ = 1.00)

The raw scores (figure 1–3) → scaled (figure 4) → softmax weights (figure 5). That final grid is the attention matrix; only then is it multiplied by V (figure 6). One transparent simplification: softmax is taken over these 6 tokens for legibility — the full model normalises over all 65.

Putting it together

Every piece of the standard formula, now unpacked:

Attention(Q, K, V) = softmax( Q · Kᵀ / √d ) · V
                     ▲        ▲       ▲     ▲
                     │        │       │     └── 6: blend the values
                     │        │       └──────── 4: variance control
                     │        └──────────────── 2,3: all pairwise similarities
                     └───────────────────────── 5: turn scores into proportions

Read left to right it’s a single sentence: compare everything to everything (2, 3), keep the numbers in a sane range (4), turn them into proportions (5), and use those proportions to blend (6) — with positions added beforehand so “where” counts as part of “what” (7).