Architectures, hyperparameters

Framing: the details other courses spare you

00:05 As you may have noticed, I'm a little bit less innovative in my lecturing them, Percy, so you're going to get PowerPoint slides rather than executable Python ones, but you should be able to find the PDFs on the website as well. So I've titled this lecture, "Everything You Didn't Want To Know About LM Architecture And Training," because we're going to get into some of the nitty-gritty details that I think most other classes would spare you the details of. Like what should my hyperparameters be? And those kinds of questions. Some minor logistics. Also, if you're doing the assignments, we are updating assignments as we find mostly minor bugs. Make sure you pull updates to the assignments as you go along. So what we're going to do, we're going to start with a quick recap of a transformer. And I'll give you two variants of a standard transformer, one that's probably coming from the standard transformer lectures that you might see in 224N. And then I'll talk about what you implement and the modern consensus variant of a transformer. And then we're going to take a much more data-driven perspective to understanding transformer architectures. So the question that we're going to ask is people have trained lots of LLMs at this point, and you can go and read all of those papers and try to understand what has changed, what has been in common. And from that, almost an evolutionary analysis, try to understand what are the things that are really important to make transformers work. So today's theme is-- the theme of the class is the best way to learn is hands-on experience. But the theme of this lecture, because we can't train all these transformers, is to learn from the experience of others. So the starting point is the original transformer.

Two transformers: the original vs. what you implement

01:46 So just as a review, hopefully, you all remember this from 224N or your other NLP classes. You've got some simple position embeddings at the bottom. You've got multi-head attention. You've got LayerNorms afterwards. You've got a residual stream going upwards. You've got an MLP, and then a softmax at the very end. And we're going to see variants to all of these different pieces until we get to basically the most modern variants of the transformer. And the latest one I'll talk about will be just a few months before. So what you implemented is not the vanilla transformer variant from the original paper. We've modified a few things. We've put the LayerNorm in front of the block. So you can see on this slide over here that there's the norm is over here right before each of these blocks in the residual stream. We've asked you to implement rotary position embeddings. The feed-forward layers use something called a SwiGLU. And then linear layers now emit these bias terms. And you might ask, why have you forced us to implement this weird variant of a transformer instead of the original transformer is all you need transformer? And so we're going to go through some of those questions. And then yesterday, I was thinking, OK, I should catch up on all the developments that have happened in architectures over the last year. And Percy warned me about this because he said, you're going to have to redo the lecture every year. And so I started looking. And I was like, all right, yeah, there's a couple of good papers recently. There's Command A, there's 2 OLMo 2 Furious. There's SmolLM and Phi-4. And then you go looking, and you're like, wow, yeah, there's Gemma 3 and Qwen 2.5 and InternLM. And then we got more. I can't even cover the screen with these guys. There's a lot of models.

The method: ~19 dense releases in a year

03:29 There were about 19 new dense model releases in the last year, many of them with minor architecture tweaks. And on the one hand, it's kind of annoying to go through all these papers and say, like, what is happening in all of these? But also, it's like actually a wealth of information because not all of them do the same thing. And you can see-- not all of you, especially in the back, can see the details of this slide. But I put together a little spreadsheet of what all these models are doing. And starting with all the way from 2017, the original transformer, all the way to 2025, what the newest models are doing. And we'll talk about this as we go. But you see certain kinds of architecture changes being explored. So here, on this column, is position embeddings. People used to do all sorts of stuff, like absolute, relative, RoPE There was a alibi phase for some people. But then, now, starting around 2023, everyone just does RoPE. So you can see this convergent evolution almost of neural architectures. And we're going to talk about all of these different kinds of things. So the parts that I'll cover-- so this is a preview of the three major sections of this lecture. And if I have time, I'm also going to talk about different attention variants at the end. The first thing is going to be architecture variations. That's what I'm going to talk about. So activations, feed-forwards, attention variants, position embeddings, all of those things. And then having nailed down the architecture, what do we have to do? Well, we have to pick hyperparameters. How big do we make the hidden dimension? How big do we make the inner projection layer inside of MLP? What do we do about the number of dimensions? How many vocab elements? Those are all important things that you have to choose when you're actually training your language

05:08 model. And you don't want to just pick these out of a hat. You want to select them in some fairly intelligent way. So we're going to start with architecture variations. And the two things that I'll mention right here-- and I'll go back to them as I talk-- the first one is, there's not that much consensus in a lot of the choices. There's been convergent evolution in the last few years, what I'll call LLaMA-like architectures at the very bottom here. But people do all sorts of things. They swap between LayerNorm and RMS Norm.

Pre-norm vs. post-norm

05:40 They do serial versus parallel layers. There's one choice that basically everyone does since the very first GPT. And I'll talk about that in a bit. But there's lots of different variations that we can learn from here. The big one, I've already talked about this guy in 224N. So if you remember that lecture, this will be review for you rather than being totally new. I think the one thing basically everyone agrees on and agreed on almost from the very start, is the use of pre-norm versus post-norm. That terminology will get a little bit more confusing. But the original transformer paper did this thing on the left over here where you had your residual stream in the gray. And in addition to the residual stream, you had these LayerNorms after every subcomponent. So you would do your multi-head attention. You would add back to the residual stream. And then you would LayerNorm that. And then you would do the same thing with your fully connected layer. And then you would LayerNorm that. And very, very early on, people realized that moving this LayerNorm to the front of this non-residual part, so this block on the right did much better in many different ways. And basically, almost all modern LMs that I use this kind of pre-norm. There have been some new innovations recently that I'll touch on in two slides. But lots of models have moved to this. The one exception is OPT-350M, which I'm guessing, they messed that one up. And that was orphaned when they were training. That was a fun find in my survey of architectures. So this pre versus post-norm thing, if you look into why it was originally developed, the arguments were that if you wanted to use this post norm

07:24 stuff, it was much less stable. And so you would have to do some careful learning rate warm up style things to make it train in a stable way. And so if you look at some of the earlier papers arguing for this pre-norm approach, Salazar and Nguyen and also this Xiong 2020 paper, you almost always see this comparison of, hey, if we use pre-norm and we do some other stability-inducing tricks, then we can remove warmup and these systems work just as well, if not better, than the post-normal LayerNorm with careful warmup-type approaches. And you see this in a machine translation setting here. You see this as well on the right on various other tasks, especially using BERT which was trained with post-norm. So there were many arguments about why this was helpful. There were arguments about gradient attenuation across layers. Like if you do pre-norm, then the gradient sizes will remain constant. Whereas if you did post-norm without warm-up, then it would blow up in this orange way. It's a reasonable argument. But I think a maybe more closer to modern intuition would be this argument that pre-norm is just a more stable architecture to train. And so some of the earlier work by Salazar and Nguyen identified all these loss spikes that if you were training with pre-norm, in blue here, you would see a lot more loss spikes and the training would be kind of unstable as you were training. So you see the gradient norm here is spiking and generally higher than the one with pre-norm. And so, today, you see pre-norm and other LayerNorm tricks being used essentially as a stability-inducing aids for using training large neural networks.

The “double norm”

09:11 And so this brings us to one new fairly recent innovation. I think this didn't exist when I gave this lecture last year, which is this variant that I don't think really has a great name, but I'm just going to call it the double norm for the moment here. So this is the original figure that I showed you at the very beginning. And we know that putting LayerNorms in the residual stream is bad. But actually, someone in 224N this year asked, well, but why do you have to put the LayerNorm in the front? Why can't you put it after the feed-forward network? And of course, you can. And not only that, recent people have gone around and just add the LayerNorm after the blocks as well. And so Grok and Gemma 2 both take this approach of LayerNorms both in front and after. OLMo 2 does only the LayerNorm after the feed-forward in the multi-head attention. And so this is actually an interesting change. Pre-norm has just been dominant and the only thing for a while. But things have been changed up a little bit. So now, there's a new variant. And there's been some evaluations of this kind of approach. People have argued it's a little bit more stable and nicer to train on these larger models. By the way, feel free to stop me and ask me questions as well. I have a tendency to keep going if no one stops me. So, yes? Why is LayerNorm in the residual bad? Why is LayerNorm in the residual bad? That's a good question. I don't think I can give you a this is the proof of why it's bad. I think one intuitive argument for why this might be bad is that the residual gives you this identity connection all the way from almost the top of the network,

