Learning Transformers via Question and Answer
Behind all the GPT-x and other Large Language Models (LLMs), there’s only one thing that is common between all these LLMs and is also the reason to make a significant leap in the field of NLP (or NLU) after RNNs and LSTMs. Before Transformers, people were dependent on RNNs, bi-directional LSTMs, and also the Encoder-Decoder architecture that still exists in the transformer's architecture. What has changed is the addition of some stacked attention layers in the Encoder as well as the Decoder. There are quite a number of places where you can learn about transformers including the Stanford CS 224n videos on YouTube. What I intend to do here is to ask as well as answer certain questions about the transformers that will help anyone answer why transformers exist today, what is their benefit in comparison to the previous architectures, and even for any kind of interview questions on transformers, attention or any deep learning interviews. So, let's get started.
I have NOT used ChatGPT to answer these questions but watched these videos and the book called Learning Transformers with Hugging Face to create questions and answers for transformers.
Question 1 — What were the challenges with the previous architectures for sequence-to-sequence problems like encoder-decoder using the RNN?
Answer —There are two major reasons:
Linear Interaction Distance: As you can see from the above figure 1, that RNNs have issue of interaction between the tokens that could be linear distance apart from each other. This means as the pair of words are far apart, say at two extremes of the sentence, RNNs capture weak interaction between such a pair. Technically or mathematically, this is due to the problem of back propagating the gradients that either keep vanishing or exploding. Even though LSTM tries to solve this problem but not perfectly. So, it is hard to learn the long-distance dependencies in RNNs or even LSTMs.
Lack of Parallelizability: This problem is very much related to the previous issue. We cannot compute the future hidden states until the past RNN hidden states have already been computed. This may not be an issue in smaller datasets but for large datasets this is a bottleneck as we cannot parallelize the computation across time (not layers). From the figure 3, you’d notice that the number of steps required to compute a hidden state at step T, depends on the previous T-1 steps.
It is natural to think that one has to move sequentially in a sentence to understand the context and semantic meaning of the sentence, but this is true during inference actually. Basically, think of two scenarios.
While learning a new language we may not necessarily see the previous words to current word. Rather, the full sentence is at our disposal. And if you recall your language learning experience, you’d notice that we do end up seeing the sentence back and forth and entire sentence at once in order to understand the meaning and the context of the sentence. Well, this is nothing but what happens during the training of any language model. So, until the era of RNNs or even LSTMs, we kept moving sequentially even while training and hence no parallelizability.
While completing a sentence in new or even old language it is natural to predict one word at a time and the current word depends on the previous words predicted. What is sequential in true sense is the inference but not the training time which is what the most time-consuming process for LLM.
Question 2— Why do we attention or how does it solve the problems of long-distance interaction and parallelizability as mentioned above?
Answer — With attention as a building block in both the encoder and the decoder, one can look at the entire sentence (at once!) to get the current word context. This solves both the problem of long-distance interaction and the problem of parallelizability as each hidden state can be computed independently. From below figure 3, one can notice that each word is using the entire sentence at once to get the context information and hence every other word can do the same independently. Note that the parallelizability offered by is across time (sequence length) but not across the layers. Hence attention brings in the following:
parallelizability — ability to look at entire context at once. The trick is to treat the sentence as a set of words than a sequence.
resolving the long-distance interaction problem by looking at entire context and not just nearby local words.
Question 3 — What is the interpretation of key, query and value in the attention.
Answer — This was inspired from the Information Retrieval(IR) world. Suppose we are visiting a supermarket to buy some ingredients to prepare a dish. Each of the ingredients can be interpreted as queries, the labels mentioned on the items in the supermarket as the keys, and when there’s a match, one can take the item(value) from the shelf. In simple cases explained in any video or blog, you’d notice the key, query and values are same as the original embedding of the word.
One straightforward reason to take key, query and value is to be able to incorporate the attention from all the words in the current word as shown in the figure 5. Notice that the SoftMax function is used in the interim step which should be remembered for the upcoming questions.
In general, and in practice, this triplet of key, query and value are usually the linear transformation of their original embedding, i.e.
key(i_th word) = K.W_embedding(i)
query(i_th word) = Q.W_embedding(i)
value(i_th word) = V.W_embedding(i)This means for explanation people assume an Identity matrix for all the three cases. However, in practice they are different matrices. The linear transformation is required as this allows a generic equation in Figure 5. and hence incorporate different meanings of a word matching against the other meanings of the other words in the sentence.
Question 4— Why do we need multi-head attention?
Answer — From the previous two questions it should be clear why attention and self-attention are required. Multi-head attention is enabling parallelization as the single step is done via, say 12(equal to the number of heads) parallel steps. And each sub-attention can capture multiple type of association between the pair separately in parallel, giving stronger reason for multi-head attention. For e.g., one part of multi-head attention captures the Noun-Verb relationship, while the other captures the gender association between the same pair.
From the above figure, one can notice how do we compute multi head attention. For the query matrix (for Linear transformation) used for self-attention, there are smaller (in width but same in height or rows) query matrices used for the same embedding or data matrix. This gives us different query vectors each capturing one type of attention or association between the same pair of words in the sentence.
Question 5— Why do we need to scale the dot product between the query and key vectors before the SoftMax operation?
Answer — I have deliberately kept a hint in the question as one could just skip the SoftMax part in the question and still the answer would be same. This suggest the hint is in the question itself. The query and key vectors are in d dimension which could be large enough to magnify the dot product. But this reason makes the question trivial and hence there is more detail to it. If you pay attention to the attention mechanism, you’d notice that after the attention layer, one adds a Batch Norm layer as well. Now when we are doing scaling in the batch normalization step, why do we need one more scaling in the attention layer! This makes it non-trivial now.
One can remember two main reasons to this question.
- Controlling the magnitude of the dot product: One is to prevent the magnitude of the dot product to become too large with the dimension. Without scaling, when the dimension of the vectors is large, the dot product can also be large, leading to large values in the SoftMax function. This can result in very small gradients during the backward pass of the training process, making the optimization more difficult. It can cause the SoftMax function to produce probabilities(see alphas in figure 5) that are close to 0 or 1, resulting in gradients that are close to 0 during training. By scaling the dot product by the square root of the dimension, the values fed into the SoftMax function are more controlled. It prevents the exponentiated values from growing too large, keeping the probabilities more balanced and the gradients more informative during backpropagation.
- Balancing Attention Heads: In the case of multi-head attention, scaling helps to balance the contributions of different attention heads. Each attention head computes its own scaled dot product (see figure 6), and scaling ensures that the dot products are of similar magnitudes across all attention heads. This balance prevents certain attention heads from dominating the attention weights and ensures a more equitable distribution of attention across the sequence.
The short answer as to why not just batch norm instead of scaling during the attention layer/computation is to make sure attention itself is computed with equal distribution to all the attention heads in case of multi head attention and hence not letting some probabilities to go to 0 or 1 giving huge/zero importance to certain aspect of that context or certain part/context of the sequence.
Question 6 — Does attention layer adds any non-linearity in the neural network? If yes how else how to add the non-linearity when incorporating the attention layer?
Answer —Yes, it does (due to SoftMax operation involved after linear transformation of the embedding vectors) but this nonlinearity only help in getting the right context for the word. However, in order for the neural network to learn complex relationships in the dataset, we need nonlinearity in terms of an activation function. There’s no activation function used on top of the output of the attention layer. Hence, we add a Feed Forward layer explicitly in the transformer architecture to bring nonlinearity.
As you can see from the above figure, Feed Forward layer is added after the Attention layer. This ensures that we add nonlinearity in the architecture.
Question 7 — What is the use or advantage of a residual layer or skip connections?
Answer — From the figure 8 of the transformer architecture, one can observe the addition of skip connections that are fed into the Add&Norm layer directrly without going via the multi-head attention layer. There are two reasons to this. One is logical and other is more technical using gradient in its explanation.
- Logical reason is to know what the difference in the learning due to the previous layer, multi-head attention in this case. Due to skip connection or a residual connection, x(2) = x(1) + L(x(1)) in comparison to when skip connection is not added, x(2) = L(x(1)). One can see in the case of skip connection, It is easy to interpret what is the contribution of previous layer, L when both the original input and the transformed input due to the layer are present.
- Technical reason is to train the model better by making sure even if gradients saturate due to sigmoid function or vanish due to ReLu function with all the values being negative, there’s still nonzero gradient to keep the learning going. In general, skip connections are used here to resolve the vanishing or exploding gradient problem as a regularization step.
Question 8 — What happens during the decoder layers to incorporate the attention mechanism?
Question 8 (alternative) — What are key, value and query in multi-head corss attention used in the decoder?
Answer — During the decoder step, we use something called as multi-head cross attention that uses each of the hidden state output from the encoder. In case of this cross attention, the cross part is due to the key and values coming from the encoder’s hidden states while the query is from the decoder’s input. Encoder’s output in the form of hidden states at each step is interpreted as some kind of a memory that the decoder refers to. So, the keys are providing some indices in that memory that eventually refers to the values inside that same memory. This cross-attention step is what gives better translation results in comparison to the traditional Neural Machine Translation (NMT way until 2017's) of doing machine translation.
From the figure 9 above, one can see the operation is same except the fact that the Key and Value vectors are coming from the Encoder’s Hidden States H, while the Query vectors Z are coming from the decoder’s input.
Question 9 — How to prevent leaks going into the decoder?
Answer — Leaks going into the decoder means that the future tokens in the sequence are not available during the prediction of current token. While the sentence or its context is being understood in the encoder layers, it is natural to see the entire sentence, but the decoder task is to predict a new sentence, say in a different language for machine translation or answering a question for a question-answer decoder only architecture like that of OpenAI’s GPT-x. So how does one prevent leaks in the decoder that could otherwise be a normal step in the encoder layers! Leaks are prevented in the decoder layer(s) by using a technique called masking and the layer called as Masked Multi-Head Self-Attention. Just like the Self Attention and Multi-Head Attention covered above in question 2,3 and 4, decoder also applies this similar layer after the embedding step. The only difference is not allowing the decoder to view the entire sentence but only the previous words from the current word being predicted. Because decoder’s task is to use the encoder and previous words as well to predict the next word. Imagine you are doing machine translation or building a question answer model architecture where one has to incorporate a decoder layer. One does not have the entire sentence available but entire question or entire original sentence is definitely available from the encoder. Entire question or the entire sentence in the original language is used via the multi-head cross attention discussed in the previous answer. But the entire sentence in the destination language cannot be available as that would be like seeing the future to predict the future!
Also note that every layer in the decoder has the masking enabled meaning every layer just sees the previous words generated by the decoder.
Question 10 — What is the issue with using self-attention?
Question 10(alternative) — What is the time complexity in T and d, where T is the length of the sequence and d is the dimension of the embedding?
Answer — The issue with self-attention is the time complexity itself which is square in the sequence length T and linear in the dimension d.
T(self-attention) = O(T^2 d)If you’ve noticed LLMs putting a limit on the number of tokens being input in the prompt is primarily due to this reason. Because of quadratic complexity in the sequence length, it is difficult to use larger documents to train as well as predict. There are different strategies and papers showing how to get rid of this quadratic self-attention complexity.
- Wang et.al. 2020 suggested the Lin-former where they projected the original sentence in T dimension to a lower dimension k such that k squared is equal to T and the entire self-attention proceeds as usual in the projected dimension. So, this makes it linear in T and hence the name as Lin-former compared to Transformer!
- Zaheer et. al. 2021 suggested Big Bird, some family of interactions than looking at all the pair of interactions. Local windows, Looking at everything and random interactions.
All these are good suggestions but in general the problem remains due to the sequence length.
I hope this question answer way of understanding the transformers and various attention mechanisms will help the purpose. Please do leave any other question you have about transformers or encoder-decoders or attention mechanisms itself.
