I keep seeing the "head" in attention models (transformers). Aside from the mathematical formula, could anyone please share the intuition behind the idea "head"?
2 Answers
I like to think about it in the context of progression of attention mechanisms in neural networks.
Early attention mechanisms were implemented as explicit sliding windows over the encoding sequence. For example, Graves 2013 implemented it as the average of several sliding Gaussians. So, it was a local attention mechanism to attend to a learnable window over the data.
Then you have Luong and Bahdanau attention, which implemented similar mechanisms, but separate from the network type and as a softmax function. This is essentially single-head attention. Now it's not a window, but a score and weighted average over all of the encoding input. The idea of the K, Q, V flows naturally from that, but in that case the Q is a hidden state of the decoder, and K/V the hidden state of the encoder. This is more flexible and can attend to more arbitrary windows over the sequence.
Transformers flow naturally from these mechanisms, but now K, Q, V all come from the same sequence, except for the second layer of the decoder, which acts the same as Luong or Bahdanau attention. In addition, they first split the K, Q, V along their last dimension (the model parameter dimension) to create several heads that can attend to different parts of the sequence. This is much better at attending to long term and varied parts of the sequence to concatenate an overall context for each Query. This, of course, is very advantageous for natural language translation, but as we are seeing, it is also great for covering the image space in vision problems. Intuitively, this makes sense as I view the Transformer as almost a superset of a convolution network, where the filters for each partition of an image attend to various parts of the image and this is learned.
Obviously, there are many papers I'm sure I haven't read and steps in between these.
A common intuition is viewing attention as probabilistic retrieval: given a query, you want to retrieve some information from some states (values) given some keys that describe the values.
With single-head attention, you get a weighted average of some hidden states with everything that is contained in there. Once you introduce, the head-specific projection, the model has much more flexibility. The information you use for retrieving (in the query-key comparison) can be of a different kind than the value. One head can retrieve a piece of information from some states, another head can retrieve a different piece of information. The head-specific projections also allow throwing away what is no longer relevant for the next layers.