Mixture of experts

Why this became a required lecture

00:05 So, we'll get started. Today, we're going to cover a mixture of experts. Last year, this was kind of a fun bonus lecture that I threw together. But this year, thanks to, you know, lots of people doing, this has become a much more critical lecture. So I've added a lot of the recent developments and at the end we'll try to walk through DeepSeek V3 and try to understand like what are all the sort of components that make up a state-of-the-art open source system or at least on the architecture side what that looks like. So mixture of experts is how a lot of you know the most modern high performance systems today are built and deployed. So there was the funny Nvidia leak of GPT-4 actually being potentially revealed as GPT-MoE-1.8T. But more you know broadly others like Grok and DeepSeek and Llama 4 now have all adopted a mixture of experts architecture and it seems like at this point in 2025 that the advantage of mixtures of experts over dense architectures is very much clear right almost all compute scales training a mixture of experts model if you do it well is going to give you benefits over a dense model and so everyone seems to be doing it in both the east and the west. And so this will be an important thing to understand if you're trying to build sort of the best model that you can for the flops that you have. So mixture of experts is very simple. It's a very terribly named concept. I think you hear mixture of experts and you think oh there must be

What an MoE actually is

01:43 experts specialized for different domains and they're like doing different things. Like there's a coding expert and like an English expert and a other languages expert. It is very far from that mental model. A mixture of experts is a type of fancy architecture that has several subcomponents called experts that are activated sparsely. And in particular when you think about mixture of experts you should be thinking about the MLPs. This is where all the action is right. So a MoE architecture and a nonoe architecture are going to be similar in almost all of its components except for one and that is you know if you look at this slide over here you know this is the components of a standard transformer you got your self attention you got your FFN if you zoom in you know in a dense model the feed forward component just sort of you know is there it's one big block in a sparse model what you would do is that you would take this FFN and you would split it up or you would copy it depending on how you're going to be setting up your multiple copies, let's say, of your FFN, your fully connected networks, and you're going to have a router that picks some smaller number of those, you know, in each forward pass or at each inference, right? So, this is the basic idea behind, and we're going to replace this one big feed forward on the left side with a selector layer and many smaller ones. And what's the advantage of this thing? Well, if it's sparssely activated, that is, let's say only picks one expert and an expert is the same size as your dense FFN, then the flops between the left side and the right side, the dense model and the model, they have the same flops, right? They're doing the same matrix multiplies

03:19 as you do your forward pass. So you have more parameters without affecting your flops. And if you're a believer that what matters is having more parameters to, for example, memorize facts about the world, well, you know, this is a great architecture. So you can kind of see the intuition behind. Hopefully that's all very clear. And you might wonder okay so it makes sense that you can get you know more parameters per flops but does that translate to actually better performance for the models that you're training. And there's been, I think at this point,

The evidence: same FLOPs, more parameters, lower loss

03:52 many, many, many papers showing that at the same flop count, at the same training amount of flops, you get better performance out of a mixture of experts, than out of a dense model. So, this is a nice paper to so today I'm going to go over a couple of the classic Google papers that, you know, put this field together. And this is one of them by Fedus et al. 2022. Where they show that you know if you flops match your you know training flops so that's the same amount of compute used for training and as you increase the number of experts the training loss of your language model just keeps going down and down and down and down and down right so you know more experts better of course the experts aren't free you need to store the memory for these experts and when you do parallelism you're going to have to think about you know routing your data into 256 separate experts so that's there's going to be systems complexities but if you're only thinking about flops. This is a great chart to see because you have the same flops, but you've gotten free, you know, test loss here. And you see the same thing reflected on the right side. You know, as you train for longer and longer, the model, the switch base with 128 experts, right? The model with more experts, you know, gets better perplexity faster, right? So hopefully that that is quite clear. You might say, well, this is a 2022 paper. Is this true sort of on modern architectures on modern scales? It you know continues to very much be true. AI2 had a very nice paper OLMoE which did a whole bunch of ablations and carefully controlled comparisons into dense versus and other architectures and they sort of see exactly the same

05:28 thing. So here on the left side this is still from Fedus et al.. You see the 7x speed up from having many experts. On the right side this is the OLMoE comparison. You see the pink one is the MoE and the teal one is dense and the training loss for the dense model goes down much more slowly than the MoE. Right? So hopefully you know I have in some sense sold you on the value of MoEs and for learning this kind of new slightly new architecture. Right? So we're going to pay a price for all of this but at least at the flops level this looks very compelling. Right? So yes question the last lecture you mentioned like yeah the bias turning part because although the it's a pretty cheap computation affects our actual process pretty badly we're you know loading in and out is there right so the question was in last lecture you know I was saying even small non- flops you know negligible flops can be really big in wall clock is anything in the MoE world going to look like that and so I think one of the drawbacks of why you know that's not the standard thing that's being taught you know let's say at CS224N is because there's significant systems complexities to making this thing efficient so I'll you know get to that it's possible to make these things very efficient especially if each expert lives on a separate device so that you know you're routing thing data to different places you can be very efficient when you do that but it's not easy right so there's a lot of infrastructural concerns s and you're going to see a lot of complexities to get this thing to work. But when it does work, you know, you're you're putting all of your flops to

07:07 use. Okay. And then the last one that I wanted to show is, you know, a lot of the companies really love because you get to present plots that look very compelling like this, right? This was from the DeepSeek V2 paper. You know, on the X-axis, this is this is a little bit of slight of hand. This is only activated parameters, right? So this is only the parameters that are used you know for computation. So you ignore all the deactivated experts and the y- axis is MMLU performance, right? And we see DeepSeek v2. Wow, look, very few activated parameters, really good MMLU performance, right? And so if you're only interested in both training and inference flops, you know, activated parameters is the name of the game. You get really good performance here. And this is, you know, not just an ablation. This is a real system that someone, you know, spent a lot of money to train and deployed out in the wild, right? And we'll see this sort of pattern recur in other examples as well. Oh, was there a question? Oh, no. All right. And so the systems thing that is also a benefit is that allow us to have another axis of parallelism. So I'm going to get

Expert parallelism, and why MoEs stayed niche

08:12 into parallelism in much more detail sort of in the systems lectures when I'm going to talk about how you're going to take your model and you're going to cut it up into many small pieces and lay them out across many different devices. But I'm going to talk at a very high level. But when you have experts, there's a very natural way to parallelize at the expert level, right? So you have multiple different feed forward blocks. You can take each of these experts and you can put them on a different device, right? And because experts are sparsely activated, all you have to do is take your token and route it to the appropriate device and the computation will happen on that device, right? So it's a natural sort of cutting point to be able to shard your model into different devices. And so this is called expert parallelism. And this is another reason why MoEs are very popular, right? If you really want to parallelize really big models, this is a thing that you're going to have to do. And kind of interestingly enough, I think, you know, developed at Google and many of the frontier labs, the closed labs were doing it, but I think the open results actually came from China very frequently. Qwen and DeepSeek were doing you know a lot of work last year and it's only really recently that I think western open source groups have started to do more work so Mixtral Grok I guess Grok's not open and then now Llama is now a MoE architecture right and so here you know Llama 4 just got released right latest and greatest this is also a sparsee and I'll talk about llama 4 as well as I go through the lecture as I said before, so you know one of the kind of starting points for

