RNN and LSTM networks are widely used in sequential tasks such as next word prediction, machine translation, text generation, and more. However one of the major challenges with the recurrent model is capturing the long-term dependency.

To overcome this limitation of RNNs, a new architecture called Transformer was introduced in the paper Attention Is All You Need. The transformer is currently the state-of-the-art model for several NLP tasks. The advent of the transformer created a major breakthrough in the field of NLP and also paved the way for new revolutionary architectures such as BERT, GPT-3, T5, and more.

The transformer model is based entirely on the attention mechanism and completely gets rid of recurrence. The transformer uses a special type of attention mechanism called self-attention. We will learn about this in detail in the upcoming sections.

Let's understand how the transformer works with a language translation task. The transformer consists of an encoder-decoder architecture. We feed the input sentence (source sentence) to the encoder. The encoder learns the representation of the input sentence and sends the representation to the decoder. The decoder receives the representation learned by the encoder as input and generates the output sentence (target sentence).

Suppose we need to convert a sentence from English to French. As shown in the following figure, we feed the English sentence as input to the encoder. The encoder learns the representation of the given English sentence and feeds the representation to the decoder. The decoder takes the encoder's representation as input and generates the French sentence as output:

Okay, but what's exactly going on here? How does the encoder and decoder in the transformer convert the English sentence (source sentence) to the French sentence (target sentence)? What's going on inside the encoder and decoder? Let's find this out in the next sections. First, we will look into the encoder in detail, and then we will see the decoder.

The transformer consists of a stack of number of encoders. The output of one encoder is sent as input to the encoder above it. As shown in the following figure, we have a stack of number of encoders. Each encoder sends its output to the encoder above it. The final encoder returns the representation of the given source sentence as output. We feed the source sentence as input to the encoder and get the representation of the source sentence as output:

Note that in the transformer paper Attention Is All You Need, the authors have used , meaning that they stacked up six encoders one above the another. However, we can try out different values of . For simplicity and better understanding, let's keep :

Okay, the question is how exactly does the encoder work? How is it generating the representation for the given source sentence (input sentence)? To understand this, let's tap into the encoder and see its components. The following figure shows the components of the encoder:

From the preceding figure, we can understand that all the encoder blocks are identical. We can also observe that each encoder block consists of two sublayers:

• Multi-head attention
• Feedforward network

Now, let's get into the details and learn how exactly these two sublayers works. To understand how multi-head attention works, first we need to understand the self-attention mechanism. So, in the next section, we will learn how the self-attention mechanism works.

Self-attention mechanism

Let's understand the self-attention mechanism with an example. Consider the following sentence:

 A dog ate the food because it was hungry

In the preceding sentence, the pronoun it could mean either dog or food. By reading the sentence, we can easily understand that the pronoun it implies the dog and not food. But how does our model understand that in the given sentence, the pronoun it implies the dog and not food? Here is where the self-attention mechanism helps us.

In the given sentence, A dog ate the food because it was hungry, first, our model computes the representation of the word A, next it computes the representation of the word dog, then it computes the representation of the word ate, and so on. While computing the representation of each word, it relates each word to all other words in the sentence to understand more about the word.

For instance, while computing the representation of the word it, our model relates the word it to all the words in the sentence to understand more about the word it.

As shown in the following figure, in order to compute the representation of the word it, our model relates the word it to all the words in the sentence. By relating the word it to all the words in the sentence, our model can understand that the word it is related to the word dog and not food. As we can observe, the line connecting the word it to dog is thicker compared to the other lines, which indicates that the word it is related to the word dog and not food in the given sentence:

Okay, but how exactly does this work? Now that we have a basic idea of what the self-attention mechanism is, let's understand more about it in detail.

Suppose our input sentence (source sentence) is I am good. First, we get the embeddings for each word in our sentence. Note that the embeddings are just the vector representation of the word and the values of the embeddings will be learned during training.

Let be the embedding of the word I, be the embedding of the word am, and be the embedding of the word good. Consider the following:

• The embedding of the word I is .
• The embedding of the word am is .
• The embedding of the word good is .

Then, we can represent our input sentence I am good using the input matrix (embedding matrix or input embedding) as shown here:

From the preceding input matrix, , we can understand that the first row of the matrix implies the embedding of the word I, the second row implies the embedding of the word am, and the third row implies the embedding of the word good. Thus, the dimension of the input matrix will be [sentence length x embedding dimension]. The number of words in our sentence (sentence length) is 3. Let the embedding dimension be 512; then, our input matrix (input embedding) dimension will be [3 x 512].

