The Transformer, One Attention Head at a Time

Why attention won

Before transformers, sequence models pushed information through recurrence, one token at a time. That made long-range dependencies expensive and training hard to parallelize. Self-attention replaced the bottleneck: every token looks at every other token in a single step, and the whole operation is a few matrix multiplies the GPU loves.

Queries, keys, and values

Each token is projected into three vectors. The **query** asks a question, the **key** advertises what a token offers, and the **value** is the content that gets mixed in. The attention weight between two tokens is the scaled dot product of one token's query with another's key, passed through a softmax.

import torch
import torch.nn.functional as F

def attention(q, k, v):
    # q, k, v: (batch, heads, seq, d_head)
    d_head = q.size(-1)
    scores = (q @ k.transpose(-2, -1)) / d_head ** 0.5
    weights = F.softmax(scores, dim=-1)
    return weights @ v

The division by the square root of `d_head` keeps the dot products from growing with dimension, which would otherwise saturate the softmax and kill the gradient.

Many heads, many questions

A single head can only attend one way. Multi-head attention runs several in parallel on lower-dimensional projections, then concatenates them — so one head can track syntax while another follows coreference.

What to remember

Attention is a soft, differentiable lookup table. Once you see it that way, positional encodings, causal masks, and cross-attention all fall out as small variations on the same theme. Swap in your own diagrams and worked examples when you make this article your own.