Visualizing transformers and attention | Talk for TNG Big Tech Day '24

I've been working on this project about visually explaining transformers. As I mentioned, I love visually explaining all sorts of things, usually in math, but I'm happy to stray into adjacent fields. The target audience here is always a little bit tricky because I know many of you are engineers and quite technical, but I don't necessarily want to assume too much background from everyone. The one thing I will assume is that you're maximally curious and eager to dig into the details of what's going on. For example, I think most people know that large language models consume a lot of computation and require substantial resources.

And anybody who has looked at the market cap of Nvidia recently will also know that they tend to be run on GPUs because they are very parallelizable. One of the things I would like you to come away from this talk with is a deep, visceral sense of what those actual computations are, what the number crunching happening is, and why it is so conducive to parallelization, which is one of the big reasons for their success. Transformers were introduced in the now very famous paper called "Attention Is All You Need," published in 2017. This paper focuses on a specific use case of machine translation.

You might imagine it being used for Google Translate and similar applications. However, since then, it has flourished in all sorts of other tasks. It's not just useful for machine translation; it's also effective for transcription. With tools like Whisper, it can work in the opposite direction by taking a piece of text and synthesizing speech. It's even beneficial for image classification.

It's a remarkably flexible framework. However, the model that you and I are going to focus on today, which is the kind that underlies chatbots, is a bit simpler than the one introduced in their paper. The model will be trained to take in a piece of text and then predict what comes next. For example, if you feed it the text, "To date, the cleverest thinker of all time was," it will predict from its vocabulary what the next word might be. Moreover, it doesn't just provide a single prediction; it assigns a probability distribution to all possible snippets of text that could follow.

In this case, there are a lot of possibilities for what it might be, so it has a relatively spread-out distribution over many different things, including some equivocation, recognizing that it's going to use words like "probably," "undoubtedly," and "arguably." If you have a model that does this—one that simply predicts what word comes next—you can turn it into something that generates new text simply by having it randomly sample from that distribution. You have it look over and weight according to the probabilities it assigned, sample a word, add that word to the input it generated, sample another word, add that word to the input it was processing, and then run this extended version through the entire process again, asking, "Okay, what do you predict will come next?" It runs through the whole process, which takes quite a lot of resources. Again, it plays this game of sampling from that distribution, adding it on, and repeating.

Now, a lot of people will ask why you would ever have it choose something that's not the most likely word if it generates that. You absolutely can do that. If you get into the notion of something called temperature, this will come from assigning the temperature equal to zero. The answer is that you get very stilted outputs that sound kind of trite. You get a little bit more creativity sometimes, or something just sounds somewhat natural if you introduce that spark of randomness into it.

Now, if you want to turn this into a chatbot—something that takes seed text and then extends it—the simplest way to start that kind of project would be to have your seed text establish an environment in which a user is interacting with some kind of hypothetical helpful AI assistant. You would take the user input that they type into your chatbot website, incorporate that into the input, suggest that what is about to be generated is what that hypothetical AI assistant would say, and then have the model begin to generate responses, one word at a time. In reality, it's one token at a time, but we'll get into that in a moment. In this case, I took a relatively simple model and had it provide some suggestions for what to do while visiting Munich. It goes through a random sampling process to generate these suggestions. I'll leave it to you to judge whether or not its suggestions seem reasonable. But it's kind of surprising that this approach actually works, I think.

Having something that simply predicts the next word can give the impression of long-range thinking. Some of this perception comes from the manner in which it performs that prediction and how extensive the associations can be. First, I want to provide a high-level overview of how data flows through the entire model. We won't go into the details just yet, but we will. So, how does the data flow through?

And then we'll dig into some of the things that I think are required to achieve the nice visual intuition we want to go for. When faced with this task, the model is given a piece of text and needs to predict what comes next. The text could potentially be quite long, but I'll just show what I can fit on the screen. The first step is to subdivide it into little pieces, which we call tokens. In the case of text, these tokens are very often words or small parts of words, and sometimes punctuation marks.

But one thing to keep in mind is that this doesn't necessarily have to be just with text. Tokens could also include snippets of an image or snippets of sound if images or sound are somehow part of that input. If you're creating a different model that performs a different task, such as transcription or image processing, the process roughly involves subdividing it into small pieces and then having some sort of vocabulary for all possible pieces. Now, in the case of text, one question that might be worth considering as we go through this is: why not just break it up into characters?

That seems like a much more natural choice for what the atomic unit of text would be: to break it into characters and then do whatever we're about to do with tokens there. I just want you to think about that. We’re going to go through this, and I’m going to return to this question. As we do, think to yourself: what could go wrong? Why wouldn’t this work well? Or maybe it would work well if we broke it into tokens instead.

Excuse me, let me break it down into characters instead. The first thing that happens is that you need to know something about the context to be able to answer that question. The very first step is to associate each of those tokens with a long list of numbers. If we think of those lists of numbers as vectors, we will refer to them as the embedding of the token. You should consider those lists as encoding the meaning of each one of those tokens in some way.

And at this phase, it's essentially just a lookup table. Each particular token is always going to be associated with the same vector every time you run it. The first thing that happens is you pass that through this all-important attention block, the thing that we're going to spend the second half of this talk exploring. This mechanism allows those vectors to communicate with each other in a specific way and share information back and forth. Therefore, a vector that initially encoded a given word might end up taking into account some of its context. For example, sometimes words can be ambiguous.

If you take the phrase "American shrew mole," the phrase "one mole of carbon dioxide," and the phrase "take a biopsy of the mole," you and I can read those and understand that the word "mole" means something fundamentally different in all three cases because of the context. So, you want to give the machine a mechanism by which it can let words interact with each other so that the meaning can be updated. But it might not just be disambiguating the meaning of a word; it might also involve capturing some sort of substance or essence from the entire context that is somehow relevant for predicting what might come next, and encoding that into a single vector. The process doesn't just go through one attention block; it flows through many different iterations.