09:50 this is some of the Chinese groups Qwen and DeepSeek have actually done some really nice work benchmarking and understanding and evaluating some of these results. So these so Qwen 1.5 was one of the first models that I knew of to have like this largecale well- tested well doumented. And what they did was they took a Qwen 1.5 dense model and they had a nice trick to upcycle it into a mixture of experts. That's a clever kind of trick to take a dense model and then turn it into a MoE. And they showed sort of significant gains at least in terms of compute efficiency while sort of m decreasing the total number of parameters relative to their sort of 7B model. DeepSeek, which is now famous, but originally when these papers were coming out were not quite as famous, did some of the, I think, really foundational work in the open-source world. A big part of this lecture is actually going to be tracing the trajectory of the DeepSeekMoE architecture. But if you look at their original DeepSeek paper, you'll see very nice papers, sorry, very nice sort of comparison showing things like what happens when you train a dense model with a particular amount of flops. What happens when you train a really naivee that doesn't do very smart routing what happens and then if you use a smarter routing called the switch sort of what happens and so you'll see all of these very carefully controlled comparisons and you see you know as you go from dense to sparse right so that's the leftmost column to the rightmost column you see all of these sort of benchmark metrics very consistently improve for a fixed amount of flops

11:28 right so this is very consistent and kind of one thing that I think almost everyone at this point has probably heard of, right, is DeepSeek V3. And that's in some sense, you know, a culmination of all this line of work. But if you had been following and you were excited about kind of this branch of neural networks and language modeling, you would have actually known about, you know, DeepSeek long before V3 got popular. And we'll see at the very end of this lecture actually DeepSeek V3 is not very different from the very earliest DeepSeek models. Architecturally, you know, they had kind of nailed it way back when they were training these sort of much smaller two billion parameter models. They really just kind of got the engineering right to get something that is actually really quite remarkably good, which is their V3 model. Okay, so now I think you know I have spent quite a few minutes trying to really hype up MoEs and they really are I think worth hyping up. They're very good. But I think there's a question of why haven't they been more popular, right? Why isn't it the standard thing we teach in, you know,

Why MoEs have not been more popular

12:33 NLP and language modeling classes? It's just that they're very complex and they're very messy and I'm hoping that they'll get simplified over the next few years, but they still remain pretty nasty. So one of the things is you know the infrastructure is very complex and the biggest advantages of really happen when you're doing multi-node training like when you have to split up your models anyway then it starts to make sense to shard experts across different models that's a very natural thing to do but until you get to that point maybe are not quite as good right so some of the earlier Google papers really talk about this trade-off where they say actually when you get these really big models that you have to split up then experts become uniquely good. There's also other things that are really tricky. If you think about it carefully, right, this decision of which expert you route tokens to is a very difficult thing to learn, right? In deep learning, we really like differentiable objectives, right? Very smooth things that we can take gradients of. Routing decisions are not differentiable because we have to pick and commit to a particular expert. So if we're doing that, you know, we're going to have a very tricky optimization problem. And the training objectives to make that work is either heuristic and or unstable, right? And so we're going to have to really carefully engineer those guys to get them to work, right? So those are two reasons why you don't really want to maybe do this normally. So what do MoEs look like? As I started this lecture with, you know, the classices that you should think of is you take, you know, the densely connected layers, the FFNs, and you split them up or you, you know, copy them and you have sparse routing decisions among them. Of course, you

14:13 could do the same kind of idea. You could have a sparsely routed attention layer. And some people have done this. There's been a couple papers and a couple releases that have taken this approach. But it is actually quite rare to see this in the major model releases. I think I've seen people talking on the internet saying like this approach is actually really much even more unstable and very difficult to really train consistently. It's sort of I haven't really seen the ablations to back that out but certainly there haven't really been many people training those kinds of models with attentions. So now you know I've told you about the basic architecture, right? It's really simple. It's just you have a router of some kind and you route and then you have different MLPs. So what are the things that might vary across different choices? You might ask how do we route right the routing function is an obviously important choice. How many experts and how big should the experts be? That's another choice. And then the final one is how would we train this router? Right? This non-differentiable objective that seems very difficult to train. So those are very important design questions and we're going to go through each one hopefully covering the design space of all these things. Okay. Any questions before I get into to each one of these different subcomponents here. Good. Okay. So if you're interested in just kind of understanding a broad overview of MoEs at least circa 2022 there's a really nice sort of survey or a review paper by Fedus et al. In 2022 that covers a lot of these and many of my figures you

Three routing families

15:51 know are credited to that paper. If we're thinking about how we're going to route or essentially match tokens to experts, right, this is the core component of because whate does is, you know, tokens are going to be coming in, right? You have your sequence that you're processing and those sequences are going to be assigned to experts, right? Not all experts will process every token. That's the whole point of a sparsely routed. And so you can ask how are these routing decisions made? So you can sort of have three different kinds of choices. You can have token choice where each token is going to have a sort of routing sort of preference for different experts and I will choose the top K experts for each token or I can have expert choice where each expert is going to sort of have a rank preference over tokens and then I'm going to choose the top K tokens for each expert. This has a really nice benefit of being balanced over experts. And then the last one is sort of you could solve some sort of complicated optimization problem to make sure that the mapping between experts and tokens is somehow balanced. Right? This is global assignment. And just to you know give you a bit of a teaser here almost all the MoEs do token choice top K. In the early days of people tried many different things spanning this whole spectrum of design space of token routers. If you look at the big releases they have all converged to basically one class of routing mechanisms which is token choice top K. So each token is going to rank order experts by affinity and then there's going to be kind of a top k choice for

17:28 each one of this and OLMoE, which I'll keep referring to throughout this lecture because they have a really nice series of ablations. So it's really nice to teach off of have exactly this ablation. They compare a token choice routing versus an expert choice routing and they show if you look at validation loss token choice is much nicer behaved much faster in loss decay. Yes. Is this function a function of the token itself for its position? It's a function of the sort of the hidden state, right? So the token is going to get processed with all the position embeddings and so on and then the hidden state will come in and then it will be processed by the MLP. And so for the other for the other two like for the experts choosing the token and also like the next one when you say it's like more balanced across the experts are they like it's still for the current like token sequence but it's like it's forcing them to be more distributed like it's still going to be the same set of tokens but really it's about kind of the ranking selector function right in token choice I'm just going to take the top K amongst the columns like maybe the fours are even identical, right? I'm just going to take the top K amongst the columns. In expert choice, I'm going to take top K amongst the rows, right? And top K amongst the columns is kind of nice because you might be able to say, oh, I can define a scoring function such that the score is how well each token gets processed by each expert. And token choice will route me to the best expert, right, for that token. So, that makes sense from processing. But expert choice has the benefit that each expert gets exactly the same number of tokens. And so now you might like if you're putting different experts on different

19:04 devices, you've got balanced utilization. So there's different trade-offs at play as you think about routing. Yes. How does a token know which expert is the best? Good. Yes. So the question was how does each token know which expert is good? That is exactly the role of the router. And I'll give you the router equation, but to give you a bit of a not really a spoiler, but you know the routers are much more lightweight than you think. So you know your token, let's say, is represented by vector X. That's your like hidden you know residual stream coming in. So now X is going to get multiplied by you know W a matrix and then you'll just take you know a sigmoid or something and that's the score. So it's really just a vector inner product almost like an attention operation in a way. Yes. Right. So the choice of so the question was is K1 here? So K is actually a hyperparameter and different will choose different things. I will talk about this again, but to give you the high level intuition, the initial argument that the earliest MoE papers made was that K should be greater than two because that way you get some exploration, right? If you're doing K equals 1, maybe you're just always exploiting the best arm and you'll never know about the potential other things you could do. But if K is two, then maybe that second arm can tell you a little bit of exploration information. So, you know, K equals 2 was the canonical choice. And K equals 2 actually continues to be very popular. That would be like double the flops. Like that's right. That's right. So that would double the flops. And so when people talk about they usually say things like x number of activated parameters and that would account for the fact that you're you know put