10:49 all the way to the bottom. And so if you're trying to train really deep networks, this makes gradient propagation very easy. So there's lots of arguments about how LSTMs and these other kinds of state space models have difficulty propagating gradients backwards. An identity connection does not have any such problems. And so putting LayerNorms in the middle might mess with that kind of gradient behavior. And that of course, you see back here. This is exactly the kind of plot you expect to see if that's happening. Cool. The other thing that people now do

LayerNorm → RMSNorm, and FLOPs vs. runtime

11:23 is in the original transformer, people did LayerNorm. And so LayerNorm is this equation over here. What you do is you have the activations x coming in. You subtract the empirical mean. So that's the average of the x's up top. And then you divide by the standard or the variance plus a little fudge factor epsilon. And then you square root that. So that you can roughly think of it as a standard deviation. So that's going to standardize your activations x. You're going to scale it up by a gamma. That's a learnable parameter. And then shift it by a beta. So this makes sense. You're going to normalize your activations. And then you're going to shift them around to whatever point you want. And many models use this LayerNorm thing. And it worked quite well. But many models have now moved on to RMS Norm. And this is one of the consensus changes. Basically, all the models have switched to using RMS Norm. And now, what do you do? You just drop the mean adjustment. So you don't subtract the mean. You don't add a bias term. And many notable models do this. The LLaMA-family PaLM, Chinchilla, T5's, they've all moved to RMS Norm. And what's the reason for this? One reason is that it doesn't really make a difference. It turns out if you train models with RMS Norm does just as well as training with LayerNorm. And so there's a simplification argument. But really, I think the argument that's often given in these papers, and I think it's good to appreciate the details of this argument, is that going to RMS Norm is-- it's faster and just as good. So in what way is it faster? Well, if I don't subtract the mean, it's fewer operations. If I don't have to add that bias term beta back, it's fewer parameters that I have to load from memory back

13:02 into my compute units. So I don't have to retrieve this state. And some of you might be thinking, but wait, you told me in 224N that nothing but matrix multiplies matter for the purpose of runtime. And this is not a matrix multiply. And so I shouldn't care about any of this. And that's a reasonable perspective to take. If you think about the number of the percentage of flops that is taken up by different operations in a transformer, this table, there's a nice paper by Ivanov et al 2023. I think the title is Data Movement Is All You Need or something that does profiling of all the different components of a transformer. And you see that tensor contractions, which are like matrix multiplies, that's like 99.8% of the flops that happen in a transformer. And so saving 0.17% of your flops doesn't seem like a huge win. But I think one of the things that's important for architecture design now is to not just think about flops. Because flops are important, but that's not the only resource that you have to think about. It's also that you have to think carefully about memory movement. And so even though tensor contraction, so this is things like matrix multiplies. That's like 99.8% of the flops. If you have things like the softmax operation or LayerNorms, all these normalization operations that happen in a transformer, there are 0.17% of the flops-- actually, they're 25% of the runtime. And a big reason for that is because these normalization operations still incur a lot of memory movement overhead. And so it does actually matter to try to optimize some of these lower level things

14:46 because it's not just about flops, it's also about memory movement. I'm going to emphasize this quite a bit more as I get into the systems lecture. When we talk about GPU architectures, it's going to become very, very, very important to think about memory, not just about flops. And so this is one of the reasons why RMS Norm has now become much more popular. And so I went back and looked at some of the earlier RMS Norm papers. I think the sad thing is that there aren't quite as many papers published by industry labs with big, nice ablations. And so many of the ablations that I'll show you are going to be from a couple of years back. But Narang et al in 2020 had this very nice ablation showing-- here's the vanilla transformer, here's the RMS Norm version. And you see the exact thing I told you. The number of steps per second that you can do in a vanilla transformer, 3.5 per second. With RMS Norm, you get 3.68. Not a huge gain, but that's, in some sense, for free. And you get a final loss that's lower than the vanilla transformer. So that's great. In some sense, we've gotten runtime improvements and we've gotten in fact, at least

Dropping bias terms

15:53 in this case, loss improvements. And so that's a win-win for us. The final thing that I'll say, which is very much in line with this RMS Norm thing in terms of theme, is that most modern transformers do not have bias terms. So the original transformer, if you look at the FFN, will look something like this. You have your inputs x. You're going to do a linear layer with a bias term. And then you ReLU it. And then you'll have a second linear layer wrapping around it. But most implementations-- if they're not gated units, which I'll talk about in a moment-- look actually something like this. They just drop the bias terms. And you can just make this argument from basically the same kinds of underlying principles. They perform just as well. Matrix multiplies are apparently all that you need to get these guys to work. And the other thing, which is maybe more subtle is actually optimization stability. I don't quite have the deepest understanding of why the bias terms are particularly bad for stability, but there's been really clear empirical observations that people have made that basically dropping these bias terms often stabilizes the training of these largest neural networks. And so now, a lot of the implementations now omit bias terms entirely and train only on these pure matrix multiply kind of settings. So that's the LayerNorm bit. And so there is two things that you should think of. This is nice because the story is pretty clear. Everyone does something, and so you should just know this. Basically, everyone does pre-norm. Or at least, they do the LayerNorms outside of the residual stream. That's kind of the iron rule. You get nicer gradient propagation.

17:34 You get much more stable training. It just doesn't make sense to do it the other way. Most people or almost everybody does RMS Norm. In practice, it works almost as well, has fewer parameters to move around. And this idea of dropping bias terms just broadly applies. A lot of these models just don't have bias terms in most places. I think the one exception to this RMS Norm one, as I was reading yesterday, is, I think, Cohere both Command A and R+ use LayerNorm. Not quite sure why. Any questions on the LayerNorm, RMS Norm and bias terms stuff before I move on? Yes, question? Do you think there are some long term lessons you can take away from these details that are more future-proofed potentially, or do you think these are-- yeah. So the question was, is there something more future-proof? And I think it's hard to have the biggest picture. In many ways, deep learning has been very empirical and bottom up rather than top down. But I do think there's some generalizable lessons that you could draw from here. I think the lesson of have very direct identity map residual connections is a story and a lesson that has played out in many, many different kinds of architectures, not just kinds of architectures. The effectiveness of LayerNorm we'll see, once again, later on in this lecture has been very effective. And so not letting your activations drift in scale is another thing that I think generally has been very effective for training stability. Those two seem like fairly generalizable lessons. We will also see the systems concerns come into play again. So this is another generalizable lesson of thinking really carefully about the impact of your architecture on the systems

Activations and gating

19:17 components of your design. So now, there's this other component, which is the activations. And there is a whole big zoo of activations-- ReLU, GeLU, Swish, ELU, GeGLU. And then, I mean, these aren't activations. There are different kinds of MLPs-- GeGLU, ReGLU, SeLU, SwiGLU, and LiGLU. And yeah, I think this is exactly the kind thing that I didn't originally want to learn. When I got into doing deep learning, I was like, I don't care about activations. It's going to train anyway. But it really does matter. Unfortunately, for both you and me, that SwiGLU and other GLU variants just consistently work well. And so I will explain those to you. And you should think about them carefully because they do work and internalize that. So I think the ReLU, and maybe the GeLU, you all should already know. The ReLU, you learn in some of the most basic deep learning classes. You just take the max of 0. And in the case of the MLP, you've got your-- I've dropped the bias terms here. x dot W1, you take the ReLU, and then you do W2. Fairly easy. A GeLU is a Gaussian error linear unit. This one multiplies the linear with a CDF of a Gaussian. And so it's basically going to be like the ReLU, but with a little bit of a bump here. Hopefully, you can see that over here. This is not just flat at the very bottom. This makes things a little bit more differentiable, which may or may not help. And the GPT family of models-- 1, 2, 3, and GPT-J and so on--