Now, from the input matrix, , we create three new matrices: a query matrix, , key matrix, , and value matrix, . Wait. What are these three new matrices? And why do we need them? They are used in the self-attention mechanism. We will see how exactly these three matrices are used in a while.

Okay, how we can create the query, key, and value matrices? To create these, we introduce three new weight matrices, called . We create the query, , key, , and value, , matrices by multiplying the input matrix, , by , , and , respectively.

Note that the weight matrices, ,, and , are randomly initialized and their optimal values will be learned during training. As we learn the optimal weights, we will obtain more accurate query, key, and value matrices.

As shown in the following figure, multiplying the input matrix, , by the weight matrices, ,, and , we obtain the query, key, and value matrices:

From the preceding figure, we can understand the following:

• The first row in the query, key, and value matrices – , , and – implies the query, key, and value vectors of the word I.
• The second row in the query, key, and value matrices – , , and – implies the query, key, and value vectors of the word am.
• The third row in the query, key, and value matrices – , , and – implies the query, key, and value vectors of the word good.

Note that the dimensionality of the query, key, value vectors are 64. Thus, the dimension of our query, key, and value matrices is [sentence length x 64]. Since we have three words in the sentence, the dimensions of the query, key, and value matrices are [3 x 64].

But still, the ultimate question is why are we computing this? What is the use of query, key, and value matrices? How is this going to help us? This is exactly what we will discuss in detail in the next section.

Understanding the self-attention mechanism

We learned how to compute query, , key, , and value, , matrices and we also learned that they are obtained from the input matrix, . Now, let's see how the query, key, and value matrices are used in the self-attention mechanism.

We learned that in order to compute a representation of a word, the self-attention mechanism relates the word to all the words in the given sentence. Consider the sentence I am good. To compute the representation of the word I, we relate the word I to all the words in the sentence, as shown in the following figure:

But why do we need to do this? Understanding how a word is related to all the words in the sentence helps us to learn better representation. Now, let's learn how the self-attention mechanism relates a word to all the words in the sentence using the query, key, and value matrices. The self-attention mechanism includes four steps; let's take a look at them one by one.

Step 1

The first step in the self-attention mechanism is to compute the dot product between the query matrix, , and the key matrix, :

The following shows the result of the dot product between the query matrix, , and the key matrix, :

But what is the use of computing the dot product between the query and key matrices? What exactly does signify? Let's understand this by looking at the result of in detail.

Let's look into the first row of the matrix as shown in the following figure. We can observe that we are computing the dot product between the query vector (I), and all the key vectors – (I), (am), and (good). Computing the dot product between two vectors tells us how similar they are.

Thus, computing the dot product between the query vector () and the key vectors tells us how similar the query vector (I) is to all the key vectors – (I), (am), and (good). By looking at the first row of the matrix, we can understand that the word I is more related to itself than the words am and good since the dot product value is higher for compared to and :

Now, let's look into the second row of the matrix . As shown in the following figure, we can observe that we are computing the dot product between the query vector, (am), and all the key vectors – (I), (am), and (good). This tells us how similar the query vector (am) is to the key vectors – (I), (am), and (good).

By looking at the second row of the matrix, we can understand that the word am is more related to itself than the words I and good since the dot product value is higher for compared to and :

Similarly, let's look into the third row of the matrix. As shown in the following figure, we can observe that we are computing the dot product between the query vector (good) and all the key vectors – (I), (am), and (good). This tells us how similar the query vector (good) is to all the key vectors – (I), (am), and (good).

By looking at the third row of the matrix, we can understand that the word good is more related to itself than the words I and am in the sentence since the dot product value is higher for compared to and :

Thus, we can say that computing the dot product between the query matrix, , and the key matrix, , essentially gives us the similarity score, which helps us to understand how similar each word in the sentence is to all other words.

Step 2

The next step in the self-attention mechanism is to divide the matrix by the square root of the dimension of the key vector. But why do we have to do that? This is useful in obtaining stable gradients.

Let be the dimension of the key vector. Then, we divide by . The dimension of the key vector is 64. So, taking the square root of it, we will obtain 8. Hence, we divide Q.K^T by 8 as shown in the following figure:

Step 3

By looking at the preceding similarity scores, we can understand that they are in the unnormalized form. So, we normalize them using the softmax function. Applying softmax function helps in bringing the score to the range of 0 to 1 and the sum of the scores equals to 1, as shown in the following figure:

We can call the preceding matrix a score matrix. With the help of these scores, we can understand how each word in the sentence is related to all the words in the sentence. For instance, look at the first row in the preceding score matrix; it tells us that the word I is related to itself by 90%, to the word am by 7%, and to the word good by 3%.

Step 4

Okay, what's next? We computed the dot product between the query and key matrices, obtained the scores, and then normalized the scores with the softmax function. Now, the final step in the self-attention mechanism is to compute the attention matrix, .

The attention matrix contains the attention values for each word in the sentence. We can compute the attention matrix, , by multiplying the score matrix, , by the value matrix, , as shown in the following figure:

Say we have the following:

The attention matrix, , is computed by taking the sum of the value vectors weighted by the scores. Let's understand this by looking at it row by row. First, let's see how for the first row, , the self-attention of the word I is computed:

From the preceding figure, we can understand that , the self-attention of the word I is computed as the sum of the value vectors weighted by the scores. Thus, the value of will contain 90% of the values from the value vector (I), 7% of the values from the value vector (am), and 3% of values from the value vector (good).

But how is this useful? To answer this question, let's take a little detour to the example sentence we saw earlier, A dog ate the food because it was hungry. Here, the word it indicates dog. To compute the self-attention of the word it, we follow the same preceding steps. Suppose we have the following:

From the preceding figure, we can understand that the self-attention value of the word it contains 100% of the values from the value vector (dog). This helps the model to understand that the word it actually refers to dog and not food. Thus, by using a self-attention mechanism, we can understand how a word is related to all other words in the sentence.

Now, coming back to our example, , the self-attention of the word am is computed as the sum of the value vectors weighted by the scores, as shown in the following figure:

As we can observe from the preceding figure, the value of will contain 2.5% of the values from the value vector (I), 95% of the values from the value vector (am), and 2.5% of the values from the value vector (good).

Similarly, , the self-attention of the word good is computed as the sum of the value vectors weighted by the scores, as shown in the following figure:

This implies that the value of will contain 21% of the values from the value vector v₁(I), 3% of the values from the value vector v₂(am), and 76% of values from the value vector v₃(good).

Thus, the attention matrix, , consists of self-attention values of all the words in the sentence and it is computed as follows:

To get a better understanding of the self-attention mechanism, the steps involved are summarized as follows:

1. First, we compute the dot product between the query matrix and the key matrix, , and get the similarity scores.
2. Next, we divide by the square root of the dimension of the key vector, .
3. Then, we apply the softmax function to normalize the scores and obtain the score matrix, .
4. At the end, we compute the attention matrix, , by multiplying the score matrix by the value matrix, .

The self-attention mechanism is graphically shown as follows:

The self-attention mechanism is also called scaled dot product attention, since here we are computing the dot product (between the query and key vectors) and scaling the values (with ).

Now that we have understood how the self-attention mechanism works, in the next section, we will learn about the multi-head attention mechanism.

Multi-head attention mechanism

Instead of having a single attention head, we can use multiple attention heads. That is, in the previous section, we learned how to compute the attention matrix, . Instead of computing a single attention matrix, , we can compute multiple attention matrices. But what is the use of computing multiple attention matrices?

Let's understand this with an example. Consider the phrase All is well. Say we need to compute the self-attention of the word well. After computing the similarity score, suppose we have the following:

As we can observe from the preceding figure, the self-attention value of the word well is the sum of the value vectors weighted by the scores. If you look at the preceding figure closely, the attention value of the actual word well is dominated by the other word All. That is, since we are multiplying the value vector of the word All by 0.6 and the value vector of the actual word well by only 0.4, it implies that will contain 60% of the values from the value vector of the word All and only 40% of the values from the value vector of the actual word well. Thus, here the attention value of the actual word well is dominated by the other word All.

This will be useful only in circumstances where the meaning of the actual word is ambiguous. That is, consider the following sentence:

A dog ate the food because it was hungry

Say we are computing the self-attention for the word it. After computing the similarity score, suppose we have the following:

As we can observe from the preceding equation, here the attention value of the word it is just the value vector of the word dog. Here, the attention value of the actual word it is dominated by the word dog. But this is fine here since the meaning of the word it is ambiguous as it may refer to either dog or food.

