Deep Learning,
Week 9

Every network so far took numbers as input. Words are not numbers, so the first step of any language model is to turn each word (or sub-word token) into a vector. The way this is done decides what the rest of the network can know about meaning.

1.1One-hot vectors know nothing

The simplest code is one-hot: a vector as long as the vocabulary, with a 1 in the word's position and 0 elsewhere. Every pair of one-hot vectors is orthogonal, so "king" is as far from "queen" as from "apple", and the vectors are enormous. An embedding instead gives each word a short, dense, learned vector. The embedding layer is nothing more than a lookup table, which is exactly a one-hot vector multiplied by a weight matrix:

import torch, torch.nn as nn
torch.manual_seed(0)
vocab = ["king", "queen", "man", "woman", "apple"]
idx = {w: i for i, w in enumerate(vocab)}
onehot = torch.eye(len(vocab))
print("one-hot dot products (king.queen, king.apple):", (onehot[idx["king"]] @ onehot[idx["queen"]]).item(), (onehot[idx["king"]] @ onehot[idx["apple"]]).item(), " <- every pair is orthogonal: no notion of similarity")

emb = nn.Embedding(len(vocab), 4)
ids = torch.tensor([idx["king"], idx["apple"]])
lookup = emb(ids)                                           # the embedding layer is a table lookup ...
matmul = onehot[ids] @ emb.weight                           # ... which equals one-hot times the weight matrix
print("embedding lookup equals one-hot @ weight matrix:", torch.allclose(lookup, matmul))
V, d = 50_000, 300
print(f"a vocabulary of {V:,} words with {d}-dimensional embeddings has {V * d:,} parameters; a one-hot vector would have {V:,} dimensions")

one-hot dot products (king.queen, king.apple): 0.0 0.0  <- every pair is orthogonal: no notion of similarity
embedding lookup equals one-hot @ weight matrix: True
a vocabulary of 50,000 words with 300-dimensional embeddings has 15,000,000 parameters; a one-hot vector would have 50,000 dimensions

The dot product of one-hot king and queen is 0, as for king and apple: no similarity is encoded. The lookup equals the one-hot product (True), and a 50,000-word vocabulary with 300-dimensional embeddings has 15 million parameters, a table that training adjusts so that useful words end up close together.

1.2Static embeddings: Word2Vec and GloVe

Word2Vec (skip-gram / CBOW) learns embeddings by predicting a word from its neighbours (or neighbours from a word); GloVe factorises a global co-occurrence matrix. Both rest on one idea: "you shall know a word by the company it keeps". The result is one fixed vector per word. To see the mechanism with nothing hidden, build embeddings the classical way: count which words appear near each other in a tiny templated corpus, convert counts to PPMI (positive pointwise mutual information), and compress with an SVD to 8 dimensions.

import numpy as np
from collections import Counter

templates = [("the {} wore a crown and he ruled the kingdom", ["king", "prince"]),
             ("the {} wore a crown and she ruled the kingdom", ["queen", "princess"]),
             ("the {} went to the market and he sold bread", ["man", "boy"]),
             ("the {} went to the market and she sold bread", ["woman", "girl"]),
             ("we sat on the {} of the river near the water", ["bank"]),
             ("she kept her money at the {} to get a loan", ["bank"])]
corpus = [t.format(w).split() for t, ws in templates for w in ws for _ in range(20)]
vocab = sorted({w for s in corpus for w in s}); idx = {w: i for i, w in enumerate(vocab)}
C = np.zeros((len(vocab), len(vocab)))
for s in corpus:                                                   # co-occurrence counts within a window of 6 words either side
    for i, w in enumerate(s):
        for j in range(max(0, i - 6), min(len(s), i + 7)):
            if i != j: C[idx[w], idx[s[j]]] += 1
total = C.sum(); pw = C.sum(1, keepdims=True) / total; pc = C.sum(0, keepdims=True) / total
with np.errstate(divide="ignore"):
    ppmi = np.maximum(np.log(C / total / (pw * pc)), 0)            # positive pointwise mutual information
U, S, _ = np.linalg.svd(ppmi)
k = 8
E = U[:, :k] * np.sqrt(S[:k]); E /= np.linalg.norm(E, axis=1, keepdims=True)         # 8-dimensional word vectors, unit length
stop = {"the", "a", "of", "at", "to", "on", "we", "she", "he", "and", "her", "near", "get", "kept", "sat", "went"}
sim = lambda a, b: float(E[idx[a]] @ E[idx[b]])
nearest = lambda v, n, skip=(): [vocab[i] for i in np.argsort(-(E @ v)) if vocab[i] not in skip and vocab[i] not in stop][:n]

print("cosine similarities:  king~prince", round(sim("king", "prince"), 3), " king~queen", round(sim("king", "queen"), 3), " king~woman", round(sim("king", "woman"), 3), " king~bread", round(sim("king", "bread"), 3))
q = E[idx["king"]] - E[idx["man"]] + E[idx["woman"]]
print("king - man + woman is closest to:", nearest(q, 3, skip=("king", "man", "woman")))
print("\nthe polysemy problem: ONE vector for 'bank' although the text uses it in two senses")
print("  nearest words to 'bank':", nearest(E[idx["bank"]], 8, skip=("bank",)))