20:55 all use the GeLU. And the original transformer and some of the older models used the ReLU. And really, almost all the modern models have switched to the gated linear units like SwiGLU and the GeGLU and others. And really, I think the Google folks really pushed for this, like PaLM, and T5 and others. But since it's been tried and true, basically, almost all the models post 2023 use a gated linear unit. And so going back to that earlier question of what generalizable architecture things can we learn from this lecture? There are some things that have really consistently been very useful-- residual connections, LayerNorms. Gating is yet another one. And so this is another place where gating just appears and is a very good way of doing things. So, originally, this is our fully connected layer, right here. This is with a ReLU. Now, instead of doing just linear and a ReLU, what I'm going to do, is I'm going to gate the output here with entrywise linear term. So x dot V is going to give me a vector. And I'm going to multiply that entrywise with my original inside term of the MLP. And then I'm going to multiply the whole thing with W2. So the way to think about this, is I've gated the hidden part of the MLP. So I've got my original activation that takes my inputs and puts it into the hidden space. And then I'm going to gate that with x dot V. And then I'm going to project that back into the hidden dimensionality using W2. So there's this gating operation that happens entrywise. And that's really the basic thing that's happening here.

22:38 And this is the GLU plus the ReLUs, the ReGLU. And then we have an extra parameter that we've added here for the gating. This is V. And so when someone says something like, oh, it's a GeGLU fully-- there's nothing to laugh about that. There's the GeGLU fully-connected layer, what I've got here is, I've got the GeLU for the non-linearity. And I've still got the exact same gating here of x dot V. And this is the architecture that was used by many of the Google models like T5v1.1, Gemma 2, Gemma 3. And then another variants, there's a SwiGLU, and this has been very, very popular. Swish is x times the sigmoid. And this is the non-linearity. And sigmoid is like this. And x is like this. So it will look just like the Gaussian error unit. And then you do the same thing here. You have a gating over the swish. And then you get a fully connected layer here. Yes? I have a question. Below a certain negative value, the swish function and also the GeLU function, it's not monotonically increasing. And in fact, it's decreasing. And a lot of the argument about how gradient descent works in introductory machine learning is that, OK, you want to do gradient descent. But here, it seems like you would go in the opposite direction if you use GeLU or swish or their gated versions. So, I'm not-- So the question was, this isn't monotonically decreasing. There's a bit on the very left of this 0 here that's flipping in the derivative. And isn't that going to be a problem? I think intuitively, you could have argued that this would be a problem. You might trap a bunch of activations at zeros.

24:20 I think, in practice, if you look at neural network optimization dynamics, what's actually happening is often you're throwing very high learning rates with momentum into the optimizer. And so you're not really going to converge to the 0 point. These activations are going to be all over the place. And so in practice, I don't think this little tiny negative piece is really an effect that's going to be huge for the model, if that makes sense. And then going back to this, the SwiGLU is basically most models today. Like the LLaMA-family PaLM, OLMo-- and I'll show you the big table later. But you'll see that the SwiGLU is very, very popular. And one thing to note-- I'll talk about this again in the hyperparameters part-- is now, remember, I've added this V term, this extra parameter. And so I want to think about how to size. This extra parameter. And what people do is gated models usually make this hidden size. The basically output dimensionality of W slightly smaller by a factor of 2/3 in order to make sure that the total number of parameters of this whole thing remains the same as the non-gated counterparts. And that's a convention thing that most people do. If you don't quite understand what that is, I'll go back over that again later. But you can just keep in mind that basically for the gated linear units, you just make everything a little bit smaller to make sure things remain parameter matched. So, yes, question. This may be obvious or have gone over in the past. One of the benefits of ReLU is like, it's very easily differentiable by the input, but if you have the derivative of the CDF of the Gaussian,

26:05 you have a squared with x. Does that not really slow things down? That's a very good question. I'm not 100% sure what the internal CUDA implementation of the SwiGLU or the GeGLU is. I think it's entirely possible that internally, they might be implemented with lookup tables. Go ahead. I mean, what really matters is the memory pressure here. And it will be the exact same because you're reading the same amount of elements. So the extra compute is negligible on [INAUDIBLE]. That's probably a better argument that basically flops why this is negligible anyway. And actually, the memory calculus is the same, so. Cool. So, do gated linear units work? I will have more modern evidence for this as well. But I thought I should take you straight to the horse's mouth, Noam Shazeer's original paper, where he evaluates all these GLU variants. And this is somewhat older stuff. So you're seeing CoLA and SST-2 performance. But you do see basically that the GLU variants consistently perform better. GLU is 84.20, 84.12, 84.36, 84.67. And wow, it's 2020s. They even give you the standard deviations so you can figure out how significant those results are. And they, in fact, are significant. And so this is some nice evidence to see here. There was also the Narang et al in 2020 paper, which is a very nice paper, studying all sorts of architecture variants, I think in the context of T5 style models. And once again, you see that the gated linear unit variants consistently achieve lower losses

27:47 than their counterparts. You see that the bolded lines are exactly at the GLU variants. And this pattern has basically held up. So for gating and activations, there are lots of variants across different models. But the gated linear unit has become basically widespread and dominant, and I think for good reason. Of course, the GLU isn't necessary for a good model. It's important to separate the two. Just because it's probably the slightly better and everyone does it doesn't mean it's necessary. And you do see examples of very high performance models. not using a GLU. Like GPT-3 is one example. A more recent one, Nemotron 340B uses a squared ReLU, which I had not seen before. And Falcon 2 11B uses a ReLU. Both of those are relatively high performance models. So you can see that it's not really necessary. And so evidence does point towards consistent gains from SwiGLU and GeGLU. And that's why we ask you to implement exactly that variant. Cool. The final thing that I want to talk about for architectures,

Serial vs. parallel blocks

28:57 and this is one final-- major, I want to say, variation that we've seen. Normally, the transformer block is serial. In the sense that for each block, the outputs come in from the bottom. And then you do your attention, and then you pass the result of that computation forward. And then you do your MLP, and then you pass that computation forward. And so this is inherently serial. You do attention, and then MLP. But of course, this might have certain parallelism constraints. If you want to parallelize this over gigantic sets of GPUs, it might be harder to do so if you have the serial connection. The systems concerns might also be more difficult. You might get lower utilization from your GPUs. And so a few models have done this thing that I'll call parallel layers. Where, basically, instead of having serial computation of attention and then MLP, they will do them both at the same time. So you will get your x from your previous layer. You will compute both the MLP and the attention side by side, and then you will add them together into the residual stream. And then that will be your output. And this was pioneered by GPT-J, which was this open source replication effort. And the folks at Google doing PaLM were bold enough to do this at the really big scale. And many others have followed since. So if you're implementing this right, you can share a lot of stuff like the Layer Norms. And the matrix multiplies can get fused together. And you can get some systems efficiencies out of that. It hasn't been quite as popular since then, at least in the last year. I think most of the models that we've seen have been serial layers rather than parallel ones.

30:36 I think the only exceptions to this are like Cohere Command A, Command R+, and a Falcon 2 11B. So now, I think we have the ability to go back to this big, hard to see chart, and then see what I was pointing at the very beginning. So this column here, you don't really need to be able to read any of the text. Because I think the colors will tell you everything you need to see. This check mark here, this is basically pre versus post-norm. The only two models I really know of in the early days that did post norm, this is the original transformer and GPT and BERT, if you want to include that into this table. And then almost everybody else-- I think, basically, everyone else has done pre-norm. The only other non checked boxes here are models that are proprietary, and I don't have details for. This column here, on the leftmost thing, this is RMS Norm versus LayerNorm. The gray boxes are the LayerNorm. The blue ones are RMS Norm. Basically, most people have converged to RMS Norm as I said. This column next to it is serial and parallel layers. Once again, most people do serial, but you see other variants. What I'm going to talk about next is going to be position embeddings, and that will be more interesting in a moment here. Any questions about any of this architecture stuff before I move on? Hopefully, that gives you a bit of an overview of at least the major variations in architectures that we see. Is the serial layer more computationally efficient than parallel layer? So the question was whether serial is more efficient than parallel. It should be actually the reverse. That parallel is more efficient than serial. And that's why you're willing to do this. So in some sense, you might expect serial to be more expressive because you're