So, the next thing it passes through has kind of an overly fancy name. It's called a multilayer perceptron. It's the simplest kind of neural network. When you break it apart, it looks like a couple of matrix multiplications with a little nonlinear something sprinkled in between. I won't go into too many details on this particular block, but the one thing I'll say is that the majority of model parameters actually live in this multilayer perceptron.

Despite the title "Attention Is All You Need," in terms of counting parameters, it's about one-third of what you need. If you want to ask, "Okay, why? What's going on with these?" it's a little hard to say. The question of interpreting any of this is a bit of a challenge, but there was some interesting research published by a couple of interpretability researchers from DeepMind in December, where they studied where facts, such as associating athletes with their sports, reside inside a model.

If you have it autocomplete a phrase like, "Michael Jordan plays the sport of blank," and it successfully predicts that "basketball" is the most likely answer, that answer was nowhere in the context. It couldn't just have come from passing information back and forth between tokens. It indicates that somewhere inside the parameters of the model, which it learned during the training process, it somehow encoded the fact that Michael Jordan plays basketball. They conducted a study where they explored various associations between athletes and sports, trying to probe the network to figure out where that information likely resided. The answer seemed to be within these multilayer perceptrons.

One thought to consider is that, insofar as prediction requires context, attention blocks become relevant. When prediction relies on general knowledge of the world, perceptrons provide additional capacity to store more information. This is not to say that perceptrons cannot exist within the attention blocks; rather, they simply offer significantly more capacity for storage. Essentially, these two elements interact, and you oscillate between multiple iterations of them. You will pass through one attention block and then through multiple others.

A layer perceptron is repeated over and over. The aspect that makes deep learning "deep" is this repetition of a certain kind of operation many times. What happens is that you flow the data through the network, sometimes on the order of a hundred times. For example, GPT-3 has 96 different layers. At the very end, you focus on the last vector of that sequence. Initially, that last vector may be encoding the word "was," but the hope is that by this point, it has absorbed so much meaning from its context and gained extensive knowledge from the weights of the network. Consequently, it is capable of making a prediction such that you can apply a certain simple operation to that last vector, which produces a probability distribution over all possible tokens that might come next.

It might feel a little bit weird that you're just doing it on the last vector because, in the rest of it, you have this big sequence of very context-rich vectors. However, what happens to be the case is that, for the purpose of training, it's very efficient if all of those other vectors are simultaneously being used to predict what comes next at various different points for subsequences of the text that you're putting in. That way, one training example can act like many. Now, I know you understand this, but just to keep it in the back of your mind, the whole framework behind any kind of machine learning is that, rather than telling the network what to do in a very designed sense—where you're going and fine-tuning, "Okay, I want you to do this with this token and that with this token"—you instead provide a very general framework. That framework is characterized by a bunch of tunable parameters, which are just real numbers that you can change.

Often, there are hundreds of billions of them. You show it many different examples of the inputs and outputs that you might want it to have. In this case, just random snippets of text, say, bold from the internet, followed by the actual word that follows. As long as you have some kind of procedure that will indicate how it's doing on a particular example—either saying, "Oh, you did poorly on that example; you need to change your weights this way," or "You did well on that example; you don't need to change too much."

Then, if you run this through many, many different examples, presumably some kind of behavior emerges within that flexible structure. The reason that's important to keep in mind is that understanding what's going on is extremely challenging because it is an entirely separate question from the design of the network itself and understanding what the computations are. What you and I are going to do is mostly focus on what those computations are and what the framework entails. We will use some motivating examples of what we think it might do to help learn what that structure is and make it more memorable. However, the actual question of what it's doing in a deep way remains an unsolved one to this day.

The specifics, if you're curious about how this training will look, are as follows: at a high level, you have a cost associated with each example. For instance, let's say it's auto-completing a phrase, such as "TNG Technology Consulting is a ___." From actual snippets of text pulled from the internet, one of the examples might include the word "leading." You would then evaluate the probability assigned to "leading" as the actual next word. After that, you take the negative logarithm of that probability, and that defines the cost.

You can kind of see the graph here. Essentially, what it means is that if it assigns a probability closer to 1, the cost is very close to 0. However, if it assigns a very low probability—indicating that it wasn't predicting the actual word—it increases very steeply as you approach that. This is just the cost for one example. In principle, the cost for the entire network will be a sum or an average of the costs over the many trillions of examples that you provide.

And so the visual you might have in your mind is that there's a kind of surface where what you and I can graph and show on a screen is limited to a function that maybe has two dimensions of inputs—two different numbers as inputs—and then that costs as a third dimension for that output. This would create a surface, and in principle, you want to find a low point on it where there's a low cost. However, the network doesn't just have two different parameters as the input; it has hundreds of billions of parameters. So it's going to be some super high-dimensional cost surface.

Still, the thought that you have in your mind is that the process of learning involves tweaking those parameters iteratively so that you're taking little steps downhill on that surface. It's not going to be in two dimensions; it's going to be in some exceedingly high number of dimensions. But still, the idea is that you're slowly tweaking it in some way to try to find a minimum on this cost surface. What that means is you really have no idea what behavior is going to emerge with those parameters, because it's whatever emerges on this unbelievably complicated surface and into which local minimum it happened to fall.

So that's something to keep in the back of your mind. Now, one of the visuals that we can try to reach for this comes from something that is very relevant to parsing everything that happens with transformers, but was not at all new to transformers. It existed in natural language processing for a while, which is this very first step where you're associating words with vectors in some way. On the one hand, the reason you have to do this is that if you're doing deep learning—indeed, in most categories of machine learning—it's very, very helpful if your inputs and outputs are both real numbers.