20:41 putting in two MLPs. Yes. So when k is greater than one like even time do we combine the outputs of the different experts into? Yes. The question was when k is one do the outputs get combined? That's right. Like if you if you look at I guess like look at the attention diagram over there you know you got the router it's routed to two MLPs up top and then they get combined together right after right so that's exactly right so in that case you can just like a simple average so the question was how does the

Top-k routing in detail

21:14 aggregation happen it's just the sum right so I'm going to go over the variance very common variance that people do and really in some ways All you need to know is top K in order to actually implement a high performance. But I'll give you the other variants because they're natural things you might think of. Top k routing is what is used in most token choice top routing topk routing. So how that works is you know you have your residual stream inputs x. That will go into a router and as I said a router is really kind of like the attention operation. There's like a linear inner product and then a softmax and then you pick the top k most highly activated experts and then those outputs are gated. Depending on the implementation you might weight the outputs based on this router weight or you might not and then you will just output the weighted average or just a straight sum depending on how your implementation works. And so a lot of the MoE papers and methods use top k switch transformer GShard grock Mixtral, Qwen and all the DeepSeek variants use

Hashing, RL, and linear-assignment routers

22:22 different top k variants. Maybe a very surprising fact and this should really make you think about what's going on with there are a lot of results that show that actually you don't even need a smart router at all. You can actually just use a hashing function at the very bottom to map these X's onto your experts. And even if you're doing hashing, so no semantic information at all, you will still get gains from a hashing based, which is pretty wild. Some of the earliest work on MoEs, I think, had the very smart idea and in many ways the right idea if you're thinking about this top down of using RL to learn the routing behavior, right? Of course, you know, the choice of where to route to is a discrete decision, and RL is great for learning discrete decisions. Why don't you use RL to learn routing? It was used in some of the earliest work on mixture of experts. As far as I know, basically no one does this now. The compute cost to do this is too prohibitive and you already have stability issues. You might not want to do that. There have been a couple of papers that have explored things like solving linear assignment problems or optimal transport style problems. They're very elegant but once again the cost of doing this is much higher than the benefits that it gives you I think in practice and it hasn't really been adopted but there's a lot of really interesting things that people are doing like this to try to improve the routing so now I can you know point at this slide and really talk through how routing works in detail so this is the kind of top K routing that almost everyone has converged to now this is the router that's used in

24:03 DeepSeek v1 to v2 Qwen and Grok do almost exactly this. There's a instead of having a softmax directly at the bottom here they do a soft DeepSeek V3 mixrol DBRX don't have a softmax at the bottom but they'll softmax the G of its but it's a very minor difference. So let's walk through what's going on here and try to reason about you know the behavior of this. So what's happening here is at the very bottom we've got our inputs. This is our u_t^l input and I would like to take this sort of residual stream input and process it through you know my first thing I'm going to do is I have to figure out which experts are going to be activated. Now how am I going to do that? Well how I'm going to do that is very similar to attention. I'm gonna take my U, which is my residual stream input, and I'm gonna take the inner products with the E of I's. These are kind of learned vectors that are for each expert that tells the expert I'm an expert, you know, that points in this direction, right? And so I'm computing in this inner product here expert and input affinity and I'm computing a softmax to determine for each, you know, token what are, you know, the best experts, right? So I normalize this is S of T. Now I take the S of I of T and I go through a top K function. I only select the K best weights and then I use this as my gate. So I zero out everything else and I take the weighted average of each of the experts outputs. And then I add that to my original residual stream and then I return that. Right? So this is hopefully very

25:45 familiar to kind of what you're all very familiar with in terms of how you know transformer works with only the difference of this top k routing piece. Is that clear kind of to everyone how this thing works? Good. Excellent. So in some sense the mechanics of the forward process of the routing is very simple. What is kind of mystifying is that fact that you can learn this very well. Right? This is in some sense a fairly complicated set of things to have to learn to do well by a model. Yes. So we're using softmax here. Previously one of the benefits of softmax is that it's going to push you pretty extremely to choosing a singular max. It's not a hard max, but it was so I'm having trouble thinking of the intuition of putting the softmax basically on top of like combining it with the top K where you're getting multiple and then you're using something that's going to push you towards choosing just one thing. Yeah. I mean I think maybe one way of thinking about the softmax is you know it the whole purpose of this is just to make it so that when I average my experts later it kind of sums to one. Don't think of the softmax as like a softmax operation even though that's literally the name. Really the softmax operation is a normalized to one operation and the normalize to one operation is going to make that a weighted average up top. The other thing that's very important is you know you might think why can't I just get rid of the top K why don't I just use the softmax here and just you know gate all the experts well then you immediately lose the systems efficiency aspect of this right you have to have top k during training otherwise

27:30 you pay the training cost of all capital n of your experts right this is the key thing about like we have to do all this gymnastics to make sure that both at training time and inference time we have a sparse number of activated ated experts. That's why we go through the top K, right? Okay. Yes. From the back. Yeah. So, because you're doing softmax first and then the top K get the weights, you no longer have to guarantee. So, the question was, yeah, so the question was if you softmax first, you no longer sum to one. And yes, that's absolutely right. You no longer sum to one. And in some ways, like there's no requirement that you have to sum to one cuz you know the next layer can magnify it back up. You know there's layer norms everywhere. It's not as if it has to sum to one. But I think that is the reason why some of the other architectures basically move the location of the softmax. There's a kind of aesthetic choice about whether you really want that, you know, weight to be normalized to one or not. Yes. Yeah. So I was wondering like how does the E vector here relates to the weight of the feed forward? Okay. So the question was whether the whether and how the E vectors relate to the feed forward. They're not really tied in any way. The E vectors are just learned vectors for the just think of the E as parameters for the router, right? They're just separate objects from the FFN. Yeah, I was just wondering how does compared to sampling from the soft. Great. The question was about how does it compare to sampling from the softmax? You can sample from the softmax and some methods actually do a kind of soft sampling from the softmax. Specifically, one of

29:08 the Google papers has a procedure where they take the top element of the softmax and then they randomly sample the second element proportional to the remainder of the softmax. And that gives you more exploration, which is good, but the drawback of that is that if you don't sample at test time, now you've got a train test mismatch. Okay. Yes. Why not just reormalize after the top K? Why not just reormalize after K was the question. Is that right? In some models do that. Some models do bring normalize off to the top K, but that's a kind of a choice like some architectures don't do that, some architectures do. It doesn't actually matter because the scale can be basically adjusted post hop, right? So there's no reason why it has to sum to one after the G operation. Cool. Oh, sorry. Yes, the bias term is U there up there. Yeah. So the first term of the sum if g is approximating probability vector could be seen as an expectation of the function fn right plus you. So so ff actually this is not an expectation of ffn because each ffn is a different fn so this is not actually an expectation and the gates are sparse. So this is like a weighted selection operation over k different or actually capital n different ffns and then the utl at the very end there you know if you remember the transformer that's the residual stream right so I'm adding back the inputs because I want sort of a identity connection throughout okay oh there's another question why does the router have such a basic parameterization like what happens if

30:48 you put more weights into your router function right the question was why is the router so basic seems like if you're going to have experts, it seems important to route to the right experts. So why don't you do that? I think, you know, there have been some ablations in some of the earlier Google papers on having like MLP routers and like more sophisticated things. I think the sort of complex answer here is that the systems concerns sort of weigh heavily. If you're using a lot of flops to make routing decisions, you know, you have to pay for those flops and so you have to get performance improvements in just the routing, you know, and I think the one other thing to appreciate here is that there are really big limits to how well you can route because the learning process for this routing thing is actually pretty dicey, right? Because how are you going to get gradients for which routers are good or bad? Well, the only thing you have is if you have top two, then you can compare the two things that you have and you can push the gradients into S of T because your G is a weight and then the S of T might inform your inner products. But that's a very indirect way to be