Position embeddings and RoPE

32:15 composing two computations rather than just adding them together. But the benefit of parallel, in theory, is that if you write the right kinds of fused kernels, a lot of these operations can be done in parallel, or the computation is shared across the different parallel parts. So, cool. So the last thing I want to talk about in architecture land-- I think this is the last thing-- is variations in position embeddings. And I think this one's interesting because in the first few years of LMs land, there were a lot of different things that people were trying. Sine embeddings were from the original transformer. You should have learned this in 224N. There's sine and cosine positions. Many others did absolute embeddings, like the GPTs and OPT, all basically just added a position learned position vector to the embedding. Some others, like T5 and Gopher, did various kinds of relative embeddings that add vectors to the attention computation. And then, I think most models have converged to RoPE, which is relative position embeddings. And this, I think, actually started in GPT-J, once again, another open source contribution. It has really rapidly been picked up by most of the models. And so the high level thought process behind RoPE is that the thing that matters is relative positions of these vectors. And so if I have an embedding f of x of i, where x is the word I'm trying to embed, and i is my position, then I should be able to write things down in this way. So there should exist the f, such that f of x, i, and f of y, j, if I take the inner product of these embeddings,

33:57 then, I can write this down as some different function g, which is a function of the two words and the difference in their positions. So this is a definition that enforces basically position invariance or absolute position invariance. So you only pay attention to how far apart these two words are. And so you can do a brief check and see, OK, what happens with sines while you get these cross terms that are not relative? So you do still leak absolute position information. Absolute positions like it's in the name. It's not a relative position embedding. And relative embeddings, well, it is relative, but it's not an inner product. So it violates this constraint. And so RoPE is this kind of clever observation that we do know one thing, that is invariant to absolute things, which is rotations. And so we're going to exploit that structure to come up with our position embeddings. We know that inner products are invariant to arbitrary rotation. So we're going to leverage that. So on the left, this is the starting point. Let's say my embedding for the word "we" is this arrow over here. And my embedding for the word "know" is this other arrow over here. Now, I want to embed this sequence, "we know that." And I look at the word "we" and "know." So how do I do that? Well, we assume position 0. So I'm not going to rotate that guy at all. "Know" is in position one, so I'm going to rotate him by one unit of rotation. And so now, I have this embedding for "we know." And now, let's say I want to embed this sequence, "Of course,

35:34 we know." Now, "we" and "know" have the same relative positioning to each other. And so let's look at what happens. We get shifted by two positions. I rotate "we" by-- I start in this vertical position, and I rotate them twice, 1 and 2. And then I rotate "know" by three positions because it's 1, 2, 3-- sorry, 0, 1, 2, 3rd position. And so, now, if you look at these two arrows, they have the same relative angle right. So their inner products are preserved. And so this is the nice fun idea about RoPE. You just rotate the vectors. And the rotation angle is determined by the position of each word and rotations-- the inner products don't care about relative rotations. And so these inner products are only going to look at the difference in distance. Now, it's easy to think about in 2D because rotations are kind of obvious in 2D. There's only one way to rotate a vector. But in high dimensional spaces where we operate, it's not obvious at all how we are going to do this rotation. So the RoPE folks came up with, in some ways, the simplest but also effective way of doing this. And the way to do it is you take your high dimensional vector. In this case, D. And I'm just going to cut it up into blocks of two dimensions. And every two dimension is going to be rotated by some theta. So there's going to be a rotation speed. And I'm going to rotate the pairs of dimensions. And so now, every pair of dimensions is encoding all of these relative positions. And much like in sine and cosine embeddings, I'm going to pick some set of thetas such that some embeddings are rotated quickly, and others are rotated much more slowly.

37:13 So they can capture both high frequency information or close by information and very far away, lower frequency positioning information. And the actual RoPE math here is, if you're going to think about rotations, it's just going to be multiplying with various sine and cosine rotation matrices. Hopefully, you remember this from linear algebra and trig. And so you can think about this as an operation where you multiply your embedding vectors with these block 2 by 2 block matrices. And there's no additive or cross terms that appear here. This is all purely relative. One thing that is different, if you're used to absolute position embeddings or sine and cosine embeddings here is that the RoPE is going to operate at the actual attention layer. You're not going to add position embeddings at the bottom. Whenever these attention computations are going to be done, you're going to intervene on that layer. And then that's going to give you your position information. And so I pulled this from I think the LLaMA implementation of RoPE. You've got the initial normal attention stuff at the very top like query keys and values. These are your normal linear projections. And then, you're going to come up with cosine and sine angles. These are rotation angles telling you how much to rotate different blocks of the query and key. And so you take your query and your key, and you're going to rotate them by the cosines and sines. And now, you've gotten rotated query and rotated key. And that's going to be what's going to go into the rest of your attention computation. So you don't do this at the bottom. You do it whenever you generate your queries and keys. Hopefully that's clear. That's really critical to enforcing

38:56 this relative positioning only information. OK? Good. So one of the things I want to highlight is that RoPE is actually one of the things that it seems like everyone has converged on. I went through all 19 of those papers over the weekend. And basically, all of them now use RoPE for various different reasons. The reason that RoPE has now many different algorithms for extrapolating context length, and that's an important part of the modern productionized language model. But also, it seems to be empirically quite effective, even at fairly small scales and small context lengths. So it's kind of worn out on this-- what's it called-- position embedding battle. Any questions before I move on to some of the hyperparameter stuff? Yes. Is the rate of rotation is consistent across all of these models? I don't think they're all the same. There's some variation in the thetas. Yes. Are the thetas for each pair, are those hyperparameters or are they trained? The thetas that determine the rotation angles, they're not hyperparameters. Much like in the sines and cosines here, there's a schedule to the rotation angles that are determined. And it's in the same intuition as the sines and cosines. You want to cover different frequency ranges in order to get higher or lower frequency information. Yes. The rotations create any difficulty with training, I wonder like this, angular rotations. The rotations themselves don't really create any issues.

Hyperparameters that don't vary

40:33 Because one way of thinking about a rotation is that it's just a matrix multiply. Since thetas are fixed and the m's here are fixed, this is really just a fixed matrix that multiplies your vector. And so in that sense, it's not really an issue. If you were learning the thetas, then maybe you have issues because you're maybe differentiating through trig functions. But you're not doing that here, so. OK, cool. So now, I think we go even one more level into the details here. And we're going to talk about hyperparameters. I feel like when you have to-- you're dropped in, and you're asked to train a new language model. There's a lot of questions you have about hyperparameters, because there's quite a few of them. And one of the things that I've realized is that actually, only a few of these really get changed across different successful models. There's actually like fairly clear rules of thumb and fairly clear guidelines that people seem to be following. So there are some things like how much bigger should the feed-forward size be? Or how many heads should I have? Or what should my vocab size be? And so we'll talk about each of those things, and we'll try to constrain the space of hyperparameters that people have. So the starting point, we're going to look at a simple feed-forward layer. Just with the bias, let's say. This is a ReLU version of it. And so there's two hyperparameters here. There's d model, which is the dimensionality of x. That's the input coming into your MLP. And then you've got dff. So this is the feed-forward dimension. This is the output hidden dimension of your MLP. And from there, you're going to project back onto d model. So what should dff be? In general.