cosine similarities:  king~prince 1.0  king~queen 0.997  king~woman 0.374  king~bread 0.316
king - man + woman is closest to: ['queen', 'princess', 'kingdom']

the polysemy problem: ONE vector for 'bank' although the text uses it in two senses
  nearest words to 'bank': ['river', 'money', 'water', 'loan', 'market', 'crown', 'ruled', 'queen']

Read it carefully. The corpus is so regular that king and prince have identical contexts (similarity 1.000), and king and queen are 0.997 similar: the shared "royal" context dominates, and gender is a small direction. Yet that small direction is consistent, so the arithmetic king $-$ man $+$ woman lands closest to queen (then princess), the famous analogy, produced by nothing but counting. (A real Word2Vec on billions of words works the same way and finds many such regularities, with plenty of exceptions: the analogy trick is a striking property, not a guarantee.)

1.3Contextual embeddings

The remedy is to make a word's vector depend on the sentence it appears in. Contextual embeddings (ELMo, BERT, GPT) do exactly that: the vector for "bank" is computed fresh from the whole sentence. The mechanism that lets one token's vector absorb its neighbours is attention, the subject of the rest of this week. A first taste (one untrained attention layer, so the numbers are arbitrary but the point is not): the same input vector for "bank" produces different outputs in the two sentences.

import torch
torch.manual_seed(0)
d = 6
vocab = ["the", "river", "money", "bank", "was", "deep", "closed"]
E = {w: torch.randn(d) for w in vocab}
Wq, Wk, Wv = (torch.randn(d, d) * 0.7 for _ in range(3))

def contextual(sentence):                                                      # one untrained self-attention layer
    x = torch.stack([E[w] for w in sentence])
    return torch.softmax((x @ Wq) @ (x @ Wk).T / d ** 0.5, -1) @ (x @ Wv)

s1, s2 = "the river bank was deep".split(), "the money bank was closed".split()
v1, v2 = contextual(s1)[s1.index("bank")], contextual(s2)[s2.index("bank")]
print("static embedding of 'bank' is one fixed vector:", E["bank"].numpy().round(2))
print("self-attention output for 'bank' in the river sentence:", v1.numpy().round(2))
print("self-attention output for 'bank' in the money sentence:", v2.numpy().round(2))
print("distance between the two contextual vectors:", round((v1 - v2).norm().item(), 3), "(0 for a static embedding)")
print("same word, same input vector; the output now depends on the neighbours, which is the idea behind contextual embeddings")

static embedding of 'bank' is one fixed vector: [-0.89 -1.51  0.37  1.46  0.94  0.77]
self-attention output for 'bank' in the river sentence: [ 0.12  0.04  0.07 -0.15  0.2   1.09]
self-attention output for 'bank' in the money sentence: [ 0.22  1.62 -0.44  0.97  0.03 -0.92]
distance between the two contextual vectors: 2.841 (0 for a static embedding)
same word, same input vector; the output now depends on the neighbours, which is the idea behind contextual embeddings

The static embedding is a single fixed vector; after one attention layer the two "bank" vectors are 2.84 apart. Training would arrange for the river sentence to move "bank" toward the water sense and the money sentence toward the finance sense.

1.4Chapter summary

• One-hot: orthogonal, huge, no similarity. Embedding: a learned lookup table (= one-hot times a weight matrix), 50,000 x 300 = 15M parameters.
• Static (Word2Vec, GloVe): one vector per word; analogies like king $-$ man $+$ woman $\approx$ queen emerge from co-occurrence counts, but polysemy is lost.
• Contextual (ELMo, BERT, GPT): the vector depends on the sentence, which attention makes possible.

Week 8 ended on a measured problem: an encoder-decoder squeezes the whole input into one context vector, and copy accuracy fell from 0.997 to 0.15 as the input grew from 4 to 32 symbols. The fix, attention (Bahdanau et al., 2014), is one idea: do not throw the encoder's states away; let the decoder look back at all of them, and choose where to look at each step.

2.1The human analogy

A person translating "John loves Mary" into Hindi does not memorise the sentence and then write; while writing "John" they look at "John", while writing "loves" they look at "loves". Attention is selective focus: for each output, a weighting over the input saying how much each part matters right now.

2.2Queries, keys and values

The cleanest way to think about it is a soft dictionary lookup, in the notes' cookbook analogy. The query is the recipe you are looking for ("what am I looking for?"), the keys are the recipe titles ("what can I offer for matching?") and the values are the recipes themselves ("what do I actually carry?"). A normal dictionary returns the value of the one key that matches exactly. Attention returns a blend: it scores the query against every key, turns the scores into weights that sum to 1 (softmax), and returns the weighted average of the values. The recipe is:

• Score the query against every key (a dot product: large when they point the same way).
• Softmax the scores into weights: non-negative and summing to 1.
• Average the values with those weights.

2.3The fix, measured

Give the Week 8 copy task an attention decoder: at every step the decoder's state is the query, the encoder states are both keys and values, and the context it uses is the weighted average. Everything else (task, hidden size 32, training length, optimiser) is identical to Week 8, so the comparison is clean:

import math, torch, torch.nn as nn
torch.set_num_threads(1)
V, BOS = 12, 10