Fine-grained and shared experts

31:52 learning your affinity. So even if you make it complex, there's no guarantee that you're going to really learn the optimal router, right? Great. Okay. So I think the one of the great innovations of the DeepSeek and which was very quickly adopted by all the other sort of Chinese uhe releases is this idea of both a shared expert and a fine grained expert. And so the basic structure that was sort of originally proposed is to take your dense architecture and kind of copy the experts over right. So in this case, you know, you're going to have, let's say, if you have two, if you have top two routing, you know, you're going to have twice the activated parameters of your original dense model, right? So you take your and you copy it over and you activate K equals 2. So this is kind of what you might think of as like the vanilla or like the basic that you might start with. People realize fairly quickly that having lots of experts is good. And the logical sort of next step beyond having lots of experts is good is I want lots of experts but I don't want to pay the parameter cost for having lots of experts. And so DeepSsee basically argued that the right thing to do then was to cut the expert up into smaller pieces, right? So remember last lecture I was you know telling you about oh the kind of golden rule in some sense is to have you know your hidden layer and then you multiply that by four and that will give you kind of your projection layer right so now what you would do is you would instead of multiplying by let's say four you might multiply by two right so now you have smaller matrices you have more fine grained experts you can have twice as many of them right and you can

33:30 kind of take that logic much more to the extreme you can like you know quadruple or multiply by eight and you can keep decreasing the size of your sort of projection dimension there that's fine grained experts and there's you know drawbacks I'll I'll talk about later it's not it doesn't come for free so you have to be very careful about how you structure these things and then the other thing that you know has been sort of studied and noted is maybe it's helpful to have at least some MLP that can capture shared structure right like maybe there's just like processing that always needs to happen no matter which token you're processing in that case it seems like kind of a waste to do all this routing work and to have all these like, you know, parameters spread out everywhere when we can just have one shared or one or a few shared experts, you know, whose job it is to handle all of this like shared processing that's needed. And so they're shared experts. And so this setup of using fine grained experts plus shared experts originally came out in DeepSeek although I think the original inspiration came from DeepSpeed and Qwen and others. So almost all of the open releases since DeepSeek have adopted some sets of these innovations because it's it's quite clear that especially fine grained experts is just really useful. That's a kind of no-brainer at this point to do. One of the things I really like about reading DeepSeek papers is that they do ablations. You know, it's not like a whatever sales tech report. You know, they actually care about whether or not their methods work. And so they have this lovely ablation in the DeepSeek paper where

35:10 they show, you know, the blue bar over here, this is GShard. This is a very basic vanilla implementation of you know, you can have one shared expert, that's the orange bar, and that gives you a big boost on some tasks and no boosts on others. You can have fine grained experts, that's the green and orange, you know, bars, and you get further boosts from that. And if you compare the blue to the orange, you know, composing all these differences give you quite the big boost over others. And so we can see that, you know, more experts and shared experts generally seem to help. Okay. Yes. Question. Like when it says seven out of something, does that mean it's doing like top seven? Yes. Sorry, I should have I should have explained that. That's right. So X out of Y means X activated out of Y total routed experts. That's right. Yeah. And so you can kind of see the pattern here as well of as you increase the number of experts, you also often increase the number of activated experts. Especially if you're doing fine grained experts, it flops wise it's free, right? Because you know each expert is now smaller. Good. Okay. So has you know basically corroborating evidence that shows really nicely that these things work. So, the bottom one I think I'll start with because it's more decisive, shows, you know, fine- grained experts going from 8 to 32 to 64 fine grained experts mirroring in some sense the DeepSeek ablations. And you see very clear trends and losses and other kinds of metrics that you see improvements going from 8 to 32 to 64, right? Fine grain experts is great. Shared experts, which is, purple versus teal at the very top, you actually don't see really any gains, at

36:51 least in the OLMoE setup. So they actually end up going with no shared experts. Even though the DeepSeek paper seemed to show more gain. So that one actually is maybe more mixed given this sort of follow-up or this you know third party replication of these kinds of ideas. So at this point you might be wondering you know what are common configurations? I think I'm going to, you know, take the page out of, you know, last lectures playbook of looking at a lot of the recent releases, you know, looking at what people do and trying to talk a little bit about the patterns that have arisen. So some of the early Google papers, so GShard, Switch Transformer, ST-MoE, some of them had really large numbers of routed experts. And there was a lot of like really interesting stuff going on in those papers. I'd encourage you to read them. Some of them happened in LSTMs and other kinds of architectures. Regardless you know very quickly I think there was like kind of a period of like 8 to 16 experts like Mixtral DBRX Grok with two active experts. Those worked reasonably well but then kind of DeepSeek or DeepSeek you know v1 comes out. That has kind of the prototypical configuration I told you about fine grained experts 64 of them six actively routed two shared experts. And each sort of expert is sort of one/4 the size of a of a normally sized expert. Take that last column with a grain of salt because I had to sort of back them out from like config files and things like that. So, I'm not 100% sure about the exact ratios here. So we've then got essentially Qwen 1.5, DeepSeek V3, MiniMax. These are, you know, Chinese. They follow essentially

38:30 in the same footsteps as DeepSeek v1. The specific numbers are different but in the in the sense that they use you know fine grained experts and they often have shared experts they're very similar to kind of this original DeepSeek configuration OLMoE MiniMax and llama are very recent they definitely do all this like fine grained expert stuff and llama 4 also uses a shared expert and you kind of see sort of variations in configuration but you see what's basically shared which is this fine grained experts ID idea and especially for the big models like llama 4 and DeepSeek very large numbers of routed experts or sorry not routed like total experts. Yes. So can you explain what the ratios are the ratio is expend is representing roughly like how much each export is sliced relative to having just the standard dense configuration. So in terms of hyperparameters, you know that if you're following the rule of thumb, your hidden dimension and sort of your projection from in your MLP should be about 1 to four or one to 2.6 if you're doing a gated network, right? And so by looking at the hidden layers of these architectures, you can kind of see how many times they sliced up that that original feed forward size. So if like for those experts, does that mean that like still increasing their group like the factor. That's right. Yeah. So, so you know, you can think of this as roughly, you know, they have, you know, 16 normally sized experts. Oh, okay. And so they, you know, they're they're of course having more parameters than the dens equivalent. They have six routed so they have eight total active experts

40:09 at any time each that are quarter sized. And so you should think of them as like roughly double the flops, right, of a of a dense equivalent. So some arithmetic, but hopefully the math is clear and consistent hopefully. Yes. Like the ratios like one are like so for some of the exotic ratios I'm not quite sure why they're that way but they are very precisely whole numbers when you when you take the ratios between the FFNs and the implied hyperparameters and so I think those are exactly the split counts of like how much they were sliced but I'm not sure why they have one over 14 I mean like does it do you ever like project to like smaller dimension because like that ratio is so small in the MLP. So yeah. So yeah. Oh, that's why you're asking like do they do they down project? Yeah, that's right. In some of them they are actually smaller. I don't remember which models in particular, but in some of them I do remember they were actually down project. Yes. What is the intuition for wanting more than one shared expert? Yeah, I mean I it does kind of seem like there was a period where some of the Chinese LM companies tried many shared experts and then you know people have come back to zero or one. And if you look at the OLMoE ablations, it's not quite clear that even one shared expert is decisively useful. I think the original motivation was that then you have equally sized, you know, experts like these are both one quarter sized experts and now you have eight active experts total and so you can keep the sizes consistent. Otherwise, I don't really see a particular justification for why it should be two smaller one versus one larger

