Understanding Multi-Headed Self-Attention

What is it used for?

MHSA is a core component of the Transformer architecture. The idea is that as the individual tokens, represented by Embeddings, come in, we need to find a contextual representation for each of them.

That is, the different meanings of stoned, in They stoned the witch. and This guy got completely stoned. need to be represented by different vectors, otherwise, we would be losing important semantic information.

How does it work?

High-level abstraction

The self-attention part

In order to achieve this change, we need to take into account the surrounding vectors in a learned manner. So, what we do is the following:

  1. We create a table which can describe the relations between every pair of tokens by putting the full input sequence on both axes, like so:
Hello world !
Hello
world
!
  1. Each vector will then ask the rest: ‘What do I interact with?’ and every vector, including itself, will say what does it interact with. What this actually means is that we will have 2 matrices, query (What do I interact with?), and key (I interact with these things!), which are going to multiply the vectors and create new vectors which will take their spots at the table, like so:
key(Hello) key(world) key(!)
query(Hello)
query(world)
query(!)

Note that queries take the rows, while keys take the columns.

  1. Once this is done, we fill the cells with dot products of their corresponding row and column vectors. This can be understood as seeing if there is any compatibility between the ‘What do I interact with?’ and ‘I interact with these things!’ of the vectors.
key(Hello) key(world) key(!)
query(Hello) 11 16 19
query(world) 14 17 20
query(!) 15 18 21

I am obviously making the numbers up

  1. However, we also need to divide them by the square root of the dimensionality of the qkv vectors, in order to keep the variance roughly around 1, independently of the model size. The resulting numbers are called the raw attention scores.
key(Hello) key(world) key(!)
query(Hello) \(\frac{11}{\sqrt{\text{d\_qkv}}}\) \(\frac{16}{\sqrt{\text{d\_qkv}}}\) \(\frac{19}{\sqrt{\text{d\_qkv}}}\)
query(world) \(\frac{14}{\sqrt{\text{d\_qkv}}}\) \(\frac{17}{\sqrt{\text{d\_qkv}}}\) \(\frac{20}{\sqrt{\text{d\_qkv}}}\)
query(!) \(\frac{15}{\sqrt{\text{d\_qkv}}}\) \(\frac{18}{\sqrt{\text{d\_qkv}}}\) \(\frac{21}{\sqrt{\text{d\_qkv}}}\)

I am obviously making the numbers up

  1. And since the dot product is not restricted to anything and the numbers can vary an awful lot, we normalise them using Softmax row by row, which is going to return nice probability distributions of what interacts with what else. These numbers are called the (normalised) attention scores.
key(Hello) key(world) key(!)
query(Hello) 0.0 0.4 0.6
query(world) 0.9 0.05 0.05
query(!) 0.8 0.1 0.1

Didn’t actually apply softmax, still making up numbers

  1. Once we have those, what we will do is that we will take the value matrices and we will use them to get yet another representation by multiplying the original embedding vectors. [That is, we will have value(Hello), value(world), value(!).] We will then multiply these representations in their corresponding columns by the (normalised) attention scores, like so:
key(Hello) key(world) key(!)
query(Hello) \(0.0 \cdot value(Hello)\) \(0.4 \cdot value(world)\) \(0.6 \cdot value(!)\)
query(world) \(0.9 \cdot value(Hello)\) \(0.05 \cdot value(world)\) \(0.05 \cdot value(!)\)
query(!) \(0.8 \cdot value(Hello)\) \(0.1 \cdot value(world)\) \(0.1 \cdot value(!)\)

The resulting numbers are called attention weights

  1. What we then do inside a transformer block is that for each row, we take the original embedding and sum it with whatever is in that row, like so: \[\text{updated\_Hello} = \text{Hello} + 0.0 \cdot value(Hello)+ 0.4 \cdot value(world) + 0.6 \cdot value(!)\] updated_Hello is the final contextual representation, Hello is the original input embedding, however, this is not technically part of the core attention mechanism, rather part of the transformer architecture, and it is therefore not done per head.

The multi-headed part