class AttnSeq2Seq(nn.Module):
    def __init__(self, hidden):
        super().__init__()
        self.emb = nn.Embedding(V, 32)
        self.encoder = nn.GRU(32, hidden, batch_first=True)
        self.cell = nn.GRUCell(32 + hidden, hidden)
        self.out = nn.Linear(2 * hidden, V)
    def attend(self, query, enc):                                     # query: (B, h)   enc: (B, T, h) = the keys AND the values
        scores = torch.bmm(enc, query.unsqueeze(2)).squeeze(2) / math.sqrt(enc.size(2))
        weights = torch.softmax(scores, dim=1)                        # (B, T): where to look in the input
        return torch.bmm(weights.unsqueeze(1), enc).squeeze(1)        # weighted average of ALL encoder states
    def step(self, tok, state, enc):
        state = self.cell(torch.cat([self.emb(tok), self.attend(state, enc)], 1), state)
        return self.out(torch.cat([state, self.attend(state, enc)], 1)), state
    def forward(self, src, tgt_in):                                   # teacher forcing
        enc, last = self.encoder(self.emb(src)); state, outs = last[0], []
        for t in range(tgt_in.size(1)):
            o, state = self.step(tgt_in[:, t], state, enc); outs.append(o)
        return torch.stack(outs, 1)
    @torch.no_grad()
    def generate(self, src, length):                                  # free running
        enc, last = self.encoder(self.emb(src)); state = last[0]; tok = torch.full((src.size(0),), BOS); res = []
        for _ in range(length):
            o, state = self.step(tok, state, enc); tok = o.argmax(-1); res.append(tok)
        return torch.stack(res, 1)

def evaluate(length, hidden=32, seed=0, steps=2500):
    torch.manual_seed(seed); model = AttnSeq2Seq(hidden); opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    batch = lambda B: (lambda s: (s, torch.cat([torch.full((B, 1), BOS), s[:, :-1]], 1)))(torch.randint(0, 10, (B, length)))
    for _ in range(steps):
        src, tgt_in = batch(64)
        loss = nn.functional.cross_entropy(model(src, tgt_in).reshape(-1, V), src.reshape(-1))
        opt.zero_grad(); loss.backward(); nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
    src, tgt_in = batch(1000)
    with torch.no_grad():
        forced = (model(src, tgt_in).argmax(-1) == src).float().mean().item(); free = model.generate(src, length)
    return forced, (free == src).float().mean().item(), (free == src).all(1).float().mean().item()

print("copy task WITH attention, hidden size 32 (the same task, size and training as Week 8's plain encoder-decoder)")
print("sequence length   teacher-forced acc   free-running acc   whole sequence exact")
for length in (4, 8, 16, 32):
    a, b, c = evaluate(length)
    print(f"{length:10d}      {a:14.3f}      {b:14.3f}      {c:14.3f}")