Numbers are things that you can change continuously, which allows you to perform calculus and work with concepts like gradients. Converting words into something expressed in this continuous realm, such as a vector space, is a very natural first step in this process. What's particularly interesting is to observe the patterns that emerge as you train the network and see what kinds of embeddings it discovers for the task at hand, whether that involves next-word prediction or some other activity that presumably requires an understanding of language in some way. One of the notable outcomes is that words with very similar meanings tend to cluster near each other. For instance, if you examine all the nearest neighbors for a word like "tower" in a simple word-to-vec model, you will find many other words that have similar tall, construction, or building-related connotations.

And these are, in some sense, a cluster that represents a loose idea represented at that point in space. But it gets a lot more interesting than that, because it's not just a matter of where the words are and how they cluster; the directions in the space might also encode different sorts of meanings that aren't necessarily captured in a particular word, but rather reflect a more generic idea. There is a paper from 2013 by some Google researchers titled "Efficient Estimation of Word Representations in Vector Space." They included a number of very memorable examples in that paper.

But one of the examples they provided was how, if you take the embedding of "woman" and subtract the embedding of "man," you can think of it as a sort of difference vector in that space. Then, you add that vector to the embedding of "man." Following that, you take the resulting difference and add it to the embedding of "king." After that, you perform a search.

Are the nearest vectors at that point in the space? What you find is that it's actually quite close to the embedding of "queen," which is very fun. It's a very mathematical way to basically play the analogy game that would be on old SATs or things like that. It is fun. I tried to reproduce this, though, by downloading a certain simple word-to-vector model.

And it kind of worked. But actually, the nearest neighbor to that point was the word "king" itself. "Queen" was the second nearest neighbor. It sits a little farther off, which kind of makes sense because the way that "queen" as a word appears in the training data is not just as a feminine version of "king." You have things like the band "Queen" or the term "drag queen."

But where it worked really well is if I played around with family relationships. So, you know, the difference between uncle and aunt is quite similar to the difference between man and woman and other things like that. The upshot is that somehow during this training process with this abstract gradient descent notion, it's as if the model learned to associate one specific direction in this space with the notion of gender in the sense that adding a little vector in that direction to a word can take you from the masculine embedding to a feminine embedding. That's very cool. And there are all sorts of other fun examples they included.

I think one that stood out to me was if you take Japan associated with sushi and then ask it to associate Germany with something, it would naturally put bratwurst in there. You can just play this game with subtraction. Now, it's very hard to visualize these things, right? Because we can draw three-dimensional value vectors, which represent lists of three numbers. But these word vectors, in principle, are very high-dimensional.

And you can understand why it is helpful for them to be high-dimensional. If you want to encode concepts with distinct directions in the space, then you would like to have many distinct directions to work with at your disposal. In this case, to take the numbers from GPT-3 as an example, the vectors that we're embedding our tokens into have 12,288 coordinates. So, in principle, this embedding space is 12,000 dimensions—very, very big. However, I also want to explain why it is even bigger than you might think.

Like weird things start to happen in higher dimensions. To get a sense of why this might be helpful, you should keep in mind that we don't just want these to encode an individual word. It's not just the various meanings of a word that might be relevant. What we want to happen as it flows through the network is for these to have the capacity to absorb whatever meaning is relevant from the context. So here, if I take a little bit from the famous Robert Frost poem, "The Road Not Taken."

When we look at that last part of the passage, "Two roads diverged in a wood, and I took the one less traveled by," the word "one" would start its life in the transformer simply being embedded into a generic embedding of "one" that knows nothing other than maybe it's a number, a pronoun, and things like that. Once it has the ability to soak in context, if you want to predict what comes next, it would be helpful to somehow know that it's one of two roads and to know the antecedent of that pronoun. It might also be helpful if you let it flow through the network and get it to ruminate on it more to somehow encode the more high-level, inarticulable, poetic ideas that it symbolizes choice and what it's a choice between. Those very high-level ideas, if you can embed them in the vector, might be relevant to predicting the next token. Anyone who has interacted with these large language models often gets the feeling that they really do seem to understand something at a deeper level that's not just as simple as defining one word versus another.

And when you think about it from that point, 12,000 dimensions doesn't actually feel like all that much. On the one hand, okay, that's a large number of different directions you could work with. But if we want to encapsulate every possible concept with a distinct direction in this space, it doesn't seem like that much to work with. I'll give two quizzes here. One is a relatively easy and uninteresting one.

And the second one is dramatically more interesting. So, quiz number one is to ask: If you have a bunch of vectors, how many vectors can you fit into an n-dimensional space so that every pair is not equal? This is dramatically more interesting. Quiz number one could instead ask: If you have a bunch of vectors, how many vectors can you fit into an n-dimensional space so that every pair is 90 degrees apart? This would be relevant if you want distinct directions in the space to correspond to distinct ideas. If the vectors are perpendicular, it helps you avoid unwanted overlap. For example, if you have one direction that corresponds to the idea that the words being spoken are in German, and another direction corresponding to the idea that the general tone of speech is one of skepticism, you wouldn’t want those to align. If they did, it would mean that any time you’re speaking German, it would accidentally imply that the speech is skeptical in some way.

Maybe that one's not a bad idea. In principle, you want more independent directions. Now with this, the answer for an n-dimensional space is n. This is actually kind of what defines the notion of dimension, at least in an abstract sense, for a vector space. Now, a fact that I didn't know, which is a purely mathematical fact and very fun and fascinating, is related to a quiz I encountered while talking to some people who conduct interpretability research into large language models. The quiz is about how many vectors you can fit into an n-dimensional space, not such that every pair is perpendicular, but so that every pair is kind of almost perpendicular. Let’s say you give a little buffer where they can be somewhere between 88 degrees and 92 degrees.