42:12 These things are going to be up projections. You're going to have more hidden units than there were inputs. But how much bigger? Well, there is actually just like a consensus. Almost everybody that uses ReLU style MLPs are going to pick dff is equal to 4 times d model. I will show you some empirical evidence for why this is a sane number later. But as far as I can tell, there's no law of nature that says you have to pick 4. This is a convention that has really held up. Now, there are a few exceptions to this rule. Remember that the GLU variants are going to scale this down by a factor of 2/3. And if you scale it down by a factor of 2/3, you're going to have roughly the same number of parameters. You can do a little bit of math. And if you scale the GLU variants down by a factor of 2/3, you'll come to the conclusion that the way to do that is to set dff equal to 8 over 3d model. That's going to be the number that you end up at. And you guys convince yourself that will give you the same number of parameters. And that's the ratio that you would get if you started with a ratio of 4. So if you look at many of the models, they actually do follow this rule of thumb. PaLM, for example-- PaLM, Mistral, and LLaMA are slightly larger. These are GLU models, but they don't follow this 2.6 rule. But if you look at, for example, LLaMA-1, Qwen, DeepSeek, Yi, and T5, they all roughly follow this 2.6 ish rule. And I can put up the big table of LMs that I made later with hyperparameters. Many, many, many of them fall into this roughly 2.6 range.

43:52 And that's the standard parameterization of a GLU unit. I'll go through one other exception. I really like this exception because I think, in many ways, big large language model training is a game of copying hyperparameters from other people. And so we don't learn very much. It's very conservative. But T5, I really like. Because in some sense, it's really bold. And I think Google people actually do some pretty bold stuff. And so if you look at the 11 billion parameter T5 model, they have a pretty incredible setting. Their hidden dim is 1,024. But their dff, their up-projected dimension is 65,000. And so that's going to give you a 64 times multiplier on the ratio of dff to d model. And of course, you compare to this, where PaLM is like a factor of 4. And everyone else is much smaller. This is a very large difference. And there's some other recent examples of using much bigger multipliers. Gemma 2 follows in these footsteps and does a factor of 8. And I'll talk a little bit about this exception later. Of course, T5 was a totally fine model. So this should tell you it is possible to train a model with such a much larger ratio. So one of the things that I think is quantitative evidence, I saw that 4x multiplier. And I thought, is that really the right thing to do, or is there some more quantitative experiment someone's done to convince me that is a good idea? So one of the figures from Jared Kaplan's scaling law paper-- and most people know this paper for the scaling law component.

45:32 But actually, there's also some really useful hyperparameter components to this paper. You'll actually see that they do exactly this thing that I'm talking about, the dff to d model ratio. And they plot essentially how much the loss increases as you vary this. And you see that there's a sweet spot. This is a ratio of 1, 2, 3, 4 and then up to 10 or so here. And so there's a pretty wide basin here, anywhere between 1 to maybe up to 10, where you can pick whatever feed-forward ratio you want. And it'll be roughly optimal. And 4 is not too far off from your optimal choices over here. It's 1, 2, 3, 4. It's like right here or maybe right here. So that's a pretty reasonable choice. So what can we learn from all of this hyperparameter stuff? I think a lot of the evidence points towards-- you can pick the same defaults of, if you're not using a GLU, you can multiply by 4. If you're using a GLU, you can use roughly 2.66. And they can work pretty well for mostly all the modern LMs. T5, once again, does show that you don't have to follow these rules. You can be a rule breaker and do whatever you'd like. There's no hyperparameter choice written in stone. You can get reasonable LMs at many other hyperparameters. That said, I think the really funny epilogue to this story is that T5 has a follow up model called T5v1.1 that's improved. And it uses a much more standard 2.5 multiplier on GeGLU. So you can read between the lines and say maybe they looked at the original T5 and said, actually, maybe we want to walk back that 64 times multiplier and pick a more standard one.

47:14 And they did end up with a better model. So, cool. What is the relationship between the ratio and the model efficiency on [INAUDIBLE]? So I think that's a good question. So the question was, what's the ratio-- or sorry, what's the relationship between this ratio that I'm talking about here and generally the impact on the model? And so if we go all the way back here, the ratio is controlling essentially how wide the hidden part of this MLP is. And so the original justification in the T5 paper for picking 64 was to say, actually, we can get bigger and fatter matrix multiplies if we make that dimension really, really large. And while that is kind of a true statement, the wider it is, you're getting more parallel computation, so to speak, rather than serial computation. So you're spending your flops and your parameters in a slightly different way than if you made your hidden units bigger, which would let you pass more information, or using more units, which would give you more serial computation. So you're spending your parameters and your flops in a slightly suboptimal way from expressive power, but you might get systems gains if your matrices are wide enough. Excellent. So another thing that is a surprising-- or maybe not surprising-- consensus hyperparameter is the ratio between the model dimension and the head dimension times the number of heads. So I clipped this from 224N. But really, the basically canonical choice

48:54 is to pick things so that the dimension d-- that's a hidden dimension-- and if you have multiple heads, you're just going to split up the number of dimensions each head gets. So you're going to keep the dimensions fixed as you add more heads. And you don't have to do that. As you add more heads, you could just keep the same number of dimensions per head, and you could just let the attention part take more and more parameters. You could do that. That's an option that you have. But most models, once again, do follow this guideline. We see GPT-3, T5, LaMDA, PaLM, and LLaMA2, they all have a ratio of one or almost exactly one. T5 is the one exception that breaks this rule. They tried the big ratio of 16. But otherwise, it is all fairly following this consensus. There's been a couple of papers that have argued against this 1 to 1 ratio. There's a notable one by-- I don't know how to pronounce this-- Bhojanapalli et al 2020, who have argued that if you have more and more heads, they're going to have lower and lower rank. And if you have very few dimensions per head, that's going to start affecting the expressiveness of the attention operation. But in practice, it doesn't really seem like we see too many significant low rank bottlenecks in practice. And most of the models with this ratio of 1 seem to do just fine. This is really a parameter that's generally been held constant by most of the models that we've seen. If I have time, I'll talk a little bit about different optimizations that people have made on this multi-head component. But hyperparameter wise, things have stayed fairly similar. I think one of the big ones in terms of hyperparameters is the aspect ratio.

50:37 So we can think about deep networks. We can have more and more layers, or we can have wide networks. And generally, if you want one knob to control the width, that would be the hidden dimension of the residual stream. That would control essentially the width of almost all the operations at once. And so this seems like a pretty critical thing to tune. You might think that deeper networks are smarter and more expressive, or wider networks are more efficient. There is generally a sweet spot of ratios that people have picked. There have been outliers. Some of the early models used much smaller ratios here. So what that means is that they were much wider than they were deep. And then some models have gone really deep where they had way more-- sorry, the other way around-- really wide, where they had way more d model than n layer. And there's been generally a sweet spot of saying, we want about 128 hidden dimensions per layer. And that has been generally stuck to by a lot of the GPT-3 and LLaMA variant models. And I'll talk a little bit about evidence for that in a second. There's considerations about aspect ratio that are quite important. They will control the amount of parallelism that we can do. So if you're doing something called pipeline parallel, what you're often going to do is you're going to take your different layers and you're going to cut them up and you're going to put them on different devices or different blocks of devices because you'll parallelize within each layer as well. And so there's going to be certain kinds of constraints that you're going to put on your model. And also, if you have really wide models, then you can do something called tensor parallel where