copy task WITH attention, hidden size 32 (the same task, size and training as Week 8's plain encoder-decoder)
sequence length   teacher-forced acc   free-running acc   whole sequence exact
         4               1.000               1.000               1.000
         8               1.000               1.000               1.000
        16               0.999               0.995               0.982
        32               0.988               0.894               0.738

Free-running accuracy at length 16 goes from 0.30 to 0.995, and at length 32 from 0.15 to 0.894; lengths 4 and 8 are perfect. No extra capacity was added to the hidden state: the decoder simply stopped depending on a single vector. (Length 32 is still not perfect, 0.894 free-running with 74% of whole sequences exact, so the model is not magic; this is a small network after a short training run.) This is the entire motivation for the architecture of the next chapters.

2.4Attention pooling

The weighted average of values is sometimes called attention pooling. Because the weights are non-negative and sum to 1, the output is a convex combination of the values: it always lies inside the range of the values, and the weights are a probability distribution over the positions, readable as "where the model looked". That readability is a large part of attention's appeal.

2.5Chapter summary

• Attention lets the decoder look at all encoder states: query vs keys $\to$ softmax weights $\to$ weighted average of values.
• Measured on the Week 8 copy task: length 16 accuracy 0.30 $\to$ 0.995, length 32 0.15 $\to$ 0.894 with no bigger hidden state.
• The output is a convex combination of the values; the weights are a distribution over positions.

The attention used in Transformers is a specific, simple recipe called scaled dot-product attention. It is the formula to know cold for the end term.

• $Q\in\mathbb{R}^{n_q\times d_k}$: one query row per position that is asking.
• $K\in\mathbb{R}^{n_k\times d_k}$: one key row per position that can be attended to (same width $d_k$ as the queries, so a dot product exists).
• $V\in\mathbb{R}^{n_k\times d_v}$: one value row per key (any width $d_v$).
• $QK^\top\in\mathbb{R}^{n_q\times n_k}$: every query against every key; softmax is applied row by row; the result times $V$ has shape $n_q\times d_v$.

3.1A worked example

The notes' example: $Q=[1,0]$; keys $[1,0],[0,1],[1,1]$; values 10, 20, 30; $d_k=2$. The scores are $q\cdot k/\sqrt2=[0.707,\,0,\,0.707]$ and the softmax gives weights $e^{0.707}/(2e^{0.707}+1)=0.401$, $1/(2e^{0.707}+1)=0.198$, 0.401:

import numpy as np
Q = np.array([[1.0, 0.0]])
K = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
V = np.array([[10.0], [20.0], [30.0]])
d_k = K.shape[1]
scores = Q @ K.T / np.sqrt(d_k)
weights = np.exp(scores) / np.exp(scores).sum()
print("scores  q.k / sqrt(2):", scores.round(3).tolist())
print("softmax weights      :", weights.round(3).tolist(), " sum =", weights.sum().round(3))
print("output = weights @ V :", (weights @ V).round(2).tolist())
print("the query [1,0] matches keys 1 and 3 equally, so they get equal weight 0.401; key 2 (orthogonal) gets 0.198")

scores  q.k / sqrt(2): [[0.707, 0.0, 0.707]]
softmax weights      : [[0.401, 0.198, 0.401]]  sum = 1.0
output = weights @ V : [[20.0]]
the query [1,0] matches keys 1 and 3 equally, so they get equal weight 0.401; key 2 (orthogonal) gets 0.198

The output is $0.401\times10+0.198\times20+0.401\times30=\mathbf{20.0}$, matching the notes. Keys 1 and 3 each overlap the query (a dot product of 1) and get equal weight; key 2 is orthogonal and gets less, but not zero: softmax never gives exactly zero.

3.2The same in code, and its properties

import torch, torch.nn.functional as F
torch.manual_seed(0)
n_q, n_k, d_k, d_v = 3, 5, 4, 6
Q, K, V = torch.randn(n_q, d_k), torch.randn(n_k, d_k), torch.randn(n_k, d_v)

scores = Q @ K.T / d_k ** 0.5                       # (n_q, n_k): how well each query matches each key
weights = torch.softmax(scores, dim=-1)             # each row: a probability distribution over the keys
out = weights @ V                                   # (n_q, d_v): a weighted average of the values
print("shapes: Q", tuple(Q.shape), "K", tuple(K.shape), "V", tuple(V.shape), "-> scores", tuple(scores.shape), "-> output", tuple(out.shape))
print("matches F.scaled_dot_product_attention:", torch.allclose(out, F.scaled_dot_product_attention(Q, K, V), atol=1e-6))
print("each row of weights sums to 1:", weights.sum(-1).tolist())
print("weights are non-negative:", bool((weights >= 0).all()))
inside = bool(((out >= V.min(0).values - 1e-6) & (out <= V.max(0).values + 1e-6)).all())
print("every output lies between the smallest and largest value (a convex combination of the rows of V):", inside)
print("changing the value of a key with weight ~0 barely changes the output; the attention weights say WHERE the model looked")

shapes: Q (3, 4) K (5, 4) V (5, 6) -> scores (3, 5) -> output (3, 6)
matches F.scaled_dot_product_attention: True
each row of weights sums to 1: [1.0, 1.0, 1.0]
weights are non-negative: True
every output lies between the smallest and largest value (a convex combination of the rows of V): True
changing the value of a key with weight ~0 barely changes the output; the attention weights say WHERE the model looked

Four properties to quote: (1) each row of weights sums to 1; (2) all weights are non-negative; (3) so each output is a convex combination of the value rows, lying within their range; (4) the whole computation is two matrix multiplications and a softmax, so it parallelises completely over positions, unlike an RNN which must go step by step.

3.3Why divide by $\sqrt{d_k}$?

If the entries of $q$ and $k$ are independent with mean 0 and variance 1, then $q\cdot k=\sum_{i=1}^{d_k}q_ik_i$ is a sum of $d_k$ independent terms with variance 1, so it has variance $d_k$ and standard deviation $\sqrt{d_k}$. Without the division the scores grow with the key dimension, and softmax of large numbers is nearly one-hot: one weight near 1, the rest near 0. There the softmax is flat, its gradient is almost zero, and learning stalls. Dividing by $\sqrt{d_k}$ restores a standard deviation of 1. Measure it:

import torch
torch.manual_seed(0)
n_keys, trials = 10, 20000
print("   d_k   std(q.k)   std(q.k / sqrt(d_k))   mean max-weight: unscaled   scaled   |   softmax gradient norm: unscaled   scaled")
for d in (4, 16, 64, 256, 1024):
    q = torch.randn(trials, 1, d); k = torch.randn(trials, n_keys, d)          # entries with mean 0 and variance 1
    raw = (q @ k.transpose(1, 2)).squeeze(1); scaled = raw / d ** 0.5
    stats = []
    for s in (raw, scaled):
        s = s.clone().requires_grad_()
        w = torch.softmax(s, -1)
        (w[:, 0]).sum().backward()                                              # how strongly does the first logit move the first weight?
        stats.append((w.max(-1).values.mean().item(), s.grad.norm(dim=-1).mean().item()))
    print(f"{d:6d}   {raw.std().item():7.2f}   {scaled.std().item():12.2f}          {stats[0][0]:21.2f}   {stats[1][0]:6.2f}   |   {stats[0][1]:25.4f}   {stats[1][1]:6.4f}")
print("a dot product of d independent terms has variance d, so without the division the softmax saturates (one weight near 1) and its gradient vanishes")

   d_k   std(q.k)   std(q.k / sqrt(d_k))   mean max-weight: unscaled   scaled   |   softmax gradient norm: unscaled   scaled
     4      2.00           1.00                           0.49     0.30   |                      0.0732   0.0882
    16      4.00           1.00                           0.73     0.32   |                      0.0461   0.0884
    64      8.00           1.00                           0.87     0.32   |                      0.0251   0.0886
   256     15.98           1.00                           0.93     0.32   |                      0.0125   0.0882
  1024     31.97           1.00                           0.97     0.32   |                      0.0067   0.0887
a dot product of d independent terms has variance d, so without the division the softmax saturates (one weight near 1) and its gradient vanishes

The standard deviation of $q\cdot k$ is exactly $\sqrt{d_k}$ (2, 4, 8, 16, 32 in the second column, from $d_k=4$ to 1024); after scaling it is 1 at every size. Unscaled, the average largest weight over 10 keys climbs from 0.49 to 0.97 and the softmax gradient norm falls from 0.073 to 0.007, a 10-fold collapse; scaled, the largest weight stays at about 0.32 and the gradient stays at 0.088 whatever $d_k$ is. That is the whole reason for the $1/\sqrt{d_k}$: it keeps the softmax in its responsive range.

3.4Chapter summary

• $\mathrm{softmax}(QK^\top/\sqrt{d_k})V$: weights are a distribution over keys; output is a convex combination of values. Worked example: 20.0.
• $q\cdot k$ has variance $d_k$; dividing by $\sqrt{d_k}$ prevents saturation (largest weight 0.97 $\to$ 0.32 and gradient 0.007 $\to$ 0.088 at $d_k=1024$).
• Two matrix products and a softmax: fully parallel over positions.

The recipe of the last chapter does not say where $Q$, $K$ and $V$ come from. The answer defines the two kinds of attention and decides whether a layer can peek at the future.

4.1Self-attention and cross-attention

• Self-attention: $Q$, $K$, $V$ are all computed from the same sequence (each from the input by its own learned matrix: $Q=XW^Q$, $K=XW^K$, $V=XW^V$). Every token asks "which other tokens of this sentence matter to me?". In "The animal didn't cross the street because it was tired", self-attention lets "it" draw on "animal". Used in the Transformer encoder (and, masked, the decoder).
• Cross-attention: $Q$ comes from one sequence (the decoder's current state) while $K$ and $V$ come from another (the encoder output). "While generating this word, where in the input should I look?" This is the Week 8 seq2seq attention, and it is used in the Transformer decoder.

import torch, torch.nn as nn
torch.manual_seed(0)
d, heads = 8, 2
mha = nn.MultiheadAttention(d, heads, batch_first=True)
encoder_out = torch.randn(1, 6, d)                     # 6 source tokens
decoder_in = torch.randn(1, 4, d)                      # 4 target tokens produced so far

out, w = mha(encoder_out, encoder_out, encoder_out, average_attn_weights=True)
print("SELF-attention  (Q, K, V all from the source)  : output", tuple(out.shape), " weights", tuple(w.shape), "= (batch, 6 queries, 6 keys)")
out, w = mha(decoder_in, encoder_out, encoder_out, average_attn_weights=True)
print("CROSS-attention (Q from decoder; K, V from the encoder): output", tuple(out.shape), " weights", tuple(w.shape), "= (batch, 4 queries, 6 keys)")
print("so in cross-attention each of the 4 decoder positions chooses among the 6 encoder positions; the output has the decoder's length")

SELF-attention  (Q, K, V all from the source)  : output (1, 6, 8)  weights (1, 6, 6) = (batch, 6 queries, 6 keys)
CROSS-attention (Q from decoder; K, V from the encoder): output (1, 4, 8)  weights (1, 4, 6) = (batch, 4 queries, 6 keys)
so in cross-attention each of the 4 decoder positions chooses among the 6 encoder positions; the output has the decoder's length

The weight matrices show the difference: self-attention is $6\times6$ (each of 6 source tokens looks at all 6); cross-attention is $4\times6$ (each of 4 decoder tokens chooses among the 6 encoder tokens), and the output has the query's length, 4. The number of keys and values only has to match each other, not the number of queries.

4.2Masked attention: no peeking at the future

When a model generates text one token at a time, position $t$ must not see tokens after $t$: in training the whole target sentence is given at once, so without protection the model would just copy the next word. The fix is a causal mask: before the softmax, set the score of every future position to $-\infty$. Then $e^{-\infty}=0$, the weight is exactly zero, and the remaining weights are renormalised to sum to 1.

import torch
torch.manual_seed(0)
n, d = 5, 4
x = torch.randn(n, d)
scores = x @ x.T / d ** 0.5

causal = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)          # True above the diagonal = the FUTURE
weights = torch.softmax(scores.masked_fill(causal, float("-inf")), dim=-1)    # -inf before the softmax = weight exactly 0 after it
print("causal attention weights (row = the position that is looking, column = what it looks at):")
print(weights.numpy().round(2))
print("rows still sum to 1:", weights.sum(-1).tolist())

def attend(x, mask):
    return torch.softmax((x @ x.T / d ** 0.5).masked_fill(mask, float("-inf")), -1) @ x
x2 = x.clone(); x2[3:] = torch.randn(2, d)                                    # change the FUTURE tokens 3 and 4
print("\nwith the causal mask the outputs at positions 0, 1, 2 do not change when tokens 3, 4 change:",
      torch.allclose(attend(x, causal)[:3], attend(x2, causal)[:3]))
nomask = torch.zeros(n, n, dtype=torch.bool)
print("without the mask they do change (the model could read the answer):", not torch.allclose(attend(x, nomask)[:3], attend(x2, nomask)[:3]))

pad = torch.tensor([False, False, False, True, True])                          # padding mask: positions 3 and 4 are padding
w = torch.softmax(scores.masked_fill(pad[None, :], float("-inf")), -1)
print("padding mask: weight given to the padding positions =", w[:, 3:].sum().item())

causal attention weights (row = the position that is looking, column = what it looks at):
[[1.   0.   0.   0.   0.  ]
 [0.05 0.95 0.   0.   0.  ]
 [0.13 0.25 0.63 0.   0.  ]
 [0.01 0.01 0.09 0.89 0.  ]
 [0.01 0.01 0.11 0.3  0.58]]
rows still sum to 1: [1.0, 0.9999999403953552, 1.0, 1.0, 0.9999998807907104]

with the causal mask the outputs at positions 0, 1, 2 do not change when tokens 3, 4 change: True
without the mask they do change (the model could read the answer): True
padding mask: weight given to the padding positions = 0.0

The matrix is lower-triangular: token 1 can only attend to itself (weight 1.00), token 5 to all five. The test is the proof that matters: with the mask, changing tokens 3 and 4 leaves the outputs at positions 0, 1, 2 unchanged (True); without it they change, i.e. information leaks from the future. The notes' mistake-list is right: masking in the decoder is causal (autoregressive), not only for padding. A separate padding mask (the last lines) stops tokens attending to the filler positions added to make a batch rectangular; its weight on padding is exactly 0.0. The two masks are combined when both apply.

4.3Self-attention cannot tell the order

There is a catch hidden in the formula. Attention treats its inputs as a set: reordering the tokens just reorders the outputs. Take "dog bites man" and "man bites dog", with the same weights, and look at the output for "bites":

import torch
torch.manual_seed(0)
d = 6
emb = {w: torch.randn(d) for w in ("dog", "bites", "man")}
Wq, Wk, Wv = (torch.randn(d, d) * 0.5 for _ in range(3))

def self_attention(x):
    Q, K, V = x @ Wq, x @ Wk, x @ Wv
    return torch.softmax(Q @ K.T / d ** 0.5, -1) @ V

a = torch.stack([emb[w] for w in ("dog", "bites", "man")])                    # "dog bites man"
b = torch.stack([emb[w] for w in ("man", "bites", "dog")])                    # "man bites dog"
oa, ob = self_attention(a), self_attention(b)
print("output for 'bites' in 'dog bites man':", oa[1].numpy().round(3))
print("output for 'bites' in 'man bites dog':", ob[1].numpy().round(3))
print("identical without position information:", torch.allclose(oa[1], ob[1], atol=1e-6), " (self-attention sees a SET of words, not a sequence)")

pos = torch.randn(3, d)                                                        # any distinct vector per position will do for this test
oa2, ob2 = self_attention(a + pos), self_attention(b + pos)
print("after adding a position vector to each token: identical?", torch.allclose(oa2[1], ob2[1], atol=1e-6), "-> the two sentences are now told apart")

output for 'bites' in 'dog bites man': [-0.72   0.958  1.767  0.24   1.691  0.924]
output for 'bites' in 'man bites dog': [-0.72   0.958  1.767  0.24   1.691  0.924]
identical without position information: True  (self-attention sees a SET of words, not a sequence)
after adding a position vector to each token: identical? False -> the two sentences are now told apart

The outputs for "bites" are identical, so a Transformer built only from attention cannot distinguish the two sentences, unlike an RNN, whose order is built into the recurrence. The cure is to add positional information to the token vectors (positional encoding, Week 10); the last line shows that once each position carries its own vector, the two sentences are told apart. It is precisely why the notes say self-attention is permutation-invariant "and that is why we add positional encoding".

4.4Chapter summary

• Self-attention: $Q,K,V$ from one sequence. Cross-attention: $Q$ from the decoder, $K,V$ from the encoder; output length = number of queries.
• Causal mask: future scores $=-\infty$ before softmax, weight exactly 0 (verified: past outputs unchanged when the future changes). Padding mask is separate.
• Attention alone is order-blind (the "bites" outputs were identical): add positional encoding.

One attention pass produces one weighting per query: one way of deciding what is relevant. A word may need several at once: its grammatical subject, the previous word, a distant noun it refers to. Multi-head attention runs several attentions in parallel, each in its own smaller space, and combines them.

5.1The formula

With model width $d$ and $h$ heads each head has width $d_k=d_v=d/h$ (8 heads on $d=512$ gives 64). Each head has its own projections $W_i^Q,W_i^K,W_i^V$, so it can learn its own notion of "relevant". The head outputs are concatenated (not averaged, a point the notes stress), restoring width $h\cdot d/h=d$, and a final matrix $W^O$ mixes them. Write it by hand and compare with PyTorch:

import torch, torch.nn as nn
torch.manual_seed(0)
d, h, n = 8, 2, 5
dk = d // h                                                    # each head works in a d/h = 4 dimensional subspace
mha = nn.MultiheadAttention(d, h, batch_first=True)
x = torch.randn(1, n, d)

Wq, Wk, Wv = mha.in_proj_weight.chunk(3, 0); bq, bk, bv = mha.in_proj_bias.chunk(3, 0)
Q, K, V = x[0] @ Wq.T + bq, x[0] @ Wk.T + bk, x[0] @ Wv.T + bv                # one big projection each, then split into heads
heads = []
for i in range(h):
    sl = slice(i * dk, (i + 1) * dk)
    w = torch.softmax(Q[:, sl] @ K[:, sl].T / dk ** 0.5, -1)                  # each head has its OWN attention pattern
    heads.append(w @ V[:, sl])
concat = torch.cat(heads, dim=-1)                                              # CONCATENATE the heads (not average) ...
out = concat @ mha.out_proj.weight.T + mha.out_proj.bias                       # ... and mix them with W_O
print("hand-written multi-head attention matches nn.MultiheadAttention:", torch.allclose(out, mha(x, x, x)[0][0], atol=1e-5))

count = lambda m: sum(p.numel() for p in m.parameters())
for heads_n in (1, 2, 4, 8):
    print(f"d_model=64, heads={heads_n}: parameters {count(nn.MultiheadAttention(64, heads_n))} = 4*64*64 + 4*64 = {4 * 64 * 64 + 4 * 64}   (head size {64 // heads_n})")
print("more heads do NOT add parameters: they split the same d_model into smaller subspaces")

hand-written multi-head attention matches nn.MultiheadAttention: True
d_model=64, heads=1: parameters 16640 = 4*64*64 + 4*64 = 16640   (head size 64)
d_model=64, heads=2: parameters 16640 = 4*64*64 + 4*64 = 16640   (head size 32)
d_model=64, heads=4: parameters 16640 = 4*64*64 + 4*64 = 16640   (head size 16)
d_model=64, heads=8: parameters 16640 = 4*64*64 + 4*64 = 16640   (head size 8)
more heads do NOT add parameters: they split the same d_model into smaller subspaces

The hand-written version matches nn.MultiheadAttention. And the parameter table has the most surprising fact of the chapter: with $d=64$, going from 1 head to 8 heads does not change the parameter count (16,640 each, $=4d^2+4d$: the four matrices $W^Q,W^K,W^V,W^O$ plus biases). More heads do not add weights; they cut the same width into smaller subspaces, so each head sees a lower-dimensional view and the total cost stays about the same as one wide head.

5.2The price: quadratic cost

Every token attends to every token, so the score matrix is $n\times n$ and the work is $O(n^2d)$, quadratic in the sequence length (the notes' Example 4). Doubling the length quadruples the cost:

print("self-attention builds an n x n score matrix per head")
print("      n     score entries    memory (float32)    n^2 * d operations (d = 512)")
for n in (512, 2048, 8192, 32768, 131072):
    entries = n * n
    mem = entries * 4
    unit = ("MB", mem / 1e6) if mem < 1e9 else ("GB", mem / 1e9)
    print(f"{n:7d}   {entries:14.2e}    {unit[1]:8.1f} {unit[0]}        {n * n * 512:.2e}")
print("doubling the length quadruples the cost; this is why long context is expensive and why efficient attention is an active research area")

self-attention builds an n x n score matrix per head
      n     score entries    memory (float32)    n^2 * d operations (d = 512)
    512         2.62e+05         1.0 MB        1.34e+08
   2048         4.19e+06        16.8 MB        2.15e+09
   8192         6.71e+07       268.4 MB        3.44e+10
  32768         1.07e+09         4.3 GB        5.50e+11
 131072         1.72e+10        68.7 GB        8.80e+12
doubling the length quadruples the cost; this is why long context is expensive and why efficient attention is an active research area

At $n=512,d=512$ it is 134 million multiply-adds per layer (the notes' figure); at $n=2048$, 2.1 billion. The score matrix alone, for a single head, is 1 MB at 512 tokens and 4.3 GB at 32,768. RNNs have cost $O(nd^2)$, linear in length but sequential; attention is quadratic but parallel, which on GPUs is a very good trade for the lengths used until recently, and a real constraint for books and long videos (hence "efficient attention" research).

5.3Chapter summary

• $\mathrm{head}_i=\mathrm{Attention}(QW_i^Q,KW_i^K,VW_i^V)$; $\mathrm{MultiHead}=\mathrm{Concat}(\text{heads})W^O$; head width $d/h$; concatenated, not averaged.
• Parameters $4d^2+4d$ for any number of heads (16,640 at $d=64$ for $h=1,2,4,8$): heads split the width, they do not add weights.
• Cost $O(n^2d)$: 134M at $n=512$, 2.1B at 2048; score matrix 4.3 GB at 32,768 tokens.

What is likely to be asked from Week 9, with 30-second answers, deeper explanations, a lecture-or-extra tag on each question, and code you can run.

Week 9 questions are about why attention exists, what its formula means, and how its variants differ. Each answer starts with a 30-second version. Questions are tagged From the notes or Extra.

6.1Interview heat-map: what is most likely

6.2Embeddings and motivation

one-hot dot products (king.queen, king.apple): 0.0 0.0  <- every pair is orthogonal: no notion of similarity
embedding lookup equals one-hot @ weight matrix: True
a vocabulary of 50,000 words with 300-dimensional embeddings has 15,000,000 parameters; a one-hot vector would have 50,000 dimensions

cosine similarities:  king~prince 1.0  king~queen 0.997  king~woman 0.374  king~bread 0.316
king - man + woman is closest to: ['queen', 'princess', 'kingdom']

the polysemy problem: ONE vector for 'bank' although the text uses it in two senses
  nearest words to 'bank': ['river', 'money', 'water', 'loan', 'market', 'crown', 'ruled', 'queen']

distance between the two contextual vectors: 2.841 (0 for a static embedding)
same word, same input vector; the output now depends on the neighbours, which is the idea behind contextual embeddings

sequence length   teacher-forced acc   free-running acc   whole sequence exact
         4               1.000               1.000               1.000
         8               1.000               1.000               1.000
        16               0.999               0.995               0.982
        32               0.988               0.894               0.738

6.3The attention formula

scores  q.k / sqrt(2): [[0.707, 0.0, 0.707]]
softmax weights      : [[0.401, 0.198, 0.401]]  sum = 1.0
output = weights @ V : [[20.0]]
the query [1,0] matches keys 1 and 3 equally, so they get equal weight 0.401; key 2 (orthogonal) gets 0.198

shapes: Q (3, 4) K (5, 4) V (5, 6) -> scores (3, 5) -> output (3, 6)
matches F.scaled_dot_product_attention: True
each row of weights sums to 1: [1.0, 1.0, 1.0]
weights are non-negative: True
every output lies between the smallest and largest value (a convex combination of the rows of V): True
changing the value of a key with weight ~0 barely changes the output; the attention weights say WHERE the model looked

   d_k   std(q.k)   std(q.k / sqrt(d_k))   mean max-weight: unscaled   scaled   |   softmax gradient norm: unscaled   scaled
     4      2.00           1.00                           0.49     0.30   |                      0.0732   0.0882
    16      4.00           1.00                           0.73     0.32   |                      0.0461   0.0884
    64      8.00           1.00                           0.87     0.32   |                      0.0251   0.0886
   256     15.98           1.00                           0.93     0.32   |                      0.0125   0.0882
  1024     31.97           1.00                           0.97     0.32   |                      0.0067   0.0887
a dot product of d independent terms has variance d, so without the division the softmax saturates (one weight near 1) and its gradient vanishes

6.4Variants and costs

SELF-attention  (Q, K, V all from the source)  : output (1, 6, 8)  weights (1, 6, 6) = (batch, 6 queries, 6 keys)
CROSS-attention (Q from decoder; K, V from the encoder): output (1, 4, 8)  weights (1, 4, 6) = (batch, 4 queries, 6 keys)
so in cross-attention each of the 4 decoder positions chooses among the 6 encoder positions; the output has the decoder's length

causal attention weights (row = the position that is looking, column = what it looks at):
[[1.   0.   0.   0.   0.  ]
 [0.05 0.95 0.   0.   0.  ]
 [0.13 0.25 0.63 0.   0.  ]
 [0.01 0.01 0.09 0.89 0.  ]
 [0.01 0.01 0.11 0.3  0.58]]
rows still sum to 1: [1.0, 0.9999999403953552, 1.0, 1.0, 0.9999998807907104]

with the causal mask the outputs at positions 0, 1, 2 do not change when tokens 3, 4 change: True
without the mask they do change (the model could read the answer): True

hand-written multi-head attention matches nn.MultiheadAttention: True
d_model=64, heads=1: parameters 16640 = 4*64*64 + 4*64 = 16640   (head size 64)
d_model=64, heads=2: parameters 16640 = 4*64*64 + 4*64 = 16640   (head size 32)
d_model=64, heads=4: parameters 16640 = 4*64*64 + 4*64 = 16640   (head size 16)
d_model=64, heads=8: parameters 16640 = 4*64*64 + 4*64 = 16640   (head size 8)
more heads do NOT add parameters: they split the same d_model into smaller subspaces

output for 'bites' in 'dog bites man': [-0.72   0.958  1.767  0.24   1.691  0.924]
output for 'bites' in 'man bites dog': [-0.72   0.958  1.767  0.24   1.691  0.924]
identical without position information: True  (self-attention sees a SET of words, not a sequence)
after adding a position vector to each token: identical? False -> the two sentences are now told apart

self-attention builds an n x n score matrix per head
      n     score entries    memory (float32)    n^2 * d operations (d = 512)
    512         2.62e+05         1.0 MB        1.34e+08
   2048         4.19e+06        16.8 MB        2.15e+09
   8192         6.71e+07       268.4 MB        3.44e+10
  32768         1.07e+09         4.3 GB        5.50e+11
 131072         1.72e+10        68.7 GB        8.80e+12
doubling the length quadruples the cost; this is why long context is expensive and why efficient attention is an active research area

Week 9 on one page: chapter glance, the formulas, the pairs to tell apart, a shapes table, a checklist and corrections to the source notes.

7.1Week 9 at a glance

7.2Shapes

7.3Key terms

7.4The pairs you must tell apart

7.5Last-minute checklist

7.6Corrections and additions to the source notes