If you think about it, when you're transitioning from two-dimensional to three-dimensional intuition and you play with it in your mind, it can seem like it just doesn't fit anymore. The little wiggle room doesn't provide much flexibility. As it turns out, as you scale things up, the way that the answer to this question grows is exponential in the number of dimensions. I've tried experimenting with this a bit, and I would actually love it if any of you out there want to provide a numerical answer to this in the case of, say, a thousand dimensions or ten thousand dimensions.

When I was playing with hundred-dimensional vectors, I could definitely fit hundreds of thousands comfortably, and they would be quite close to each other in this almost orthogonal way. It gets a little computationally expensive if you try to do it with a thousand-dimensional vector and, you know, hundreds of millions of them. But it's shocking, actually. You can fit a very large number, even though 88 degrees and 92 degrees, to our eye, look almost orthogonal. This might explain a couple of things. One, it might clarify how these models can actually accommodate quite a few ideas in a seemingly small space.

Compared to the number of concepts in the world, 12,000 isn't that big, which is also why they scale surprisingly well. You don't really experience the effect of exponential growth until you reach a certain dimension. As you go from, say, 10 to 100 dimensions, it's not just that you're gaining 100 times as many directions at your disposal; you're obtaining many more than that.

You're getting many more than that. As you move from 100 to 1,000 dimensions, you experience a sort of super-linear increase. Now, that's just something to keep in the back of your mind. Every time I'm trying to motivate something by saying, "Oh, imagine some direction in this space encodes such and such idea," even though that's a somewhat made-up example for the sake of illustrating the structure, it's perhaps not that implausible for certain very specific concepts to correspond to certain specific directions in the space. With all that said, let's get into the meat of the matter, which is what's going on inside this attention block. What is the specific mechanism by which information is transferred back and forth between these vectors?

And once you do, you can have a sense of what the number crunching that happens inside a large language model is, and why it is as parallelizable as it is. The way I’m going to do this, as I said, is by telling a little bit of a lie, where we’re going to imagine something that we might want it to learn, something we would hope that it learns, but this isn’t to say that it’s necessarily what actually happens in production models. Let’s say we have some kind of phrase like, “A fluffy blue creature roamed the verdant forest.” We want to think about how the nouns in that phrase get updated by the adjectives. What I mean by that is that when you initially encode all of these words, the vectors associated with them don’t have any notion of their context from the initial encoding. For example, the vector associated with the word "fluffy" is just some abstract notion of fluffiness, not tied to anything.

The vector associated with the creature is just a very generic notion of a creature that won't be any more specific than that. However, we want to ensure that these vectors can communicate with each other, incorporating the relevance of fluffiness and blue nits into that creature concept. It’s important to note that it's not just about encoding the particular word; they also encode a bit of extra information, which is the position that the word appears in the context. This positional information is also integrated into the vector.

And this is a very important idea because one of the philosophies here is to keep everything parallelizable. This means that if sequential relationships are relevant—if something is right next to another—you don't want that to depend on the way the data is processed, such as going through it step by step. Instead, it should be entirely inferable from the state of the vector itself, regardless of the order in which things are processed. In principle, each vector encodes not only the meaning of the word but also its position. So, at the moment, when we envision it with these little diagrams, you would want that after the attention process is complete and after we carry out any computations, the vector associated with a token representing "creature" would encode the much more specific notion of a "fluffy blue creature." Similarly, the vector associated with the position of "forest" would now encode a much more specific notion of a forest that has some greenness to it.

And the constraint that we have is not a hard constraint, but a helpful design choice if you're doing any kind of deep learning. It is important to ensure that all the operations you're working with are expressed in the language of matrix multiplication. This way, you can effectively ask questions about which matrices are associated with each other and how they should update each other. By using matrices filled with tunable parameters, the training process, which relies on gradient descent and a method called the backpropagation algorithm, can function very efficiently. This efficiency is partly due to the fact that matrix multiplication can be performed very quickly on GPUs in parallel. Additionally, this is because it is a linear operation, and when you are performing calculus with it, the derivatives are straightforward and simple.

There's no added complexity that emerges. So generally, try to say how I can somehow express this notion of letting one thing influence another, only using the language of multiplying matrices, where those matrices have tunable weights in them that can, hopefully, learn to do what we want them to do. There are three matrices that are going to come up. They're called the query, the key, and the value matrices. We will go through each one of them, one by one.

To set the stage, the query matrix works as if each word is asking a question. You might imagine the word "creature," or any other term, asking, "Are there any adjectives sitting in front of me?" Right now, it can't really see its context; it just knows its position and that it's a noun. You also want to give the adjectives that precede it the capacity to respond, numerically or mathematically, by answering, "Yes, I'm an adjective."

I'm in that position. Now, of course, they don't literally ask. Instead, the way this looks is to associate the embedding for that word and creature with a new vector. That's what we're going to call the query vector associated with that token. This query vector is commonly much smaller than the embedding.

So, if that embedding dimension was like 12,000 dimensions, a common choice for keys and queries would be around 64 or 128 dimensions. The way you generate this query vector involves using a matrix. You take a large matrix, multiply it by the embedding, and it should be able to recognize that this embedding corresponds to a noun and its specific position. The query that it produces should encode the idea of searching for an adjective in certain positions. The visual you might have in your mind is that if the embedding vector exists in a very large space, pointing in some direction and encoding all that information, we are mapping it down to a smaller space to ask the question.

And this isn't just something you do to one of the words; you do it to every single one of the words. You multiply them all by that same query matrix. What that means is that every time it hits a noun, the query it produces is effectively asking, "Hey, are there any adjectives in such and such position sitting in front of it?" You might naturally ask, what is it doing to all the words that aren't nouns?

And who knows? Maybe it's trying to do something else in parallel for the moment, given that we're just creating a toy example to motivate the structure. I can say, let's just focus on the nouns and then discuss all the other things that you might want a mechanism like this to do later, and whether that would be accomplished in parallel in the same matrix or not. So, you apply this to all of the different embeddings to get the sequence of query vectors. At the same time, you also multiply them by a distinct matrix, which we're going to call the key matrix.

