Self-attention in one page

Self-attention lets every token in a sequence look at every other token and decide how much each one matters for its own new representation. It is the core of every Transformer.

Queries, keys and values

Each token embedding $x_i$ is projected three ways with learned matrices:

• Query: what this token is looking for.
• Key: what this token offers.
• Value: what this token passes on if it is chosen.

Scaled dot-product attention

Stacking all tokens into matrices $Q, K, V$:

Row $i$ of the softmax is a probability distribution over all tokens: how much token $i$ attends to each one.

Why divide by √dₖ?

If the entries of $q$ and $k$ have unit variance, their dot product has variance $d_k$. Large scores push softmax into regions with near-zero gradients. Dividing by $\sqrt{d_k}$ keeps the variance at 1.

The data flow

flowchart LR
  X[Token embeddings X] --> Q[Q = XW_Q]
  X --> K[K = XW_K]
  X --> V[V = XW_V]
  Q --> S[Scores QKᵀ / √dₖ]
  K --> S
  S --> P[softmax]
  P --> O[Output = PV]
  V --> O

In NumPy

import numpy as np

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

def self_attention(X, Wq, Wk, Wv):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    scores = Q @ K.T / np.sqrt(K.shape[-1])
    weights = softmax(scores)          # (tokens, tokens)
    return weights @ V, weights

Going further

• Multi-head: run $h$ attentions in parallel on smaller projections and concatenate.
• Causal mask: set future positions to $-\infty$ before the softmax for decoder models.
• Cost: $O(n^2 d)$ in sequence length $n$, which is why long contexts are expensive.