Module 5 cheat sheet: Transformer & Attention

Attention lets every word look at every other word at once, so the model understands context and trains in parallel.

Key ideas

RNNs read word by word: slow, and they forget far-away words. S17 →
Each token makes a query (what I need), key (what I offer) and value (what I carry). S17 →
Dividing by √dₖ stops scores getting too large, so softmax does not saturate. S18 →
Masks hide padding and, in decoders, future tokens. S18 →
Multi-head = several smaller attentions in parallel, each learning a different relation, for about the same cost. S18 →
Long context: attention cost grows with n², the KV cache with n; models can miss facts in the middle. S19 →

The formula

Attention(Q, K, V) = softmax( Q·Kᵀ / √dₖ ) · V

Q·Kᵀ: how well each query matches each key (a score table).
/ √dₖ: scale the scores down (dₖ = size of each key vector).
softmax: each row becomes weights that add up to 1.
· V: each token gets a weighted mix of the values.

Making attention cheaper

Method	Idea	Exact?
Sliding window	Each token looks only at nearby tokens	No
Sparse attention	Skip most token pairs, keep a pattern	No
FlashAttention	Same maths, smarter use of GPU memory	Yes
GQA / MQA	Heads share keys and values → smaller KV cache	Small change

Numbers to remember

Double the context → about 4× attention work (n²).
KV cache, Llama 3.1 8B (FP16): about 128 KB per token, so ~1 GB for 8,000 tokens.
Original transformer: 8 heads, model size 512 (64 per head).

Code: attention in NumPy

import numpy as np
X = np.random.randn(4, 8)                      # 4 tokens, 8 numbers each
Wq, Wk, Wv = [np.random.randn(8, 4) for _ in range(3)]
Q, K, V = X @ Wq, X @ Wk, X @ Wv
scores = Q @ K.T / np.sqrt(K.shape[1])
w = np.exp(scores - scores.max(1, keepdims=True))
w /= w.sum(1, keepdims=True)                   # softmax per row
out = w @ V                                    # new vector for each token

Common mistakes

Forgetting the √dₖ scaling: training becomes unstable.
Forgetting the causal mask in a decoder: the model “cheats” by seeing the answer.
Treating attention maps as a full explanation of the model’s reasoning.
Assuming a long context window means the model uses every part equally well.

Created by Ajeetkumar · EasyLearnSpace