52:20 you slice up the matrices, and then you distribute those on GPUs. And one thing that we'll learn in, I think, 1, 2, 3, 4 or 5 lectures is that these different parallelism paradigms are going to have different constraints. You need really fast networking for tensor parallel. And you can maybe get away with slower networking or higher latency networking for pipeline parallel. And so your networking constraints might, in turn, drive some of these width-depth considerations. But setting that aside, you might abstractly ask, what is the impact of aspect ratio on model performance? And once again, Kaplan et al have a really nice visual aid showing how aspect ratio impacts performance. And so this is three different scales-- 50 million, 274 million, and 1.5 billion parameters. And the x-axis is aspect ratio. Y-axis is loss difference in percentage change. And you see, that around 100, which is, once again, I told you was around the consensus choice of hyperparameters, is the minimum across different scales. So this is backed by some of this large scale hyperparameter data that's been published by Kaplan et al and roughly matches that intuition. And a really nice thing here, is it seems to be the case that aspect ratio optima does not shift too much across several orders of magnitude here. So if this holds up even more, that's very good news. You can keep training on one fixed aspect ratio. One thing I will know that is quite an interesting result is Yi Tay and others at Google had this very interesting paper

53:59 studying the impact of depth versus width both upstream and downstream. And one of the things that they found was that if you're looking at losses, then, it doesn't really matter. Parameter is the only thing that matters. Deeper models don't help you. But the story is less clear if you're looking at downstream accuracy. At the time, they were looking at fine tuned superglue accuracy. They were arguing that for the same amount of flops, deeper models might be better. So I'll just leave it at that. There's not quite as much follow up to this work, at least in the open, that I've seen. But downstream performance may actually be slightly different in terms of the aspect ratio considerations here. Cool. The final thing that I want to talk about in this very low level hyperparameter world is what are the vocabulary sizes that you might want to pick. And in general, vocabulary sizes have been trending upwards. And I think a big part of why is because LLMs are being deployed out in the wild. They're becoming more useful services. And when that happens, you're going to interact with people speaking different languages, people using emojis, all sorts of other kinds of almost modalities or languages than what you might expect. And so I think some of the earlier models, and especially monolingual models, ranged around in the 30,000 to 50,000 token vocabulary range. And you can see this in GPTs, the early LLaMAs. But if you look at the multilingual, or I would call production systems that have come out, they've all been shifting towards the 100,000 to 250,000

55:39 range for their vocabulary sizes. And I looked at Command A, which is one of Cohere's models. They're a company that emphasizes a lot of multilingual stuff. You see very large vocab sizes from them. Even with GPT-4 and many others that have copied the GPT-4 tokenizer are going to be around 100k tokens. And so that's the standard that a lot of people are operating at, roughly at 100k to 200k token size. And I think there's been work showing that as models get bigger, these models can, in some sense,

Regularization: weight decay vs. dropout

56:12 handle more and more, or make good use of more and more vocab elements. And so you might see increasing trends to token counts as models get scaled up or more and more data is used to train them. Cool. So the last thing, this is no longer specific hyperparameters, but two other things that you might need to do before you set your model to run, which is dropout and other kinds of regularization. And I think this one was really interesting to me when I was originally doing the research for putting this lecture together. If you think about pre-training, pre-training is about the furthest place that you might think of from regularization. Because pre-training, you do usually one epoch. You can't even go through all of your data because you have too much of it. So you're going to do one epoch training, and you're almost certainly not overfitting the data in that one pass that you're doing. And so you might think, we don't need regularization for pre-training. Let's just set your optimizer loose. It's all about minimizing loss. And this is really good arguments for why you shouldn't need to regularize. But then, if you look at what people do, the story is actually kind of mixed. And this story actually is maybe even more mixed than what has turned out to be. But early days, people did a lot of dropout. And then there's a lot of weight decay that also seems to be happening. And these days, I think, a lot of the people have stopped publishing details on precisely they're training hyperparameters.

57:50 But dropout has gone out of fashion. But weight decay has really been something that a lot of people continue to do. And why is that? That's a really odd thing to be doing. So I'll give you a moment to just think about this state of affairs. If you're training a really large neural network for one pass on SGD, on vast amounts of data, why would you use weight decay when you're doing that? So maybe some of the answer. But I think that's a kind of interesting thing to think about. It's very intuition violating, at least for me. So the reason is because it's not to control overfitting. In the sense that if you look at weight decay, different amounts of weight decay don't really seem to change the ratio of training loss to validation loss. So you can train with different amounts of weight decay if you train for long enough or you control your hyperparameters appropriately, you'll end up with the same train to val loss gap. So overfitting, nothing's happening here, even with zero weight decay. But what is interesting is that the weight decay seems to be interacting somewhat in a strange way with the learning rate schedules of the optimizers. And so what's happening is that if you look at a constant learning rate, so this is a model trained on constant learning rate. And then you suddenly decrease the learning rate to near zero. So you see this drop off as you decrease the learning rate. And then let's look at different kinds of weight decay that you could do.

59:28 And what happens is with weight decay, the model is not training very well at this high learning rate. And then when you decrease the learning rate, it'll very rapidly drop off. And when you look at cosine learning rate decay, what happens is that the models with high weight decay start out very slow. But then as they cool down, that is, their learning rate decreases, they very rapidly optimize. And so there's some very complex interaction happening here between the optimizer and the weight decay and some of implicit acceleration that happens near the tail end of training that ends up giving you better models. And so the answer to the question I posed you is you don't weight decay because you want to regularize the model, which is kind of what it was designed for. Your weight decaying in order to get actually better training losses. And you end up doing that because of the various learning dynamics at the tail end of training, as you decrease your learning rates to 0. It's a very interesting and complex. And in some ways, troubling thing to be doing with language models. But now, you see why, if you look at a lot of the reports, you'll see we use weight decay. This is why that ends up happening. Cool. So putting all that together, so there are certain things that I think are just no-brainer. So if you're picking various hyperparameters for your model. You don't really need to think too deeply about them in the sense that they've been validated and basically everyone else does them. So this is things like the hidden size of MLP, the head dimensions of your multi-head attention,

61:07 your aspect ratio, and your choice of regularization through weight decay. All of those, there's fairly good, I think, consensus evidence of how to pick most of these hyperparameters. And those defaults roughly give you the kinds of things that we suggest in the assignment. So you can follow along, and they'll roughly give you something similar to this. So any questions about the hyperparameter piece? Yes. Is there a reason why dropouts gone out of fashion? That's a good question. I don't think I've seen-- the question was, why did dropout go out of fashion? I haven't quite seen a deep analysis of why dropout is or isn't helpful. I haven't seen any result that, for example, shows that it helps for training loss. And both this paper argues and logic would dictate, there's not really a training overfitting issue with these models that can't even do one epoch over their training data. Yes. Multilingual vocabularies actually contribute to improve performance in one language. So, [INAUDIBLE]. So the question was, do multilingual vocabularies contribute to improving performance in one language? When you say one language, you mean do multilingual or larger vocabularies help performance in English? Is that the right question? Yeah. So I think in your high resource language, the impact is less. So if you're only thinking about English language modeling, you can get away with smaller vocabularies. This much is kind of true. But the place where larger vocabularies is really helpful is when you're starting to get at,

62:48 I wouldn't say the tail of your distribution, but when you get to languages that are more minority. And one great example of this, if you look at any of the Cohere announcements about their models or their tokenizers, they basically always argue that because of the way they have larger vocabularies and the way they train their tokenizer, non-English and low resource languages, they are packed into much fewer tokens. And so people using those pay much lower cost at inference time, which is a great benefit. Yes, question? [INAUDIBLE] if weight decay doesn't have a significant impact on the val loss, why do we care about the training dynamics or the favorable [INAUDIBLE] dynamics? So the question was, if it doesn't have an impact on val loss, why do we care about training dynamics? The goal is still, I want to get good training loss. This is the game that we're playing. And the surprising thing about weight decay is that somehow, it gets us better training losses. I think the intuitive thing that makes sense is, you do weight decay. It gives you better val losses. But that's not what happens. What it's getting you is better training losses, which are also the same as val losses. Yes. Are there differences in the architecture hyperparameter choices people make as they move towards multimodal architectures, if they're doing images as well as text? So the question was about multimodal models. That is a great question. My survey of multimodal models is very incomplete. What I can say is, a lot of the academic and open work that I've seen, they do what you might call

