Mixture of Experts LLM - MoE explained in simple terms

Key Takeaways

hello Community let's have fun today mixture of expert system here we go imagine you would have eight different mistal 7B llms and you want to combine them to a mixture of experts how would you do this now we know that here our mixture of expert M build up on the observation that LMS can be decomposed into smaller specialized expert system that only focus here on this distinct aspects of the input data so it is sensitive to the input data which enables a more efficient computation and a research allocation especially if you do parallel GPU Computing so AI models now with a mixture of expert system can significantly reduce the computational cost of our large language model but please notice that the mixture of expert system goes with sparsely activated a lot of people only talk about mix of expert but it is sparsely activated mix of expert and I will explain this in a second yes you all noce at the very beginning of 2017 they emerged we had here publication by Google brain this sply gated mixture of experts layer for outrageously large newal Network and you know there is a rumor that even gp4 maybe is a mixture of expert layer now from Mistral we have some torrent information that they are exploring here the mixture of expert performance of multiple mistal 7B models so it all started in 2017 and here you have it the very beginning you have a mixture of expert layer here in blue embedded within a recurrent language model so what we have we have an input a token comes in and then we have a intelligence and this intelligence is the gating Network and it decides to which expert system sent this token off to maybe only to one expert and maybe to two expert system this was the very beginning in 2017 so if we look now here at the original publication of our sparsely gated mixture of expert layer we scroll down and you see here this is the very simple beginning of MOA layers and and we of course we are interested to computers so here we go yes you can read this is beautiful but you notice what is now really interesting is here the intelligence of the gating Network the switching off input tokens to specific expert system and we go here with the 27 literature so as I told you the simplest way is the softmax gating function you have a simple choice of nonsparse gating function this was from 1994 my good in the old Millennium you multiply the input by a trainable weight Matrix and you have a l function in the training and you have the soft Max function and then you just optimize the complete system this is it remember soft Max sums up to one no negative you notice then you have a noisy topk gating we add two components to the softmax gating Network sparcity and Noise so before taking here the soft mix function we add a tunable Gau noise and keep only the top K value setting the rest to for example minus infinity forms a sparcity creat sometimes theoretically discontinues in the output of gting function there's no Pro no problem in practice helps with load balancing beautiful this is the way and then we have here even in the 2070 version we have we have to train here the gating Network to perform its task of assigning the tokens to specific experts they say hey we trained the gating network but simple back propagation along with the rest of the model if we choose here K greater than one so two expert system the gate values for the top key expert have nonzero derivatives with respect to the weights of the gating Network this type of occasionally sensitive behavior is already described in 2013 refers back to the noisy reac defies and the gradients also back propagate through the gating Network to its input we integrate here of course our gating Network this is it then you have here some problem data parallelism model parallelism yeah we want here the mixture of expert system to have here the perfect parallelism structure for multiple GPU notes of course but then you see network bandwidth balancing here utilization this is more or less it from 2017 previous dat of order and their model imagine consisted still in 2017 we had nothing of of a GPT system they still had two stacked lsdm layers and my youngest viewer will have no idea what it is but never mind we just look here for the mixture of expert functionality now we go to Mega Blocks we jump here to 2022 just about yeah November a year ago we had Mega blogs efficient Sports training within the mixture of expert system and it is rather easy what they did here they addressed the limitation of the classical Mo Moe system and they reformulated here the computation in terms of block spars mathematical operation and develop veloped here the infrastructure compute so new block sparse GPU kernels for the optimization of the matrix multiplication of block Spar Matrix structures or tensor structures this is all there is to it that's it if you want to have here an overview of Mega Blocks you have here incoming tokens now instead of this uh gate that we have here the intelligence now they call it root they have now probability assigned to it and then you have the permutation to computation and it's done great let me focus on the router because I think the router is really a crucial element in our mixture of X but the routa is implemented as also as a trainable neuron Network and the routa determines how to distribute singular tokens to various experts based on the current state of the input data and of the exper of each expert system so what we do we start with a token representation now each token a piece of data or a segment of a sentence or whatever you have is you notice represented as a high dimensional Vector normally 768 Dimension or 1024 Dimension here importance of your tokenizer do not underestimate a bad tokenizer can crush your model even if you do anything else right tokenizer and token structure are important for this great token Vector representation this Vector encapsulate the features of the token may include semantic meaning the context you notice from our Transformer structure and then the toker the router takes now the token's vector representation as an input and now calculates a set of scores one score for each of the available expert and these scores represent the Ro's assessment of how relevant each expert is for processing that specific special token and this process involves a learned a transformation followed by a very simple soft soft Max operation and you know why because the scores are non negative and they sum up to one so we exactly see here the maximum or the top two Max allocations great or after scoring the root of selects the top K expert system with the highest score for each token K is a small number sometimes one but you can go here for at least k equal 2 which would send each token to the two most relevant expert system beautiful now I already received question about your okay yes you have to train the router so the router is trained jointly with the rest of the mixt of expert model here and it simply learns like a new network the optimal scoring and the distribution strategy based on the feedback from the overall model performance so what it does the training process adjusts the routs parameter to minimize the loss function of the entire model as you know from training any newal network the the loss function the same happens here so which includes measuring how well the combined effort of all expert perform on a given task so this is quite time intensive so watch out to have the perfect training data sets for this you train the router you train the complete system and you achieve here the mixture of expert benefit now you know that in Computing now this mix of expert efficiently we have two problems we have a problem of how to dynamically root this things and we have a load inbalance computation now to address here this problems here the Autos here of the Mega Blocks developed an approach for rooting and computation based on sparse Primitives I will give you an example what this means in a minute because this is important spars Primitives is the building block here so to say our approach now drops tokens this is beautiful and Maps efficient to Modern gpus enabling a speed up of up to 40% beautiful so Mega PLS accelerate here their process 40% yeah they showed that the compute here in uh mixture of expert layer can be expressed as a block sparse operation to accommodate imbalance assignment of token to experts and they use this formulation to train here and they call it now a dropless a no token drop the dropless mixture of exper system and if you see this little D it's not for differentiation it is for dropless beautiful and then of course if they have now the block sparse matrix multiplication operations they develop now of course the hardware compute s they optimize the GPU Kels here for specifically the block spse Matrix operation and you know there are different methodologies in mathematics how how to optimally compute the product of one or two Matrix and if you have a block Spa Matrix there's also there are multiple ways to optimize the mathematical operation here in calculating this in the GPU chel beautiful that's all you need to know so here we go first summary we enhance the number of model parameters of our llm without a corresponding rise and a computational expense we have our experts system and only the input data that are specific to an expert go to this single or maybe k equal 22 expert system great now I was asking here a simple explanation yes ban this is for you now how can you imagine this in very simple terms imagine you have a large group of students in the classroom and every student has a question as it is normally but in the corner you see here our teachers and each teacher is specialized in a particular subject and there are I don't know 50 students and only one two three four five six seven teachers now if all students rush to the teacher you have a bottleneck so how we do this easy the idea is to send each student to the teacher that is best suited to answer their question but of course each teacher can only attend a certain number of student at a time now imagine this is a science class and this here is our professor of mathematics and the rest of the teacher only I don't know baseball or basketball or something else so all the people want to go to the mathematics teacher and nobody's interested in baseball so you see we have an extreme bottleneck so how do we cope with this and this is exactly how the our experts here how we have here the mixture of experts that all question questions from here the students are answered yes yes yes the example of of course if too many students go to only one teacher some students won't get the answer this is what we call a token dropping in the mixt of expert system but we also want to ensure that all students get help all all questions are answered so you might have to limit the number of students each teacher can see but this also means that some other teachers will be idle all the basketball and volleyball teachers are idle while here our mathematic Professor is above 100% capacity so we waste computational resources so we generate a second mathematical teacher to have here to reduce your the bottleneck function great now what does mega blog do Mega blog solve this by efficiently organizing which student goes to which teacher yes yes yes and it does this by dynamically adjusting the flow of students and we do this SE mathematics with a block sparse method now if you don't know what a block Spar is I will show you in the next slide but Aken to creating flexible group sciences that can change as needed so we ensure that an efficient and balanc distribution of students to teacher is happening so what mega block does in this example it reformulates here a complete mixture of expert computation using now this block BL sparse Matrix operation notice that it's not only a sparse Matrix operation but a block sparse and I will show you this this allows for the handling of imbalance assignment of tokens to expert without the need to drop any token so we have we achieve our goal no token is dropped mega block does this with a dynamic routing yes of course and of course we adopt this here to our GPU cornner so we have the optimal parallelization happening so in the best case each teacher here is a GPU and we have parallel parallel parall parallel GPU computations great so we are fast and efficient what is a FAL Matrix Wikipedia if you have almost only zeros in this Matrix you have a sparse Matrix so here only non nine non-zero elements 26 seral elements this is a sparse Matrix you can calculate the sparity and today density of it beautiful in mathematics you know we have sparse input dense sparse output dense input and dense output and sparse Den input a lot of different possibilities but what we do here is with a block Spar formulation we achieve that no token is Left Behind so what we do we apply now this block Spar computation for each mixture of expert layer to avoid it that any token is dropped yes yes L we can do this yes computation is zero yeah if you go from a mathematic side you know sparse Matrix is where a large Matrix where almost everything is with zero this is a spar Matrix a block Spar structure is now a block Spar Matrix is a special type of a spar Matrix where non-serial elements are clustered into blocks in this structure of a block is now more efficient than a fully Spar Matrix when we calculate matrix multiplication for example so give you an example you are here in your school and you have four blocks and a block is here playing field for example and here they play soccer and here they play something else and maybe they play volleyball and they play whatever so you have build your sports blocks here and the blocks here Focus here on one particular semantic topic let's call it in this way or execute here only one game so you group them together and if you think now as this as a matrix you see this is beautiful because now you have clustered everything together that belongs together and now our matrix multiplication will become very fast very efficient if we also optimize the GPU cornels for this particular Matrix multiplication that's it yeah B if you read this here Sports tournament with many teams yes yes yes and you see exactly what I mean okay not just sport you have your science club you have the art Club you have the chess club you see what's happening we are building blocks for specific tasks Mega Blocks approach here is exactly this yes blocks pass Matrix says yes matrix multiplication you know this or another example now the size can vary of these blocks think about it you have in a school and you have now the science club The Art Club and the chess club but different amount of people would like to go to the art club and the chess club so in a traditional setup here in a nonblock sparse approach each club would meet need in a large Hall and occupy here three equally sized hallways or rooms if we use here the magga block approach we adjust the rooms so if you have less students we take a smaller room a smaller block and have a lot of students we take a bigger block so our matrix multiplication you see is becoming even more efficient and it is dynamically adjusted here also to the input data beautiful this is all we want said okay this was 2022 what about here the advancement in 2023 well you know it yes we have a mixture of expert system here with instruction tuning and here you have Google University Berkeley MIT Amherst and University of Texas and they have here July 2023 instruction tuning of mixture of expert system so now we are in the fine tuning and you know instruction tuning is a very special fine tuning with we have an instruction data set where we have the instruction and a tuple with an output defined we have here we Define here for the model how we want that the model responds to instruction given and with this training data set we can train now our mixture of expert system yes yes yes emphasize necessity of an instruction tuning stage for a mixture of expert models to support the performance of den models you know dance model is where we do not have sparse Matrix calculation on Downstream task or held out task now you know switch Transformer have a video on switch Transformer so similar to the switch Transformer architecture the aors replace here the feet forward component of every other Transformer layer with a mixture of expert layer this mixture of expert layer consist of a collection of independent Feit forward expert that are now our experts and then like in 2017 you have now a let's call it intelligence gating functionality that uses here softmax activation function to model a probability distribution over the expert so you know exactly now to which expert to send your specific token each M Mo layer is a learnable gating network is trained to use it input to activate the best two expert for each token of an input sequence k equal 2 here you Noti this inuring inference the Learned gating Network dynamically picks here the best two experts for each token great if you want to learn more about here there's a multitude of instruction tuning methodologies here for llms and for mixture of expert system I would recommend you have a look at this uh scientific publication from October 9 2023 and here you have a beautiful amount of instruction tuning explained all the different data sets and everything that you can use also here to train your mixture of expert system beautiful I hope this was a first touch on Moe systems what are they when did they happen how how did they evolve what methodologies do we have to train them to fine-tune them for instruction fine-tuning for example and I hope this gives you a first idea what are mixture of expert system for the upcoming development in our large language model and in our vision language model it would be great to see you in my next video