Google Gemma-2: Technical Report Deep Dive

Google just released Gemma 2, the second iteration of their open weight models. This family contains three models including two 9 and 27 billion versions. But Google is only releasing the bigger two models at the moment. Now, on benchmarks, this seems to be best in class, where the 9 billion version is able to outperform llama3.8 billion and the 27 billion version is very close to the 70 billion version of llama3 on academic benchmarks. Now, personally I'm not a big fan of academic benchmarks, but Google has been testing this model on lmsys Chatbot arena for a while under a pseudo name and this seems to be doing pretty good in the overall ranking. This is sitting at number 11 and I'm talking about the bigger 27 billion version which is right below Gemini 1.5 Flash and their scores are pretty close, which is pretty impressive for a model of its size. The llama3.70 billion instruct is at number 12, llama370 billion is at number 12, while the smaller Gemini 2 is at number 9. Now, there is probably a reason why it's doing really good on the LMSENSE benchmark when people are using these models. That will become clear when we look at the technical report. Google also released a highly detailed technical report titled Gemma 2 Improving Open Language Models at a Practical Size. It's a lot more detailed compared to some of the other technical reports that we have seen. And a section on knowledge distillation of training the smaller models is very interesting. Now, I read the technical report but wanted Claude 3.5 sonnet to help me with some understanding. So I provided that and said create a mind map of the paper. So here's the mind map that it came up with. It talks about introduction. So it's a lightweight state of the art open weight models which ranges from 2 to 27 billion parameters. It also talks about the architecture, so it's a pretty good summary of the main points. It's a decoder only transformer and there are some innovations that they have introduced when it comes to the architecture. But for me I was really interested in the training data. So there are three different data sets of varying size that are used to train different models, which is pretty interesting to see. And they use this approach called knowledge distillation for the smaller models, which could really become a defective way of training smaller models from bigger models if these results hold true and we're going to spend quite some time understanding how knowledge distillation from a bigger model to a smaller model works. Now they have the same tokenizer which they use for the the first series of Gemma and Gemini model that has a pretty high vocabulary size of 256,000 tokens. And that's why it's pretty good potentially at multilingual tasks, although it's specifically trained on English language. Now when it comes to training they are doing pre training and then supervised fine tuning with behavioral cloning. That will become evident when we talk about knowledge distillation. And they used a policy for RLHF along with merging multiple different iterations of the models to see whether that improves the performance. So there are a lot of details. We are going to look at some of the stuff in this video. So it's going to be very different than my normal videos. One thing which I'm going to highlight is the way they have used the LMSYS training data or the LMSYS prompts in the training that could have an impact on why these models are able to perform better on the LMSYS arena chatbot leaderboard. Okay, so one of the major issues with training these large language models is that even the smaller models require a lot of data to fine tune them. So for example, the Llama 3 family is fine tuned up to 15 trillion tokens which is pretty huge depending on it gives Only less than 1% improvement compared to the state of the art models which I think in that case they are talking about the Mistral family. Now even with this huge amount of data, there is evidence that shows that these models are still under trained. So there has to be a better way of training these models if you don't really have access to a lot of data. Now that's where this knowledge distillation from a bigger model to a smaller model comes in. So you have a teacher model and you use the probability distribution of that teacher model to train your student model. And they have adopted that methodology for the smaller 9N 2 billion parameter models and they do pretty good. Now for the 27 billion model, they basically trained this model from scratch. So not they're not doing any knowledge destination. Okay, so I've been talking about about knowledge destination. So let's understand what exactly this is and how it works. So the normal way of training these models is to to use them as next word prediction models. So it's a single architecture which predicts the next token. Now the way it works is you have your language model and a sequence comes in. So for example, as sequence is the quick brown fox and the model has to make a prediction for the next token based on its training data. What it has seen there is a distribution of tokens, probabilities. So for example, in this case jumps has the higher probability to be the next word, right? So the model picks that word with highest probability, make that as a predicted token and then compares it with the ground truth or what in reality is supposed to be the token that is supposed to predict, right? So you have the predicted token, the ground truth token, and then you calculate the loss. That loss is used to update the model weights, right? So that's how we do pre training of these large language models. But there's another way. So instead of one model you can use two different models. In this case you have a teacher model which is a much bigger model, something like GPT4 or Gemini 1.5 Pro or even Cloudsonnet. And then you have the student model. That's the model that you want to fine tune or train, right? Now in this case what you do is you take your sequence, you pass it through the teacher model. The teacher model, will it use its token distribution or the probabilities that it has to pick the next most probable token, right? Then you do the same thing by looking at the probability distribution of the student model. But rather than computing the loss function for the student model based on its prediction, what you want to do is you want to use KL divergence. And the way KL divergence work is it compares the distributions from two different, let's say models. So you want the probability distribution of student model to be very close to the teacher model and you update the model weights based on the KL divergence. Now to understand KL divergence, here's a quick example. So for example, here the KL divergence of the teacher model. So this is the probability distribution that you want to adopt. And here's the probability distribution for next token of the student model. Now if you simply change the probability distribution of the student model to get closer to the distribution of teacher model. So teacher distribution, you will reduce the KL divergence. And during the training process this is your objective. So you will keep changing the weights of the student model in a way that it will try to adopt the distribution that is seeing from the teacher model. Now the good thing about this is that indirectly you can distill knowledge from this bigger teacher model to the smaller student model. And that's the technique that Google is using for the Gemma 2 models. Now interestingly enough, they do this both during pre training as well as the fine tuning stage when it comes to the smaller models. Now this is a way in which you can really improve the performance of these models without needing trillions and trillions of tokens. Now, I suspect that they're using the 27 billion model, probably as a teacher model or maybe even using the bigger Gemini models, but I think that detail is not present in the paper. So in terms of the training processes, where it's similar for all the models, you have the supervised, you have pre training, then supervised fine tuning. Then on top of it they have a policy model which does reinforcement learning from human feedback. And, and then they are doing model merging. So they basically train multiple versions of the models and then merge their weights together by averaging the model weights. And on top of that they have a safety layer which basically filter information based on whether it's unsafe outputs or there are duplicates. So here's how the prompt template looks like. So if you are to use this model, you want to make sure that this is the prompt template you follow. This is for the single turn conversation. So you have special tokens, then the role is user, then end of the turn, then you use the same tokens as the beginning of and you have the model role, then end of the tokens and you put end of the sequence. So if you were to use it for a second turn, so this is basically you just append that. So you need to make sure that there is this end of sequence token that will indicate the end of the conversation. So in the beginning of the video we said that Gemini 2 seems to be doing really well on Chatbot arena and this might be explained with the help of what they have used for training. So in the post training they talk about the data set that they used. So they say we extended the training the post training Data from Gemini 1.5 with a mixture of internal and external public data. In particular, we use the prompts but not the answer. Some Adamsis chat 1 million and all of our data go through the filter process. Now keep in mind, they took the prompts that are available on LMSIS chat and use those as inputs in their training process, but they didn't use the responses since they used that data during the supervised fine tuning process. So I assumed that the teacher model was used to generate responses for those prompts which are in the Adamsis chart data set. Now the benefit of this approach is that it will avoid biases from predetermined responses and encourage the model of creativity and flexibility. Since we are using the teacher model to generate responses for those inputs that are in the Syschart data set and we they're doing RLHF. So there is a tendency that the model might adopt to, especially like the teacher model might adopt to responses which are going to be favored by human annotators. And we are doing knowledge destination on top of it. So during the knowledge destination process, the student model will basically clone the behavior of the teacher model. Right. So it will start generating responses very similar to the teacher model, especially if you run it for a lot longer with a lot of details. Right. And as a result that it probably is going to learn how to answer these prompts that are present in the chatbot arena. And that could potentially explain why the responses from even a smaller models are liked more compared to the other open weight models. Now that's pure speculation on my part. It may or may not be actually happening, but this is something that I think will require a lot more experimentations. Now they also did ablation studies, which is very important. Especially they were trying to compare models which are trained from scratch versus which are trained using this knowledge distillation process. And this seems to show that for the smaller models, so for example, the 2.6 billion token, when it's trained from scratch on benchmarks, it's only able to get up to 60, which is the average score. But with knowledge destination process, it is able to get to 67.8, which is a substantial improvement. So if this holds true, it could be a way of actually training smaller models. And not only that it improves the score and benchmarks, but even the perplexity of models that are distilled from bigger models is much better compared to the models that are trained from scratch. So that's one very important, I think finding the second one is changing sliding window size. Right? So this show that even if you keep changing the sliding window size, the effect on perplexity is minimal. This means that during inference time you can change the sliding window size. You can get much better inference speed with a minimal reduction in performance. So this is going to be huge if it, if these results actually hold true. If you want to use this model, the easiest way is to use in Google AI Studio. So under the models now you have Gemma 2, you can interact with the model here. It's also available on Hugging Face. So the weights are available both for the 9 billion and 27 billion version. And I believe even the quantized versions are available now. So I'll create a subsequent video and in which we're going to test the model. I'll show you how to integrate it with your own code basis if that's something you're interested in, make sure to subscribe to the channel. Thanks for watching and as always, see you in the next one.