Training a router you cannot differentiate

41:47 one. Okay, cool. So then hopefully you know you get a sense of how the routing works for a lot of these and how it's all set up. The forward pass hopefully you fully understand. Now we need to think about training and training is pretty gnarly right and the major challenge I foreshadowed earlier right when we train we cannot turn on all the experts because if we do that then we pay the full flops cost of all the experts right having a model that's like I don't know 256 times more expensive to train is a total no-go right so we need train-time sparsity but sparse gating decisions are obviously not differentiable we now have a kind of annoying RLish problem. And so we could do any of these things like RL to optimize gating policies. We could do you know bandit inspired things of doing randomization to do exploration. Or you know we can just have some heruristics that try to balance things out right like put some loss terms in there and hope things work out. You know having gone through deep learning classes of many kinds you can kind of guess internally which one people use in practice. And I'll talk about each one of these three in turn. Okay, so RL I think is one of the earliest things that people tried. It's probably the most principle thing that you can do in this space, right? You have a you know non-differentiable routing decision. Well, think of that as a policy, throw RL at it and then solve the problem. Unfortunately it's not better than a lot of the other things that you can do. There is a paper by

43:24 Clark et al. In 2020 who were exploring various like scaling related questions in and they do have an RL baseline that you know I was able to dig up. But unfortunately it's not really that much better than say using hashing for decisions and they were you know they were really interested in benchmarking this thing on the left called S-BASE which is like a linear assignment kind of a method and that thing you know handily beats you know doing RL and I think In practice, the gradient variances and complexity means that it's pretty finicky to use and no one, you know, at scale has really used an RL based approach to optimize these gating decisions as far as I know. A thing that has been done much more at scale is stochastic approximations of various kinds. So what they might do is they might add a bit of you know perturbations. So here is an example of one from Shazeer in 2017. This is one of the early papers where they're still going to do kind of top k routing. So they're going to keep the top k elements of this h of x operation and they're going to softmax that to get the gate. But what we're going to do to get this you know h of x operation is kind of the following. So what we're going to do is we're going to have our original sort of linear you know affinity. This is identical to what we were doing before. We were basically just computing, you know, our inputs X and, you know, a sort of learned weight for each gate. And so this part's the same, but I'm actually now going to jitter it a little bit. I'm going to add a normal and then I'm going to pick sort of a W noise scale that's learned.

45:06 And this thing is going to control how much noise to inject into this process. And you can kind of think of this as a stochastic exploration policy. And by manipulating W noise in particular ways like sort of annealing it down or doing various things I can control the exploration exploitation trade-offs that this is going to have right and so this is going to give you one solution to the to the explore exploit dilemma and especially if you're noising things up each expert might randomly get you know some other tokens that it wasn't expecting to get. So it'll lead to experts that are less specialized but maybe a little bit more robust. And so that that seems generally quite nice. Of course the stochasticity also means that you don't get as much specialization and that leads to loss of efficiency. And you know there's another approach that people have done where they sort of multiply the router logits or sorry they add yeah have a multiplicative perturbation to the to the router logits with the goal of getting less brittle experts. But this sort of jitter process was kind of removed in some of the later papers because they found it just didn't work as well as some of the heuristic loss based approaches. And so this was an approach that was tried in a couple this kind of stochastic routing tricks were tried in a couple of the early Google papers. But I think that has generally been abandoned by a lot of the people training these. Okay. So yes, for the stochastic like what problem does that solve? Because we're still taking the top K. So we still can't differentiate backwards, right? Well, if you think of this, so

46:44 the question was we still can't differentiate because we're taking the top K. But if you kind of change the your interpretation of the problem a little bit, if you think about a bandit problem, right, it has the same structure as this where you know you pull a bandit arm and you don't see any of the other arms. So you can't, you know, really allocate your resources efficiently. If you pull some of the other MoEs at random, now you've got enough data to be able to do some optimization. And so this jittering is very similar to in spirit to this kind of like epsilon greedy style exploration thing where you're randomly pulling some of the other arms with some probability where the probability itself depends on how confident you are about this routing decision. So that's kind of the intuition and then of course you know that's going to give you some way of getting some signal back. Okay. So, the thing that in practice people have ended up with is you know we don't do any of that. We don't do you know RL we don't do stochastic exploration. But we rely on really another mechanism to sort of keep things reasonable. So if we're doing top two routing right technically speaking we do get some signal in the gradient descent process because we can compare the top two you know experts that we did evaluate. And so it's possible to do you know some optimization but when we do you know ignore if we if we drop all the other constraints the big issue that arises is you just end up sort of picking one expert all the time and that expert is good at everything and all the other experts are terrible right you end up in this local minimum where you've routed all of your tokens to one experts all the time so really the key game becomes then how do we get out of that

Load-balancing losses

48:22 local minimum and loss balancing or like balancing losses is really the key trick to get out of this. And this is this is kind of important to understand because this is the loss that mostly everyone actually uses to train the right. So if you were zoning out earlier, you know, you probably should make sure to pay attention to these this particular set of equations here. So this is originally from the switch transformer from Fedus et al. In 2022 and they add this particular loss where what they're going to do is they're going to you know loop all over each of the experts and they're going to take you know the you could think of this as an inner product between the vector f and the vector p. And so what are these vectors? Well f is for each of the experts this is the fraction of the tokens that were allocated to expert I. So you can think of this as kind of a probability vector that's telling me you know what fraction of my tokens in my batch or in my you know whatever the unit is here did I route to expert per I. Now P of I is the fraction of the router probability that was allocated to expert I. So the router probability is kind of the original sort of softmaxed routing decision that I was sort of intending to send. Right. So this is kind of measuring P of I is what was sort of the intended probability from the router and then F of I was what was the actual sort of like you know what was the actual routing decision made by the top K method. And one thing that's kind of interesting to look at here is let's say we take the derivative of that loss with respect to P of I. So you know this is a

50:02 linear function with respect to P of I and you'll see that the strongest downweighting action happens on the sort of biggest experts with the biggest allocations. Right? So the it's actually in fact proportional to the amount of tokens that you get. So you're going to be pushed downwards sort of more strongly if you got more tokens. And so this is kind of the basic behavior of this loss. And you know almost everybody uses this kind of F.P kind of a of a trick to try to balance tokens across different units. So the basic unit that you might want to balance over initially is batches. You might want each batch to get allocated evenly to experts but you might actually have other kinds of balancing that you might want to do. And DeepSeek does exactly this kind of thing. I'll talk about all the variants that they've thrown in but you know the first thing is per expert balancing per batch. So each batch they want to make sure experts get an even number of tokens and you know this is from the DeepSeek paper and hopefully this looks you know very familiar to you. This is exactly the same you know F.P inner product structure as you saw before you know P of I is defined a little bit differently that's S of I of T you know but that should be familiar from earlier as well that's the softmax pre-top K right so hopefully this looks all pretty good to you the other thing you might want though is you know you might want to balance across experts that's all well and good but you might also want to think about the systems concerns right because you're going to shard your experts onto different devices and you might want to balance per device right and so you might have another loss that's essentially the same structure,