I'll label it as WK, and you produce a sequence of keys. It's entirely analogous. Those keys are also going to be smaller vectors in that same smaller space. The way to think about it is that they're kind of answering a given question. The keys and the queries are trained together so that if a query vector is picking up on a pattern—like looking for adjectives before it—the key matrix should be designed in such a way that when there are adjectives sitting in a position before it, it outputs something that corresponds to the answer "yes."

We'll see what that means in just a moment for them to correspond. But like everything we observe, this key matrix is full of parameters that are learned during the training process. Whether this is actually what it does is a separate question, but the hope is that you're providing the model with the capacity to perform such tasks. If you have all of these keys, along with all of these queries, and you've trained everything, you're essentially relying on a large amount of data to somehow yield the outcomes you desire. You might consider those key vectors in a similar way.

You've taken something from that very big embedding space and mapped it down into that key query space. You should think of that smaller space as really being the same one, where the keys and the queries reside in the same space. What we mean by saying that the key is answering the query is that they align in this space, pointing in the same direction. For instance, if you take whatever query was produced by a creature asking, "Hey, are there any adjectives sitting before this position four?" then the key produced by the words "blue" and "fluffy" should point in a similar direction. This way, you can recognize that they correspond by measuring whether they point in the same direction. So, how do we measure that things point in the same direction?

I mean, there are a number of different ways, but certainly in the context of machine learning, where everything is straightforward if you can just multiply things and add them up without doing anything more complicated than that, the easiest tool in the toolkit is the dot product. This dot product is going to be positive whenever two vectors align with each other. It's going to be zero whenever they are perpendicular to each other, which we can think of as meaning unrelated. It's negative whenever they point in opposite directions. So, after producing all these keys and queries, if you take each pair of them and calculate the dot product between every possible pair, this will give you a grid of values that indicates which words are relevant for updating the meanings of which other words.

So the way I might like to visualize this is by thinking of all of those points inside the grid as literal dots, where the big dots correspond to large dot products and small dots correspond to smaller dot products. In our particular example, with adjectives and nouns, that pattern would look something like this. For instance, if we zoom in and examine the relationships, we can see that "fluffy" and "blue" each provide a very large association when related to the word "creature," while everything else is small in comparison. That's going to be something we can work with.

And the question is: what are we going to be able to do with this if you pick up on which words are relevant to which others? Now, the thing that we're going to want to do—there's a certain terminology for this, by the way—is that whenever I'm using the term "attend," it means that the tokens on the left attend to the token on the top. In this case, the adjectives attend to the creature. What you want to do is take some kind of weighted sum according to how much each of these vectors on the left attends to one of them at the top. We will take weighted sums along the column in such a way that we want the larger values to matter more, while the smaller values, such as those associated with less relevant words, will matter less.

But at the moment, these aren't weights. These are just outputs of a dot product that may be very large positive numbers or large negative numbers. If we want them to behave like weights, we need them all to be between zero and one, and for them to all add up to one. Once we have it in that format, we'll take a weighted sum. In a couple of minutes, you'll see what it is that we're taking a weighted sum of.

But just knowing that we want them to act like weights, the tool in the belt for this, another very common function that comes up in machine learning contexts, is something called the softmax. It has a very elegant little formulation. You essentially exponentiate everything and then normalize it. But all you really need to think about is that it takes an arbitrary list of numbers; they could be positive, could be negative, or even infinity. The output is going to be a list of numbers where each one is between 0 and 1, and if you add all of them up, they total 1. The larger numbers from the input will correspond to larger probabilities, essentially larger weights in that output, while the smaller numbers from the input will end up closer to 0.

The reason we call this "softmax" is that it acts somewhat like a maximizing function. If one of the numbers in the input is notably larger than the others, then the output distribution has almost all of its weight towards that particular number. Therefore, if you were to sample from that distribution, it would effectively be the same as just selecting the maximizing input. However, it is "softer" than a maximum in the sense that as you continuously change the inputs, you achieve a smooth transition in the output. This smoothness is beneficial if you want to perform calculus with it, which is essential for gradient descent operations happening under the hood. We take this list of numbers, which we want to act as weights (currently arbitrary values between negative infinity and positive infinity), and we run them through the softmax function. We do this for every single column in this diagram. Each column corresponds to one of the queries, and by applying the softmax to all the dot products, we obtain a sense of weights that indicate which other words will be relevant for updating the word at the top of that column. There is a slight nuance that I am not illustrating in the diagram here: before applying the softmax, you divide by the square root of the dimension for that key-query space. However, that detail doesn't really matter in this context.

It's something that's helpful for numerical stability. But if you're curious, that's a little detail missing here. So once we have this grid, this is properly what I'm going to call the attention pattern associated with this. Well, it's actually not an attention block, but just ahead of that attention block. And this attention pattern, again, I kind of like to visualize with a bunch of dots, just giving you a sense of which things are most relevant to which other things.

But there's one tiny nuance that we need to address before we actually use this to make any updates. It has to do with something I referenced a little bit earlier: during training, there isn't just one prediction that takes place, but actually quite a few different predictions as it processes a given input. For example, if it's taking in the phrase "the fluffy boo creature roamed the verdant forest" and predicting the next word, as it all flows through the transformer, what comes out on the other end is going to be a prediction at every single step along the way for every different subsequence of tokens. Again, this might not seem super relevant in the case of something like a chatbot, where it's generating text.

Why are you predicting the next word? You already know what the next word is. However, it is very relevant for training because it means that if you input this particular data, rather than having just one specific training example, which is whatever actually follows this, you get a whole collection of training examples where the same weights are used to make predictions along the way. As a result, all of the weights can be updated by the various predictions made for that single piece of processing. Additionally, the context can become quite large.

