The Transformer, from the dot product up

1 Why drop recurrence? "Attention is all you need"

In Session 14 we used a recurrent neural network (RNN) — and its gated cousins the LSTM and GRU — to process sequences. An RNN reads a sequence x₁, x₂, …, xₙ one step at a time, carrying a hidden state hₜ forward:

hₜ = f(hₜ₋₁, xₜ)   # step t depends on step t-1

That little recurrence — "step t needs step t−1" — is elegant, but it causes two serious problems for long sequences.

Problem 1 — recurrence is inherently sequential (slow)

Because hₜ depends on hₜ₋₁, you cannot compute step 5 before step 4. The work forms a chain of length n, so even on a GPU with thousands of cores, an RNN's forward pass takes O(n) sequential steps — the hardware sits idle waiting. Modern GPUs love doing thousands of independent multiplications at once; recurrence refuses to give them that.

Problem 2 — long-range dependencies get forgotten

To connect word 1 with word 100, the signal must survive 99 hops through the hidden state. Each hop multiplies gradients by a factor; do that 99 times and the gradient either shrinks to nothing (vanishing gradient) or blows up (exploding gradient). LSTMs soften this with gates, but the fundamental issue remains: the path length between two positions grows with their distance.

Read that table carefully — it is the entire argument for the Transformer. Self-attention makes the path between any two tokens length 1 (every token looks directly at every other), and it does the work as one big matrix multiply with no sequential dependency. The price is the O(n²) term: every token attends to every other, so cost grows with the square of the sequence length. For typical sentence lengths that trade is a bargain.

2 Self-attention & Q/K/V

Self-attention takes a sequence of token vectors and produces a new sequence of the same length, where each output vector is a weighted blend of all the input vectors. The weights are learned from the content itself. Here is the machinery, built up piece by piece.

Step 1 — three roles per token: Query, Key, Value

Start with the input as a matrix X of shape (n, d_model): n tokens, each a d_model-dimensional vector (the embedding from the previous layer). From X we create three new matrices by multiplying with three learned weight matrices:

Q = X · WQ      # queries  — "what am I looking for?"
K = X · WK      # keys     — "what do I offer / contain?"
V = X · WV      # values   — "what information do I pass on?"

• Query (Q) — each token's "search request." Shape (n, d_k).
• Key (K) — each token's "advertised label." Shape (n, d_k).
• Value (V) — each token's "payload" to be mixed into others. Shape (n, d_v).

The matrices WQ, WK, WV are the only learned parameters here — training decides what each token should ask for, advertise, and carry.

Step 2 — scores: how well does each query match each key?

We measure match with a dot product. The dot product of two vectors is large when they point the same way (similar) and small/negative when they don't. Compute every query against every key in one shot:

scores = Q · Kᵀ     # shape (n, n): scores[i][j] = how much token i attends to token j

Step 3 — scale by √d_k

Here's a subtle but important detail. If d_k (the dimension of each query/key) is large, the dot products grow large in magnitude too — roughly proportional to d_k. Large inputs push the upcoming softmax into a region where its gradients are tiny (it saturates), which stalls learning. The fix is to divide by √d_k to keep the variance around 1:

scaled = (Q · Kᵀ) / √d_k

Step 4 — softmax into attention weights

Apply softmax along each row so each query's scores become a probability distribution over the keys (all positive, each row sums to 1). Row i now says "how much should token i draw from each other token."

Step 5 — weighted sum of values

Multiply those weights by V to blend the payloads. Putting all five steps together gives the celebrated scaled dot-product attention formula:

Q = [[1, 0],     K = [[1, 0],     V = [[10, 0],
     [0, 1],          [0, 1],          [ 0,10],
     [1, 1]]          [1, 1]]          [ 5, 5]]

q₁·k₁ = (1)(1)+(0)(0) = 1
q₁·k₂ = (1)(0)+(0)(1) = 0
q₁·k₃ = (1)(1)+(0)(1) = 1
scores = [1, 0, 1]

weights ≈ [0.401, 0.198, 0.401]   # sums to 1.0 ✓

out₁ = 0.401·[10,0] + 0.198·[0,10] + 0.401·[5,5]
     = [4.01, 0] + [0, 1.98] + [2.005, 2.005]
     ≈ [6.02, 3.98]

3 Multi-head attention

A single attention computation can only learn one way of relating tokens. But language has many simultaneous relationships (subject–verb agreement, coreference, adjacency, semantic similarity…). Multi-head attention runs several attention functions in parallel, each in its own learned subspace, then combines the results.

The recipe

Pick h heads (the original paper used h = 8). Split the model dimension among them: d_k = d_v = d_model / h. Give each head i its own projection matrices WQᵢ, WKᵢ, WVᵢ:

headᵢ = Attention(X·WQᵢ, X·WKᵢ, X·WVᵢ)     # each head: shape (n, d_v)

MultiHead(X) = Concat(head₁, …, headₕ) · Wᴼ   # concat then project back to d_model

Two moves to notice:

• Concatenate the h head outputs side by side. Since each head produces d_v = d_model/h features, the concatenation is back to width d_model.
• Project with a final learned matrix Wᴼ of shape (d_model, d_model), which lets the heads' outputs mix into one coherent vector.

4 Positional encoding

Here's the catch with attention: it is permutation-equivariant — shuffle the input tokens and the outputs just shuffle the same way. The formula softmax(QKᵀ/√d_k)V contains no notion of position; token 3 and token 30 are treated identically. But word order carries meaning, so we must inject it.