51:39 but instead of summing, you know, which tokens go to which experts, you might measure which tokens go to which devices, right? And that's going to be a different f that's measured over the device groups rather than over each expert. And so now you can set up a different loss to balance over devices. You optimize this. You're naturally going to try to learn routing functions that make sure each GPU or each TPU, what have you, have an even number of tokens leading to even utilization, right? And that would be great from a systems perspective. So, basically, everyone does, you know, kind of this kind of a thing. And so, DeepSeek V3, actually kind of innovates a little bit. This is this is kind of cool, and I don't think I've seen this before. It's one of the first things in the MoE world that doesn't actually come from Google really. Which is that they have gotten rid of this per expert balancing term. They've gotten rid of this entirely and instead what they now do is they basically take their softmax scores and they add a little fudge factor B of I where B of I is a little fudge factor score for each expert. Right? So expert I you know might get upweighted or downweighted. So, if an expert isn't getting enough tokens, you know, it's going to be given a higher B of and then that's going to allow it to grab more tokens. And the way that this works is the way that this works is that they're going to learn BOFi through a really simple online gradient scheme, online learning. And so they're going to measure at each batch, you know, what are each of the experts getting, like are they getting an even number of tokens? And if they're not

53:19 getting enough tokens, they add sort of gamma, some learning rate to B of I, sort of making it higher. If they're if they're getting too many tokens, they're going to subtract gamma, making that expert slightly less attractive, right? So they're just learning little, you know, offsets for each of the S of I. And notice here, you know, you're only using the B of I to make the routing decisions. You're not actually sending it over as part of your gating weights, right? That's a that's a sort of somewhat important thing to do. So they call this auxiliary loss free balancing. If you go and read the DeepSeek V3 paper, which all of you should because it's a it's a really nice paper, they'll make a big deal about how this makes training so stable, so great. So wonderful. And then of course you like keep reading the section and they're like actually but we decided that you know for each sequence maybe we still want to be balanced and this doesn't work well enough so we've added the you know the heuristic loss back. So they do have something called a complimentary sequence-wise auxiliary loss that you know is basically exactly the auxiliary loss that they decided they needed because what they wanted to do was to balance load balance the experts at a per sequence level rather than a per batch level. I'm not sure why they do this particular thing rather than any other sort of you know B of style trick but that's just kind of what they do in DeepSeek v3. So it's not fully auxiliary lossfree as they'd like you to believe. Okay. Oh yes. Question. This is a bit of an unfair question, but if we did not have to worry about systems optimizations, do you think the performance of this model would be a lot better or would it stay roughly the same? If we did not think about systems optimization, would the performance of this model be better or stay the same? When you say this model, what do you

54:59 mean? DeepSeek V3 or like just in general like this model never. So are you saying like if we ignore the systems concerns do we think are still good? Is that kind of one way of asking that question? Like would the performance on downstream pass for example be better than what we have right now? Yeah. So I think I didn't have to balance this like I must set roughly equal number of tokens for every expert or Yeah. Yeah. That's right. That's right. Well, I think actually per expert balancing this term, right? This is not a systems concern. So, you still want to do this because if you don't do this, what you'll find and actually there's, you know, I'm going to keep referring to the OLMoE paper because they have so many ablations. They have a really nice ablation where they get rid of exactly this. And what they find is basically early on in training the model just picks like one or two experts and all the other experts are dead. Like the router never sends anything to them. So, you're just wasting memory at that point, right? So now you've just lost performance for free. You've effectively gotten a smaller model. And so even if you ignore all the other like device balancing parallelism concerns, you've just gotten a worse model because you didn't properly allocate your experts, right? It's the same way as like you want to use all your parameters, right? You would like to effectively use your parameters. You want to do expert balancing. Sorry, say device. What does device refer to? Yeah, actually so normally this would refer to like GPU or TPU. There is a subtlety. I'll talk about this maybe in the very last or second to last slide. There are more sophisticated and cool versions of this where you try to balance things to minimize communication costs as well. And so there's you know broader notions

56:37 of device like you know one rack or whatever else but here it usually refers to like GPU. Yes. Going back to the fact that like hashing as a routing algorithm seems to improve performance like is there intuition for that because that's effectively just like randomly choosing a like one of the few forward members to send it through right so like why does having multiple copies of that I guess each of which get less data why does that make performance better yes the question was why does hashing do anything at all I don't have the you really precise intuition for this. But you can make arguments either two ways. One is, you know, even if you're hashing, the same tokens are going to go to the same, you know, or the same kinds of, you know, sequences are going to go to the same expert every time, right? And so each expert will still get some deterministic subset of the inputs. And so there's some specialization that can still occur. It's just non-semantic or, you know, non-learned. And if you're a distribution Zipfian, like the word 'the', might dominate one expert, you know, and so you might still get actually semantic specialization where like one expert is effectively dominated by like very frequent things, but like a random routing function probably wouldn't be a day like a pure random thing that's not dependent on input. Yeah, I would bet that that would be really terrible. Yes, I have never run or seen that, but yes, I think that would be that would be horrible. Good. Yes. Yeah. So for like during LM like you have many layers right many transform I think in the lecture you mentioned that each expert okay so like you do like ad like 32 layers like 64

58:16 experts that's like a lot of GPUs or I wonder if like experts are bundled together on like a single GPU is that the question was like won't you need lots of GPUs if you have lots of layers and lots of experts. Yeah if you if you exclusively give a GPU to a single expert. Yes, that would be that would be kind of crazy. But you would kind of shard things so that each GP would hold, you know, enough of these units to, you know, effectively use memory, right? The name of the game in parallelism is you always want to use up all of your memory because that's one of your resources, right? You don't want to parallelize more than you have to. Cool. Okay. Excellent. Oh, okay. I did put the ablation in here. Yeah. So, this is exactly what happens to the question of what happens if you don't do, you know, expert balancing loss. I think the great picture to see is this bottom left one. If you don't do load balancing, you know, what are the tokens assigned to which expert? You see the pink and the yellow expert. They just like kind of take over. They take up, you know, about 50% of the tokens. All the other experts are dead. They do nothing, right? And so you've wasted, you know, the majority of your experts at this point. You know, six out of eight of your experts. And you've created a two experte unintentionally. And you know that gives you know worse losses up seen up on the top right the teal lines. Of course maybe that's still better than the dense model because at least you've got two experts going. But you could have done better right counterfactually speaking. Okay. So I won't go quite as deep as I could into the system side because I haven't really started to cover you know the core systems concepts necessary for you to deeply appreciate a

The systems side: all-to-all, block-sparse kernels, token dropping

59:54 lot of the parallelism concerns like you know basically the hierarchy of communication speeds in a data center and so on. But really as I said before you know one thing to keep in mind is just how nicely can fit into devices. You know the thing that people say is expert parallel you know that involves sending or putting one or a few you know experts onto each device and what happens when you know are basically processing a token. Well, you would hit the router and after the router, you now have picked few experts. And so now you would have a collective communication call like all to all communication dispatch that would send the tokens to the relevant devices. You know, the feed fors would compute you know their outputs and then you would return the tokens to sort of where they belong or you would you know combine I guess multiple experts and so you would need another sort of collective communication call. And so if your feed forward computations are sort of big and beefy enough, you can kind of pay for the cost of basically doing this expert parallelism. And the one of the thing that's nice about this is that it's another form of parallelism in your in your toolkit. So you've got on the right side, you know, you know, data parallelism, model parallelism of, you know, two or three different kinds and then you've got expert parallelism and you can combine all of them to come up with sort of ways of trading off all the resources you have. So the communication speed, the amount of data that you have, your batch size, and your, number of experts and your memory. So, I'm not going to go into too much detail about how specifically this is going to help, but keep in mind that this gives you

