The Key to Modern AI: How I Finally Understood Self-Attention (With PyTorch)

You've heard of the paper "Attention Is All You Need" and how it started the current AI boom that we're in. Well, once you understand attention, you'll have unlocked the most important thing you need to understand AI today. I spent the last month delving deep into self-attention specifically, just so I could wrap my head around it. Today, I'm going to walk you through three levels of understanding. Finally, we'll explore some code to show how this actually works in PyTorch.

Understanding attention is one of the most difficult concepts to grasp, but it unlocks all the cutting-edge work that's going on today. At a high level, it really just means the meaning of the other words in your context window. Think of an ambiguous word like "light." On its own, it can have many different meanings, but given the context around the word, its meaning can change. When a large language model starts, the word "light" is ambiguous and is represented by the average of its many possible meanings.

Then, by using the context of the sentence, the LLM can update its meaning of the word "light" to capture a richer meaning. Self-attention helps the model understand how much each word should pay attention to every other word in the sequence. Right now, we are just focused on the word "light," but this process applies to every word in our context window. At a high level, just remember that attention is all about understanding the context and using that context to update the meaning of every token in your sequence. If you want to think of the process like this: in self-attention, we take our input sequence, which consists of all the words in our context window, we pass it through the self-attention mechanism, and on the other side, we receive the input enriched through that self-attention mechanism.

And now, on the other side, we get the input enriched with the context that self-attention has picked up. The way we compute this enriched meaning is through self-attention, which involves three high-level steps. First, we compute the attention score, which essentially measures how much a specific token should care about the other tokens in its context window. In other words, it assesses how much attention a given token should pay to the other tokens. Next, we perform some mathematical operations to normalize that attention score into attention weights.

Again, they mean the same thing. How much should I care about the other tokens in this sequence? How much attention should I pay to the other tokens in my sequence? The third step is to update the meaning. We do this by taking the attention weights and multiplying them by the meanings of all the other words in the sequence, such that we get a weighted sum with every contextual word contributing its amount of weight to the new definition.

And that gives us the input enriched with context. So now we have these embeddings of the tokens with the context added through the process of self-attention. The key takeaway is that what self-attention does is add this context to our input embeddings, allowing us to capture the meaning as the words flow through our network. All right, now let's walk through an example of this process using a new sentence. When we compute self-attention, it's always within the bounds of our context window.

In our example here, the context window is this eight-word sentence. I am going to illustrate this process for just the word "lightly" here, to keep with our example of ambiguous words and their meanings being updated. However, this process of attention will be carried out for every individual word. I'll hide our process over here in the corner. So, this sentence is our input sequence.

It is our context for this computation of attention. We then compute the attention score by taking every word and comparing it to each other word in the sequence. For our example, again, we're just going to use the word "lightly" here. So, to compute the attention score, we take the word "lightly" and we basically ask every word in the sequence, "Hey, are you similar to me? Do you contribute to my definition?" They all return a score, and these scores are just a set of numbers—one for every token in the context window. These scores then get turned into weights that will sum to one, which are just a normalized version of the attention score. We then take these attention weights to figure out how much each word in our context will contribute to the updated definition. If these are the attention weights for the word "lightly," what we do is take each word, and each word contributes its weighted amount of meaning to the new definition.

In this example, the thickness of the arrow represents how much weight each word is contributing to the new meaning, which is again based on the attention weights. Now, the word "lightly" in the output of our attention mechanism has this amount of meaning added to it. Conceptually, this is how we compute attention and add context to a word, allowing the model to capture a richer and more nuanced meaning.

All right, before we go ahead, I just want to run through some quick prerequisites that'll help with your understanding. If you remember this from high school math, this is a two-dimensional grid. We have the x-axis and the y-axis, and this point here, (3, 2). This point can be represented as this dot here on the grid, or as this coordinate (3, 2) over here. It's basically capturing the same information, which is that we have three on the x-axis and two on the y-axis.

Now, a way I like to visualize is to have these axes as horizontal number lines. This still captures the same information in point 3-2. The reason I'm talking about this is that words in large language models are represented as high-dimensional numbers with over 12,000 different dimensions. I find it helpful to visualize high-dimensional points like this, where there are just different weights or different values on these various axes. It still captures the same meaning.

