Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
deep learning machine learning arxiv explained neural networks ai artificial intelligence paper attention transformer attention mechanism google google brain shazeer trillion trillion parameter language model gpt3 gpt-3 gpt 3 t5 sharding mesh mtf mesh tensorflow query key value feed forward experts routing mixture of experts sparse sparse experts data parallelism model parallelism expert parallelism trillion parameters perplexity scaling flops bfloat16
#ai #technology #switchtransformer

Scale is the next frontier for AI. Google Brain uses sparsity and hard routing to massively increase a model's parameters, while keeping the FLOPs per forward pass constant. The Switch Transformer compares favorably to its dense counterparts in terms of speed and sample efficiency and breaks the next magic number: One Trillion Parameters.

OUTLINE:
0:00 - Intro & Overview
4:30 - Performance Gains from Scale
8:30 - Switch Transformer Architecture
17:00 - Model-, Data- and Expert-Parallelism
25:30 - Experimental Results
29:00 - Stabilizing Training
32:20 - Distillation into Dense Models
33:30 - Final Comments

Paper: https://arxiv.org/abs/2101.03961
Codebase T5: https://github.com/google-research/text-to-text-transfer-transformer

Abstract:
In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) defies this and instead selects different parameters for each incoming example. The result is a sparsely-activated model -- with outrageous numbers of parameters -- but a constant computational cost. However, despite several notable successes of MoE, widespread adoption has been hindered by complexity, communication costs and training instability -- we address these with the Switch Transformer. We simplify the MoE routing algorithm and design intuitive improved models with reduced communication and computational costs. Our proposed training techniques help wrangle the instabilities and we show large sparse models may be trained, for the first time, with lower precision (bfloat16) formats. We design models based off T5-Base and T5-Large to obtain up to 7x increases in pre-training speed with the same computational resources. These improvements extend into multilingual settings where we measure gains over the mT5-Base version across all 101 languages. Finally, we advance the current scale of language models by pre-training up to trillion parameter models on the "Colossal Clean Crawled Corpus" and achieve a 4x speedup over the T5-XXL model.

Authors: William Fedus, Barret Zoph, Noam Shazeer

Links:
TabNine Code Completion (Referral): http://bit.ly/tabnine-yannick
YouTube: https://www.youtube.com/c/yannickilcher
Twitter: https://twitter.com/ykilcher
Discord: https://discord.gg/4H8xxDF
BitChute: https://www.bitchute.com/channel/yannic-kilcher
Minds: https://www.minds.com/ykilcher
Parler: https://parler.com/profile/YannicKilcher
LinkedIn: https://www.linkedin.com/in/yannic-kilcher-488534136/
BiliBili: https://space.bilibili.com/1824646584

If you want to support me, the best thing to do is to share out the content :)

