Self-attention in transformers

Each token asks a question (query) and every token offers a label (key) and content (value).

Model answer.

The idea: each token's new representation is a relevance-weighted mix of all tokens — e.g. in "The animal didn't cross the street because it was tired", "it" can draw heavily on "animal".

Queries, keys, values
For an input matrix X (n tokens × d_model), three learned projections produce:

• Q = XW_Q — what each token is looking for,
• K = XW_K — what each token offers to be matched against,
• V = XW_V — the information each token passes on.

Formula
Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V

• QKᵀ is an n × n matrix of similarity scores between every pair of tokens.
• Softmax over each row turns scores into weights that sum to 1.
• Multiplying by V gives each token a weighted average of value vectors.

Why divide by √d_k?
If the components of q and k are independent with mean 0 and variance 1, their dot product has variance d_k. For large d_k the scores become large in magnitude, pushing softmax towards a near one-hot output where gradients are tiny. Dividing by √d_k brings the variance back to about 1 and keeps training stable.

Multi-head attention
Several heads run in parallel with separate, smaller projections; outputs are concatenated and projected, letting heads capture different relationships.

Cost

• Computing QKᵀ for n tokens is O(n² · d) time and needs O(n²) memory for the score matrix per head.
• This quadratic growth limits long contexts. Remedies include sparse or sliding-window attention and FlashAttention, which computes exact attention in tiles without storing the full n × n matrix.

Positional information
Attention alone is permutation-equivariant, so word order would be lost. Transformers add positional encodings (fixed sinusoidal or learned embeddings) or use relative schemes such as rotary position embeddings (RoPE).

Masking: decoder (GPT-style) models apply a causal mask so each position attends only to earlier positions.

Follow-up: "Why did transformers replace RNNs?" Positions are processed in parallel during training, and any two tokens connect in one attention step rather than a long recurrent chain.