Attention & Transformers

Start here — what this is really about

Forget the formulas for a moment. Attention is one small idea: to understand any one piece of an input, look at the other pieces and decide how much each one matters.

Read this sentence: "The animal didn't cross the street because it was too tired." What does it refer to — the animal or the street? You knew instantly: the animal. Your brain looked back at the surrounding words and decided animal was the relevant one. That act — scanning the neighbours and weighting them by relevance — is exactly what an attention layer does, with numbers instead of intuition.

The problem attention solves — meaning depends on context

A word, or a patch of an image, rarely means anything on its own. The word bank is a riverbank or a place for money depending on the words around it. A patch of brown fur only means dog once you notice the ear and the leash nearby.

So a model needs a way to let each element gather information from the rest of the input before it decides what that element represents. The naive fixes are bad: reading strictly left-to-right (like older RNNs) forgets things far away, and looking only at fixed neighbours (like a small convolution) can't connect a word to another word ten positions away.

A token is what we call one such element. When every token looks at every other token within the same input, we call it self-attention — and it is the engine inside the Transformer.

Attention is a weighted average you learn

Here is the mechanical heart of it, stripped bare. Attention produces, for each token, a weighted average of other tokens' information — where the weights say how much to listen to each one.

Two pieces make this work, and the rest of the page just builds them out:

1. A way to score how relevant each token is to the one doing the looking.
2. A way to turn those scores into weights that sum to 1, then take the weighted average.

The score in step 1 is a dot product (the same operation from Linear Algebra): two vectors that point the same way score high, so similar tokens pay more attention to each other.

Queries, Keys, and Values — three roles for every token

To score relevance and then fetch information, each token plays three different roles. The model learns to project every token vector into three smaller vectors:

• The Query $\mathbf{q}$ — what am I looking for?
• The Key $\mathbf{k}$ — what do I contain, so others can find me?
• The Value $\mathbf{v}$ — what do I actually hand over if you pick me?

Scaled dot-product attention — the whole formula

Now stack every token's query into a matrix $Q$, every key into $K$, every value into $V$, and do all the lookups at once. This single equation is the core of every Transformer:

Decoding every symbol, ELI5 — and why each part is there:

• $Q$, $K$, $V$ — the stacks of all queries, all keys, all values (one row per token). Bundling them into matrices lets us score every token against every token in one matrix multiply.
• $QK^\top$ — every query dotted with every key. The superscript $\top$ means transpose (flip the matrix on its side) so the shapes line up for multiplication. The result is an $n \times n$ grid of relevance scores: entry $(i, j)$ is how much token $i$ cares about token $j$. This grid is the whole point — it is the attention pattern.
• $\sqrt{d_k}$ — the square root of $d_k$, the length of each key vector. Dividing by it keeps the scores from getting huge. Why it matters: without this, big vectors produce big dot products, the softmax saturates, and gradients vanish so the layer stops learning (full reason in the Depth panel below).
• $\text{softmax}(\cdot)$ — squashes each row of scores into positive weights that sum to 1, so every row becomes a proper weighted-average recipe.
• multiply by $V$ — take that weighted average of the value vectors. Why it matters: this is the actual output — each token's new representation, blended from the tokens it decided were relevant.

Softmax is the step that turns raw scores into a clean weighting:

Decoding it, ELI5:

• $\mathbf{z}$ — one row of raw scores (for one token, its scores against all tokens).
• $e^{z_i}$ — $e$ (about 2.718) raised to each score. Exponentiating makes everything positive and stretches gaps, so a slightly higher score wins a lot more weight.
• $\sum_j e^{z_j}$ — the $\Sigma$ ("sum") of all those exponentials; dividing by it forces the whole row to add up to exactly 1.
• Why it matters: it converts "how relevant" into "what fraction of my attention," and it's smooth (differentiable), so gradient descent can train through it.

Let's actually run it. This cell implements scaled dot-product self-attention in pure NumPy on a tiny 4-token sequence — no deep-learning library needed — and confirms each attention row sums to 1.

Multi-head attention — several viewpoints at once

One set of Q/K/V projections can only track one kind of relationship at a time. Language and images have many at once: grammar, meaning, position, object parts. So we run several attention operations in parallel, each with its own projections, and combine them.