61:31 another sort of tool in your expert toolkit. Another thing that is also you know useful is let's say you have multiple experts on a single device you know you might hope that because the computations are sparse like let's say you know token one this first token you know gets multiplied to export zero the second one is expert one and this third one's expert two so this is really three matrix multiplies that are small and sparse and you might hope that modern GPUs can sort of take advantage of these kinds of you know complex these kinds of sparse matrix multiplications. And that's exactly right. So if you know lay out your sort of experts correctly and the weights are sort of fused in the right way then modern sort of sparse matrix multiply sort of engines can sort of effectively make sure that you're not wasting any flops in doing this one big matrix multiply. So, so modern libraries like MegaBlocks can basically take advantage of this you know device level sort of sparsity support to do multiple expert computations sort of all at once. So this is yet another advantage that you get with. So one fun side thing which maybe isn't mysterious to you all anymore because you've sort of grown up in the era of GPT-4. But when the GPT-4 API first came out, it was kind of mysterious to me because when you set the temperature to zero, you know, you kind of got different responses even though it was supposed to be deterministic. And lots of people speculated about why would that be. I'm not saying this is the answer to

63:10 that reason. But there is actually a interesting source of randomness in right. So in MoEs, think about you know what happens. You're going to route your tokens to experts, right? And experts live in different devices. It could be that you know you have a lot of examples. You're going to batch of course batch your queries when you're processing them. And so if you've batched your queries, these tokens are going to get routed into different experts. So imagine you've got, you know, this batch to process and you've got a bunch of experts, but for whatever reason, this batch really loves expert number three. Like all the tokens go to expert number three. So now what happens? Well, the device for expert number three doesn't have enough memory to load all of those tokens. And then what happens is what people call token dropping. And this happens at training time as well. You often have what's called a load factor where you're sort of controlling the maximum number of allowed tokens. And if the router just allocates too many tokens to an expert, you just drop those tokens off either for systems reasons or because you're just worried that that expert is going to take over at least in the training time. And so now this token has gotten dropped and it's not going to get anything at all. Like the MLP is just going to do a zero computation and the residual connection is just going to pass things straight forward. And then you're going to return an output. And so if your token got dropped, you're going to get a different result than if your token didn't get dropped. And so based on, you know, who else is in your batch, can induce stochasticity both at training time and inference time, which is like kind of an interesting thing that you don't normally think about because you almost never think about like cross batch effects when doing inference. Okay, so that's kind of the

Stability and fine-tuning

64:48 the main bits of you know the main basic components of building a MoE. One fun side thing. If you were to actually go out tomorrow and trying to train a MoE I think the system side will make you a little bit sad, but the other thing that would make you sad is probably the stability side of things. So kind of have this property that sometimes they'll just kind of blow up on you if you try to fine-tune them. They're very difficult to fine-tune and they'll sometimes blow up on you. And so you know Barret Zoph and others really studied they had a whole paper on basically trying to make more stable and there's a paper which is the one I'm I'm referencing here whose entire purpose is to stabilize training and there's a couple tricks that I'll I'll mention that I think are relevant and that people do. The first one is you know if you're doing the router softmax so this goes back to last lecture about stability right like what did I say about stability well the thing to be afraid of is the softmaxes right the softmax is always where you want to be afraid and so for the they do all the computations in float 32 for the router computations just to be safe right and sometimes they also add the you know an auxiliary z-loss so hopefully you remember that it was just last lecture you know you do log of the sum of the exponentiated you know values in the softmax and you square that and you add that as an extra loss right so this is going to keep the normalizer values near one which is nice for stability so this is actually one of the places where z-loss was used earlier before it got sort of more popular for training models you can kind of see the effects here if you look at the losses I think the center the second

66:30 plot here is maybe a great one you know if you remove the z-loss from your router routing function you see these like giant loss spikes in your validation loss where you know the model just kind of goes a little bit crazy for a couple iterations and then gets kind of pulled back. Of course, it like still trains okay, but you are better off having the Z loss than not having a Z-loss. There is a pretty noticeable gap in the validation loss by the end here, right? Other things that can happen, people, you know, of course you want to fine-tune your like also RLHF your if you're going to, you know, ship and release it. But this turns out to be, kind of problematic. Some of the earlier work, you know, when people were starting to do, this was back in kind of the BERT and T5 era. So there was a lot of fine-tuning going on. And you know, one of the things that people saw was, you know, actually there's a lot of overfitting that happens if you were kind of doing sparse models. You see this big gap between train and val, right? This blue and orange line. Whereas the dense model, this green and red line, has a smaller train test gap. And so there was a lot of worries about overfitting because you have these like gigantic parameter models that you're fine-tuning on small data. One of the solutions that was proposed at the time, I don't think this is very popular, as far as I understand, is to architect yours such that not every layer is layer, but you like let's say alternate dense layers and layers. Then you can just fine-tune the dense layers and then that will that will still be fine, right? That behaves just like a dense model. So that was fine. Another solution the one that we saw in the DeepSeekMoE paper is just kind of use a lot of data like if

68:10 overfitting is a problem you know we have access to lots and lots of SFT data just shovel all of those guys in so in the case of DeepSeek they use 1.4 million training examples. Then maybe you're you're not quite as worried about these overfitting concerns. The last thing I'll end with which is a is a trick in the toolkit that I that people have done and seen is upcycling. And so this idea is to take a dense model like the one over here and then you take your MLP and you make a bunch of

Upcycling a dense model

68:43 copies of it. And then you maybe perturb it and then you have your router that's initialized from scratch and then you just pretend this is and then you train it from that point on right you just initialize the MoE from a dense model and this is a trick that's kind of called upcycling and you know people have shown that if you can get it to work it is a very cost-effective way of getting right and the MoE is great for inference because not every MLP is going to be active or at inference time right So, so you're going to you might effectively get a much larger parameter model without doing the training of a much larger parameter model. And several people have succeeded at this. MiniCPM, which I'll mention again in the scaling wall lecture, but this is a Chinese open LLM that basically tried to build really good small language models and they succeeded at taking a dense model and upcycling it into that you can see that their numbers get significantly better in the last two rows, right? So the dense models to the MoE, they get a pretty non-trivial bump in performance. Qwen, I mentioned at the start of this lecture, one of their earliest attempts at a MoE was taking one of their dense models and then building upcycled. And they got, you know, fairly significant performance gains relative to sort of smaller models at the time. Like they got models on par with their 7B models with a 2.7 billion parameter active model. So to wrap up I want to sort of walk through the DeepSeek architecture at the very end here. And hopefully

DeepSeek MoE v1 to V3, plus MLA and MTP

70:20 this will give you a sense of you know the first thing I want to do is I want you to understand the DeepSeek V3 architecture setup and all the changes that they did. Because that's an example of a modern high performance open source system. I also want you to maybe appreciate that architectures don't change that much. DeepSeek v1 or you know DeepSeek v1 is you know it's not that new it's like maybe a year and a half or something maybe two years old and they basically nailed the architecture at that point right so I want to see I want you to see what they changed from that very ear the earliest attempt to their big training run so this is the very first starting point this is DeepSeek I'm calling it v1 but actually you know probably the right way to refer to it is DeepSeek it's a 16 billion parameter model with 2.8 of those parameters active. And you've seen already this diagram over here. This is the shared two shared plus 64 fine grained experts of which four of them are active at a time or maybe about six of them are active at a time. Sorry. And the routing you know you've already seen this I presented this in the in the middle of the lecture here. This is the very standard top K routing where the softmax is at the bottom before the top K selection. And for balancing right at training time all they do is to add this auxiliary loss balancing term right both the expert and device level balancing terms right so hopefully you know you remember those from earlier so that's the v1 and then they saw how sort of effective their model was so I guess to add some more context right DeepSeek originally had a dense model and then they had a MoE model and

