Multi-Head Attention

Multi-head attention lets a network examine the same set of tokens through several different “views” at once. Rather than forcing one attention pattern to capture every relationship, it gives multiple small attention mechanisms room to specialize.

How the heads work

For each token representation, learned linear layers create a query (what this token is looking for), a key (what it offers), and a value (the information it can contribute). Each head computes scaled dot-product attention: it compares queries with keys, converts the scores into weights with softmax, and takes a weighted mixture of values. A head might learn to link a token to a nearby modifier, while another tracks a distant dependency. Their outputs are concatenated and passed through another learned projection, allowing the model to combine these separate views.

Masks and transformer blocks

Attention weights are restricted with masks when needed:

A causal mask prevents a token from seeing later tokens, which is essential for autoregressive prediction.
A padding mask prevents meaningless padding positions from affecting real tokens.

In a transformer, multi-head attention sits inside a residual block alongside LayerNorm and a feed-forward network. The residual path preserves the incoming representation, while normalization keeps activations and gradients stable. In PyTorch, this mechanism appears as torch.nn.MultiheadAttention; modern transformer implementations commonly use optimized attention kernels rather than constructing the full score matrix directly.

Why it matters in training

Multiple heads increase representational flexibility without making each individual comparison excessively large. Scaling query–key dot products by the square root of their dimension prevents softmax scores from becoming too extreme, which would produce near-one-hot weights and weak gradients. The central cost is attention’s quadratic memory and compute in sequence length: every position can compare with every other position. Too few heads can limit useful relationship patterns; too many tiny heads can waste capacity. If masking is wrong, a model can leak future information during training, show deceptively low loss, and fail when asked to predict without access to that future.

Multi-head attention runs several attention computations in parallel, each with separate learned query, key, and value projections. Each head can focus on different relationships or representation subspaces; their outputs are concatenated and projected into a single representation. It matters because it lets a transformer model multiple dependencies simultaneously, increasing representational capacity without requiring a single attention map to capture every relevant relationship.

Imagine several people reading the same sentence, each asked to look for a different kind of clue. One notices who is doing an action, another watches for time words, and another spots which earlier word “it” refers to. Together, they build a richer understanding than any one reader could alone.

Multi-head attention gives an AI system several of these “readers” at once. Each attention head looks across the same words or pieces of information, but can focus on different relationships. This helps the system connect ideas near and far apart—useful for understanding language, generating text, and working with other kinds of data.