This is a point in five-dimensional space. However, if you extrapolate this out, imagine a word having a point in a 12,000-dimensional space. Now, every word would share that same 12,000-dimensional space, but each would have its own unique point within it to capture the meaning of that word. For example, if we consider the word "light," our five-dimensional space, for illustrative purposes, will have different meanings along the axes. Instead of having arbitrary axis names, we can conceptualize the axes in this 12,000-dimensional space of words.

In our example, we can think of this five-dimensional space as having different meanings associated with each axis. When we move up and down the weight axis, we change how much the word "light" relates to the meaning of weights or weightiness. The same applies to color; we adjust how much the word "light" relates to a color. While the dimensions in a large language model aren't always interpretable by humans in this way, for our example, we will use these five dimensions to help us understand the word in question, which is "light." The word "light" is ambiguous.

And so a large language model, in the beginning, will have some average meaning of all the different things this ambiguous word could mean. Then, by using self-attention, the model will update the meaning of the word "light" by looking at the other words in the sentence. In this case, we have "bright"; "mood" might be a little bit related; it definitely doesn't mean anything related to "volume"; "color" could be impacted; and "weight" not at all. That's why context is important—it helps us frame or derive richer meanings for the words in our window. Now that we understand, at a high level, what self-attention is trying to do, we take this input sequence, compute an attention score, normalize that into an attention weight, and then update the meaning using those attention weights from the previous step. This process gives us an input that's enriched with the context of all the words around it.

So, if we start with our input sequence here, this sentence represents our entire context window. The scope of all the words around it is included within this context. Therefore, the attention scope is limited to the context window. We first need to compute the attention scores.

And now, the way we do this and the way we're able to train large language models is that we have to have a set of weights that we can train during the large language model training process. These sets of weights capture different relationships, different meanings, and various concepts that we humans can't always interpret or understand. If you remember earlier, the attention score starts by having the word in question—so in this case, "lightly"—take itself and look at what it needs to know from these other tokens. Remember, for every token in the sequence, or in our case, we're just using words, and for every word in our context window, each word is going to ask, "What can you tell me about myself?" This is what the query matrix is. The query matrix is just a set of weights that helps each token know what question to ask about itself. Now, how does it know what question to ask about itself? Well, let's just wave our hands and say it's magic, right? This occurs through the training process.

But you can think of this matrix as just a way for each token to know what question to ask. To say, "Hey, I'm the word 'lightly.' What can you tell me about myself? Or can you even answer a question about me?" That's what the query matrix is.

And so, if we're just computing attention for this single word right now, we take the vector from this query matrix just for the word "lightly" in this example. Then, we perform the dot product with what's called the key matrix. The dot product helps us measure the similarity to determine how well a given key answers a specific query. So, remember, in this case, the query matrix is asking, "What information can I get from the other tokens? What information do you provide from the other tokens?"

It helps the word "lightly" ask the question, "What other words are relevant to me?" The key matrix is what answers that query. Each of these tokens has a vector in the key matrix that indicates its relevance or the answers it can provide. When you perform the dot product between the query vector of the word "lightly" and the key matrix, it determines what information can be derived from the other tokens.

And then the key matrix indicates that for each token, this is the information I can provide. When you perform the dot product between those two, if the question is being answered by the key, we obtain a higher attention weight. This process is repeated for every token in the sequence. In our case, we consider the word lightly and use its query vector to ask the question: what information can I obtain from other tokens, or what information do I want from other tokens? We calculate the dot product between the values in the key matrix for every token. Essentially, the key matrix states, this is the information that I can provide about my token.

And we do the dot product there to get the attention scores. When we have a question that is asked of the key matrix and is similar—similar in the sense that they are pointing in roughly the same direction, or as close to the same direction as possible in a 12,000-dimensional space—we get a higher attention score. The attention scores are normalized into attention weights using a mathematical function called softmax. You can think of softmax as jumps, which allows us to achieve a higher attention score.

The attention scores are normalized into attention weights using a mathematical function called softmax. You can think of softmax as a mechanism that helps normalize these scores so that every score sums to one, allowing us to interpret these scores as probabilities or weights. It also contributes to numerical stability, making the training process more efficient. Once we have our attention weights, we use these weights for our word embeddings and multiply them by the value matrix. Inside the value matrix, each of these tokens has its own vector, which represents the information it can contribute to other words. When we apply the weights to the query vector, we are indicating how much each word in this sentence should contribute to its meaning. We then obtain the weighted sum of all that new information. Initially, all these matrices start out as random values.