It can be thousands of different tokens. This can mean the difference between a single training example and thousands of training examples for one particular run. Now, if this is going to work—if you want to achieve training efficiency and speed up by processing everything in parallel—it means that even though it's processing the full block of text, for example, if it's saying, "the fluffy boot creature blank," and it's predicting what verb or word might come next, it has to somehow ignore the fact that it already knows the word that comes next, as that is included in the input. Given that the attention mechanism allows all the vectors to communicate with each other, if you let them interact without restriction, each one along the way would have the answer revealed to it. Therefore, what you have to do is take all of the positions that correspond to later tokens influencing an earlier token and do something to zero them out.

So, in this case, everything where there's a later token, like the word "roamed," sitting before an earlier token, like the word "creature," is something that we want to ensure is set to zero. By the way, if you see diagrams like this in other sources, there will be different conventions regarding whether people place the queries at the top or the keys at the top. As a result, you might see this masking look transposed from its current position. However, the important thing to remember is that you don't want later tokens to influence earlier ones; you must set those to zero.

Now, if you simply set them to zero, the problem is that your columns wouldn't be normalized anymore. Since we want to use them as weights for a weighted sum shortly, that would complicate matters. You might suggest just setting them to zero and then renormalizing the columns. However, a more elegant approach before applying the data is to use Softmax to set all the values inside that masked region to negative infinity.

And this feels a little bit weird if you haven't spent a lot of time with softmaxes, but the effect of that is that any component that was negative infinity, after you pass it through the softmax, goes to zero. Then automatically, because the output of a softmax always adds up to one, the columns will be normalized in this way. This gives you what you would call a masked attention pattern, or sometimes people will call it causal attention or causal self-attention, because you don't want later elements to influence earlier ones. Now, one thing you might consider at this point is that creating this attention pattern, which determines which tokens are relevant to updating which other tokens, grows quadratically in complexity as you increase the context size. Context size refers to how much text is being incorporated into predictions. Therefore, this pattern that it produces expands significantly with the size of that context. This is why it can be very non-trivial to simply take a traditional transformer and scale up the context.

So anytime you have something like the newer large language models with very large context windows, they have to be doing something a little bit cleverer. They have to have some variation on the original attention mechanism. Another thing you might pause and ponder here is that, even though it is a very large pattern you’re going to get when you have very large context sizes, as you go through the process of some large language model producing new text one token at a time, you might notice that there will be a lot of redundancy in each of those computations. Remember, in principle, each one of those tokens being generated goes through the whole process where all the data flows through the network and does everything that we are talking about here for each one of those tokens. However, because a lot of them are the same when generating new text, there is a lot of redundancy, and clever caching can be implemented to take advantage of this, so you are not faced with quadratic speedup at inference time.

It's something that's much more relevant during training or during the very first pass when it's making its very first token. Another thing you might want to reflect on is, earlier when I said, "Why don't we just break up our input into characters? Why aren't tokens just characters?" Because that feels like a much more natural, atomic unit of what text is. There are a couple of answers here.

One of the issues is that the context becomes much larger. Given that scaling with context involves scaling quadratically, this can bloat your network very quickly. Another challenge is that as you embed these tokens and aim for them to convey some notion of meaning, processing character by character requires it to be handled by the first few layers of the network before any meaning of the relevant words can be conveyed. In contrast, using tokens allows you to jump directly to the meaning of words from the very first step. However, there's a balancing act to consider; you can't make the tokens too large, as they need to appear sufficiently in the training data for effective learning.

If your tokens were like paragraphs, each one would only show up at most once in the training data, and you would hardly be able to learn anything from it. So the pattern that people have found, which seems to work quite well, is something called byte pair encoding. If you want to learn more about it, you can look it up. But, in the back of your mind, just think of it as words, pieces of words, punctuation marks, and things like that. Given that you have figured out, or the model has presumably figured out, which words need to be relevant to updating which other words, the question is how that actually happens. How is it, for example, that the embedding of a word like "fluffy" communicates with the embedding of "creature" in order to update it to point in a more specific direction that encodes the notion of a fluffy creature?

The simplest way to do this is to introduce yet another matrix, which we will call the value matrix. This matrix will operate in such a way that if you multiply it by a given adjective, the resulting vector will be such that when you add it to the embedding of a noun, it points that noun in a more specific direction. This direction will also encode the relevant meaning of the adjective. For instance, if you multiply the value matrix by an adjective like "fluffy," it may add a certain "fluffiness" direction to the word "creature." Conversely, if you multiply it by the adjective "blue," it will produce a new direction that, when added to the previous result, points to an even more specific meaning: a "fluffy blue creature."

You could do this. This would actually work perfectly fine. However, if you look at the number of parameters involved here, you'll see that the key and query matrices were mapping to smaller spaces. If we refer to the numbers from GPT-3, each of these ends up having about 1.5 million parameters. With this high-dimensional embedding space, if the value is mapping from the embedding space to the embedding space itself, it involves almost 100 times as many parameters, which could be acceptable. However, it is a bit more elegant to work with, especially in how this all plays out in the context of multi-headed attention, which we'll discuss shortly. If that value map is actually broken up a little bit and involves fewer parameters—where it first maps down to a smaller space and then maps back up to the larger space—it can enhance the overall efficiency. Ultimately, it's designed to take in these 12-dimensional inputs and produce other 12-dimensional outputs, but for the linear algebra enthusiasts among you, you might consider it a low-rank transformation.

At a high level, it's a method that allocates fewer parameters of the network to this task. When you have this value map, you will multiply it by each of the embeddings in the same way that you multiplied the key map to obtain several keys. Associated with each position of that key, you have a specific value. What we will do then is take weighted sums along each column, weighted according to the attention pattern, and the sums are taken from these values. For example, in the column associated with the word "creature," almost all of the weights after the softmax function are going to be zero, or something that is effectively zero.

