An intuitive explanation of Self Attention | by Saketh Kotamraju | Towards Data Science

In this article, I am going to explain everything you need to know about self-attention.
What do transformer neural networks contain that make them so much more powerful and better performing than regular recurrent neural networks?
Answer: They use the Multi-headed Self Attention block to give the word embedding inputs more contextual information. But what exactly does the self-attention block do to the word embeddings that make these transformer models so much more powerful?
That is the focus of this article. The main purpose of this post is to describe both the intuition behind each part of the self-attention block and the mathematical operations involved in it.
This article does not aim to explain the overall structure of transformer neural networks. It also does not also describe the difference between self-attention and regular attention.
We know that word embeddings are vectors that represent the semantic meaning of a word. Words with similar meanings may have similar embeddings. However, in a sentence, the individual meanings of the words do not represent their meanings in the sentence. For example, if I had the phrase, Bank of a river, the embeddings of Bank and river individually mean completely different things, but they have a strong correlation in the sentence. Word embeddings without self-attention do not possess this sense of contextual information, so given the phrase above, a language model would have a low chance of predicting river. In order to address this problem, the self-attention block was proposed in the paper Attention is all you need as part of the original transformer architecture.
A self-attention module works by comparing every word in the sentence to every other word in the sentence, including itself, and reweighing the word embeddings of each word to include contextual relevance. It takes in n word embeddings without context and returns n word embeddings with contextual information. For example, in the phrase, Bank of the river, Bank would be compared with Bank, of, the, and river, and as Bank is compared with those four words, its word embedding would be reweighted to include the relevance of the words to its own meaning in the sentence accordingly.
The Self-attention block is comprised of three steps/parts:
Dot product similarity to find alignment scoresNormalization of the scores to get the weightsReweighing of the original embeddings using the weights
Dot product similarity to find alignment scores
Normalization of the scores to get the weights
Reweighing of the original embeddings using the weights
When the input word embeddings are passed into the self-attention module and find out who they should pay attention to, they need a function to find how similar they are to the other words in the sentence. This is where dot product similarity comes into play.
For the following explanation, I will use the following example:
Words: v1 v2 v3 v4 
where v1, v2, v3, and v4 are the word embeddings of words in a sentence.
In our example, we will give self-attention to v3, and compare it with v1, v2, v3, and v4. The reason why we also have to compare each word to itself (e.g. v3 to v3), is so that the model can learn which parts of the semantic meaning of a word it should pay attention to. It may seem unusual that a word may have to pay attention to itself. However, sometimes a word can have more than one meaning represented in its embedding (like the example with Bank), and unless we also include itself in the dot product similarity, the model won’t be able to learn which parts of the word’s meaning it should pay attention to. Simply put, words usually have more than one meaning, and in order for the model to recognize which meaning of the word it should pay attention to, we must compare words to themselves by including them in the dot product similarity. In our example, where v3 is compared to every other word and itself, that looks like this:
(v1*v3) = S31(v2*v3) = S32(v3*v3) = S33(v4*v3) = S34
We will take the dot product of v1 and v3, v2 and v3, v3 and v3, and v4 and v3 to determine the alignment score of each pair of embeddings. The alignment scores, S31, S32, S33, and S34, will tell us how similar the semantic meanings of v3 and each of the four compared words. The higher the alignment score, the more similar the semantic meanings of the words, and the more attention v3 will have to pay attention to the other word in the pair. Note that each of the alignment scores is a single number, not a vector or matrix of numbers.
I’ve only done this process for one word: v3. In reality, this process will be done simultaneously for all four words using vectorization and linear algebra.
Take the word embedding of King and Queen.
King = [0.99, 0.01, 0.02]Queen = [0.97, 0.03, 0.02]
The alignment score of this pair would be:
(0.99 * 0.97) + (0.01 * 0.03) + (0.02 * 0.02) = 0.961
Now take the word embeddings of King and Dog
King = [0.99, 0.01, 0.02]Dog = [0.01, 0.02, 0.02]
The alignment score of this pair would be
(0.99 * 0.01) + (0.01 * 0.02) + (0.02 * 0.02) = 0.0105 
Note that these word embeddings are made up and that the real word embeddings of dog, king, and queen are likely much larger in size. Also note that the first index, where king and queen had a high number, represents the royalty of the particular word.
As you can see the alignment score of King and Queen is much greater than King and Dog. This is because the semantic meaning of the king and dog are dissimilar so they are not important to each other in a sentence. Index 1 represents the royalty of the word, so King and Queen are similar in that regard.
In short, if the semantic meaning of both words differs, then the alignment score won’t be high, the words will have a low correspondence, and each word will pay low attention to each other so the original word embeddings of each word won’t change too much. If the semantic meaning of two words is similar in any regard, then the alignment score will be high, the words will have a high correspondence, and the two words will pay a high amount of attention to each other so the original word embeddings of each word will undergo a significant change.
Normalizing the values of the outputs of any Neural Network through an activation function is a common procedure. We will normalize the alignment scores using a SoftMax activation function to obtain the weights that we will apply to the original word embeddings. The SoftMax function will make each of the alignment scores a probabilistic distribution so they all add up to one. In our example,
S31 -> SoftMax -> W31S32 -> SoftMax -> W32S33 -> SoftMax -> W33S34 -> SoftMax -> W34
all four alignment scores will be SoftMaxed to obtain the final weights that have to be applied to the original word embeddings in order to create the final contextualized embedding.
Note that when I say ‘weights’ in the paragraph above, I am not referring to the parameters that the model will learn.
Now the final reweighing process occurs. In order to determine the total amount of attention the word(v3) should pay to the other words, we multiply the weights with their respective original word embeddings and then add all of these values together to get the final word embedding for the word being compared(v3); in this case,
W31 * v1 = Y31W32 * v2 = Y32W33 * v3 = Y33W34 * v4 = Y34Y31 + Y32 + Y33 + Y34 = Y3
Y3 is the final reweighted word embedding vector of v3. Remember that Y31, Y32, Y33, and Y34 are all individually vectors, so we must add them together to get one final vector.
Remember that the process described above will be done for every other word in the sentence as well.
Intuitively, all of the processes above make sense except for one thing: the fact that because there are no weights to learn, so the alignment scores and self-attention weights will essentially be pre-determined and the model won’t be able to learn any deeper connections between two words. This is why we will introduce weights in the self-attention block. But where? In order for the model to learn the deepest connections between words, we will introduce the weights at three locations: the input word embeddings, the dot product similarity comparison, and the final step of reweighing the word embeddings. In addition, one more benefit of introducing the words at these locations is that the shapes/dimensions of the vectors multiplied do not change. Remember that if we have a vector of 1 x k shape, and we multiply it by a matrix of k x k shape, this should result in a vector of 1 x k shape. Therefore, even if we introduce weights in the locations described, nothing should change in dimensionality or shape.
This is why we must introduce weights at the places where we have the original word embedding vectors, v1, v2, v3, and v4, because the shape of these vectors is all the same, and even if we introduce a k x k matrix of weights, the shapes of these will still be 1 x k. The three locations where we use the original embedding vectors is the input word embeddings, the dot product similarity when we compare each word to other words, and during the final reweighing process when we multiply the normalized weights times the original word embeddings. With weights, our calculations will look as follows:
(v1 * Mk) * (v3 * Mq) = S31(v2 * Mk) * (v3 * Mq) = S32(v3 * Mk) * (v3 * Mq) = S33(v4 * Mk) * (v3 * Mq) = S34S31 -> SoftMax -> W31S32 -> SoftMax -> W32S33 -> SoftMax -> W33S34 -> SoftMax -> W34(v1 * Mv) * W31 = Y31(v2 * Mv) * W32 = Y32(v3 * Mv) * W33 = Y33(v4 * Mv) * W34 = Y34Y31 + Y32 + Y33 + Y34 = Y3
As you can see, wherever we use the original word embeddings v1, v2, v3, and v4, we multiply those vectors with the corresponding weight Matrixes. Mk, Mq, and Mv are simply the key, query, and value matrices/weights that the model will learn. Remember that in the calculations above, I only did the self-attention operation for one word, v3. In reality, this will happen to all of the words simultaneously through vectorization and some linear algebra.
Note that in the diagram above, when I multiply the key, query, and value matrices with their respective embedding vectors, I do a matrix multiplication, not a dot product.
Although it may seem reasonable that one self-attention block is enough for a word to obtain contextual relevance, this is not the case. Often, a word will have to pay attention to multiple other words, and only one self-attention block may not be enough attention for multiple words. This is especially evident in examples where the input text is very large (e.g. for a text summarization task). A pair of words that have contextual relevance, therefore, will sometimes not get enough attention that will make an observable change in their respective embeddings.
In order to fix this problem, we will use a multi-headed attention block. The Multi headed attention block expands the model’s ability to focus on different positions in the input text.
A multi-headed attention block is essentially the same thing as a regular self-attention block, but instead of just one attention block, the multi-headed attention block will contain multiple self-attention blocks that operate in parallel. These self-attention blocks will not share any weights; the only thing they will share is the same input word embeddings. The number of self-attention blocks in a multi-headed attention block is a hyperparameter of the model. Suppose that we choose to have n self-attention blocks. The consequence of this is that after all the individual calculations are done by each of the self-attention blocks, we will have n embeddings for each word. In order to fix this problem, the multi-headed attention block concatenates these embeddings and finally passes them through a dense layer. Remember that we want the shape and number of the inputs to equal the shape and number of the outputs. Because we concatenated the embeddings, we keep the number of inputs equal to the number of outputs. Because we passed the concatenated outputs through a dense layer, we were able to control their shape, making sure that it’s the same as the shape of the inputs.
The self-attention block takes in word embeddings of words in a sentence as an input, and returns the same number of word embeddings but with context. It accomplishes this through a series of key, query, and value weight matrices. The multi-headed attention block consists of multiple self-attention blocks that operate in parallel and share no weights. After each of the self-attention blocks return new contextualized word embeddings, the multi-headed attention block concatenates these new embeddings together and passes them through a dense layer in order to control the shape of the output. This allows for a word embedding to receive enough attention that will make an observable change in its embedding.
I hope you found this content easy and intuitive to understand. If you think that I need to elaborate further or clarify, drop a comment below.
Attention Is All You Need (arxiv.org)
Rasa
Illustrated: Self-Attention (towardsdatascience.com)