64:26 like shallow or later fusion or earlier fusion of the modalities. And the way that works is you bolt the vision modality onto an existing language model. In those cases, the hyperparameter and architecture choices are fixed. One thing I will note, and I will talk about this in just a few slides, is that the multimodal models pioneered some pretty interesting techniques in stabilizing language model training. And that's been a really big theme, and I'll talk a little bit about those. So what is different is often when you bolt on this new vision piece and you retrain with that, that's a big shock to the model.

Stability tricks

65:00 And so you have to think carefully about how to stabilize that training process. And those innovations have actually seeped back into pure text language model training. Cool. So I went back through and I looked through all these new papers. And as I was trying to think about, OK, what's been new in the last year? And what new architecture and related things have happened? Actually, the core architecture hasn't changed much. But I think the one thing that stood out as being very emphasized in a lot of the releases has been what I would call stability tricks. And so these are things where you would like to train your model in much more stable ways. And as you make bigger and bigger models, you train for longer and longer. These kinds of issues start to appear more and more. So I've taken this from the OLMo 2 paper. And actually, that paper is a great set of academic results on LLM training stability. And one thing they start with is this figure. And you look at this blue curve over here. And you look at this, L2 norm of the gradient graph. And this is terrifying graph to look at. Your loss curve kind of seems to be behaving OK. But you've got some bad spikes every now and then. And you open up your gradient norm, and it's this horrible plot where you've got spikes everywhere, where your norms are completely blowing up. And if you're training models like this, you're going to have a really tough time getting it to converge reasonably. At some point, it's going to hit-- gradient norm explodes, and you can't do anything.

66:39 And your training is done so you can't train any further. And so there's been a lot of emphasis basically trying to turn this blue curve into something that looks a lot like the orange curve. And of course, this loss is higher, but ignore that fact because I think they just switched data sets in between these two training runs. But this orange curve has nice low gradient norms throughout. And that's really the kind of plot that you would much rather see. And so you might ask, where do stability issues arise in transformers? And of course, they can arise basically everywhere. But if you look at the kind of interventions that people are making, there's really one place that really stands out as the kind of problem child. And that's the softmaxes. And it can be a problem because you're going to be taking exponentials. And those can be numerically badly behaved. You're also dividing two numbers. And so you might have a division by 0. So for many different reasons, this softmax piece is a part that you might have lots of issues with. And so, actually, one more thing I want to talk about. So where are the softmaxes in a transformer? Well, there's one at the very end. So you've got to be careful about that output softmax. And also, there's softmaxes in your self-attention. So there's two softmaxes that we're going to think a little bit about. And for each one, I'm going to mention stability intervention that has generally seemed to be effective. So the first one is called the z-loss. And it might desire to cite a paper that's older. I've gone back to Devlin in 2014, where in a machine translation paper, their goal was to try to make sure that this normalizer was

68:22 near 1. So if you look at p of x, that's the output softmax over here. The output softmax is two terms. You exponentiate your logits, and then you divide by the normalizer z. And the z is just summing up the values across all the vocab. And so if you want this z of x, you want to train the network to have a z of x close to one. Well, then you can rewrite your loss, and you can add a little second term here to try to force log of z of xi to be close to 0. So you're going to end up with an auxiliary loss term that's alpha log squared z of xi right. You can see that derivation on the right here. And this is in some sense what people often call the z-loss. I think Jacob Devlin and others did this for machine translation for totally different reasons than what it's used for today. But this was, I think, the first instance of this in language modeling land was PaLM, who used this, as they called it, auxiliary loss of z-loss 10 to the negative 4 log squared z to basically encourage the softmax normalizer to behave nicely. And you can reason through the behavior of this regularizer. If it succeeds and it forces log of z of x to always be 0, then the log and the exponential cancels, and you've basically just got U of r of x. And that's a good place to be. That's a nice numerically stable operation. So all of these problematic operations go away. And so you can think of the softmax as being well-behaved when z of x is close to 1. Or log of z is close to 0. And PaLM, in some sense, is very much a pioneer. Because they did this z loss trick, and many others

70:02 didn't really do it for a long time, or at least the ones that had open papers. But then there was a kind of sequence of papers that have done this. Baichuan 2 is actually the earliest follow up that I know of. And then DCLM and OLMo 2 and now several others have basically picked up on z-loss. So a very nice, convenient intervention for improving stability. And then, the other trick that we see, so that was how to stabilize the output softmax. But we've got another softmax we've got to deal with. The other softmax we have to deal with is in the attention operation. And so this is from an NVIDIA paper. I forgot to put the citation marker. But here, this is a block diagram of how attention works. You got your layer norm at the beginning. You got your QKVs. Ignore this for the moment. You might multiply your Qs and your Ks. You'll softmax it. You'll multiply the V, and then you'll project it. And then that's going to give you your fully connected and your output. So if you ignore this little piece over here, this looks just like your normal multi-head attention operation. So what's the difference here? So several folks came up with this idea or this approach called the QK norm, where you take the queries and the keys, and you pass them through a LayerNorm layer before you take their inner product for the softmax operation. And this is a very different kind of approach to controlling the behavior of the softmax. Here, you're not controlling the normalizer z. Instead, you're controlling the inputs to the softmax to be bounded in size. And that's going to naturally control the bad behaviors

71:45 of the softmax. And as I said before, this is originally an innovation from the vision and multimodal model community. Dehghani in 2023, this was a paper on training, very large vision transformers. And then Chameleon and Idefics from Hugging Face use these tricks for their multi-modal training components. And then it got picked up by several others like Gemma 2, DCLM, OLMo 2, all basically uses this kind of techniques in order to stabilize their training. And I think I'm allowed to add one joke per lecture. And so this is the one I'm going to go with here. I think one of the things that really has stood out in terms of stability interventions has been just how strikingly effective LayerNorms are. So we've seen going from LayerNorms just in the pre part of the block to both the beginning and the end of the non residual component. And now, we've also thrown it into the Q and the K component. At least in terms of improving stability, LayerNorms have been shockingly effective without affecting performance too much. The last trick that I'll note, I think this one has been not quite as frequently used, which is to soft-cap the logits that go into the softmax. So the other approach that you can take, so QK norm is in some sense a very heavy handed intervention because we're going to operate over the entire vector. But one thing you could do is after you take the inner products for self-attention, you could pass them through a soft maximum operation. So you can pass them through this equation over here.

73:29 So you have your logits as your input. Divide it by the soft cap. Multiply by the soft cap. What does that do? Well, if your logits start exceeding the soft cap by a lot, the tanh is going to clip them off to one. And so you're going to have a maximum value of soft cap over here. So this is going to control, in some sense, soft clipping of the logits and Gemma 2, and I think OLMo 2 also do this. It hasn't been I think quite as popular otherwise. And I think the other evidence against this, the NVIDIA folks that I mentioned earlier, did actually quite a few different stability improving interventions. And what they find is you have your baseline model over here. This is the perplexity of the baseline model, 11.19. Soft capping makes it worse. QK norm actually makes it better because you can use more aggressive learning rates and push the optimizer further. Cool. So that's the end of the stability improving intervention stuff. Does anyone have any questions? I think that's been kind of the new development over the last year. Yes. So for the QKV norm, I understand that during training, you will have the LayerNorm being applied. At inference time, it's the layer norm still being kept? Yes. So the question was, at inference time, do you still use the norm? And the answer is yes. Because the layer norm has learned parameters. Like the whole action of the layer norm is it takes an activation, normalizes it to a unit, and then scales them to some size. If you take that out, that's a huge change to the model. It will have no idea what to do with those unnormalized activations. Cool.

Attention variants: MQA/GQA and sliding windows