And then, through the process of training a large language model, they each begin to learn together to answer and ask different questions and capture that meaning in the value matrix. We can't control what answers or questions they're asking and what information they're capturing; it is simply done through the training process. That is part of the magic of neural networks and large language models. Once we take the attention weights and multiply them by the value matrix, we obtain a new vector for the word "lightly" that has a richer and more nuanced contextual meaning.

And again, under the hood, what really happened was the word "lightly" started over here as one point in our 12,000-dimensional space, and it became a different point entirely over here that captured more nuanced and richer meaning. All right, so now let's head over to actually implement this in code to show you an example. We are going to import some libraries in this Jupyter notebook, and then we will create actual embeddings for our dummy example. Our batch size is one because we are only going to have one training sample. Our context length is eight, so we have eight words in our context window.

And then each word is assigned five dimensions. If you recall, we're going to take our inputs, compute an attention score, then normalize that into an attention weight, update the meaning, and generate an output that is our context window enriched with the sentence's context. So here is our input text. We're just going to create our embeddings. What we have here is our PyTorch tensor with a batch size of one; we have one training sample.

We have eight different rows here. Each row has five dimensions. You could think of each row as one token or one word in our sentence. In our case, we're using the words as individual tokens, but in reality, tokens are more like subwords. Thus, each row corresponds to one word in our input text.

And this is the five-dimensional point that represents the word "the." This is the five-dimensional point that represents the narrator. And here is our five-dimensional point that represents the word "lightly." Now, these actual numbers here are just randomly generated, created for illustrative purposes.

If we wanted to use real embeddings, we'd have to download some sort of pre-trained language model and use that, which we can do in another video. But for this example, let's just keep it light—no pun intended—by using random numbers here to illustrate how this works. So we have our input embeddings. Now we need to compute the attention score. Remember, to compute the attention score, we take the query vector of every token in our context window and multiply that by the key matrix.

That will give us our attention score. Now we initialize our query, key, and value matrices. We take our token dimensions as our input dimension, which is five for our five-dimensional tokens. Then, we project it onto a four-dimensional space.

Typically, you will have multiple attention heads in a transformer block. We're just using one, but in general, these weight matrices are smaller than the actual context window. You would have multiple attention heads, which end up learning different aspects of the context itself. Here, we can see our 5 by 4 output matrix for our query, key, and value matrices.

And so, from here, it's time to compute our attention scores, right? Except instead of doing just one word in the sequence, we're going to compute the attention scores for the entire sequence at once. Rather than taking the vector here, we will obtain the query vectors for every word by multiplying our input by the query matrix. We take our entire input and multiply it by the query matrix, which gives us all the query vectors. Now that we have all the query vectors and all the key vectors, we compute the dot product, which gives us our unnormalized attention score for every token. Next, we compute the attention weights by normalizing. Remember, we have our attention scores, and we will run them through softmax to convert them into attention weights. We compute the attention weights by taking the softmax of our attention scores and normalizing by the square root of our embedding size.

Again, this normalization by the square root is simply to help with numerical stability and training, which is outside the scope of this video. Now we have our attention weights, which look like this. These weights indicate how much each token should pay attention to every other token in the sequence. If we take one of these rows and sum it, we should get one, and we do. These are now our normalized attention weights.

And so the next step is with the attention weights. We multiply the weights by the value matrix to obtain all of our updated vectors. These attention weights are assigned to every word in the sequence. For example, the first row here corresponds to the word "the," and the second row corresponds to the word "narrator." We take the weights for each of these words and multiply them by the value matrix to get the updated context vectors.

Now that we have our context, we take the weights for each one of these and multiply them by the value matrix to get the updated context vectors. We now have our context vectors, which, as you may have noticed, are four-dimensional vectors for each word. This is because our attention head has a size of four. These context vectors are then fed into a multilayer perceptron, which is a neural network that learns from this enriched meaning rather than taking the raw input meaning. This process is repeated multiple times in parallel for multiple attention heads and multiple transformer blocks.

But that is how you do self-attention. So, at a high level, we took our input, performed self-attention, and obtained an input that was now enriched with context. That is what is meant by "Attention is All You Need." You essentially take the context of all the words around a given word and create a new, richer meaning for each word in your input sequence. With that, I would like to gently suggest that you like and subscribe to the channel. I hope you enjoyed this video. Please comment below to let me know what you think or if you have any additional questions. I really appreciate it. Thank you so much! I'll see you in the next one. Thank you.