Week 15: The Attention Mechanism

1. Queries, Keys & Values

Attention reframes Week 14's "relevance score" idea using three learned projections of each token's embedding: a query (what this token is looking for), a key (what this token offers, for others to match against), and a value (the actual content this token contributes once it's deemed relevant).

import numpy as np

d_model = 8    # embedding dimension
seq_len = 4

X = np.random.randn(seq_len, d_model)   # 4 tokens, each an 8-dim embedding

W_q = np.random.randn(d_model, d_model) * 0.1
W_k = np.random.randn(d_model, d_model) * 0.1
W_v = np.random.randn(d_model, d_model) * 0.1

Q = X @ W_q   # queries -- one per token
K = X @ W_k   # keys    -- one per token
V = X @ W_v   # values  -- one per token

A useful analogy: think of it as a search engine. The query is your search text, the keys are every document's indexed title, and the values are the actual document contents. Every token in a sequence simultaneously plays all three roles — it issues a query looking for relevant context, offers a key so other tokens can find it, and carries a value that gets pulled in wherever it's deemed relevant.

2. Scaled Dot-Product Attention, Step by Step

With Q, K, and V in hand, attention is exactly three operations you already know from Weeks 2 and 14: a dot product for relevance, softmax to normalize, and a weighted sum.

def softmax(z):
    shifted = z - np.max(z, axis=-1, keepdims=True)
    exp_z = np.exp(shifted)
    return exp_z / np.sum(exp_z, axis=-1, keepdims=True)

def attention(Q, K, V):
    d_k = K.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)   # (seq_len, seq_len) -- every query vs. every key
    weights = softmax(scores)         # each row sums to 1 -- a distribution over "which tokens matter"
    return weights @ V                # weighted sum of values, per query

output = attention(Q, K, V)   # shape: (seq_len, d_model) -- one context-aware vector per token

The division by sqrt(d_k) is the "scaled" part — without it, dot products grow larger as the embedding dimension grows, pushing softmax's inputs into a range where its gradient becomes tiny (the same saturation problem sigmoid has, from Week 11). Scaling by the square root of the key dimension keeps the scores in a well-behaved range regardless of d_model. The output for each token is a weighted blend of every value in the sequence, weighted by how relevant that token's query found each key — exactly the "look back at everything, weighted by relevance" idea from Week 14, made precise.

3. Multi-Head Attention & Why Multiple Heads Help

A single attention computation can only capture one notion of "relevance" at a time. Multi-head attention runs several smaller attention computations in parallel — each with its own learned Q/K/V projections — and concatenates their outputs, letting different heads specialize in different kinds of relationships (one head might track grammatical subject-verb agreement, another might track coreference between a pronoun and the noun it refers to).

def multi_head_attention(X, num_heads, d_model):
    d_k = d_model // num_heads
    outputs = []

    for _ in range(num_heads):
        W_q = np.random.randn(d_model, d_k) * 0.1
        W_k = np.random.randn(d_model, d_k) * 0.1
        W_v = np.random.randn(d_model, d_k) * 0.1
        Q, K, V = X @ W_q, X @ W_k, X @ W_v
        outputs.append(attention(Q, K, V))   # each head's own (seq_len, d_k) output

    concatenated = np.concatenate(outputs, axis=-1)   # back to (seq_len, d_model)
    W_o = np.random.randn(d_model, d_model) * 0.1
    return concatenated @ W_o   # final learned mixing of all heads' information

Splitting d_model across num_heads smaller heads (rather than running num_heads full-size attentions) keeps the total computation comparable to a single large attention operation, while still giving the model multiple independent "perspectives" on the same sequence — a genuinely richer representation for roughly the same cost.

4. Self-Attention vs. Cross-Attention, and Attention Masks

When queries, keys and values all come from the same sequence, it's called self-attention — each token attending to every other token in the same input. When queries come from one sequence but keys/values come from a different one (e.g. a decoder attending back to an encoder's output), it's cross-attention — this is precisely the mechanism from Week 14's seq2seq fix, and you'll see it again in Week 16's encoder-decoder architecture.

# A causal mask prevents a token from attending to future tokens --
# essential for a model that generates text left-to-right, one token at a time
seq_len = 4
mask = np.triu(np.ones((seq_len, seq_len)), k=1).astype(bool)   # True above the diagonal

def masked_attention(Q, K, V, mask):
    d_k = K.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)
    scores = np.where(mask, -np.inf, scores)   # forbidden positions get -infinity
    weights = softmax(scores)                   # softmax(-inf) -> 0, effectively removing them
    return weights @ V

Setting forbidden scores to -infinity before softmax is the key trick: softmax turns -inf into a probability of exactly 0, so those positions contribute nothing to the weighted sum — a clean way to enforce "don't look ahead" without changing the attention mechanism's core math. This causal mask is exactly what makes the decoder-only architecture you'll build in Week 16 generate text one token at a time without "cheating" by seeing the answer.

5. Computational Cost of Attention

Computing Q @ K.T produces a (seq_len, seq_len) matrix of scores — meaning attention's compute and memory cost grows quadratically with sequence length. Doubling the input length quadruples the attention computation, which is precisely why "context window" (Week 17) is both an architectural and a cost constraint, and why serving very long contexts (Week 25) is expensive.

6. Hands-on Exercise

7. Knowledge Check

Four quick questions. Expand each to check your answer.