Better and Faster LLMs via Multi-token Prediction

Key Takeaways

hi everyone today I'm going to be summarizing or give you kind of my takeaways and my thoughts on this new and exciting paper from glow kill atal and they are proposing this multi-token prediction as a way to train More Sample efficient large language models and we will touch on that and what that means and there are a ton of interesting results we will cover some of the more interesting ones and what this could potentially mean for the future of llms right so we know that large language models have different challenges and one of those being the sample efficiency but also if you are going to propose an approach on our approach to do language modeling right so the traditional way that we do that today is doing this next token prediction if you're going to propose a method right you need to be able to do experiments that assesses the robustness of that approach and this is one of the reasons I like this paper and I wanted to cover it I think it's a great example of how to perform experiments and how to think out of the box in terms of how to test the robustness of an approach there are a lot of exciting results and what this paper proposes here is they're trying to basically train large language models to predict multiple future tokens at once and they are claiming that this again results in higher sample efficiency and specifically this particular training Paradigm here or approach is asking the model to predict the following end tokens using and independent heads operating on top of a sheared model trunk and this is a Transformer um model trunk that we're we're using here so we're going to get into those details as we go through the paper but the idea is that we have these output heads that are independent and right they have the ability to predict multip tokens as opposed to just predicting one and by predicting you know training this mod to be able to predict multiple tokens we can improve right results on some benchmarks but also it shows that it can create speed up of up to 3x I believe right here they report that um so there's a lot of details here that I'm going to skip because I'm going to be focusing more on the important parts now as for results right they have a 13 billion parameter model that they trained with this approach and it solves 12% more problems on human eval and 177% more on mbpp bench mark than comparable next token molds right so they are actually comparing with these traditional way of how we uh this Auto regressive language model so Lama GPT and so on and so forth but something interesting about the size selection here actually with this approach what they showed is that and you can see here in the chart I'm jumping to this right away uh as you increase the model size you can see here we go from 0.3 billion to 13 billion you can see how it significantly improves the results right and on this particular Benchmark now these benchmarks I should have said that they are more around code generation right and the trying to assess the code generation capabilities of llms right they're very popular and they are used heavily in in llm research and so that that's kind of exciting to see and what we will do is will jump on or touch on a few things that are kind of interesting in this paper now if you want to understand a little bit on the actual multi-token prediction approach right they they mentioned that during training the model predicts four future tokens at once right so you're getting that at once now when you're doing inference you essentially what you do here and and they mention it here they employ only one only the next token output head right but they have the option also to use the other tree head as well and this gives you a speed up inference time right so this goes up to 3x I believe that's what they report that's exciting I think the result because you get all of these benefits in in one shot now I think I would say before I go into all the different aspects of this paper that this is a very promising approach now the results that you see are more centered around uh you know these common benchmarks now how does this translate to real world capabilities and that's something we were we're going to discuss and and the implication of that now there's a whole section here on the methodology you can read more about that you know the laws this is very common I'm not going to go through the mat this is just for review um but they also have right uh the the kind of changes that they made on this particular model um and you can see here that they end up with this multi token prediction cross entropy loss and they Define it and you can see that this architecture in particular uses a Shear trans from a trunk um which produces the hidden representation right from the observe context you can see context is is common with language model um and then it has n independent output heads implemented in terms of the Transformer layers and a share on embedding Matrix so you can you can look at it there's a summary of that here in this particular figure and this will give you it even has a little bit of code here on exactly what is happening here and during this training process now something that they highlight here is this memory efficient implementation right so there is a bottleneck here and the bottleneck is that the logic vectors right uh become the GPU memory usage bottleneck right you have a huge vocabulary and so on so those dimensions are kind of increasing what they have done is they are proposing a way to kind of deal with this and you can read more about how they do it here but the way I understood what they're doing in this paper is that you know they figured out a way to leverage this shared trunk right so they have a shared trunk here and they sequentially compute the forward and backward pass of each independent output head so you you might be asking how do they deal with these output heads there are a couple of them you in this case they're referring to four as a as a number and we will get into why they actually chose four because they show in the experiments that's kind of um kind of an ideal number for this setup you can see here that they that the way this is done is sequentially right 1 2 tree you know we have a loss goes back into the head goes back in sheer trunk and then is it goes in a sequential way right that forward and backward pass um and it's accumulating gradients at the trunk right the trunk is really tying things up together okay and and you can see how they manage to reduce GPU memory utilization here and for inference right they're using I do think yeah they mention it here they're using cell speculative decoding method um for the next token prediction and the multi- token prediction which leads to again it leads to being able to utilize right all that capability that you're getting from training on multiple tokens at inference time so they report some results on that I believe the main number that they share is that they get up to 3x speed up on when you're using right as it mentioned in in the in the first figure here when they use the other tree heads as opposed to just using one the standard setting there are a lot of experiments I want to highlight a few of these before uh finalizing the video here which were interesting to me again I am coming from a application developer how might this be interesting for someone like me that develops applications on top of large language models I'm also a researcher but I'm trying to just gain insights from the research that's happening and the trends and see how I can translate that into production ready systems and so what I will do is I will highlight a few of the results that are interesting and but there is a lot of results in this paper and this is remarkable because you don't don't typically see papers today highlighting so many different aspects of these models or assessing different aspects of the model so that's really cool and then you can see that different insights are coming out so what's interesting here this first figure here you can see the results on end token prediction models on mbpp again we're using this Benchmark here and by model size so this shows again what the trend is as the model scales right you get this kind of Boos in performance on both of these code generation benchmark so that shows consistency there and you want to see that right you want to see consistent and stable results when you're developing new approaches for doing language modeling specifically okay so that's the benefit of scale uh with model size faster inference again it touches on the really self speculative decoding and they observe a speed up of 3.0x here on code with an average of 2.5 accepted tokens you can take a look at the table here I won't go through it but you can just take a look at if you want interested in this particular part I think this is the part that we are most interested in because again when you think about oh you're doing multiple token prediction you are considering more future tokens the first thing you think is okay this is going to be an inefficient approach it's going to be computationally intensive uh but it turns out that at inference time you can leverage those heads can do prediction of four tokens at once this is the reason why I think I got excited about it when I saw it the first time this table just introduces or presents a few results uh they have like bite level models as well you can see um you can get this one consistently at n equal 8 consistently gets better performance compared to the other ones but the ones I'm interested to see is the actual a model strained on on on 200 billion tokens and one on one trillion tokens which is now more common right um and these ones have vocabulary of 32k tokens they try different end values so you can see that for using for and this is why they use four as the common example across the paper that it provided the best results right and this setting here so this is for searching for the optimal n and you can see how n equal 4 usually provided the better results there is obviously an exception here with this particular Benchmark which is another code generation uh type of Benchmark you can see that this for this one with Nal 6 actually provided good results but it seems that this one is also acceptable in terms of the performance theem all right so one interesting aspect of this paper that I really liked as I read through it is the fact that again they focus in different aspects of language models and the different challenges and right now there's a lot of conversation about these models having constraint to understand you know history right understand context right and and understand not only local patterns but understand global patterns as well and that's very important because there's a lot of semantics that might be hidden and that are important right as the language MO is generating tokens that might be missed in that generation process so by considering future tokens you might be able to you know gain and be able to build a system that can understand more global patterns and learn global patterns um so that's what they kind of talk about here in this particular section so you can read more about that and kind of the results that they're reporting I thought this was an interesting one uh because again they mention it at the beginning here as one of the motivations right not only do they want to train a system that can have good reasoning capabilities and so on right and can acquire language and World Knowledge and so on this is what we're interested in but at the same time we also want something that can you know deal with these issues of just focusing on local patterns and know we know also that next token prediction based on some kind of experiments that have been reported that they Overlook hard decisions as well right so there's a lot of these challenges while these systems are already very capable there are all of these little things that we kind of sort of need to look at more closely and this paper does a good job at that this particular one is training on multiple Epoch so there's also another kind of idea as well on how to effectively train large languish models if you have been following the research around llms there's also this idea that you can train llms more optimally if you consider doing more EPO and a few Labs a few research lab have shown that you can do so effectively so the question is can this multi- toen prediction setting be helpful and kind of is suitable when you're implementing multiple EO it turns out that yes they do that and they do show that they have some you know they get some performance improvements when they test it out in their different benchmarks that they're using okay so they do show that as well not only do we want that learn Global pattern not only do we want the systems to be good at inference time not only do we want the systems to be efficient at training time not only do we want these systems to be able to used in different ways like training using multiple Epoch and also fine-tuning this molds right because so far there's our conversations about you know where might fine tuning fit but a lot of like developers and companies are using fine tune mods so if this was made available can we finetune a multi- toen predictor right the model that has been trained using this new setting of multi token prediction so they do show that and a couple of results uh stand out here um you can see here they make a comparison of fine uning performance on code contest so code contest is another it's an excellent Benchmark as well it's actually one of the harder benchmarks um but they show that consistently if they're using the setting of n equal 4 how that outperforms just the normal setting where you're just doing next token prediction that's basically what this one is saying and I think this particular sentence actually summarizes that also they wanted to test because so far we have spoken about the coding benchmarks but we want to talk also about the natural language benchmarks specifically and as I read through the paper it got really interesting because they show that you know it regresses right when you're using this multi-token training with you know the 7B models it doesn't improve performance on Choice tasks or these multiple choice tasks and something they highlight here is that you know some of these benchmarks might not be suitable they're not there might not be like you know leveraging generative capabilities of language model as well so we probably don't think that this is the right way to test the performance of these models and I appreciate that kind of transparency and the fact that the authors actually reported that result here and and what they do is they try to use something like summarization so specifically what they highlighted here was the use of the abstractive Tex rization task and they chose you know a couple of benchmarks for that and they show in this setting that you know this has a lot of potential right so n equal 4 setting here you can see how this one is performing a bit better um and more stable results right as you know the trading tokens increases um there's also some other theoretical stuff as well that might be interesting like the induction capability of the multi-t token prediction models U there's some more ablations on the synthetic data as well again the induction capability they discuss more that here there's also another polinomial arithmetic task that's another complex task so to speak that you can test models with and again they're showing really good results with n equal 4 setting um but you can check it out and and and they have settings like in domain and out of domain but I will leave that out of this video because I want to keep it as short as possible and then again the algorithmic reasoning which is something they highlighted that these models have they're better algorith reasoning I think it makes sense uh because I think doing this multi-token prediction might have good benefits for tasks that require these models to look at longer patterns these you know the semantics the long-term setting and they also go into like why it works some speculation about it I think that this was an interesting one as well according to the author's intuition uh is that multi- token prediction mitigates the distributional discrepancy between training time teacher forcing and inference time out aggressive generation they also talk about this idea of Choice points you can take a look at this as well so Choice points might be you know uh areas that the model you know struggles with right or maybe these next token prediction are ignoring this particular challenging transitions of sequences and these particular models are able to handle that well or promote Better Learning of these particular transition so they talk about an example here you can see the ground Ro 1 2 3 4 5 a b so this particular transition here might be a little bit harder to catch with just normal you know standard next token prediction but if you have a multi- token prediction setting you might be able to capture those because there's again what I mentioned was it's it's able to promote that the ability to learn that because it leverages it correlates tree a and 5 to C so tree a and 5 to C okay in the data again because it's looking at Future tokens and they go into like kind of the the ways associated with it and so forth right but I think this is the overall summary of that and they believe that the quality of text Generations depends on picking the right decisions at Choice points those Choice points are really important and that in token prediction lose losses promote those right promote learning those really enjoyed this paper I think the main takeaway I would say without spending too much more time here is that you have a potential alternative way of training these large language models right and getting all these different benefits right without any kind of additional computational overhead right if that makes sense so it's like you get better inference you get better results and you get all these other benefits and you can see how they looked at for instance one one one of the are we will have to take a little bit uh a closer look at is this idea of learning global patterns because I think that's where this particular approach might be very very um very interesting to apply the other stuff like multiple EPO fine tuning and so on I think those are kind of standard experiments that we will uh also I think have more discussions around as we progress with these ideas and how to improve large language model training and so on but I think those were kind of the areas that I wanted to highlight um and obviously faster inference as well but overall the idea that this model can be more sample efficient even while using you know more feature tokens I think that's a a really Co idea and that's what this paper essentially proposed um so that will be it for this video hopefully you like the style that I'm presenting this again as I mentioned I am a researcher but I'm also building applications of on top of large language models and helping companies do that so I regularly try to keep up with you know what are the latest ideas what could be the latest trends that we can bring into like in the production setting and I think this paper presented a few things that highlight or maybe emphasize on where kind of the field is moving towards consider leaving a like I'll appreciate that and let me know in the comments if you have any questions if you have any other paper ideas that you want me to give you my takes on as well I'm happy to do that I already got some requests from some of you um and yeah I I totally appreciate that so And subscribe to the channel if you think these are interesting that will kind of give me a signal as to whether this is interesting for this community thanks and have a good day