The fix — add a position signal to each embedding

Before the first attention layer, we add a positional encoding vector PE(pos) to each token's embedding:

inputₜ = Embedding(tokenₜ) + PE(pos = t)

Both have width d_model, so they add elementwise. Now each vector carries "what I mean" + "where I sit."

Sinusoidal positional encoding (the original choice)

Vaswani et al. used fixed sine and cosine waves of different frequencies. For position pos and dimension index i (out of d_model):

PE(pos, 2i)   = sin( pos / 10000^(2i/d_model) )
PE(pos, 2i+1) = cos( pos / 10000^(2i/d_model) )

So even dimensions get sines, odd dimensions get cosines, and the wavelength grows as you move to higher dimensions: low dimensions wiggle fast (encode fine position), high dimensions wiggle slowly (encode coarse position). It's like representing a number in many "bases" at once — the same trick as the hands of a clock (seconds, minutes, hours) telling you position at multiple scales.

5 The full Transformer block

We now assemble the parts. A Transformer has an encoder stack and a decoder stack, each made of N identical blocks (the paper used N = 6). Each block has two or three sublayers, and every sublayer is wrapped in the same two helpers.

The two wrappers: residual connection + LayerNorm

• Residual connection — add the sublayer's input back to its output: x + Sublayer(x). This gives gradients a direct highway back through the network (recall the same skip-connection idea from ResNets), so very deep stacks still train. The sublayer only has to learn a correction to its input, not rebuild it.
• Layer normalisation (LayerNorm) — renormalise each token's vector to mean 0, variance 1 (then scale/shift with learned parameters). It keeps activations stable across the deep stack. So every sublayer computes LayerNorm(x + Sublayer(x)).

The feed-forward sublayer

After attention mixes information across tokens, each block has a position-wise feed-forward network (FFN) that processes each token independently: two linear layers with a non-linearity (ReLU/GELU) between, typically expanding to d_ff = 2048 then back to d_model = 512:

FFN(x) = max(0, x·W₁ + b₁) · W₂ + b₂     # same weights applied to every position

Encoder block vs decoder block

Masked attention — no peeking at the future

The decoder generates output one token at a time, so when predicting token t it must not see tokens t+1, t+2, … (they don't exist yet at inference, and letting it cheat at training would be useless). We enforce this with a causal mask: before the softmax, set the scores for all future positions to −∞, so softmax drives their weights to 0.

scores = (Q·Kᵀ)/√d_k
scores = scores + mask        # mask[i][j] = -inf if j > i, else 0
weights = softmax(scores)      # future positions now get weight 0

The whole pipeline, end to end

The encoder reads the whole input and builds rich representations. The decoder generates the output autoregressively: its masked self-attention looks at what it has produced so far, its cross-attention pulls relevant information from the encoder, the FFN refines, and a final linear layer + softmax gives a probability over the vocabulary for the next token.

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    # Q, K, V shapes: (batch, heads, seq_len, d_k)
    d_k = Q.size(-1)

    # 1) scores: (batch, heads, seq_len, seq_len)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / d_k ** 0.5

    # 2) causal / padding mask: set masked positions to -inf
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))

    # 3) softmax over the last dim (keys) → attention weights
    weights = F.softmax(scores, dim=-1)

    # 4) weighted sum of values
    output = torch.matmul(weights, V)          # (batch, heads, seq_len, d_k)
    return output, weights


class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, h):
        super().__init__()
        assert d_model % h == 0
        self.h, self.d_k = h, d_model // h
        self.W_q = torch.nn.Linear(d_model, d_model)
        self.W_k = torch.nn.Linear(d_model, d_model)
        self.W_v = torch.nn.Linear(d_model, d_model)
        self.W_o = torch.nn.Linear(d_model, d_model)   # final projection

    def forward(self, x, mask=None):
        B, n, _ = x.shape
        # project, then split into h heads → (B, h, n, d_k)
        split = lambda t: t.view(B, n, self.h, self.d_k).transpose(1, 2)
        Q, K, V = split(self.W_q(x)), split(self.W_k(x)), split(self.W_v(x))

        out, _ = scaled_dot_product_attention(Q, K, V, mask)

        # concat heads back to (B, n, d_model), then final projection
        out = out.transpose(1, 2).contiguous().view(B, n, -1)
        return self.W_o(out)

★ Putting it all together

We replaced the slow, forgetful recurrence of Session 14 with pure attention — and built the architecture that powers modern AI. Here's the one-paragraph story tying the five topics together:

Quick self-check

📚 References & Further Reading

Class material

• SST Deep Learning handout (Session 15) — your course handout for this session, covering the Transformer architecture, self-attention, and multi-head attention.

Papers, docs & deep dives

• Attention Is All You Need (Vaswani et al., 2017) — the original paper that introduced the Transformer; read it to see the exact formulas and architecture firsthand.
• The Illustrated Transformer (Jay Alammar) — the best visual, step-by-step walkthrough of Q/K/V, multi-head attention, and the full stack.
• The Annotated Transformer (Harvard NLP) — the paper re-implemented line-by-line in runnable PyTorch; perfect for connecting the math to real code.
• GenAI Session 1 — Introduction to GenAI — our sibling course's intuitive, math-free treatment of attention and transformers; a gentle complement to today's deeper dive.