Attention Is All You Need

Lesson 4

Self-attention

Every token queries every other token in the same sequence

Self-attention, every token attends to every other token

Vaswani et al. (2017) asked whether we can capture sequence relationships without recurrence or convolution. Self-attention is the core of their answer: every position builds a new representation by looking at every other position in one step, using queries, keys, and values derived from the same sequence.

"Can we design a model that captures sequence relationships without recurrence or convolution?"

, the motivating question behind Attention Is All You Need. See also the interactive narrative at

Lex AI, Attention is All You Need

The answer: direct links between all positions

Self-attention addresses the RNN pain points from Lesson 1 in a different way: direct connections between positions instead of a chain of hidden states.

✓ Full parallelization

All positions compute attention at once on modern hardware (e.g. "The", "cat", "sat" in parallel), GPUs exploit thousands of cores across the sequence and batch.

✓ Direct long-range links

Any token can attend to any other in one layer, e.g. keys … are without passing through every intermediate RNN step.

✓ Short mixing paths

In one attention layer, information can move between any two positions in a single mixing step (vs O(n) steps along an RNN). Deep stacks still add depth; the comparison is to per-layer sequential recurrence.

Story so far: input Z = token + position

Before any attention math runs, the model forms one vector per position:

\mathbf{Z} = \text{token embeddings} + \text{positional encodings}

(Lessons 2–3.) Each row z_pos is d_model-dimensional and carries both what the token is and where it sits. That Z (or, after stack depth ℓ, the hidden sequence H^(ℓ)) is the input to the block that builds Q, K, V, not the same thing as Q, K, V themselves.

From Z (or H) to Q, K, V

Self-attention uses three learned linear projections on the same sequence matrix (call it X for “input to this attention sublayer”, for the first layer, X is Z up to layer norm ordering):

Q = XW^Q,\quad K = XW^K,\quad V = XW^V

W^Q, W^K, W^V each have shape (d_model, d_model) in the simplest telling; in multi-head attention, the effective head dimension is d_k = d_model / h (e.g. 512 / 8 = 64), Lesson 6 expands that split.

Row i of Q is the query at position i: “what am I looking for?” Rows of K and V are keys (“how do I label myself for matching?”) and values (“what vector do I pass if selected?”).

Why Q / K / V cannot be the input embeddings

A common confusion: “Aren’t Q, K, V just the embeddings?” No, for three separate reasons (as spelled out in the Lex AI pipeline and the paper).

1 · They change every layer

Token (+ position) embeddings are fixed at the input to the stack (then updated only as part of the full forward through layers). At each transformer layer, Q, K, V are recomputed from the current hidden representation X^(ℓ). If Q/K/V were the raw embeddings, they could not evolve with depth, yet depth is what builds richer context.

2 · They differ across heads

Multi-head attention runs h attentions in parallel. Each head h has its own W_Q^h, W_K^h, W_V^h:

Q_h = X W_Q^h,\quad K_h = X W_K^h,\quad V_h = X W_V^h

So the same token row in X becomes h different triples (Q_h, K_h, V_h), different “views” of the same content. That is impossible if Q/K/V were identical to a single embedding vector per token.

3 · Dimensionality (d_k vs d_model)

Input vectors live in d_model (e.g. 512). Each head projects to d_k = d_model / h (e.g. 64) so that h heads concatenated restore width h · d_k = d_model. Q, K, V per head live in that smaller working space for efficiency and structure; they are not the same shape or role as the raw d_model embedding row.

Representation	Role
Input Z / hidden X	Meaning + position (and deeper context after layers).
Q, K, V	Projections for matching (Q·K) and mixing (V), rebuilt per layer and per head.

Scaled dot-product scores → softmax → weights

Attention compares each query to all keys with a dot product, scales, then applies softmax so weights form a probability mass over positions (nonnegative, sum to 1). That mass selects a convex combination of value rows.

Scores (logits):

\text{scores} = \frac{Q K^\top}{\sqrt{d_k}}

Weights (row-wise softmax):

A = \text{softmax}(\text{scores})

Output:

\text{Attention}(Q, K, V) = A V

Why √d_k? Dot products grow roughly with dimension; dividing keeps scores in a range where softmax is neither flat nor one-hot saturated, stable training.

Why softmax? It turns raw scores into competition: each query allocates 100% of its attention budget across keys, exactly the “who should I listen to?” story. (See Lesson 4 for the same mechanism in the generic Q/K/V picture.)

"cat" distributes its attention across all tokens. Brighter cell = more attention. Click any token to switch.

Head Causal mask

rows = querycols = keybrighter = higher weight

Query: —

Click a row to see its attention distribution.

Scaled dot-product attention over seeded token vectors. Toggle causal mask: each token attends to itself + earlier only (upper-right goes grey) — that's what makes GPT autoregressive. (Seeded, not trained — illustrative.)

Build it yourself

Press Run to compute self-attention in NumPy. Notice that Q, K, and V all come from the same X, that's what makes it self-attention:

Python

import numpy as np

np.random.seed(0)

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

# One input sequence X drives Q, K, and V (n = 3 tokens, d_model = 4)
X = np.array([[1.0, 0.0, 1.0, 0.0],
            [0.0, 1.0, 0.0, 1.0],
            [1.0, 1.0, 0.0, 0.0]])
d_model, d_k = 4, 4
Wq = np.random.randn(d_model, d_k) * 0.5
Wk = np.random.randn(d_model, d_k) * 0.5
Wv = np.random.randn(d_model, d_k) * 0.5

Q = X @ Wq            # one query per token
K = X @ Wk            # one key   per token
V = X @ Wv            # one value per token

scores = Q @ K.T / np.sqrt(d_k)   # (n x n): every token scores every token
A = softmax(scores)               # row-wise attention weights
out = A @ V                       # weighted sum of values

print("attention weights (rows sum to 1):\n", np.round(A, 3))
print("\nrow sums:", np.round(A.sum(axis=1), 3))
print("\nself-attention output (n x d):\n", np.round(out, 3))
print("\nQ, K, V all came from the SAME X -> that's 'self'-attention")

First run loads the Python runtime (~10 MB) — takes ~5–10 seconds. Subsequent runs are instant.

What “self” means

Self-attention: Q, K, V all come from the same sequence X. Cross-attention (encoder–decoder) uses Q from one side and K, V from the other, Lesson 11.

Multi-head attention (preview)

Multi-head attention stacks h parallel attentions, each with its own projections, then concatenates heads and applies W_O. That is how the model captures many relationship types at once (syntax, coreference, long vs short range, …). Lesson 6 is dedicated to it.

Complexity

Full attention materializes an n × n attention matrix per head, O(n²) time and memory in sequence length n. That is the price of all-pairs direct links; long-context systems use approximations (sparse, linear attention, etc.).

Next: multi-headed attention in full detail.