Thus, if the value vector of other words dominates the actual word in cases as shown in the preceding example, where the actual word is ambiguous, then this dominance is useful; otherwise, it will cause an issue in understanding the right meaning of the word. So, in order to make sure that our results are accurate, instead of computing a single attention matrix, we will compute multiple attention matrices and then concatenate their results. The idea behind using multi-head attention is that instead of using a single attention head, if we use multiple attention heads, then our attention matrix will be more accurate. Let's explore this in more detail.

Let's suppose we are computing two attention matrices, and . First, let's compute the attention matrix .

We learned that to compute the attention matrix, we create three new matrices, called query, key, and value matrices. To create the query, , key, , and value, , matrices, we introduce three new weight matrices, called . We create the query, key, and value matrices by multiplying the input matrix, , by ,, and , respectively.

Now, the attention matrix can be computed as follows:

Now, let's compute the second attention matrix, .

To compute the attention matrix , we create another set of query, Q₂, key, K₂, and value, V₂, matrices. We introduce three new weight matrices, called , and we create the query, key, and value matrices by multiplying the input matrix, , by ,, and , respectively.

The attention matrix can be computed as follows:

Similarly, we can compute number of attention matrices. Suppose we have eight attention matrices, to ; then, we can just concatenate all the attention heads (attention matrices) and multiply the result by a new weight matrix, , and create the final attention matrix as shown:

Now that we have learned how the multi-attention mechanism works, we will learn about another interesting concept, called positional encoding, in the next section.

Learning position with positional encoding

Consider the input sentence I am good. In RNNs, we feed the sentence to the network word by word. That is, first the word I is passed as input, next the word am is passed, and so on. We feed the sentence word by word so that our network understands the sentence completely. But with the transformer network, we don't follow the recurrence mechanism. So, instead of feeding the sentence word by word, we feed all the words in the sentence parallel to the network. Feeding the words in parallel helps in decreasing the training time and also helps in learning the long-term dependency.

However, the problem is since we feed the words parallel to the transformer, how will it understand the meaning of the sentence if the word order is not retained? To understand the sentence, the word order (position of the words in the sentence) is important, right? Yes, the word order is very important as it helps to understand the position of each word in a sentence, which in turn helps to understand the meaning of the sentence.

So, we should give some information about the word order to the transformer so that it can understand the sentence. How can we do that? Let's explore this in more detail now.

For our given sentence, I am good, first, we get the embeddings for each word in our sentence. Let's represent the embedding dimension as . Say the embedding dimension, , is 4. Then, our input matrix dimension will be [sentence length x embedding dimension] = [3 x 4].

We represent our input sentence I am good using the input matrix (embedding matrix). Let the input matrix be the following:

Now, if we pass the preceding input matrix directly to the transformer, it cannot understand the word order. So, instead of feeding the input matrix directly to the transformer, we need to add some information indicating the word order (position of the word) so that our network can understand the meaning of the sentence. To do this, we introduce a technique called positional encoding. Positional encoding, as the name suggests, is an encoding indicating the position of the word in a sentence (word order).

The dimension of the positional encoding matrix, , is the same dimension as the input matrix . Now, before feeding the input matrix (embedding matrix) to the transformer directly, we include the positional encoding. So, we simply add the positional encoding matrix to the embedding matrix and then feed it as input to the network. So, now our input matrix will have not only the embedding of the word but also the position of the word in the sentence:

Now, the ultimate question is how exactly is the positional encoding matrix computed? The authors of the transformer paper Attention Is All You Need have used the sinusoidal function for computing the positional encoding, as shown:

In the preceding equation, implies the position of the word in a sentence, and implies the position of the embedding. Let's understand the preceding equations with an example. By using the preceding equations, we can write the following:

As we can observe from the preceding matrix, in the positional encoding, we use the sin function when is even and the cos function when is odd. Simplifying the preceding matrix, we can write the following:

We know that in our input sentence, the word I is at the 0^th position, am is at the 1^st position, and good is at the 2^nd position. Substituting the value, we can write the following:

Thus, our final positional encoding matrix, , is given as follows:

After computing the positional encoding we simply perform element-wise addition with the embedding matrix and feed the modified input matrix to the encoder.

Now, let's revisit our encoder architecture. A single encoder block is shown in the following figure. As we can observe, before feeding the input directly to the encoder, first, we get the input embedding (embedding matrix), and then we add the positional encoding to it, and then we feed it as input to the encoder:

We learned how the positional encoder works; we also learned how the multi-head attention sublayer works in the previous section. In the next section, we will learn how the feedforward network sublayer works in the encoder.

Feedforward network

The feedforward network sublayer in an encoder block is shown in the following figure:

The feedforward network consists of two dense layers with ReLU activations. The parameters of the feedforward network are the same over the different positions of the sentence and different over the encoder blocks. In the next section, we will look into another interesting component of the encoder.

Add and norm component

One more important component in our encoder is the add and norm component. It connects the input and output of a sublayer. That is, as shown in the following figure (dotted lines), we can observe that the add and norm component:

• Connects the input of the multi-head attention sublayer to its output
• Connects the input of the feedforward sublayer to its output:

The add and norm component is basically a residual connection followed by layer normalization. Layer normalization promotes faster training by preventing the values in each layer from changing heavily.

Now that we have learned about all the components of the encoder, let's put all of them together and see how the encoder works as a whole in the next section.

Putting all the encoder components together

The following figure shows the stack of two encoders; only encoder 1 is expanded to reduce the clutter:

From the preceding figure, we can understand the following:

1. First, we convert our input to an input embedding (embedding matrix), and then add the position encoding to it and feed it as input to the bottom-most encoder (encoder 1).
2. Encoder 1 takes the input and sends it to the multi-head attention sublayer, which returns the attention matrix as output.

3. We take the attention matrix and feed it as input to the next sublayer, which is the feedforward network. The feedforward network takes the attention matrix as input and returns the encoder representation as output.
4. Next, we take the output obtained from encoder 1 and feed it as input to the encoder above it (encoder 2).
5. Encoder 2 carries the same process and returns the encoder representation of the given input sentence as output.

We can stack number of encoders one above the other; the output (encoder representation) obtained from the final encoder (topmost encoder) will be the representation of the given input sentence. Let's denote the encoder representation obtained from the final encoder (in our example, it is encoder 2) as .

We take this encoder representation, , obtained from the final encoder (encoder 2) and feed it as input to the decoder. The decoder takes the encoder representation as input and tries to generate the target sentence.

Now that we have understood the encoder part of the transformer, in the next section, we will learn how a decoder works in detail.

Suppose we want to translate the English sentence (source sentence) I am good to the French sentence (target sentence) Je vais bien. To perform this translation, we feed the source sentence I am good to the encoder. The encoder learns the representation of the source sentence. In the previous section, we learned how exactly the encoder learns the representation of the source sentence. Now, we take this encoder's representation and feed it to the decoder. The decoder takes the encoder representation as input and generates the target sentence Je vais bien, as shown in the following figure:

In the encoder section, we learned that, instead of having one encoder, we can have a stack of encoders. Similar to the encoder, we can also have a stack of decoders. For simplicity, let's set . As shown in the following figure, the output of one decoder is sent as the input to the decoder above it. We can also observe that the encoder's representation of the input sentence (encoder's output) is sent to all the decoders. Thus, a decoder receives two inputs: one is from the previous decoder, and the other is the encoder's representation (encoder's output):

Okay, but how exactly does the decoder generate the target sentence? Let's explore that in more detail. At time step , the input to the decoder will be <sos>, which indicates the start of the sentence. The decoder takes <sos> as input and generates the first word in the target sentence, which is Je, as shown in the following figure:

At time step , along with the current input, the decoder takes the newly generated word from the previous time step, , and tries to generate the next word in the sentence. Thus, the decoder takes <sos> and Je (from the previous step) as input and tries to generate the next word in the target sentence, as shown in the following figure:

At time step , along with the current input, the decoder takes the newly generated word from the previous time step, , and tries to generate the next word in the sentence. Thus, the decoder takes <sos>, Je, and vais (from the previous step) as input and tries to generate the next word in the sentence, as shown in the following figure:

Similarly, on every time step, the decoder combines the newly generated word to the input and predicts the next word. Thus, at time step , the decoder takes <sos>, Je, vais, and bien as input and tries to generate the next word in the sentence, as shown in the following figure:

As we can observe from the preceding figure, once the <eos> token, which indicates the end of sentence, is generated, it implies that the decoder has completed generating the target sentence.

In the encoder section, we learned that we convert the input into an embedding matrix and add the positional encoding to it, and then feed it as input to the encoder. Similarly, here, instead of feeding the input directly to the decoder, we convert it into an embedding, add the positional encoding to it, and then feed it to the decoder.

For example, as shown in the following figure, say at time step we convert the input into an embedding (we call it an output embedding because here we are computing the embedding of the words generated by the decoder in previous time steps), add the positional encoding to it, and then send it to the decoder:

Okay, but the ultimate question is how exactly does the decoder work? What's going on inside the decoder? Let's explore this in detail. A single decoder block with all of its components is shown in the following figure:

From the preceding figure, we can observe that the decoder block is similar to the encoder and here we have three sublayers:

• Masked multi-head attention
• Multi-head attention
• Feedforward network

Similar to the encoder block, here we have multi-head attention and feedforward network sublayers. However, here we have two multi-head attention sublayers and one of them is masked. Now that we have a basic idea of the decoder, first, let's look into each component of the decoder in detail, and then we will see how the decoder works as a whole.

Masked multi-head attention

In our English-to-French translation task, say our training dataset looks like the one shown here:

By looking at the preceding dataset, we can understand that we have source and target sentences. In the previous section, we saw how the decoder predicts the target sentence word by word in each time step and that happens only during testing.

During training, since we have the right target sentence, we can just feed the whole target sentence as input to the decoder but with a small modification. We learned that the decoder takes the input <sos> as the first token, and combines the next predicted word to the input on every time step for predicting the target sentence until the <eos> token is reached. So, we can just add the <sos> token to the beginning of our target sentence and send that as an input to the decoder.

Say we are converting the English sentence I am good to the French sentence Je vais bien. We can just add the <sos> token to the beginning of the target sentence and send <sos> Je vais bien as an input to the decoder, and then the decoder predicts the output as Je vais bien <eos>, as shown in the following figure:

But how does this work? Isn't this kind of ambiguous? Why do we need to feed the entire target sentence and let the decoder predict the shifted target sentence as output? Let's explore this in more detail.

We learned that instead of feeding the input directly to the decoder, we convert it into an embedding (output embedding matrix) and add positional encoding, and then feed it to the decoder. Let's suppose the following matrix, , is obtained as a result of adding the output embedding matrix and positional encoding:

Now, we feed the preceding matrix, , to the decoder. The first layer in the decoder is the masked-multi head attention. This works similar to the multi-head attention mechanism we learned about with the encoder but in a small difference.

To perform self-attention, we create three new matrices, called query, , key, , and value, . Since we are computing multi-head attention, we create number of query, key, and value matrices. Thus, for a head, , the query, , key, , and value, , matrices can be created by multiplying by the weight matrices, , respectively.

Now, let's see how masked multi-head attention works. Our input sentence to the decoder is <sos> Je vais bien. We learned that the self-attention mechanism relates a word to all the words in the sentence to understand more about each word. But there is a small catch here. During test time, the decoder will only have the words generated until the previous step as input. For example, during testing, say at time step the decoder will only have the input words as [<sos>, Je] and it will not have any other words. So, we have to train our model in the same fashion. Thus, our attention mechanism should relate the words only until the word Je and not the other words. To do this, we can mask all the words on the right that are not predicted by our model yet.

Say we want to predict the word next to the word <sos>. In this case, the model should see only the words up to <sos>, so we mask all the words on the right of <sos>. Say we want to predict the word next to the word Je. In this case, the model should see only the words up to Je, so we mask all the words to the right of Je, and the same applies for other rows, as shown in the following figure:

Masking words like this help the self-attention mechanism to attend only to the words that would be available to the model during testing. Okay, but how exactly we can perform this masking? We know that for a head, , the attention matrix, , is computed as follows:

The first step in computing the attention matrix is computing the dot product between the query and key matrices. The following shows the result of the dot product between the query and the key matrix. Note that the values used here are arbitrary, just to get a good understanding:

The next step is to divide the matrix by the dimension of key vector . Suppose the following is the result of :

Next, we apply the softmax function to the preceding matrix and normalize the scores. But before applying the softmax function, we need to mask the values. For example, look at the first row of our matrix. To predict the word next to the word <sos>, our model should not attend all the words to the right of <sos> (as this will not be available during test time). So, we can mask all the words to the right of <sos> with :

Figure 1.49 – Masking all the words to the right of <sos> with

Now, let's take a look at the second row of the matrix. To predict the word next to the word Je, our model should not attend all the words to the right of Je (as this will not be available during test time). So, we can mask all the words to the right of Je with :

Figure 1.50 – Masking all the words to the right of Je with

Similarly, we can mask all the words to the right of vais with as shown:

Figure 1.51 – Masking all the words in the right of vais with

Now, we can apply the softmax function to the preceding matrix and multiply the result by the value matrix, , and obtain the final attention matrix, . Similarly, we can compute h number of attention matrices, concatenate them, and multiply the result by a new weight matrix, , and create the final attention matrix, , as shown:

Now, we feed this final attention matrix, , to the next sublayer in our decoder, which is another multi-head attention layer. Let's see how that works in detail in the next section.

Multi-head attention

The following figure shows the transformer model with both the encoder and decoder. As we can observe, the multi-head attention sublayer in each decoder receives two inputs: one is from the previous sublayer, masked multi-head attention, and the other is the encoder representation:

Let's represent the encoder representation by and the attention matrix obtained as a result of the masked multi-head attention sublayer by . Since here we have an interaction between the encoder and decoder, this layer is also called an encoder-decoder attention layer.

Now, let's look into the details and learn how exactly this multi-head attention layer works. The first step in the multi-head attention mechanism is creating the query, key, and value matrices. We learned that we can create the query, key, and value matrices by multiplying the input matrix by the weight matrices. But in this layer, we have two input matrices: one is (the encoder representation) and the other is (the attention matrix from the previous sublayer). So, which one should we use?

We create the query matrix, , using the attention matrix, , obtained from the previous sublayer and we create the key and value matrices using the encoder representation, R. Since we are performing the multi-head attention mechanism, for head , we do the following:

• The query matrix, , is created by multiplying the attention matrix, , by the weight matrix, .
• The key and value matrices are created by multiplying the encoder representation, , by the weight matrices, and , respectively. This is shown in the following figure:

But why do we have to do this? Why do we obtain the query matrix from and key and value matrices from ? The query matrix essentially holds the representation of our target sentence since it is obtained from and the key and value matrices hold the representation of the source sentence since it is obtained from . But how and why exactly is this useful? Let's understand this by calculating the self-attention step by step.

The first step in self-attention is to compute the dot product between the query and key matrices. The query and key matrices are shown in the following figure. As we can observe, since the query matrix is obtained from , it holds the representation of the target sentence, and since the key matrix is obtained from , it holds the representation of the input sentence. Note that the values used here are arbitrary, just to get a good understanding:

The following shows the result of the dot product between the query and key matrices:

By looking at the preceding matrix, , we can understand the following:

• From the first row of the matrix, we can observe that we are computing the dot product between the query vector (<sos>) and all the key vectors – (I), (am), and (good). Thus, the first row indicates how similar the target word <sos> is to all the words in the source sentence (I, am, and good).
• Similarly, from the second row of the matrix, we can observe that we are computing the dot product between the query vector (Je) and all the key vectors – (I), (am), and (good). Thus, the second row indicates how similar the target word Je is to all the words in the source sentence (I, am, and good).
• The same applies to all other rows. Thus, computing helps us to understand how similar our query matrix (target sentence representation) is to the key matrix (source sentence representation).

The next step in the multi-head attention matrix is to divide by . Then, we apply the softmax function and obtain the score matrix, .

Next, we multiply the score matrix by the value matrix, , that is, , and obtain the attention matrix, , as shown:

Say we have the following:

The attention matrix, , of the target sentence is computed by taking the sum of value vectors weighted by the scores. To get some clarity, let's see how the self-attention value of the word Je, , is computed:

As shown, , the self-attention of the word Je is computed as the sum of the value vectors weighted by the scores. Thus, the value of will contain 98% of the values from the value vector (I) and 2% of the values from the value vector (am). This basically helps the model to understand that the target word Je means the source word am.

Similarly, we can compute number of attention matrices, concatenate them, and multiply the result by a new weight matrix, , and create the final attention matrix, as shown:

Now, we feed this final attention matrix to the next sublayer in our decoder, which is a feedforward network. Let's see how that works in the next section.

Feedforward network

The next sublayer in the decoder is the feedforward network, as shown in the following figure:

The feedforward layer in the decoder works exactly the same as what we learned in the encoder. Next, we will look into the add and norm component.

Add and norm component

Just like we learned with the encoder, the add and norm component connects the input and output of a sublayer, as shown in the following figure:

Next, we will look at the linear and softmax layers.

Linear and softmax layers

Once the decoder learns the representation of the target sentence, we feed the output obtained from the topmost decoder to the linear and softmax layers, as shown in the following figure:

The linear layer generates the logits whose size is equal to our vocabulary size. Suppose our vocabulary consists of only the following three words:

Now, the logits returned by the linear layer will be a vector of size 3. Next, we convert the logits into a probability using the softmax function, and then the decoder outputs the word whose index has a high probability value. Let's understand this with an example.

Suppose the input to the decoder are the words <sos> and Je. Now, the decoder needs to predict the next word in our target sentence. So, we take the output of the topmost decoder and feed it to the linear layer. The linear layer generates the logits vector whose size is our vocabulary size. Let the logits returned by the linear layer be the following:

Now, we apply the softmax function to the logits generated by the linear layer and obtain the probability:

From the preceding matrix, we can understand that the probability is high at index 2. So, we look for the word that is at index 2 in our vocabulary. Since the word vais is at index 2, our decoder will predict the next word in our target sentence as vais. In this way, the decoder predicts the next word in the target sentence.

Now that we have understood all the decoder components, let's put them all together and see how they work as a whole in the next section.

Putting all the decoder components together

The following figure shows the stack of two decoders; only decoder 1 is expanded to reduce the clutter:

From the preceding figure, we can understand the following:

1. First, we convert the input to the decoder into an embedding matrix and then add the position encoding to it and feed it as input to the bottom-most decoder (decoder 1).
2. The decoder takes the input and sends it to the masked multi-head attention layer, which returns the attention matrix, , as output.
3. Now, we take the attention matrix, , and also the encoder representation, , and feed them as input to the multi-head attention layer (encoder-decoder attention layer), which again outputs the new attention matrix.
4. We take the attention matrix obtained from the encoder-decoder attention layer and feed it as input to the next sublayer, which is the feedforward network. The feedforward network takes the attention matrix as input and returns the decoder representation as output.
5. Next, we take the output obtained from decoder 1 and feed it as input to the decoder above it (decoder 2).
6. Decoder 2 carries the same process and returns the decoder representation of the target sentence as output.

We can stack number of decoders one above the other; the output (decoder representation) obtained from the final decoder will be the representation of the target sentence. Next, we feed the decoder representation of the target sentence to the linear and softmax layers and get the predicted word.

Now that we have learned how the encoder and decoder work in detail, let's put the encoder and decoder together and see how the transformer model works as a whole in the next section.

To give more clarity, the complete transformer architecture with the encoder and decoder is shown in the following figure:

In the preceding figure, Nx denotes that we can stack number of encoders and decoders. As we can observe, once we feed the input sentence (source sentence), the encoder learns the representation and sends the representation to the decoder, which in turn generates the output sentence (target sentence).

We can train the transformer network by minimizing the loss function. Okay, but what loss function should we use? We learned that the decoder predicts the probability distribution over the vocabulary and we select the word that has the highest probability as output. So, we have to minimize the difference between the predicted probability distribution and the actual probability distribution. First, how can we find the difference between the two distributions? We can use cross-entropy for that. Thus, we can define our loss function as a cross-entropy loss and try to minimize the difference between the predicted and actual probability distribution. We train the network by minimizing the loss function and we use Adam as an optimizer.

One additional point we need to note down is that to prevent overfitting, we apply dropout to the output of each sublayer and we also apply dropout to the sum of the embeddings and the positional encoding.

Thus, in this chapter, we learned how the transformer works in detail. In the next chapter, we will get started with BERT.

We started off the chapter by understanding what the transformer model is and how it uses encoder-decoder architecture. We looked into the encoder section of the transformer and learned about different sublayers used in encoders, such as multi-head attention and feedforward networks.

We learned that the self-attention mechanism relates a word to all the words in the sentence to better understand the word. To compute self-attention, we used three different matrices, called the query, key, and value matrices. Following this, we learned how to compute positional encoding and how it is used to capture the word order in a sentence. Next, we learned how the feedforward network works in the encoder and then we explored the add and norm component.

After understanding the encoder, we understood how the decoder works. We explored three sublayers used in the decoder in detail, which are the masked multi-head attention, encoder-decoder attention, and feedforward network. Following this, we understood how the transformer works with the encoder and decoder, and then at the end of the chapter, we learned how we train the network.

In the next chapter, we will learn what BERT is and how it uses the transformer for learning contextual embeddings in detail.

Let's put our newly acquired knowledge to the test. Try answering the following questions:

1. What are the steps involved in the self-attention mechanism?
2. What is scaled dot product attention?
3. How do we create the query, key, and value matrices?
4. Why do we need positional encoding?
5. What are the sublayers of the decoder?
6. What are the inputs to the encoder-decoder attention layer of the decoder?

To learn more, check out the following resources:

• Attention Is All You Need by Ashish Vaswani, Noam Shazeer, and Niki Parmar, available at https://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf
• The Illustrated Transformer blog by Jay Alammar, at http://jalammar.github.io/illustrated-transformer