Transcription
If you feed a large language model the phrase "Michael Jordan plays basketball," and have it predict what comes next, and it correctly predicts basketball, then that suggests that somewhere, within its hundreds of billions of parameters, there is information embedded about a specific person and their specific sport. And I think generally anyone who has played with one of these models has a pretty clear sense that it has memorized a lot and a lot of facts. So a reasonable question you can ask is, how exactly does that work? And where do those facts live? Last December, a number of researchers from Google DeepMind published work on this question, and they were using this specific example of matching athletes to their sports. And while the full mechanical understanding of how facts are stored is still unsolved, they did come up with some interesting partial findings, including a very general high-level conclusion that facts seem to live inside a specific part of these networks, fancifully known as multilayer perceptrons, or MLPs for short. In the last few chapters, you and I have been digging into the details behind transformers, the architecture that powers large language models, and which also powers many other forms of modern AI. In the last chapter, we focused on a piece called "attention." And the next step for you and me is to dig into the details of what happens inside these multilayer perceptrons, which make up the other large chunk of the network. The computation here is actually relatively simple, especially when compared to attention. It can basically be summarized as a pair of matrix multiplications with a little something in between. However, interpreting what those computations are doing is a very big challenge. Our main goal here is to walk through the computations and make them memorable, but I'd like to do that in the context of showing a specific example of how one of these blocks, at least in principle, can store a concrete fact. Specifically, the fact that Michael Jordan plays basketball. I should mention that the design here is inspired by a conversation I had with a researcher at DeepMind, Neil Nanda. For the most part, I'm going to assume that you've either watched the last couple of chapters, or that you have a basic sense of what a transformer is, but a refresher never hurts, so here's a quick reminder of the general flow. You and I have been studying a model that's trained to ingest a piece of text and predict what comes next. The input text is first broken down into a set of tokens, which is to say small pieces that are typically words or small pieces of words, and each token is associated with a high-dimensional vector, which is to say a long list of numbers. Then this sequence of vectors is repeatedly passed through two types of operations, attention, which allows vectors to pass information between each other, and then multilayer perceptrons, which is the thing we'll be digging into today, and there's also a certain normalization step in between. After the sequence of vectors flows through many and many different iterations of both of these blocks, in the end, the hope is that each vector has ingested enough information, either from context, or from all the other words in the input, as well as from the general knowledge that's been embedded into the model's weights through training, such that it can be used to predict what token comes next. One key idea I want you to take away is that all these vectors live in a very high-dimensional space, and when you think about this space, different directions can encode different kinds of meaning. So one classic example I love to refer to is how if you look at the embedding for the word woman and subtract the embedding for the word man, and take that little step and add it to another male name, like uncle, you land in a place that's very very close to the corresponding female name. In that sense, that particular direction encodes gender information. The idea is that many other distinct directions in this high-dimensional space might correspond to other features that the model might want to represent. In a transformer, these vectors don't just encode the meaning of a single word. And as these elements flow through the network, they ingest richer meaning based on all the context around them, as well as based on the model's knowledge. Ultimately, each one of them needs to encode something far beyond the meaning of a single word, as that needs to be sufficient to predict what comes next. We've already seen how the attention blocks allow you to ingest context, but the majority of the model's parameters actually live inside the MLP blocks, and one idea for what they might be doing is that they provide additional capacity for storing facts. As I said, the lesson here will focus on the concrete game example of how the fact that Michael Jordan plays basketball is stored. Now, this simple example will require you and me to make a few assumptions about that high-dimensional space. First, we'll assume that one direction represents the idea of the first name Michael, then an approximately orthogonal direction represents the idea of the last name Jordan, and then a third direction represents the idea of basketball. So specifically, what I mean by this is if you look into the network and pick out one of the vectors that's being processed, if its dot product with the first name Michael direction is one, then that means that the vector encodes the idea of a person with that first name. Otherwise, the dot product will be zero or negative, meaning that the vector doesn't really align with that direction. And for simplicity, let's completely ignore the very reasonable question of what it might mean if the dot product is greater than one. Similarly, the dot product with these other directions will tell you whether it represents the last name Jordan or basketball. So, let's assume that the vector is supposed to represent the full name, Michael Jordan, then its dot product with both directions should be one. Since the text "Michael Jordan" spans two different tokens, this also means that we have to assume that a prior attention block has successfully passed information to the second of these two vectors to ensure that it's able to encode both names. With all these assumptions, let's now get to the meat of the lesson. What happens inside the multilayer perceptron? You might think of this sequence of vectors flowing into the block, and remember that each vector was originally associated with one of the tokens from the input text. What will happen is that each individual vector from this sequence will undergo a short sequence of operations, which we'll unpack in a moment, and in the end, we'll get another vector of the same dimension. That other vector will be added to the original vector that flowed in, and that sum is the result that flows out. This sequence of operations is something that you can apply to each vector in the sequence, associated with each token in the input, and it all happens in parallel. In particular, the vectors don't talk to each other in this step, each one is doing its own thing. And for you and me, this actually makes things much simpler, because it means that if we understand what happens to just one of the vectors through this block, then we actually understand what happens to all of them. When I say that this block will encode the fact that Michael Jordan plays basketball, what I mean is that if a vector that encodes the first name Michael and the last name Jordan flows in, then this sequence of computations will produce something that includes the basketball direction, which will be added to the vector at that position. The first step in this process looks like multiplying this vector by a very large matrix. No surprises here, this is deep learning. And this matrix, like all the other matrices we've seen, is filled with the model's parameters that have been learned from data, which you can think of as a set of knobs and switches that are adjusted and tuned to determine the model's behavior. Now, one good way to think about the process of matrix multiplication is to imagine each row of that matrix as its own vector, and taking a set of dot products between those rows and the vector that's being processed, which I'll call E for embedding. For example, suppose the first row equals the first name Michael direction that we assume exists. This means that the first component in this output, which is the dot product here, will be one if this vector encodes the first name Michael, and zero or negative otherwise. It's more fun to take a moment to think about what it might mean if this first row is the first name Michael direction plus the last name Jordan direction. And for simplicity, let me go ahead and write that as M plus J. Then, when taking the dot product with this embedding E, things distribute really nicely, so it looks like M dot E plus J dot E. And notice how this means that the final value will be two if the vector encodes the full name Michael Jordan, and otherwise it will be one or something less than one. And that's just one row in this matrix. You can think about all the other rows in parallel, asking some other kinds of questions, exploring some other kinds of features of the vector that's being processed. Often this step also involves adding another vector to the output, which is filled with the model's parameters learned from data. And this other vector is known as the bias. For our example, I want you to imagine that the value of this bias in this first component is negative one, meaning that our final output looks like the relevant dot product, but minus one. You might very reasonably ask why I want you to assume that the model has learned this, and in a moment you'll see why it's so clean and beautiful if we have a value here that's positive if and only if the vector encodes the full name Michael Jordan, and otherwise it's zero or negative. The total number of rows in this matrix, which is sort of like the number of questions being asked, in the case of GPT-3, which we're following the numbers for, is just under 50,000. In fact, it's four times the number of dimensions in this embedding space. That's a design choice. You can make it more, you can make it less, but having clean multiples tends to be hardware friendly. As this matrix full of weights takes us into a higher dimensional space, I'll give it the shorthand W for up. I'll continue to call the vector that we're processing E, and let's call this bias vector B and put all of that back into the diagram. At this point, the problem is that this operation is purely linear, but language is a completely non-linear operation. If the input that we're measuring is high for Michael plus Jordan, then it's also necessary that it's somewhat influenced by Michael plus Phillips and also Alexis plus Jordan, even though there's no conceptual connection between them. What you really want is a simple yes or no for the full name. The next step is to pass this large intermediate vector through a very simple non-linear function. A common choice is one that takes all negative values and clamps them to zero and leaves all positive values unchanged. And continuing the deep learning tradition of exaggerated names, this very simple function is often called a rectified linear unit, or ReLU for short. And this is what the graph looks like. So, if we take our imagined example where this first input to the intermediate vector is one, if and only if the full name is Michael Jordan and zero or negative otherwise, after passing it through ReLU, you end up with a very clean value where all the zero and negative values are clipped to zero. So this output will be one for the full name Michael Jordan and zero otherwise. In other words, it directly mimics the behavior of an AND gate. Models often use a slightly modified function called JLU, which has the same basic shape, but it's a little bit smoother. But for our purposes, it's a little cleaner to just think about ReLU. And also, when you hear people referring to transformer neurons, they're talking about these values here. When you see the common neural network picture with a layer of dots and a bunch of lines connecting to the previous layer, which we had earlier in this series, that typically means passing this combination of a linear step, matrix multiplication, followed by some simple non-linear function like ReLU. You could say that this neuron is active when this value is positive and it's inactive if this value is zero. The next step looks very similar to the first step. You multiply by a very large matrix and add to it a certain bias term. In this case, the number of dimensions in the output goes back to the size of the embedding space, so I'll go ahead and call this the down projection matrix. And this time, instead of thinking about things row by row, it's better to think about them column by column. As you can see, there's another way you can hold matrix multiplication in your mind is to imagine taking each column of the matrix and multiplying it by the corresponding term in the vector that's being processed and adding all those rescaled columns together. The reason it's better to think about it this way is that the columns here have the same dimensions as the embedding space, so we can think of them as directions in that space. For example, we'll imagine that the model has learned how to transform the first column into the basketball direction that we assume exists. What this means is that when the relevant neuron at this first position is active, we will add this column to the final output. But if that neuron is inactive, and if this number is zero, then this will have no effect. And it doesn't have to be just basketball. The model can also integrate with this column and many other features that it wants to associate with something that carries the full name Michael Jordan. And at the same time, all the other columns in this matrix tell you what will be added to the output if the corresponding neuron is active. And if you have a bias in this case, it's something that you add in every time, regardless of the neuron values. You might be wondering what this does? As with all the parameter-filled objects here, it's kind of hard to say specifically. Maybe there's some accounting work that the network needs to do, but you can ignore it for now. And to make our notation more compact again, I'll call this large matrix W, and similarly I'll call the bias vector B, and put it back into our diagram. As I previewed earlier, what you do with this final result is add it to the vector that flowed into the block at this position, and that gives you this final output. For example, if the vector flowing in encodes both the first name Michael and the last name Jordan, then since this sequence of operations will result in an AND gate firing, it will add the basketball direction, and so what comes out will encode all of these elements together. And remember, this operation happens for each one of these vectors in parallel. In particular, if we take the numbers for GPT-3, this means that this block doesn't just have 50,000 neurons, but it also has 50,000 times the number of tokens in the input. And that's the entire operation, two matrix products, each with a bias term added and a simple pooling function in between. Any of you who have watched previous videos from the series will recognize this structure as the most basic type of neural network we studied there. In that example, it was trained to recognize handwritten digits. Here, in the context of a transformer for a large language model, this is one piece in a larger architecture and any attempt to interpret what exactly it's doing is closely tied to the idea of encoding information in vectors from a high-dimensional embedding space. That's the core lesson, but I want to step back and think about two different things, the first is a kind of accounting, and the second involves a thought-provoking fact about higher dimensions that I actually didn't know until I dug into transformers. In the last two chapters, you and I started to account for the total number of parameters in GPT-3 and see exactly where they were, so let's finish that game quickly here. I've already mentioned how this down projection matrix has just under 50,000 rows and each row matches the size of the embedding space, which for GPT-3 is 12,288. And multiplying these values together, we get 604 million parameters for that matrix alone, and the down projection has the same number of parameters just with the transposed shape. And so, they combined give about 1.2 billion parameters. The bias vector also accounts for a few other parameters, but they're a negligible fraction of the total, so I won't even show them. In GPT-3, this sequence of embedding vectors flows not through one block, but through 96 distinct MLP blocks, and so the total number of parameters dedicated to all of these blocks comes out to about 116 billion. And this represents about two-thirds of the total parameters in the network, and when you add that to everything we had before, for the attention blocks, embedding, and de-embedding, you actually get the advertised total of 175 billion. It's probably worth mentioning that there's another set of parameters associated with the normalization steps that were skipped in this explanation, but like the bias vector, they represent a very negligible fraction of the total. As for the second point of reflection, you might be wondering whether this central example that we've spent so much time on reflects how facts are actually stored in real large language models. It's true that the rows of that first matrix can be thought of as directions in this embedding space, and that means that the activation of each neuron tells you how well a given vector aligns with a given direction. And it's also true that the columns of that second matrix tell you what will be added to the output if that neuron is active. Both of those are just mathematical facts. However, evidence suggests that individual neurons rarely represent a single clean feature like Michael Jordan, and there might actually be a very good reason for that, which is related to an idea that's been floating around among researchers in interpretability these days known as superposition. This is a hypothesis that might help explain why models are particularly difficult to interpret and why they can scale so surprisingly well. The basic idea is that if you have an n-dimensional space and you want to represent a set of different features using vectors that are orthogonal to each other in that space, you know, in that way if you add a component in one direction, it won't affect any of the other directions, then the maximum number of vectors you can fit is just n, which is the number of dimensions. To a mathematician, this is actually the definition of dimension. But it gets interesting if you relax that constraint just a little bit and tolerate some noise. Let's say you allow these features to be represented by vectors that are not perfectly orthogonal, but are approximately orthogonal, maybe between 89 and 91 degrees. If we were in two or three dimensions, this wouldn't make any difference. And this doesn't give you any extra space to fit more vectors, which makes it more counterintuitive where the answer changes dramatically for higher dimensions. I can give you a very quick and direct example of this using some sparse Python that will generate a list of 10,000 distinct 100-dimensional vectors, each initialized randomly, and that list will contain 10,000 distinct vectors, that's 100 times the number of vectors as there are dimensions. The plot here shows the distribution of angles between pairs of these vectors. And since they started out randomly, these angles can be anything from 0 to 180 degrees, but you'll actually notice, even for random vectors, that there's a significant bias for things to be closer to 90 degrees. Then what I'll do is run a certain optimization process that repeatedly pushes all these vectors so that they try to become more orthogonal to each other. After iterating that a few different times, this is what the distribution of angles looks like. We actually have to zoom in here because all the possible angles between pairs of vectors fall within this narrow range between 89 and 91 degrees. In general, one consequence of what's known as the Johnson-Lindenstrauss lemma is that the number of vectors you can cram into a space, that are approximately orthogonal like this, grows dramatically with the number of dimensions. And this is extremely important for large language models, which might benefit from associating independent ideas with approximately orthogonal directions. And this means that it's possible to store many, many more ideas than the dimensions of the space allocated for them. This might partially explain why model performance scales so well with size. A space that has 10 times the dimensions can store well over 10 times the number of independent ideas. And this applies not only to the embedding space where the vectors flowing through the model live, but also to this vector full of neurons in the middle of the MLP that we just studied. And this means that at GPT-3 sizes, it might not be able to explore just 50,000 features, but if it instead leverages this enormous additional capacity by using near-orthogonal directions for the space, it might be able to explore many, many more features of the vector that's being processed. But if that's the case, then what that means is that individual features won't be visible when a single neuron lights up. It will have to look like a specific combination of neurons instead, that is, superposition. For any of you who are curious to learn more, the main search term here is sparse autoencoder, which is a tool that some interpretability specialists use to try to extract the true features, even if they're very superimposed across all these neurons. I'll put a link to some great anthropological posts on this topic. At this point, we haven't touched on every detail of the transformer, but we've gotten to the most important points. The main thing I want to cover in the next chapter is the training process. On the one hand, the short answer for how training works is that it's all about backpropagation, and we've covered backpropagation in a separate context in previous chapters in the series. But there's more to discuss, such as the specific cost function used in language models, the idea of fine-tuning with reinforcement learning with human feedback, and the concept of scaling laws. A quick note for the active followers among you, there are a number of non-machine learning videos that I'm excited to dive into before I get to the next chapter, so it might take a little while, but I promise it will come in due time. Thank you.