Self-Attention

Self-attention is the core mechanism of the transformer. It computes relationships between all pairs of positions in a sequence.

Queries, Keys, Values: Given input $X$ (sequence of token embeddings), compute:

Attention computation:

Step by step:

1. $QK^\top$ — dot product of every query with every key → $n \times n$ relevance scores
2. $/ \sqrt{d_k}$ — scale to prevent softmax saturation (large dot products → near-one-hot softmax)
3. Softmax — normalize each row to a probability distribution (how much each position attends to every other)
4. Multiply by $V$ — weighted sum of value vectors according to attention weights

Intuition: "soft lookup in a dictionary." Each position asks a query, compares it to all keys, and retrieves a weighted mix of values. The network learns what to look for (Q), what to advertise (K), and what to provide (V).

def scaled_dot_product_attention(q, k, v, mask):
    # q, k: (batch, ..., seq_len, d_k);  v: (batch, ..., seq_len, d_v)
    d_k = q.shape[-1]
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(d_k)
    scores = scores.masked_fill(~mask, -torch.inf)   # mask True = attendable
    return softmax(scores, dim=-1) @ v

The $1/\sqrt{d_k}$ scaling matters because dot products grow like $\sqrt{d_k}$; unscaled, large logits make the softmax peaky and resistant to gradient updates.

Complexity: $O(n^2 d)$ — quadratic in sequence length. This is the main bottleneck for long sequences.

vs. RNNs: attention gives direct connections between any two positions regardless of distance, solving the long-range dependency problem.

See also: Multi-Head Attention, Causal Masking, Positional Encoding