71:58 the MoE model was remarkably good and so when they went to v2 they went straight to the MoE, and now this is a 236 billion parameter model of which 21 of those billion parameters are active right so you need a lot of memory but your flops consumption for inferencing this model is not so bad now the architecture is identical I copied literally the same figure because the architecture is literally the same minus changes to the number of you know experts that are active and we've got now sort of some new things happening but not too many new things. So the top selector is the same. So the equation from before, this previous equation, this is identical. This this is still how they do things. But they have this very clever trick that they add on. And this is you know, I was going to say at the very beginning, you know, what's the drawback of having fine grained experts? Why can't I have, I don't know, 1024 fine- grained experts or 2048 fine-grained experts? Well, the problem is when you shard your experts very finely and you have a lot of active experts, right, you're going to have to route to those experts, right? So, your communication costs potentially grow and if you're very fragmented, you might have to send a lot of tokens to a lot of devices, right? And so, the clever thing they come up with is to say, I'm not just going to, you know, at for each batch route to the top K experts naively, which might force me to send my tokens to lots of devices. What I'm going to do is I'm going to first pick top M devices, right? So I'm going to do my, you know, normal scoring calculation, but I'm first going to sort of subset the set of allowed devices to top M, right? And once I've picked my

73:37 devices, then I'm going to pick top K for each token within each device, right? So so now I've restricted the devices. This really controls the communication cost. And now this gives you more efficient training when you're scaling up to these gigantic sizes, right? You need to start really engaging with the systems aspect of things when you're training a 236 billion parameter model. The other thing which reflects the systems concerns that are necessary at this scale is that they add a communication balancing loss. One way of thinking about things is, you know, for an expert, there's kind of inputs and outputs, right? The inputs are, you know, the token comes in and you route to your expert. And the outputs are, you know, you have to kind of bring the tokens back where they belong, right? So, if a batch belongs on this device, it has to go back where the original device was. So, we have to think about both the input communication cost and the output communication cost. And so, they add a balancing loss to try to balance out the output communication cost as well, not just the sort of input side. So that's a minor note but you can kind of see their attention to detail on trying to make sure all the different sort of systems aspects are properly taken care of. And then finally we kind of get to the you know the big DeepSeek v3 sorry that should say v3 not v2 up there 671 billion parameters of which 37 are active. You know once again you know exactly the same figure because the MoE architecture itself doesn't change. That's stayed the same since DeepSeekMoE, right? Like if it works, don't change it. They do change a couple things. Maybe they were, you know, hearing you all say, why don't you normalize to one? And so, you know, they've normalized the gate to one. They've moved kind of the softmax normalizer operation up there. But they're not actually exponentiating

75:14 sort of the sort of gating decisions. They're actually taking sigmoids which is a sort of softer sort of more nicely behaved operation you know, than the softmax. So they've got some changes here, but conceptually this is still the same as the top K routing decision, right? You hopefully see very similar things happening. And then in terms of the losses, they've gone to this auxiliary loss-free trick of this B of I being incremented or decremented based on the expert load. And then they have a sequence-wise auxiliary loss. And just to, you know, add some context, why would you want to balance different experts on a single sequence? Well, the thing that they're very concerned about is at training time, you know, it's fine to not have a sequence-wise balancing loss, but at inference time, it might be the case that someone sends you very out of distribution sequences, and that might overwhelm certain experts, right? So, at inference time, you can't control which sequences you get. So, you might want sort of stronger balancing that operates at a single sequence level rather than overall batch level. Okay. And then in the Oh, sorry. Yes. Does 3 still do like the top M devices like does it keep the B2 improvement? Yeah, they keep the top M improvement. They do not keep for example the communication loss. So so they've they've jettisoned some things but top M is a I mean it seems like a pretty clever idea they keep it. Yeah. Yeah. But it's not like they always add things. They have removed some of the things. And so in the last two or so minutes, of the class, I'm going to go over the none

76:52 parts of DeepSeek V3 because I think, you know, we're already at the point where I've explained most of DeepSeek V3. I might as well go through the steps of explaining the rest of DeepSeek V3 at this point. So, you all know kind of how that works. So, they have a clever sort of optimization for the attention piece called MLA or multi head latent attention. And you all actually already know all the ingredients that you need to understand this because at the end of last lecture you know I talked about like GQA and MHA right so those are all inference optimizations that you need in order to optimize the size of the KV cache. So the deepse folks take a different tac or different approach at optimizing this instead of reducing the number of heads they're actually going to sort of project the heads into a lower dimensional space. So you have your inputs H of T and instead of sort of generating the K's and V's directly from these H of T's, what I'm going to do is I'm going to first generate a lowdimensional C. This you can think of this as like a, you know, compressed version of H. And this C is going to be smaller and easier to cache. And I'm just going to cache these C's. And whenever I need, you know, these K's and V's, well, I can sort of up project from this KV sort of conceptually speaking. And then, you know, I can take the inner products with the Q's, right? So you can kind of see how this would be a KV cache savings if I only have to save the C instead of the higher dimensional H of T. And that's exactly the idea. So you take your H of T, you project it into a lower dimensional C and then you up project this back into the K's and V's, right? And if the C's are small, well that's you've compressed the KV cache. That's good. And then you know in terms of the computation, right? If

78:31 you're thinking about flops, well, you might think, well, this is not good because I have to multiply an extra matrix W U K, right? I didn't have this matrix before. That's an extra matrix multiply that I have to pay for. But kind of the clever thing here is remember that on the other side of K, right? I'm going to take K. Q, right, there, that Q.K is going to be an inner product in the attention operation, right? And Q itself has a projection matrix Q. And so the trick here is you can merge this W U K and this Q matrix together into one matrix. So I haven't gotten any extra matrix multiplies. I've just merged this new matrix multiply into my other one. Right? This is, you know, just associativity. I can just merge the two. They also compress the queries for memory savings during training, but really that one is not quite as necessary because it doesn't interact at all with the KV cache. I'm only going to mention this last one in passing. Because it is a subtlety, but it's kind of a clever subtlety that you realize, which is that this original trick, this sort of thing that I just described at the top is not compatible with RoPE, right? And the reason is because, you know, the RoPE matrices, you know, basically you have the Q's and the K's and you rotate each of those Q's and the K's by multiplying with a rotation matrix RQ and RK. But if you do that, then these RQs and RKs are in between the query projection and this up latent vector up projection matrix. And since I can't reorder these matrix multiplies, you know, RoPE kind of gets in the way. And they still have a solution of basically doing RoPE on non-compressed dimensions. That's kind of a side point. I think it's not quite as important. You can kind of look

80:08 at the paper if you're super interested. The other thing that they do, and this is the last thing I promise, is that they have a minor change in their loss function called MTP where they predict multiple tokens in parallel. And so what they can do is normally, right, you have your inputs, you shift them to the to the left by one. So you're predicting one token in the future and then your transformer is going to predict all those tokens, right? That's your normal transformer loss. But then what you can do is right before you make those predictions you can take you know the hidden state you can pass it to a very lightweight one layer transformer and that model can predict you know one token in the future right so now the model is not just predicting the next token it's predicting the two tokens into the future right so that hopefully all makes sense and this is just a small lightweight model that can do that you can sort of see the architecture right here the one thing that is kind of disappointing that I learned as I was sort researching for this lecture is actually they only do MTP with one token ahead. So even though they have this very complicated diagram of how they could do it for many tokens, turns out it's only done for one token. Okay, so now I'm all done. MoEs are now at the core of how you would build and deploy, you know, a really high performance large scale system. And they take advantage of kind of the sparsity idea that you don't need all of the parameters all the time. And discrete routing is the real big challenge. And this is I think one of the big reasons why MoEs didn't immediately catch on. It's very scary to have to try to optimize this top K routing decisions. But heuristics somehow seem to work, right? Like they just do. And so there's a lot of empirical

81:45 evidence now that MoEs, at least for flop-constrained settings, are just a good idea. It's cost effective. You should do it. So definitely worth learning. Thanks a lot for listening.