Multi-head attention gives each head its own small slice of the representation, so heads can specialize instead of competing.

Decoding the symbols, ELI5:

• $\text{head}_i$ — the output of the $i$-th attention operation (its own $\text{softmax}(QK^\top/\sqrt{d_k})V$ using head $i$'s own projections).
• $h$ — how many heads run in parallel (e.g. 8 or 12). More heads = more relationship types tracked at once.
• $\text{Concat}(\cdots)$ — stick the heads' output vectors side by side into one long vector.
• $W^O$ — a final learned matrix that mixes the concatenated heads back into the model's dimension. Why it matters: without this mix step the heads would stay in separate lanes; $W^O$ lets the layer combine what different heads found.

Order matters — positional encoding

Here is a subtle trap. Look again at $\text{softmax}(QK^\top/\sqrt{d_k})V$: it is a sum of weighted values. If you shuffle the tokens, you shuffle the rows and columns but compute the exact same set of weighted sums — the output for each token is unchanged. Attention is permutation-invariant.

The fix is a positional encoding: give every position its own distinctive pattern of numbers and add it to the token's vector before attention. The original Transformer used fixed sine and cosine waves of different frequencies.

The Transformer block — attention plus the plumbing

Attention is the star, but a working Transformer block wraps it with three supporting pieces, then stacks that block many times:

x = x + MultiHeadAttention(LayerNorm(x))   # look around, add the result back
x = x + MLP(LayerNorm(x))                   # think per-token, add the result back

The three helpers each earn their place:

• The x = x + ... pattern is a residual connection (from CNNs): each sub-layer only learns a correction to add, which keeps gradients healthy through dozens of layers.
• LayerNorm steadies the numbers going into each sub-layer.
• The MLP lets each token privately process what attention just gathered.

Here is a compact self-attention module in real PyTorch. It needs a GPU-class library, so it opens in Colab rather than running in your browser — but notice the shape is exactly the NumPy version above, just batched and multi-head.

import torch, torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, dim, heads=8):
        super().__init__()
        self.h, self.dk = heads, dim // heads
        self.qkv = nn.Linear(dim, dim * 3)     # builds Q, K, V in one matmul
        self.proj = nn.Linear(dim, dim)        # the W_O mixing step

    def forward(self, x):                       # x: (B, N, D)
        B, N, D = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]        # each (B, h, N, dk)
        att = (q @ k.transpose(-2, -1)) / self.dk ** 0.5   # scaled scores
        att = att.softmax(dim=-1)               # weights sum to 1 per row
        out = (att @ v).transpose(1, 2).reshape(B, N, D)   # weighted average of values
        return self.proj(out)

Why it matters for vision

The Vision Transformer (ViT) makes one clever leap: chop the image into small square patches, flatten each patch into a vector (a token), add positional encodings, and feed the sequence into an ordinary Transformer. With enough training data it matches or beats CNNs, and because it speaks the same "sequence of tokens" language as text models, it unlocks unified multimodal systems like CLIP.

The catch is cost. Building the full $n\times n$ score grid means attention has quadratic complexity in the number of tokens.

• $n$ — the number of tokens (patches or words).
• Why it matters: doubling the sequence roughly quadruples the compute and memory. This is the bottleneck for long documents and high-resolution images, and the reason for efficient variants — windowed attention (Swin), sparse and linear attention, and FlashAttention's memory-savvy implementation.

Practice — and how to make it stick

Start hands-on. This lab gives you a working self-attention in NumPy and asks you to make the pattern visible — edit the code, run it (Ctrl/Cmd+Enter), and read the heatmap.

More to try, by hand and in code:

1. By hand: given query $[2, 0]$ and keys $[2,0], [0,3], [1,1]$, compute the three dot-product scores, softmax them, and check the weights sum to 1.
2. From scratch: add real learned projections $W^Q, W^K, W^V$ (random matrices) to the NumPy cell and confirm the output shape is unchanged.
3. Multi-head: split the 4-D vectors into two 2-D heads, run attention in each, concatenate, and compare the pattern to the single-head version.
4. Build a tiny ViT: patchify CIFAR-10, add a class token plus positional embeddings, and train a few Transformer blocks.

Continue to Vision Transformers to see attention applied to images end-to-end, then Foundation Models for where this architecture leads.