The multi-headed part of the name just means that we are doing multiple of the above operations in parallel, however, just how is it parallelised is quite clever:

  1. We have the query, key and value representations, however, these often need to be quite big, so we split them into smaller chunks. That is, if these vectors have 30 columns and we want to have 10 heads, we would just split each of the matrices into 10 smaller ones, each ending up with 3 columns.
  2. Then, what we do is that we do the above process in parallel for each of the heads and concatenate the outputs into one big vector (which happens to have the same dimensions as the token representation vector).
  3. We then pass this vector through a linear transformation matrix and just like that, we have the output representation.

The idea behind this, aside from being efficient, is that it lets each of the heads specialise in something and do just that very well. We then slap those representations together and the pass through the linear transformation makes it so that the resulting vector is actually coherent, not just a bunch of small relationless representations.

The actual maths

In practice, we still have to do all of that, but we also need to worry about batches and multiple heads operating in parallel. Additionally, for efficiency reasons, we need to do everything via multi-dimensional tensor objects.

Here is the process, step by step:

  1. We start with an input tensor of dimensions \([batch\_size, n\_tokens, embedding\_dim]\).
  2. This gets multiplied by the three query, key and value matrices, each of dimensionality \([embedding\_dim, qkv\_dim]\), resulting in three new tensors, each of dimensionality \([batch\_size, n\_tokens, qkv\_dim]\). Each stores one of the 3 representations and I will be calling them queries, keys and values.
  3. The last dimension then gets split into the heads, so the tensors end up with the following dimensions \([batch\_size, n\_tokens, n\_heads, head\_dim]\) where \(n\_heads\) is the number of heads and \(head\_dim\) is the dimensionality of each head, as explained above.
  4. We then transpose the second and a third dimension, so that the dimensionality now looks like this: \([batch\_size, n\_heads, n\_tokens, head\_dim]\).
  5. Now, we have all the tensors we want in the correct format, so we need to start using them. As explained above, we need to get the dot product of the queries and keys. To do this, we transpose the third and fourth dimensions of the keys tensor and multiply it with queries. That is, we are multiplying \(queries[batch\_size, n\_heads, n\_tokens, head\_dim] \cdot keys[batch\_size, n\_heads, head\_dim, n\_tokens]\). As mentioned in the previous step the first two dimensions \([batch\_size, n\_heads]\) sort of act as a 2-level nested loop, while the remaining two are essentially doing a matrix-matrix multiplication on every iteration of the inside loop with the following dimensionality: \([n\_tokens, head\_dim] \cdot [head\_dim, n\_tokens]\). This elegantly returns a matrix of \([n\_tokens, n\_tokens]\) dimensionality, that is, the self-attention matrix explained above, in the high-level section. The final dimensionality of the attention scores tensor is therefore \([batch\_size, n\_heads, n\_tokens, n\_tokens]\).
  6. We then apply Softmax by row to this tensor divided by the square root of qkv dimensionality. (I am going to clarify, because there are a lot of dimensions floating around here, by qkv dimensionality, I mean the number dimensions of a single qkv vector in a single head. We divide by a scalar.)
  7. We then multiply the attention scores tensor by the values tensor, that is, \(attn[batch\_size, n\_heads, n\_tokens, n\_tokens] \cdot values[batch\_size, n\_heads, n\_tokens, head\_dim]\). This returns the attention weights tensor, which has dimensions\([batch\_size, n\_heads, n\_tokens, head\_dim]\).
  8. In order to extract the context vectors for the individual tokens from this tensor, we need to transpose back to the original tensor dimensionality, that is we transpose the second and third dimensions and attention weights becomes \([batch\_size, n\_tokens, n\_heads, head\_dim]\).
  9. We then combine the heads back together and the dimensionality becomes \([batch\_size, n\_tokens, qkv\_dim]\).
  10. We then multiply this by the linear projection matrix to make sure that the heads’ outputs actually work well together as a representation.
  11. And thus, we end up with a tensor of contextualised representations where the last dimension is the dimension of the contextual representation with the same dimensionality as \(embeddin\_dim\).