And the only other tokens that actually attend to it are the adjectives. What we want to do is take a weighted sum of those two adjectives while everything else is zeroed out. Once we add all of those together, that's going to give us the change we want to make to the embedding up top. This will produce some change that might give us a little name like delta E. The result that will come out of the attention block at the end will be whatever the original vector was that came in, plus this change. This change is this weighted sum along the column.

The hope is that, for this kind of example, it takes something that encodes a creature and then produces something that provides a more specific representation of that fluffy blue creature. This process is not limited to just one column; it actually applies to all of the different columns based on the attention pattern and the various values. It takes this weighted sum, producing a sequence of changes, and adds that sequence of changes to the original sequence of vectors.

And those sums are going to be what flows out of the attention block, at least almost. So far, I've just been describing what you would call a single head of attention. As it is, this is actually quite a lot to think about because each one of them is being multiplied by a key matrix, a query matrix, and a value matrix. Those keys and queries are producing this large attention pattern, and that's going through the softmax. This is a lot to consider, but this happens about 10,000 times for a single pass of the network because this is just a single head of attention.

And then each attention block often involves many different heads. It might be something on the order of a hundred; larger models could have even more. Furthermore, there are going to be many different layers of attention. So if this is just one of the heads, multi-headed attention refers to those heads operating in parallel. Again, the mantra for transformers is that they are very well architected to work on GPUs, allowing you to execute computations in a shorter amount of time, even if the amount of computation is greater.

And what I mean by multiple attention heads here is that this adjective-noun example is just one out of many different ways that you might imagine context is relevant to updating meanings. You can imagine all sorts of other things, like also in the sphere of grammar, adverbs updating the meanings of verbs, or nouns and their antecedents going together. But it doesn't even necessarily have to be grammatical. Just really anything where you want to somehow transform information that was in one part of a passage to something that's later, such that it's relevant for predicting the next word. That has to somehow happen through these key-query and value mechanisms that are asking which things are relevant to which other ones and what changes should take place when that relevance is there.

So each one of the heads inside one of these multi-headed attention blocks has its own distinct key and query matrices, which are used to produce their own distinct attention patterns. They also have their own distinct value maps, which generate their own unique sequences of values. You perform a weighted sum along the columns for this process, so it really is a lot to think about. Each output from those heads represents a proposed change to the embedding that was initially present. You add all of these outputs together with the original embedding that flowed in. There are many different components that can be added, or at least potentially added, depending on whether they are zeroed out by these attention patterns or not.

And so, in this little schematic illustration that I provided, we have an attention block and some dancing lines to indicate that they are communicating with each other. The process happening under the hood involves all of those distinct attention heads, each of which incorporates various patterns interacting with one another. As I mentioned at the beginning, there isn't just one attention block; it flows through a multilayer perceptron and then into another attention block, which again allows them to update each other, continuing this flow. The overarching idea is that all of these vectors are gaining richer and richer meanings, while the context they are drawing from is itself becoming increasingly rich.

And by the end, your hope is that it's not just the meaning of the original token anymore. It's really whatever it needs to be to predict what comes after at that point. Now, there's definitely a lot more that can be said about transformers. We could, for example, dive into what's going on with those multilayer perceptrons. But I think this is a pretty good point to step back and ask a little bit: why are these as effective as they are?

There are many different tactics that people have been trying in deep learning for a while. However, since that 2017 paper, transformers have been on the rise, being applied to more and more distinct fields. There are a couple of things one could point to here. One of them is that a significant lesson in machine learning over the last couple of decades is that scale alone matters. Simply making things bigger and providing more training data can sometimes lead to qualitative improvements in model performance.

And there are laws that people write about this. They are not mathematical laws; rather, they are more heuristic laws, but they can be phrased mathematically. These laws predict, for a given size of the model and for a certain amount of training, what the cost function will look like. If that cost function is decreasing, it typically corresponds to improvements in the model's performance that are qualitatively visible. For example, it simply looks like a chatbot that behaves better, among other things.

And if scale matters, that means that parallelizability matters. In earlier natural language processing, systems would read through text the way that you and I do. They would start from the beginning and go through to the end, updating the meaning bit by bit as they read. However, once researchers moved away from those mechanisms and retained only one— the attention mechanism— which allows elements to communicate with each other without relying on sequential processing, it enabled them to process text very differently from how you and I do. They can take in entire passages at once, allowing all of the different embeddings to interact simultaneously. This approach may seem strange to us as we try to empathize with it compared to our method, but it is extremely efficient for GPUs, enabling a user to perform a greater number of floating-point operations in a given amount of time.

The other benefit here is the fact that we are performing next token prediction. This means that you can train on a vast amount of data that does not require human labeling. By having this pre-training phase, you can provide massive amounts of data that do not necessitate a large number of human labelers to identify what is on an image or similar tasks. This allows you to accumulate a significant number of training repetitions in this very large model without being constrained by human feedback. Human feedback usually comes in later; often, the reason we refer to it as pre-training during next token prediction is that what is required to make chatbots, for example, work more effectively is some form of reinforcement with humans in the loop afterward. Nonetheless, the idea that most of the actual floating-point operations (FLOPs) in training can come from something unsupervised ties into the concept that scale, and scale alone, tends to matter. Another appealing component of transformers is that, even though we are focusing on language where tokens are short snippets of words, the approach remains effective.

The fact that you can tokenize essentially anything—just break up whatever your data type is into little pieces and then embed those as vectors—means that you can have lots of distinct data types work in conjunction with each other. If you want a model that processes not just text, but text together with an image or together with sound, it can treat all of those as equal participants in the same landscape, rather than having to create different bespoke architectures for each one and having a completely separate process for how they communicate with each other later. So, with that, I think I’ll call it an end here, say thank you, and turn things over to the Q&A. I'm more than happy to talk more about transformers if you want. I'm also happy to answer any questions about "3 Blow and Brown" if you’d like.