75:12 So I have this last bit, last few slides that I want to end with. If we go over then, we can always push this into the MOE lecture. But I think we also have a lot of content next time because I have to cover DeepSeek V3. So the last thing I want to cover is variations on the attention heads. So attention heads, I think, haven't had as much work done to them. But there have been a few, I think, important changes that you need to know about in order to understand the models that are being trained. So the one thing I'll talk about-- the first thing I'll talk about is GQA and MQA. And these aren't really critical to the training time behavior of the models, but they're very important in understanding the inference costs and inference behavior of the models. And because this is an important architecture change, I'll mention them here in addition to probably being mentioned by Percy in some of the inference lectures. The other thing that's a new development I'll mention is how the most recent models, like LLaMA 4, if you've heard of it, supports supposedly 10 million tokens of context. How does it do that? Well, it does so by messing with the attention pattern in very structured ways. So I'll talk about that as well. So GQA and MQA, if you looked at some of the larger models, like the big LLaMA models or others, you'll have heard or seen this term GQA or MQA. And I'll talk through what that means. So to set the stage, let's think about the compute that you need to do attention. So this is, once again, 224N slides here. You're going to take your XQ, your query, and your XK.

76:52 And then you're going to form your big quadratic attention matrix. And you can walk through each of these matrix multiplies. And you can convince yourself that the total number of arithmetic operations is going to be b times n times d squared. So that's going to be b is the batch dimension. n is the sequence length. And d squared is going to be the hidden dimension squared. And you can ask about the total memory accesses. And this is going to be b times n times d. And this is going to be, for example accessing just this matrix here. This XQ is going to be that size. And then the softmax is going to be b times h times n squared. And you can convince yourself of that by just thinking about the size of the softmax matrix, which is going to be batch times number of heads times all of the different softmax activations that you have. So that's n squared of them. And you've got a projection, and you've got d squared projection operations at the very end over here. And so we can take the ratio of total memory accesses and arithmetic operations. And this is going to be something that will be very important in a couple of lectures, this idea called arithmetic intensity. So we want our arithmetic intensity to be high. What that means is, we want to be doing a lot of compute for every single memory access that we do. And this is going to be because memory accesses are very expensive on a GPU, relatively speaking. And compute is relatively cheap. And so in this batch computation that I'm showing you here, the arithmetic intensity, if you take the ratio of those two things, is going to be 1 over k plus 1 over bn inverse.

78:31 And so this is going to mean that we can keep our GPUs running. Because if we have of large number of heads and we have large batch size and large sequence length, those are all going to be good large numbers. Of course, this is what happens at training time. So the issue is that inference time, we do not have these big chunky matrices to multiply together. And so that's going to really change the nature of the behavior of our algorithms. So when we're generating text, remember that we have to generate a token. And then the transformer has to read that token. And then it has to process it. And now, we can get the next token distribution. And then we do the things autoregressively one token at a time. And by doing this, we can't parallelize this generation process. We need to go step by step for every single new token. And when we do this, we're going to need to incrementally compute attention, an idea that people call the KV cache. And so what do you do? This is a lovely animation of a KV cache that's been explained. So if you can look at this figure, what you're doing is you've got a query token. A query token here is you've generated a new token. You're conditioning on it. And now, you want to ask, what information should I look up in the past, past that query token? And your query tokens are shifting from 1 through n because you're generating new tokens one at a time. You're building up this key cache over here, where, basically, I'm building up all of the past tokens keys. And the past tokens keys don't change because they only depend on things in the past. And so, incrementally, as I generate tokens, building up

80:12 all of these past keys. And each time, I can compute one new element of Q dot K. So the big attention matrix is going to be this lower triangular matrix. I'm computing one row at a time. And that row is exactly what's necessary to generate the next token. So this KV cache idea, if you've not seen this before, is this idea of saying, I'm going to generate the K's and the V's incrementally as I go, as I generate each token. And I'm only going to compute QK that's absolutely necessary to do my operations. And so once again, you can go through and do the various arithmetic components of how many flops do we do? What's the total number of memory accesses? And if you think about the KV cache, I'm only multiplying the absolute necessary keys and values. Since I'm saving all of the intermediate computations, I'm not wasting any matrix or vector multiplies. The total number of arithmetic operations remains exactly the same, bnd squared. But the memory access patterns are now different. Why is that? Because when I do this KV caching thing, I'm going to have to move various kinds of parameters in and out of memory repeatedly. Whenever I multiply with a key matrix, I'm going to have to put that into memory, and then multiply by K. And then I need to put that away, and I need to compute some activations. And so I'm repeatedly loading in different matrices. And that's going to give me a much higher total memory access of bn squared d plus nd squared. And so when you take this ratio, now, the arithmetic intensity is not so good.

81:50 You're going to get n over d plus 1 over b inverse. And so if we reason through this-- so if I want arithmetic intensity to be high, I want this thing inside to be very small. So I need really large batches. And I need n over d to be small. What does that mean? I need really short sequence lengths or really big model dimensions. And this n over d is really unfavorable because I don't want a bigger model, and I don't want a shorter sequence length. And so this is the core, in some sense, inference cost trade off that people face. You have this very bad memory access pattern where you have this one term, n over d, that's really killing you in terms of the throughput of your system. And so this motivates this thing called MQA. And the key idea here, hopefully, you see from this figure back here that, really, the part that's really bad is the keys and the values. They have this KV cache thing being built up, and there's memory moving in and out. So what you do is you can have multiple heads for the query, multiple query heads, but only one dimension or one head for the keys and values. This immensely simplifies things. Once you do this, now, you're moving much less information for the K's and the V's. And so K and V is shared. But query has many heads. And so you still have multi-head attention or multiple queries, but only single K's and V's. So that's why it's called multi-query attention. And now, when you do the same kind of arithmetic we have fewer memory accesses because we've shared the K's and the V's. And the arithmetic intensity is much, much better behaved. And so we can increase things like--

83:34 we've decreased the first term by a factor of n. So longer sequence lengths are now viable. And the second term is now divided by the number of heads. So this term is also not so terrible. So all the different terms are controlled now. And MQA can give you much better behaviors. GQA or grouped-query attention basically changes this slightly. Instead of having single query-- or sorry, multiple query and single key, you can reduce the number of keys by some multiple. And so this will let you trade off between the inference time behaviors and the expressiveness of the model. Because maybe going from multi-head all the way to multi-query is a little bit too aggressive. Some works show that GQA doesn't hurt, but multi-head attention hurts. I'm not going to get into that. I'm just going to close off with this very last thing, which I think is a really interesting development in the last few months. So back in 2019, OpenAI had this cool paper basically arguing how to build longer attention models. And they were basically arguing, well, one way to do that is to come up with sparse attention patterns. So instead of paying attention to all the sequence, I'm going to pay attention to, let's say, a local window at each chunk. And then I can have other attention patterns that are like diagonals that help propagate information across. So you can build sparse or structured attention that trades off various kinds of expressiveness versus runtime. GPT-3 uses exactly these kinds of tricks when they originally released it to get larger attention windows.

85:13 Sliding window attention is another variant of this idea, where at each layer, you only pay attention to a small region around your current position. And this also is going to control the total amount of resources that you need-- total amount of resources you need in order to do longer contexts. So your effective receptive field is now the local 1 times the layers. The final trick-- so those were the older ideas. But the way that this has kind of been modern instantiation is some of the recent papers like LLaMA 4 and Gemma and Cohere Command A have now come up with this very clever trick of basically having transformer blocks, where in this case, you have a set of four transformer blocks. The very bottom one uses full self-attention with no position embedding. So there's no RoPE, no nothing. It doesn't know about position at all. But it's full self-attention. And it only happens once every four blocks. And then the three blocks above it use sliding window attention with RoPE. And so this is actually a really clever trick to both control the systems aspect of things because the full attention only happens every now and then. And also, the length extrapolation aspect because RoPE only deals with local context windows. And anything that's really, really long range has no position embeddings at all so it could extrapolate very, very aggressively. Because you don't have to do this position extrapolation that you do with something like RoPE. So that's a really cool development that we've seen in the last couple of months. So I think we're coming up on time.

86:52 Feel free to ask me questions about architecture or hyperparameters. I'll be happy to answer questions after.