The attention mechanism is one of the key pieces of technology inside LLMs and many other prominent systems in robotics, computer vision etc. This article attempts to explain how it processes data. It was first introduced in a 2017 landmark paper called Attention is all you need
Imagine you are given a sentence :
“The student did not answer the question because it was too difficult”
What does it refer to? The student or the question?
We humans know immediately it refers to “question” by looking at the other words present in the sentence. But how will a model implement this process?
Before the attention mechanism, recurrent neural network (RNN) was the default technique used. It processed text sequentially from left to right. It was slow, had vanishing or exploding gradients and there was loss of information for long input sequences. The attention mechanism on the other hand lets each token communicate with every other token and transfer information.
To easily understand the attention mechanism, we will explain it in the context of an LLM. The goal of an LLM is to take in a piece of text and predict the next word. The first step is to tokenize the text i.e. break it into small units called tokens. Each token is then associated with a corresponding high dimensional dense vector that captures its semantic meaning called its embedding.
In this vector space where all the embeddings exist, direction might correspond with semantic meaning such that words with similar meanings appear close to each other. The embeddings of lion, zebras and elephants will be clustered as they represent the same ideas ie animals.
Embeddings come from an embedding matrix where the number of rows represent the total vocabulary size and the number of columns the embedding dimensions. Each token has its own lookup position in the matrix. Embeddings however only capture individual word meanings and does not distinguish words with different meanings.
Imagine these two sentences :
- We were stuck in a traffic jam for hours
- I love my toast with strawberry jam
The word jam in the two sentences does not convey the same idea, but the embedding table assigns a single embedding to jam. Ideally we want the embedding of jam to adapt based on its meaning in context such that in the first sentence it would be closer to traffic than and in the second sentence it will be closer to food.
Capturing context is possible thanks to the attention mechanism. It is the process by which tokens communicate with each other and incorporates contextual meaning in their general embedding transforming their semantic meaning into contextual meaning and making sure words are understood in their specific context
The attention mechanism (Q,K,V) is represented by this formula :
To understand this we will need to understand what the Q, K and V matrices represent.
We said earlier that the goal of an LLM is to take in a piece of text which is broken down into tokens and predict the next token.
In the sentence “The student did not answer the question because it was too difficult” assuming every word is a token. Each token is associated with its corresponding embedding. The embeddings are then stacked together row-wise forming a matrix which constitutes the input (X) to the attention mechanism.
To get context information i.e compute attention, tokens interact with each other. They do this via their query, key and value vectors. The input X is projected into three different spaces by multiplying it with three different learned weight matrices Wq, Wk, and Wv forming the Query(Q), Key(K) and Value(V) matrices. The entries of the different weight matrix are parameters of the model and their true behavior is learned during training.
Each row in the query matrix is a query vector representing a token. Same for the key and value matrices. The query represents what a word is looking for and the key represents what a word has to offer.
To understand how tokens relate to each other, we perform the dot product of the query and the transpose of the key matrices which gives the attention scores. The row of the dot product operation represent the word in the context asking a question and the columns represent the words being evaluated. The model computes a score for every possible token pair. A high score means the token assigns more value to that specific token. These scores tell us how much each token attends to one another, i.e. how much focus to place on the other token.
The attention scores are still just raw numbers. The ideal would be to then normalize them so they sum to 1 and produce valid attention weights with a valid probability distribution. But since the dimension of the key matrix can be extremely large, the scalar attention scores can easily vanish or explode if normalization is applied. Dividing the dot product by the square root of the dimension of the key and then normalizing yields proper attention weights.
Every row of the attention score is now a valid probability distribution and sums to 1 revealing how much focus each token places on the other.
Using these weights we compute the weighted sum of the value vectors i.e. we multiply each value by how much attention it receives to produce the new embedding.
The output is embeddings that have captured their contextual meanings. Every token embeding is deeply enriched with the contextual meaning of the sentence in which it appears..
By letting every word interact with every other word, we were able to capture the full structure and meaning of language
The attention mechanism gives every token the ability to see the whole sentence and updates its embedding.
Leave a Reply