Transformer Attention: Why does Attention need more than one head?

Published

This is a series where I answer the questions my tutoring students really ask – the ones textbooks tend to skip. I teach machine learning, statistics, and math to learners of all levels.

The question. After the queries-keys-values post, a student asked the obvious follow-up: “One set of Q, K, V already lets every token look at every token. What do extra heads add? Isn’t multi-head attention just the same computation repeated several times?”

It would be, if the heads shared their projections. They don’t – and that small difference turns one question per token into several questions asked in parallel.

One head means one question

Remember what a single head is: three learned projections, and then the score-normalize-mix routine. Each token gets one query, so it can ask one question of the sentence, and the softmax forces it to spend its entire attention budget on the answer to that one question. The weights must sum to one, so a single head is always a compromise: whatever pattern it learns, it has to cover every situation with the same distribution.

The head we built in the queries-keys-values post is a good example. Its row for “journey” came out as:

head 1, “journey”: [0.15, 0.23, 0.22, 0.13, 0.09, 0.18]

That is a mild preference for the tokens that are similar overall – “journey” itself and “starts” – spread thinly across everything. One question, one compromise answer.

A second head can ask something completely different

Now give the model a second head with its own projections. Trained heads specialize – in real language models, researchers have found heads that track subject-verb agreement, heads that point from a pronoun to its referent, heads that just look at the previous token. We cannot train anything in a blog post, so here is a hand-set second head that stands in for a trained one: its query and key read only the third coordinate of each embedding, and its value delivers the first two.

W_q2 = torch.tensor([[0.], [0.], [3.]])  # query: reads coordinate 3
W_k2 = torch.tensor([[0.], [0.], [3.]])  # key: advertises coordinate 3
W_v2 = torch.tensor([[1., 0.],
                     [0., 1.],
                     [0., 0.]])          # value: delivers coordinates 1 and 2

scores2 = (inputs @ W_q2) @ (inputs @ W_k2).T
w2 = torch.softmax(scores2, dim=-1)

print(w2[1])  # the row for "journey"
# tensor([0.603, 0.154, 0.136, 0.022, 0.006, 0.080])

head 1, “journey”: [0.15, 0.23, 0.22, 0.13, 0.09, 0.18] head 2, “journey”: [0.60, 0.15, 0.14, 0.02, 0.01, 0.08]

Same token, same sentence, two utterly different distributions. Head 2 puts 60% of its budget on “Your,” because “Your” has by far the largest third coordinate – that is the only thing this head cares about. Meanwhile head 1 keeps its broad, similarity-based mixture. The token “journey” now receives two context vectors, [0.3061, 0.8210] from head 1 and [0.4345, 0.4181] from head 2, and the model simply concatenates them into four numbers and moves on. (A final output projection, W_o, mixes the heads back down to the model width – that is the only other moving part.)

Two attention heads over the same sentence. Head 1 spreads its weights in a mild similarity pattern; head 2 concentrates on the tokens with a large third coordinate. Same input, two different questions.

One honest caveat from our session: when we initialized a second head with random projections, both heads produced nearly identical patterns, because the geometry of the input embeddings dominates untrained weights. The heads only come to ask different questions because training pushes their projections apart. Our hand-set head is a stand-in for what training produces, not a substitute for it.

One thing to try

Set the second head’s projections equal to the first head’s (W_q2 = W_query, W_k2 = W_key, W_v2 = W_value, from the queries-keys-values post) and verify that the two outputs become identical – the heads differ only because their projections differ. Then count parameters: two heads of width two versus one head of width four cost almost exactly the same to run. Multi-head attention is not more computation; it is the same computation reorganized so that no single softmax has to answer every question at once.

The code in this article is adapted from the teaching notebooks I use with my students, and it follows the progression in Sebastian Raschka’s Build a Large Language Model (From Scratch), which continues from here to grouped-query attention and the rest of the GPT architecture.

References

  1. Vaswani, A., et al. (2017). “Attention Is All You Need.” Advances in Neural Information Processing Systems 30 (NeurIPS 2017).
  2. Clark, K., Khandelwal, U., Levy, O., & Manning, C. D. (2019). “What Does BERT Look At? An Analysis of BERT’s Attention.” Proceedings of BlackboxNLP 2019.
  3. Raschka, S. (2024). Build a Large Language Model (From Scratch). Manning.

I tutor machine learning, deep learning, statistics, and math. If attention – or anything else – is not clicking for you, you can find me on Wyzant.