Inside Transformers: An In-depth Look at the Game-Changing Machine Learning Architecture — Part 4

Note: AI tools are used as an assistant in this post!

Let’s continue with the components.

Multi-Head Attention

The next idea on top of the self-attention mechanism (scaled dot-product attention) is Multi-Head Attention. We can see multiple heads in the attention mechanism as kernels in convolutional neural networks. Each head is one scaled dot-product attention, and all of these heads will work in parallel and have separate matrices. For each head (self-attention layer), use different Wq, Wv, Wk, then concatenate the results, Ai. The main transformer paper used 8 attention heads, i.e., Wq1, Wv1, Wk1, ..., Wq8, Wv8, Wk8. This allows the network to attend to different parts of the sequence differently. The following image shows the difference between one head and multiple heads.

main transformer paper

The following image shows the input and output dimensions of the multi-head attention.

source: Introduction to Deep Learning — Raschka

As shown, in the main transformer paper, the input dimension is T*de and the output dimension after concatenating the output of all attention heads will be T*dv*h = T*de, so the input and output sizes are the same. Note that T is the number of input tokens and dv = de/h = 64 where de=512, h=8 to have the output size of de = dv*h after concatenation.

For the last Linear layer after concatenation, the dimensions are as follows:

source: Introduction to Deep Learning — Raschka

The main transformer paper used do = dv*h = de to have the same output size as the input: de. So, the output of the multi-head attention would be again T*do=T*de.

To recap, the following slide shows a single head scaled dot-product attention procedure:

source: Introduction to Deep Learning — Raschka

Additionally, the following slide demonstrates the process for multiple heads, starting with input that is multiplied by various weight matrices of various heads to create Query, Key, and Value, to concatenating the outputs of various heads and sending them into the Linear layer:

source: Introduction to Deep Learning — Raschka

We have multi-head attention in both the encoder and decoder. As you can see in the transformer architecture at the beginning of this post, in the multi-head attention block in the encoder part, the same input is used to generate Query, Key, and Value. This is called self-attention because it tries to learn the interaction of the input with itself. On the other hand, the multi-head attention block in the decoder of the transformer architecture, after the masked multi-head self-attention layer in the decoder, uses the output of the encoder to generate Value and Key, and the output of the masked multi-head attention → Add&Norm to generate Query. This is called cross attention, as it tries to learn the interaction between input and output. This layer is designed to allow each position in the decoder to attend over all positions in the input sequence from the encoder.

In other words, when generating an output token, the decoder not only considers the previously generated tokens (through masked self-attention) but also takes into account the entire input sequence (through cross attention). This process is crucial to tasks like machine translation where a word’s translation often depends on multiple words or the entire context of the input sentence.

In terms of implementation, the cross-attention mechanism works similarly to the standard (self-) attention mechanism. The difference lies in what is used as queries, keys, and values:

• Queries come from the previous decoder layer (the output of the masked self-attention layer).
• Keys and values come from the output of the encoder.

This way, the decoder can focus on different parts of the source sequence for every target position.

To give a specific example, let’s consider a translation task from English to French. The encoder takes in the English sentence and generates a sequence of representations. Then, while generating the French translation, for each French word (each step in the decoder), the model uses the cross-attention to focus on different words in the English sentence. This helps it understand the context and semantics of the source sentence and produce a coherent and accurate translation.

But what is the masked version of the multi-head attention block. Masked multi-head attention is a variant of multi-head attention used in the decoder part of the Transformer model to prevent the model from “seeing” future tokens in the output sequence during training, which could lead to unrealistically good performance and problems at test time. This aligns with the autoregressive nature of sequence generation tasks like translation, summarization, etc., where you generate one token at a time.

In the standard multi-head attention mechanism, the attention scores are calculated for all pairs of tokens. In contrast, in the masked version, we add a mask to the attention scores before they’re passed through the softmax function, effectively zeroing-out positions that correspond to ‘future’ tokens. This mask ensures that when generating the output token at position i, the model only has access to output tokens at positions less than i.

To explain further, consider the sentence “The cat sat on the mat”. When the model is predicting the word “sat”, we don’t want it to have access to “on the mat”, because in a real-world scenario, it won’t have this future information. So we mask or hide this information, forcing the model to make the best prediction based on the words “The cat” only.

To perform the masking operation, an upper triangular matrix with entries set to a very large negative value (like -1e9) is created. This mask is then added to the scaled dot product of the query and key vectors. Because these large negative values are then passed through a softmax function, they become very close to zero, effectively ignoring the future tokens.

This mechanism is vital to maintain the causality property in the decoder, preventing the current output token from depending on future output tokens.

source: Introduction to Deep Learning — Raschka

The following slide shows another example for masking:

source: Introduction to Deep Learning — Raschka

In terms of learning different concepts via different heads, there is no guarantee that different attention heads in the Transformer architecture will learn different concepts. However, the hope is that by initializing the weights of the model randomly, and by having multiple heads that learn in parallel, the model will find different useful patterns to pay attention to.

Evidence suggests that different heads do indeed learn different types of attention patterns. For example, some heads may specialize in syntax (like paying attention to the previous word), others in longer-range dependencies (like matching parentheses in a sentence), or specific semantic relationships (like attending from an object to its attributes).

We can analyze this post-training by looking at the attention patterns that different heads produce on a dataset, or using probing tasks that aim to determine what kind of information is captured by different layers or heads in the model.

However, the interpretability of attention heads is still an open research question. While we can find heads that seem to correspond to certain intuitive patterns, many heads do not correspond to anything easily interpretable. Furthermore, recent research has shown that the attention weights might not be as interpretable as we’d like to think, and we should be cautious in over-interpreting them.

In terms of encouraging heads to learn different things during training, there’s ongoing research into techniques for doing this, such as using regularization terms in the loss function that encourage diversity among the heads.

Let’s continue in the next post :)

Thank you for taking the time to read my post. If you found it helpful or enjoyable, please consider giving it a like and sharing it with your friends. Your support means the world to me and helps me to continue creating valuable content for you.