If you want to support me financially (completely optional and voluntary, but a lot of people have asked for this):
SubscribeStar: https://www.subscribestar.com/yannickilcher
Patreon: https://www.patreon.com/yannickilcher
Bitcoin (BTC): bc1q49lsw3q325tr58ygf8sudx2dqfguclvngvy2cq
Ethereum (ETH): 0x7ad3513E3B8f66799f507Aa7874b1B0eBC7F85e2
Litecoin (LTC): LQW2TRyKYetVC8WjFkhpPhtpbDM4Vw7r9m
Monero (XMR): 4ACL8AGrEo5hAir8A9CeVrW8pEauWvnp1WnSDZxW7tziCDLhZAGsgzhRQABDnFy8yuM9fWJDviJPHKRjV4FWt19CJZN9D4n
 Hi there, today we'll talk about switch transformers scaling to trillion parameter models with simple and efficient sparsity by William Fetus, Barrett Zoff and Noam Shazir of Google Brain. So as you can see right off the title, we're going towards trillions of parameters. GPT-3 had 175 billion parameters. This paper claims to have a model with a trillion parameters. Now is it really five times bigger or 10 times bigger than GPT-3? That's a debatable question, because the trillion parameters are not used in the same way as in a classic transformers. They are used actually in a sparse way. That's why the word sparsity is in here. And the way they are used in sparse manner is this new architecture called the switch transformer. It's not entirely new. It's built on mixture of experts. In this paper, that's also called MOE. That has been around for a while and we're going to see what that is. Now on a high level, switch transformers takes mixture of experts to an extreme in that it is a transformer. And the feed forward layer is divided up into these experts. And the switch transformer only routes each token to one expert only. That's the sparse part. So the mixture of experts previously, they always claimed you need at least two experts in order to get a stable training signal, the switch transformer manages to get it down to a single expert. So it's like a hard routing of information to just a single endpoint per layer of each token. So in that that means you can now scale the experts and you can scale the number of parameters in the model without making the model compute more. That's a very special notion. So you can up the parameters of the model. But if a forward pass of a data point will still have the same amount of flops that it needs to forward propagate through the network. Very special architecture right here. So yeah, that's why I'm saying trillion parameters not necessarily comparable to the 175 billion parameters of something like GPT three. So how do they do it? Because previously it was claimed it was unstable. They have new ways of making the training stable, such as selective dropout, selective casting of parameters to different precisions, and a better initialization. So that's the high level overview of the paper. And we'll dive into it, we'll explore kind of what mixture of experts is and how the model works. And what turns out it's a very long paper, as you can see, when papers have a table of content. That's a lot of fun. But it's a lot of engineering as well. And we're mostly interested in the model here, what it can do, and what does it how does it sort of fit in to the big world of transformers and language models and so on. Last thing I want to say, trillion parameters is, you know, it's a catchy title that most of the paper, they don't work with trillion parameter models, they work with models in the in the order of billions of parameters. And at the end, they build a model with a trillion parameters. It doesn't do as well as their models with as their smaller models. They also, it feels like they don't put that much work into it, because it's probably also quite fuzzy and expensive. But just know, we're not going to have trillion parameter models around anytime soon. Just yet. Interesting fact, the original resonant paper also built a 1000 layer convolutional neural network. Even though the resonance we have today, you know, they are maybe 50 or 150 layers deep, they did build a 1000 layer model. So maybe compare it a bit to that one. It's just like we can do it, not necessarily we need to. So here you can see something they discover. The curve on the left is very, very known to people that are in the language model game, let's say, or in the in the let's scale up AI game. And that is as you increase the size of the model, the loss will go down. And that's loss, as I understand it. So that's test loss. I believe that is perplexity. So scaling properties, exactly that that might be perplexity or test loss on some downstream task in any way, as you scale up the model parameters, the model gets better and better and better. The interesting thing right here is twofold. First of all, I believe they do hold the data set constant. So the data set is always the same, the amount of compute you put into it, the amount of either number of steps or time is also always the same. And in this specific case, the amount of flops per forward pass is also the same. The only thing that changes is the number of parameters. Again, it's very special to have a model where you can scale up the number of parameters, yet the flops required to forward propagate stay the same. So you can see here that there is a almost unhalted decrease here, it flattens out a little bit towards the bottom, though that is not necessarily does not necessarily mean it will ever flatten out before it's you know, at zero. I will approach zero, I guess. So and you can you can see that, you know, they scale up the model quite a bit. And also, their main comparison here is the T five base. So that's the text to text transfer transformer. By the way, if you don't know what a transformer is, or what a language model is, it's best you go back to my earlier videos and look up like the GPT three paper or the attention is all you need paper, I've made videos about lots of these things, I assume that you know them. You can see right here that if you compare to number of training steps, for example, the this switch models, all of them, no matter how big they are, they provide massive gains over like something like a T five. And they also do this in time. So this paper is very much about trade offs, you do require more storage for your weights. So you have to have more memory more RAM. However, that memory can be distributed, it can be sharded, because they use this mesh TensorFlow library to implement the switch transformers. And because their model has this sparsity, they can efficiently shard the model. So you trade off more memory, which can be sharded. But what you gain is training speed, and both in terms of time and number of training steps required. So you are much more efficient. Note that this only all of this holds in this super large regime, right? We this is, they say they've also discovered the speed ups in smaller models. But you know, as far as the paper is concerned, we are talking about millions, hundreds of millions of parameters, billions of parameters, even to trillion of parameters, together with these giant corporate corpora of, of text. So that's sort of the regime we are in. And the results do not necessarily transfer down to the lower scale problems that you know, you might face with your lonely one, collab in the corner. All right, so in a transformer, you have a transformer is nothing else but a bunch of these layers right here. This is this is in itself a transformer layer in its basic form. And it consists of sort of two parts, it consists of this self attention, right here. Now, that's the standard transformer self attention. That's what was introduced in attention is all you need. And what's been used ever since in all the transformers. This one right here is a is an, as I understand it, a language model. So you know, this this is very standard. However, after the self attention, you have this feed forward layer. Now usually, what you do is you have an input sequence, and you transform that through multi head attention into another sequence right here. Okay. And then what you do is you take each of these things and feed them through a feed forward layer. And if I am, as I understand it, this feed forward layer is simply, you know, a regular feed forward layer that you would find in a neural network, and you pass them, you pass these things individually. So this here, it's a vector, you pass it through here, and boom, that becomes the next layer representation, this thing right here, you pass it through as well. Boom, that becomes this one, and so on, right? You pass them individually to get the next layer representation. So this, this part right here, the attention part, it sort of aggregates information and relates the individual items of the sequence to each other, and transforms them into, you know, a new sequence, where sort of all the every token can gather information from every other token. That's what the attention mechanism does. That's step one. In step two, every token is isolated, every token is for itself. And the feed forward layer simply determines, you know, what's given one token given token number one, what is, you know, given its representation in this layer, what is the best representation for the next layer? Okay. So that's token number one of the next layer. So the multi head attention is kind of relating tokens to each other, and the feed forward layers, they are relating layers to each other. Okay, so up here, you would have the next multi head attention layer. So you can see the feed forward layer as sort of translating from one layer to the next layer, right, getting saying, oh, you come from this layer, I'm going to translate you such that the next layer understands you. And that happens on a token by token basis. Now you can see this is it's always the same feed forward layer for all the tokens, right, the tokens are sort of treated like a batch of samples. The idea of this switch transformer and also of the earlier mixture of experts transformer is that it might not be a good idea to have only a single one, right? This is the only feed forward layer, it's the same for all the tokens, it might actually be a good idea to have a couple of them that sort of specialize in different things. So what could that be? You know, in a in a basic world, this could just be like one for nouns. And this could be a feed forward layer for verb verbs, tokens that are verbs, tokens that are adjectives, and sort of maybe here is like punctuation tokens, right? You might think, well, if you are a noun token, the next layer might want to look differently at you than if you are a punctuation token, right? So this translation from one layer to the next layer can now happen dependent on what the token represents, right? Now we we of course, first of all, we don't have these annotations. And second, it's not necessarily that you know, we want to always divide it by noun, verb, adjective punctuation. Ideally, we want to learn this routing. So we simply want to say, look, instead of just one feed forward layer, we give the model four feed forward layer, feed forward layer one, two, three, and four. And for each token, the model can decide to which of these feed forward layer it sends the token to. So here you can see this is a token. Now, you know, we are dealing with word pieces. Let's just say the word more, I was like, I was thoroughly confused by when I saw this like, huh, why does it say more parameters, but here, it's the string more, right, and the string parameters. And these are in the vocabulary, and they get an embedding vector associated with them. So that's what's going on here. Then they go through self attention, as you can see here, both go through self attention, and then each one of them is routed to one of these four experts. Now the, the one here, the one on the left and the one on the right, these are the same experts, right, they're just duplicated visually here. But these would be the same weight matrices in there. So you have four feet forward layers in this layer. And each token can be routed to any one of them. And this routing here, this is learned. So in here, you have a matrix, they call it like WR. And using WR, you simply do an inner product of WR with your input right here, let's call that H with your input H. I guess they use H for a different thing. I think they call this X again. So you do this with X. And then you get, you get H, which is your routing, and then you simply build a histogram, you normalize the histogram, I think with a softmax. And that those are your routing weights. So it's very much like another attention mechanism, except that the queries, this thing here, these are like the queries, these are sort of the queries of this attention mechanism. And this here, these are the keys and the values. So that's the keys and the values of this attention mechanism. The queries are just learned, so the queries are not dynamically generated. And the keys and values, they are not. Yeah, it's a weak analogy, but you can sort of think of it like this. So there is this routing mechanism. And it decides where a token gets goes to. Now, as you can see, the router is soft, that means there is never a one or a zero right here, there's always kind of a number in between, but they hard clip that. So they hard clip it, they just route it to the maximum, as you can see here, number two is the maximum. And they just route it to number two, they don't route it proportionally or anything. They just take argmax and they route it through, they do multiply the output by the actual number that they got out here. So if the router is unsure, then the output is less. If the router is sure, the output is more. But this hard routing is what's the key right here. And that means, you know, before, before, you'd have one feed forward layer. So any token that goes forward goes through one feed forward layer. If you do a mixture of experts in the classic sense, and you route it in a soft way, you now have four feed forward layer. So every token goes through four of these computations. So you've basically multiplied the amount of computation by four, because you've multiplied the amount of parameters by four, right, you have four times as many parameters. Now when you do this argmax routing, like the switch transformer, you have multiplied the number of parameters in your model by four, but any token will still only incur one feed forward layer. That means you keep the amount of computation that you do per forward pass the same. And that's, that's sort of the key right here. So now they can scale up massively the number of experts, while still keeping the amount of flops the same. And notably, you also don't need any data transfer in between the experts. Every expert can be can, you know, receive their tokens and then do their independent work. And you can certainly chart this across many, many machines. This is how this looks. So in this case, you have three experts and your sequences are of line of length six. So you want to sort of route each token there and there can be overflow, like every token is independently routed. So it can happen, something like this, that a, you know, a token like three token gets routed to one expert, but it only has space for two tokens. And they have some tricks like they have this capacity factor right here, or they can reroute. These are very much engineering things, which are important. But you know, they don't change the sort of final, final result. Now I want to go down here where they have a display of this sharding more like an explanation of the sharding, which I think is very illustrative. So how, what do they essentially do? If you think of many machines, you have 16 machines. So each little square here is one machine. Okay. Here are the different ways of how you can shard a model and model sharding. Now we are not going to build a machine anytime soon that can hold a trillion parameters, that's not going to happen. Okay. So you need to somehow shard the model or the data or both. And these are the different ways how you can do it. So if you use data parallelism, that is the easiest that is also directly built into things like PyTorch and so on. What you do is, so the top row shows how to model weights are split and the bottom row shows how the data is split. So how to read this is when you do data parallelism, the weights are split such that each of the 16 cores has the same weights. You see, so this, these weights right here are the same as these weights are the same. They're all the same. So this is sharded. The data is run so that you take a data set, you take a batch of data and now you distribute this data point goes here, this data point goes here, this data point goes here, and so on. You distribute the data and you do the forward propagation and at the end, you sort of gather them again, right? So you gather them together again, because you have to, you know, calculate your gradient. Okay. So that's data parallelism. The model is spread out and if you want to do an update to the model, then you need to communicate around these weights. Okay. So all these different pieces have to then communicate with each other when there's a weight update. If you do data parallelism, here is how the data split. We've already seen this. So one piece, this piece of data is split over 16 cores. So you can see like this core right here only has this little piece of the data and not all of the data. On the other hand, you can do model parallelism. In model parallelism, you can see it's exactly the other way around, namely that one core only has a little piece of model, right? And, but every core gets all of the data. So this data here, the bottom row is data, all of the data. The point here is that if you do model parallelism, that's what you do when the model itself doesn't fit, right? Over here, the model fits on your machine, but not the whole batch at the same time. Model parallelism you do when the model itself doesn't fit. What you have to do is you have to take your data, right? And you have to send it sequentially. So maybe this is the first layer, like that's layer one weights. And then you have to compute layer one, and then you have to send it to layer two, and so on. So you have to send it sequentially through the through the sharding of the model, right? Because you want to forward propagate through all of the model. This is has very, very much of a cost of communication, you can build very big models, but it comes at a cost right at the end, you get your why and you calculate your loss and you backprop again backwards through the whole thing. You can mix them, right? You can do model and data parallelism. So here you can see that the weights, so this is this is layer one weights, layer two, layer three, layer four. And here again, you have layer one, layer two, layer three, layer four, and so on. So you can mix the two in that you can have model and data parallelism, if both your model and also your data don't fit in a single machine. And you can see here that the this upper left part receives, they receive the same data, but this here receives different data, right? So you split your mini batch into four different parts. And you send the first part up here, like that's data one, you send that up here, and that goes through the model in this sequence sequential fashion. You send data to right to here and so on. So we mix the two. Now in expert and data parallelism, that's what they that's what they do in the switch transformer. So this here is the switch transformer. And this here over here will then that's the switch transformer, one trillion. So for the one trillion model, they actually need to mix all of them. But you want to add, you know, if you can, you want to avoid model parallelism, model parallelism is really the thing that kills you because of the very high communication cost. So in the switch transformer, they have expert and data parallelism. What does it mean? So the top row is how the model weights are split. And you can see the weights are split, but the different color means that they're different weights. So here are weights number one, weights, two, weights, three, weights, four, and so on. Now we've already had this over here, right? Different weights in the model parallelism case were split over different machines. However, if you look at the data, the data is also split, and the weights, they're not the same. And these are exactly these experts. So experts, this means that, you know, this piece of data here only goes to this expert, and then to the output. This piece of data right here only goes to this expert, and then to the output, right? There is no communication between the different experts, whereas here you have this super high communication. Okay? So you can see you can scale up the experts as you scale up your data, as long as each shard of data is routed to only one expert. And then of course, you can mix the expert model and data parallelism if you really if not even a single expert fits on a machine, right? If that's the case, you need to again, shard, you do model sharding on the experts. All right? So the switch transformer, as I said, this here is the switch transformer that the most of the paper is about. And now we can dive into the results. The results are pretty spectacular. They mostly compare, as I said, to t5 base and t5 large. And as you can see right here, the switch model has significantly more parameters. So 7.4 or here 26 billion parameters compared to not even a billion of t5 large, yet the number of flops is matched. So they build models where the number of flops for forward prop is matched. But the the number of parameters are higher. So you know, it is somewhat of a fair comparison, right? You have the same amount of compute done per forward prop. And now we see what does it help to just have raw again in parameters. And it turns out it helps a lot. You've probably already seen that we get these massive speed ups, massive sample efficiencies over a dense model. You've probably so this we've looked at exactly in the in the intro, they also have benchmarks on. Let's see this down here. They also have benchmarks on multilingual on multilingual data set. And you can see in every single language, the switch transformer gains on the dense transformer by quite a bit. So this is in this is log space, as you can see. And it's quite impressive, actually. And these gains are in time as well as a number of steps. So that's pretty, pretty cool. So as I as I said, the the trade off here, of course, is that you need more machines, you need to actually add more machines. And you can see this largest model that they built is this switch xxl, which is matched in flops to trans to t five xxl model, yet has many more parameters and beats the t five at log perplexity and in as I understand in downstream tasks by quite a bit. They also built this trillion parameter model. It is not as good, mainly because they, as I understand it, they just want to get to a trillion parameters. And I think I think it's you know, training isn't really easy at that size. So they scale it down, as you can see, it has less number of heads, less number of layers. But the number of experts are way up. So that's how they scale to a trillion. And the results are, you know, better than the t five xxl, which is impressive, given that it has less flops per token. However, it is still worse than the switch xxl. So the trillion parameter model, it's still you know, it's still not everything to have a lot of parameters, you actually need to do good trade offs. And here they've traded off too many parameters for you know, less number of heads and less number of layers. And that hurts again. So very, very interesting stuff right here. The last thing I want to look at is their tricks for getting this to work. So they detail three tricks for getting this to work. And they are right here, three tricks, how they can do this. And people before them have said, No, you need at least two experts, otherwise it's unstable. So they do selective precision with the large sparse models, which means that if for some of these computations, it you know, it, it pays off to do them in higher precision, you don't want to send around these flow 32 precision things, you don't want to send those from machine to machine, right? So you have your input, you have your multi head attention. And then here, again, this is whatever x prime, and then you send that to the experts. Right here are the different experts. And then you send that back. And that's why okay, now, you don't want this here is communication cost. If you were to send around float 32 vectors, that's a lot of data that you have to transmit. So you'd rather send around 16 bit precision, right as they do right here. And however, if you do 16 bit precision, you're you know, the whole machine learning part doesn't work as well. So what they do is they do as soon as it as a as soon as a vector arrives here, this is in 16 bit, they scale it up. They cast it to a 32 bit vector, they calculate using the 32 bit vector 32. And then they cast it again to a 16 bit vector to send it back. And that seems to work. So they do selective selectively casting the precision up. And also they do selective dropout that's down here. So they do expert dropout, which means they don't apply dropout to the whole network uniformly as you would do regular normally. But they say they can do a much larger dropout rate at expert layers. And that makes a bit of sense because the expert each expert is only used very sparsely. So it makes sense to up their dropout rate. Because you know, in the end, you might drop out as much signal from a sparsely used expert, if you raise the dropout rate, then you do from a densely used layer in with a smaller dropout rate. And the last thing is that they simply do better initialization. So they find if they scale down the the initial scale of the original transformer by a factor of 10, that leads to a lot more stable training. It's astounding that after so many years, still something like initialization can, you know, make or break such a model that is just insane to see. There's a lot more to this paper, they do a lot of downstream tasks. They also talk a lot about, you know, this is not only this model, they do a lot of optimizations under the hood, they use mesh tensorflow and so on. It's clear that a lot of work has gone into this. And interestingly enough, they can also distill these models. So what they can do is they can take this large model and they distill it to a model that is as big as T5 base, a dense model. So they go from a sparse large model, and they distill it into a dense model that is equivalent to T5. And they do outperform T5 if it were trained from scratch. And they gain up to something like 30%. So 30% of the gains they made from here to here, they can retain by distilling it down. They say they can distill it down way over 90-95% of the model, which is also pretty interesting and, you know, pretty cool. Because then you could sort of distribute the trained models around and people could use them. All right, so that was it for me. Hopefully check out the paper and all the experiments, downstream tasks and so on. It's a very cool paper, has a lot of cool experiments. There's code, at least TUDO code. And that was it. Thank you. I'll check out Tudotism and let you know.
 Hi there, today we'll talk about switch transformers scaling to trillion parameter models with
 simple and efficient sparsity by William Fetus, Barrett Zoff and Noam Shazir of Google Brain.
 So as you can see right off the title, we're going towards trillions of parameters.
 GPT-3 had 175 billion parameters.
 This paper claims to have a model with a trillion parameters.
 Now is it really five times bigger or 10 times bigger than GPT-3?
 That's a debatable question, because the trillion parameters are not used in the same way as
 in a classic transformers.
 They are used actually in a sparse way.
 That's why the word sparsity is in here.
 And the way they are used in sparse manner is this new architecture called the switch
 transformer.
 It's not entirely new.
 It's built on mixture of experts.
 In this paper, that's also called MOE.
 That has been around for a while and we're going to see what that is.
 Now on a high level, switch transformers takes mixture of experts to an extreme in that it
 is a transformer.
 And the feed forward layer is divided up into these experts.
 And the switch transformer only routes each token to one expert only.
 That's the sparse part.
 So the mixture of experts previously, they always claimed you need at least two experts
 in order to get a stable training signal, the switch transformer manages to get it down
 to a single expert.
 So it's like a hard routing of information to just a single endpoint per layer of each
 token.
 So in that that means you can now scale the experts and you can scale the number of parameters
 in the model without making the model compute more.
 That's a very special notion.
 So you can up the parameters of the model.
 But if a forward pass of a data point will still have the same amount of flops that it
 needs to forward propagate through the network.
 Very special architecture right here.
 So yeah, that's why I'm saying trillion parameters not necessarily comparable to the 175 billion
 parameters of something like GPT three.
 So how do they do it?
 Because previously it was claimed it was unstable.
 They have new ways of making the training stable, such as selective dropout, selective
 casting of parameters to different precisions, and a better initialization.
 So that's the high level overview of the paper.
 And we'll dive into it, we'll explore kind of what mixture of experts is and how the
 model works.
 And what turns out it's a very long paper, as you can see, when papers have a table of
 content.
 That's a lot of fun.
 But it's a lot of engineering as well.
 And we're mostly interested in the model here, what it can do, and what does it how does
 it sort of fit in to the big world of transformers and language models and so on.
 Last thing I want to say, trillion parameters is, you know, it's a catchy title that most
 of the paper, they don't work with trillion parameter models, they work with models in
 the in the order of billions of parameters. And at the end, they build a model with a
 trillion parameters.
 It doesn't do as well as their models with as their smaller models.
 They also, it feels like they don't put that much work into it, because it's probably also
 quite fuzzy and expensive.
 But just know, we're not going to have trillion parameter models around anytime soon.
 Just yet.
 Interesting fact, the original resonant paper also built a 1000 layer convolutional neural
 network.
 Even though the resonance we have today, you know, they are maybe 50 or 150 layers deep,
 they did build a 1000 layer model.
 So maybe compare it a bit to that one.
 It's just like we can do it, not necessarily we need to.
 So here you can see something they discover.
 The curve on the left is very, very known to people that are in the language model game,
 let's say, or in the in the let's scale up AI game.
 And that is as you increase the size of the model, the loss will go down.
 And that's loss, as I understand it.
 So that's test loss.
 I believe that is perplexity.
 So scaling properties, exactly that that might be perplexity or test loss on some downstream
 task in any way, as you scale up the model parameters, the model gets better and better
 and better.
 The interesting thing right here is twofold.
 First of all, I believe they do hold the data set constant.
 So the data set is always the same, the amount of compute you put into it, the amount of
 either number of steps or time is also always the same.
 And in this specific case, the amount of flops per forward pass is also the same.
 The only thing that changes is the number of parameters.
 Again, it's very special to have a model where you can scale up the number of parameters,
 yet the flops required to forward propagate stay the same.
 So you can see here that there is a almost unhalted decrease here, it flattens out a
 little bit towards the bottom, though that is not necessarily does not necessarily mean
 it will ever flatten out before it's you know, at zero.
 I will approach zero, I guess.
 So and you can you can see that, you know, they scale up the model quite a bit.
 And also, their main comparison here is the T five base.
 So that's the text to text transfer transformer.
 By the way, if you don't know what a transformer is, or what a language model is, it's best
 you go back to my earlier videos and look up like the GPT three paper or the attention
 is all you need paper, I've made videos about lots of these things, I assume that you know
 them.
 You can see right here that if you compare to number of training steps, for example,
 the this switch models, all of them, no matter how big they are, they provide massive gains
 over like something like a T five.
 And they also do this in time.
 So this paper is very much about trade offs, you do require more storage for your
 weights.
 So you have to have more memory more RAM.
 However, that memory can be distributed, it can be sharded, because they use this mesh
 TensorFlow library to implement the switch transformers.
 And because their model has this sparsity, they can efficiently shard the model.
 So you trade off more memory, which can be sharded.
 But what you gain is training speed, and both in terms of time and number of training steps
 required.
 So you are much more efficient.
 Note that this only all of this holds in this super large regime, right?
 We this is, they say they've also discovered the speed ups in smaller models.
 But you know, as far as the paper is concerned, we are talking about millions, hundreds of
 millions of parameters, billions of parameters, even to trillion of parameters, together with
 these giant corporate corpora of, of text.
 So that's sort of the regime we are in.
 And the results do not necessarily transfer down to the lower scale problems that you
 know, you might face with your lonely one, collab in the corner.
 All right, so in a transformer, you have a transformer is nothing else but a bunch of
 these layers right here.
 This is this is in itself a transformer layer in its basic form.
 And it consists of sort of two parts, it consists of this self attention, right here.
 Now, that's the standard transformer self attention.
 That's what was introduced in attention is all you need.
 And what's been used ever since in all the transformers.
 This one right here is a is an, as I understand it, a language model.
 So you know, this this is very standard.
 However, after the self attention, you have this feed forward layer.
 Now usually, what you do is you have an input sequence, and you transform that through multi
 head attention into another sequence right here.
 Okay.
 And then what you do is you take each of these things and feed them through a feed forward
 layer.
 And if I am, as I understand it, this feed forward layer is simply, you know, a regular
 feed forward layer that you would find in a neural network, and you pass them, you pass
 these things individually.
 So this here, it's a vector, you pass it through here, and boom, that becomes the next layer
 representation, this thing right here, you pass it through as well.
 Boom, that becomes this one, and so on, right?
 You pass them individually to get the next layer representation.
 So this, this part right here, the attention part, it sort of aggregates information and
 relates the individual items of the sequence to each other, and transforms them into, you
 know, a new sequence, where sort of all the every token can gather information from every
 other token.
 That's what the attention mechanism does.
 That's step one.
 In step two, every token is isolated, every token is for itself.
 And the feed forward layer simply determines, you know, what's given one token given token
 number one, what is, you know, given its representation in this layer, what is the best representation
 for the next layer?
 Okay.
 So that's token number one of the next layer.
 So the multi head attention is kind of relating tokens to each other, and the feed forward
 layers, they are relating layers to each other.
 Okay, so up here, you would have the next multi head attention layer.
 So you can see the feed forward layer as sort of translating from one layer to the next
 layer, right, getting saying, oh, you come from this layer, I'm going to translate you
 such that the next layer understands you.
 And that happens on a token by token basis.
 Now you can see this is it's always the same feed forward layer for all the tokens, right,
 the tokens are sort of treated like a batch of samples.
 The idea of this switch transformer and also of the earlier mixture of experts transformer
 is that it might not be a good idea to have only a single one, right?
 This is the only feed forward layer, it's the same for all the tokens, it might actually
 be a good idea to have a couple of them that sort of specialize in different things.
 So what could that be?
 You know, in a in a basic world, this could just be like one for nouns.
 And this could be a feed forward layer for verb verbs, tokens that are verbs, tokens
 that are adjectives, and sort of maybe here is like punctuation tokens, right?
 You might think, well, if you are a noun token, the next layer might want to look differently
 at you than if you are a punctuation token, right?
 So this translation from one layer to the next layer can now happen dependent on what
 the token represents, right?
 Now we we of course, first of all, we don't have these annotations.
 And second, it's not necessarily that you know, we want to always divide it by noun,
 verb, adjective punctuation.
 Ideally, we want to learn this routing.
 So we simply want to say, look, instead of just one feed forward layer, we give the model
 four feed forward layer, feed forward layer one, two, three, and four.
 And for each token, the model can decide to which of these feed forward layer it sends
 the token to.
 So here you can see this is a token.
 Now, you know, we are dealing with word pieces.
 Let's just say the word more, I was like, I was thoroughly confused by when I saw this
 like, huh, why does it say more parameters, but here, it's the string more, right, and
 the string parameters.
 And these are in the vocabulary, and they get an embedding vector associated with them.
 So that's what's going on here.
 Then they go through self attention, as you can see here, both go through self attention,
 and then each one of them is routed to one of these four experts.
 Now the, the one here, the one on the left and the one on the right, these are the same
 experts, right, they're just duplicated visually here.
 But these would be the same weight matrices in there.
 So you have four feet forward layers in this layer.
 And each token can be routed to any one of them.
 And this routing here, this is learned.
 So in here, you have a matrix, they call it like WR.
 And using WR, you simply do an inner product of WR with your input right here, let's call
 that H with your input H. I guess they use H for a different thing.
 I think they call this X again.
 So you do this with X. And then you get, you get H, which is your routing, and then you
 simply build a histogram, you normalize the histogram, I think with a softmax.
 And that those are your routing weights.
 So it's very much like another attention mechanism, except that the queries, this thing here,
 these are like the queries, these are sort of the queries of this attention mechanism.
 And this here, these are the keys and the values.
 So that's the keys and the values of this attention mechanism.
 The queries are just learned, so the queries are not dynamically generated.
 And the keys and values, they are not.
 Yeah, it's a weak analogy, but you can sort of think of it like this.
 So there is this routing mechanism.
 And it decides where a token gets goes to.
 Now, as you can see, the router is soft, that means there is never a one or a zero right
 here, there's always kind of a number in between, but they hard clip that.
 So they hard clip it, they just route it to the maximum, as you can see here, number two
 is the maximum.
 And they just route it to number two, they don't route it proportionally or anything.
 They just take argmax and they route it through, they do multiply the output by the actual
 number that they got out here.
 So if the router is unsure, then the output is less.
 If the router is sure, the output is more.
 But this hard routing is what's the key right here.
 And that means, you know, before, before, you'd have one feed forward layer.
 So any token that goes forward goes through one feed forward layer.
 If you do a mixture of experts in the classic sense, and you route it in a soft way, you
 now have four feed forward layer.
 So every token goes through four of these computations.
 So you've basically multiplied the amount of computation by four, because you've multiplied
 the amount of parameters by four, right, you have four times as many parameters.
 Now when you do this argmax routing, like the switch transformer, you have multiplied
 the number of parameters in your model by four, but any token will still only incur
 one feed forward layer.
 That means you keep the amount of computation that you do per forward pass the same.
 And that's, that's sort of the key right here.
 So now they can scale up massively the number of experts, while still keeping the amount
 of flops the same.
 And notably, you also don't need any data transfer in between the experts.
 Every expert can be can, you know, receive their tokens and then do their independent
 work.
 And you can certainly chart this across many, many machines.
 This is how this looks.
 So in this case, you have three experts and your sequences are of line of length six.
 So you want to sort of route each token there and there can be overflow, like every token
 is independently routed.
 So it can happen, something like this, that a, you know, a token like three token gets
 routed to one expert, but it only has space for two tokens.
 And they have some tricks like they have this capacity factor right here, or they can reroute.
 These are very much engineering things, which are important.
 But you know, they don't change the sort of final, final result.
 Now I want to go down here where they have a display of this sharding more like an explanation
 of the sharding, which I think is very illustrative.
 So how, what do they essentially do?
 If you think of many machines, you have 16 machines.
 So each little square here is one machine.
 Okay.
 Here are the different ways of how you can shard a model and model sharding.
 Now we are not going to build a machine anytime soon that can hold a trillion parameters,
 that's not going to happen.
 Okay.
 So you need to somehow shard the model or the data or both.
 And these are the different ways how you can do it.
 So if you use data parallelism, that is the easiest that is also directly built into things
 like PyTorch and so on.
 What you do is, so the top row shows how to model weights are split and the bottom row
 shows how the data is split.
 So how to read this is when you do data parallelism, the weights are split such that each of the
 16 cores has the same weights.
 You see, so this, these weights right here are the same as these weights are the same.
 They're all the same.
 So this is sharded.
 The data is run so that you take a data set, you take a batch of data and now you distribute
 this data point goes here, this data point goes here, this data point goes here, and
 so on.
 You distribute the data and you do the forward propagation and at the end, you sort of gather
 them again, right?
 So you gather them together again, because you have to, you know, calculate your gradient.
 Okay.
 So that's data parallelism.
 The model is spread out and if you want to do an update to the model, then you need to
 communicate around these weights.
 Okay.
 So all these different pieces have to then communicate with each other when there's a
 weight update.
 If you do data parallelism, here is how the data split.
 We've already seen this.
 So one piece, this piece of data is split over 16 cores.
 So you can see like this core right here only has this little piece of the data and not
 all of the data.
 On the other hand, you can do model parallelism.
 In model parallelism, you can see it's exactly the other way around, namely that one core
 only has a little piece of model, right?
 And, but every core gets all of the data.
 So this data here, the bottom row is data, all of the data.
 The point here is that if you do model parallelism, that's what you do when the model itself doesn't
 fit, right?
 Over here, the model fits on your machine, but not the whole batch at the same time.
 Model parallelism you do when the model itself doesn't fit.
 What you have to do is you have to take your data, right?
 And you have to send it sequentially.
 So maybe this is the first layer, like that's layer one weights.
 And then you have to compute layer one, and then you have to send it to layer two, and
 so on.
 So you have to send it sequentially through the through the sharding of the model, right?
 Because you want to forward propagate through all of the model.
 This is has very, very much of a cost of communication, you can build very big models, but it comes
 at a cost right at the end, you get your why and you calculate your loss and you backprop
 again backwards through the whole thing.
 You can mix them, right?
 You can do model and data parallelism.
 So here you can see that the weights, so this is this is layer one weights, layer two, layer
 three, layer four.
 And here again, you have layer one, layer two, layer three, layer four, and so on.
 So you can mix the two in that you can have model and data parallelism, if both your model
 and also your data don't fit in a single machine.
 And you can see here that the this upper left part receives, they receive the same data,
 but this here receives different data, right?
 So you split your mini batch into four different parts.
 And you send the first part up here, like that's data one, you send that up here, and
 that goes through the model in this sequence sequential fashion.
 You send data to right to here and so on.
 So we mix the two.
 Now in expert and data parallelism, that's what they that's what they do in the switch
 transformer.
 So this here is the switch transformer.
 And this here over here will then that's the switch transformer, one trillion.
 So for the one trillion model, they actually need to mix all of them.
 But you want to add, you know, if you can, you want to avoid model parallelism, model
 parallelism is really the thing that kills you because of the very high communication
 cost.
 So in the switch transformer, they have expert and data parallelism.
 What does it mean?
 So the top row is how the model weights are split.
 And you can see the weights are split, but the different color means that they're different
 weights.
 So here are weights number one, weights, two, weights, three, weights, four, and so on.
 Now we've already had this over here, right?
 Different weights in the model parallelism case were split over different machines.
 However, if you look at the data, the data is also split, and the weights, they're not
 the same.
 And these are exactly these experts.
 So experts, this means that, you know, this piece of data here only goes to this expert,
 and then to the output.
 This piece of data right here only goes to this expert, and then to the output, right?
 There is no communication between the different experts, whereas here you have this super
 high communication.
 Okay?
 So you can see you can scale up the experts as you scale up your data, as long as each
 shard of data is routed to only one expert.
 And then of course, you can mix the expert model and data parallelism if you really if
 not even a single expert fits on a machine, right?
 If that's the case, you need to again, shard, you do model sharding on the experts.
 All right?
 So the switch transformer, as I said, this here is the switch transformer that the most
 of the paper is about.
 And now we can dive into the results.
 The results are pretty spectacular.
 They mostly compare, as I said, to t5 base and t5 large.
 And as you can see right here, the switch model has significantly more parameters.
 So 7.4 or here 26 billion parameters compared to not even a billion of t5 large, yet the
 number of flops is matched.
 So they build models where the number of flops for forward prop is matched.
 But the the number of parameters are higher.
 So you know, it is somewhat of a fair comparison, right?
 You have the same amount of compute done per forward prop.
 And now we see what does it help to just have raw again in parameters.
 And it turns out it helps a lot.
 You've probably already seen that we get these massive speed ups, massive sample efficiencies
 over a dense model.
 You've probably so this we've looked at exactly in the in the intro, they also have benchmarks
 on.
 Let's see this down here.
 They also have benchmarks on multilingual on multilingual data set.
 And you can see in every single language, the switch transformer gains on the dense
 transformer by quite a bit.
 So this is in this is log space, as you can see.
 And it's quite impressive, actually.
 And these gains are in time as well as a number of steps.
 So that's pretty, pretty cool.
 So as I as I said, the the trade off here, of course, is that you need more machines,
 you need to actually add more machines.
 And you can see this largest model that they built is this switch xxl, which is matched
 in flops to trans to t five xxl model, yet has many more parameters and beats the t five
 at log perplexity and in as I understand in downstream tasks by quite a bit.
 They also built this trillion parameter model.
 It is not as good, mainly because they, as I understand it, they just want to get to
 a trillion parameters.
 And I think I think it's you know, training isn't really easy at that size.
 So they scale it down, as you can see, it has less number of heads, less number of layers.
 But the number of experts are way up.
 So that's how they scale to a trillion.
 And the results are, you know, better than the t five xxl, which is impressive, given
 that it has less flops per token.
 However, it is still worse than the switch xxl.
 So the trillion parameter model, it's still you know, it's still not everything to have
 a lot of parameters, you actually need to do good trade offs.
 And here they've traded off too many parameters for you know, less number of heads and less
 number of layers.
 And that hurts again.
 So very, very interesting stuff right here.
 The last thing I want to look at is their tricks for getting this to work.
 So they detail three tricks for getting this to work.
 And they are right here, three tricks, how they can do this.
 And people before them have said, No, you need at least two experts, otherwise it's
 unstable.
 So they do selective precision with the large sparse models, which means that if for some
 of these computations, it you know, it, it pays off to do them in higher precision, you
 don't want to send around these flow 32 precision things, you don't want to send those from
 machine to machine, right?
 So you have your input, you have your multi head attention.
 And then here, again, this is whatever x prime, and then you send that to the experts.
 Right here are the different experts.
 And then you send that back.
 And that's why okay, now, you don't want this here is communication cost.
 If you were to send around float 32 vectors, that's a lot of data that you have to transmit.
 So you'd rather send around 16 bit precision, right as they do right here.
 And however, if you do 16 bit precision, you're you know, the whole machine learning part
 doesn't work as well.
 So what they do is they do as soon as it as a as soon as a vector arrives here, this is
 in 16 bit, they scale it up.
 They cast it to a 32 bit vector, they calculate using the 32 bit vector 32.
 And then they cast it again to a 16 bit vector to send it back.
 And that seems to work.
 So they do selective selectively casting the precision up.
 And also they do selective dropout that's down here.
 So they do expert dropout, which means they don't apply dropout to the whole network uniformly
 as you would do regular normally.
 But they say they can do a much larger dropout rate at expert layers.
 And that makes a bit of sense because the expert each expert is only used very sparsely.
 So it makes sense to up their dropout rate.
 Because you know, in the end, you might drop out as much signal from a sparsely used expert,
 if you raise the dropout rate, then you do from a densely used layer in with a smaller
 dropout rate.
 And the last thing is that they simply do better initialization.
 So they find if they scale down the the initial scale of the original transformer by a factor
 of 10, that leads to a lot more stable training.
 It's astounding that after so many years, still something like initialization can, you
 know, make or break such a model that is just insane to see.
 There's a lot more to this paper, they do a lot of downstream tasks.
 They also talk a lot about, you know, this is not only this model, they do a lot of optimizations
 under the hood, they use mesh tensorflow and so on.
 It's clear that a lot of work has gone into this.
 And interestingly enough, they can also distill these models.
 So what they can do is they can take this large model and they distill it to a model
 that is as big as T5 base, a dense model.
 So they go from a sparse large model, and they distill it into a dense model that is
 equivalent to T5.
 And they do outperform T5 if it were trained from scratch.
 And they gain up to something like 30%.
 So 30% of the gains they made from here to here, they can retain by distilling it down.
 They say they can distill it down way over 90-95% of the model, which is also pretty
 interesting and, you know, pretty cool.
 Because then you could sort of distribute the trained models around and people could
 use them.
 All right, so that was it for me.
 Hopefully check out the paper and all the experiments, downstream tasks and so on.
 It's a very cool paper, has a lot of cool experiments.
 There's code, at least TUDO code.
 And that was it.
 Thank you.
 I'll check out Tudotism and let you know.