But with that, I'll call it an end. Thank you. Thank you. Thank you. Thank you.

Thank you. I've got one question. I've also seen your videos, so it seems you have been working on this topic for quite some time already. I would be interested to know how much work you have put into this topic in the last month, year, or whatever time frame. It's a lot more than I was expecting, but it's been fun.

Especially because I'm more of a math person by background, it has involved a lot of discussions with people who know more than I do on the matter. I've found it especially enjoyable to engage with the interpretability researchers in this area, as they are motivated to understand what is actually happening. From my point of view, I don’t care if that’s what’s actually going on.

I just want a good pedagogical motivation for the structure. But it's also thought-provoking in its own right, trying to explore how we can probe these models. One of the things I would love to discuss further, which we would have talked about today, is just very hard to fit in because you need all the background knowledge. For example, there is some work done to address that little fact I referenced, where the number of vectors you can fit into a space that aren't orthogonal, but almost orthogonal, scales exponentially. This is related to what they call the superposition hypothesis.

And there are very clever ways that you can try to probe what those feature directions might be, even when there are many more than the number of dimensions. There could be a whole video about that. There have been several rabbit holes like that, which may or may not actually turn into videos, but I'm quite happy to explore them. The visuals are always a labor of love, as I pour time into them to try to communicate what I want. I have a question here.

Sure! Here’s the reformatted text with correct capitalization and spelling: "Yeah, thank you for the inspirational talk. I have a question regarding transformable models. It's a bit technical. This trend of using deeper and deeper transformable models reminds me of the good old days with convolutional neural networks when they hit a barrier in terms of exploding and vanishing gradients. They started using things like residual networks to let the signal pass through without any transformations when not needed. Are there similar trends in transfer models that you're aware of?"

I'm not entirely sure I follow the question. So, you're saying that the trend where residual networks, for example, became quite helpful— is there something similar in this case? Yes. I mean, the concept of residuality is sort of baked in from the beginning with how it’s even done. The idea is that as you're going from layer to layer, rather than generating an entirely new data type, you're kind of adding to whatever was there before.

So it's not just in the sense that you're preserving dimensions so it's friendly for the machine, which seems to be one aspect that's under consideration. The people who designed it are very hardware-aware, in contrast to perhaps some other machine learning frameworks. However, there appears to be a continuation of the trend you’re referencing, where making things residual is deeply beneficial. Whether that's because it promotes training stability or if there's something more conceptual at play, I can't definitively say, but it certainly seems to be the trend. We have another question over here.

Hey, thank you for the great presentation. So, when it's a very long presentation or a lengthy talk, do you have some metrics that you keep track of, such as how many jokes should be included in this 10-minute period to keep people awake? Or do you have a team of creative thinkers, so to speak, that arranges everything in the background for you to finalize in the end? Oh, to have a team! No, I don't think I'm very good at that, actually.

Like, I'm sort of a nerd who just digs into the topic itself and says, "All right, let's just dig into the topic." The thing that should motivate the next step is centered around understanding and clarity rather than humor. There certainly isn't enough deliberate thought that goes into that component. For me, what I find most motivating in learning something is just the promise that I'll actually understand it. I think the way I try to motivate audiences with, say, a longer video, is not so much that there will be an unrelated joke in the middle, but more that there's a real promise that, hey, we're on a path to understanding here.

And insofar as that promise is something people believe, I think that's when they stick with it a little bit more. A question from behind: If I understand you correctly, you use GPUs to do a lot of transformations, but you're not really interested in the absolute values that you calculate in each transformation. Wouldn't the mix of analog computers together with digital ones resolve this issue?

Energy hungriness? It's a good question. I am not an expert on the drawbacks or advantages of analog computers, and therefore I should stop talking now. However, from a first-principles standpoint in physics, it does make a lot of sense. If you're more error-tolerant in that way, the big advantage of digitization over analog computers becomes apparent. If you were to go back to the 1940s and explain why someone should bet on one horse and not another, it would have to do with error propagation.

And so in this case, if I kind of don't care about that because everything's just a matter of soft association anyway, I'll believe it. I am as curious as you are to understand why that's not the case. My best guess for why they haven't replaced it is just that there's so much smart mind power that has gone toward optimizing the living daylights out of matrix multiplication in a digital context. It's just hard to catch up with that. Even if something is going to be developed from a first principles standpoint, there's such a head start on that one. But if you find someone with a clearer answer, please send that answer my way because I'm as curious as you are. So, we have a question in Zoom. We have one good question, yes. Thanks for trying three times to come here and actually doing it this time. My question is about tokenization and images. You briefly mentioned that text and sound are basically one-dimensional in some sense because you read from left to right or you listen through time. However, images are inherently two-dimensional, and I would argue that the pixels above and below are more important than the one at the end of the current pixel row.

So, how could you tokenize that nicely? Yeah, I guess there are two things to consider. One is tokenization, and the other is how attention will work. I mean, for tokenization, loosely speaking, you would want to think of it as little patches of the image that will serve as the tokens. Additionally, there is a notion of positional encoding that will encode not just the x position, but also the y position. This way, the vector that embeds that patch will have information about where it is positioned, which can be used for associations.

But as I understand it, the naive way of doing that would result in either way too many tokens or tokens that are way too small. If you think about the resolution of an image and the desire for tokens that appear frequently, you would likely end up with something that bloats the network significantly. Therefore, you have to be a bit clever, but in principle, it's just the patches. Regarding how they interact with each other, there are various methods where you can mask the attention so that it first processes column-wise and then row-wise, or something like that. However, as I understand it, there isn't one single way to approach it. In all cases, it does take into account your point that it's two-dimensional.

And the way that you'd see that is usually in the positional encoding. I appreciate your point that it's two-dimensional. I think that was the last question, Grant. Thank you very much, everyone, for coming.

This is delightful to actually be here.