Learning rate scheduling with TensorFlow

Key Takeaways

in our other videos we talked about the basics of fine-tuning a language model with tensorflow and as always when i refer to videos i'll link them below but still can we do better so here's the code from our model fine tuning video and while while it works we could definitely tweak a couple of things by far the most important thing is the learning rate in this video we'll talk about how to change it which will make your training much more consistently successful in fact there are two things we want to change about the default learning rate for adam the first is that it's way too high for our models so by default adam uses a learning rate of 10 to the minus 3 which is very high for training transformers we're going to start at 5 by 10 to the minus 5 which is 20 times lower than the default and secondly we don't just want a constant learning rate we can get even better performance if we decay the learning rate down to a tiny value or even to zero over the course of training so that's what this polynomial decay schedule thing is doing that name might be intimidating especially if you only vaguely remember what a polynomial is for maths class so i'll show you what that decay looks like in a second but first we need to tell the scheduler how long training is going to be so that it decays at the right speed and that's what this code here is doing so we're computing how many mini batches the model is going to see over the entire training run and to compute that we're taking the size of the training set dividing it by the batch size which gives us the number of batches per epoch and then we're multiplying that by the number of epochs to get the total number of batches it's going to see over the whole training run so once we know how many batches how many training steps we're taking we just pass all of that information to the scheduler and we're ready to go so what does the polynomial decay schedule look like with default options it's actually just a linear schedule so it looks like this it starts at our initial value which is 5 by 10 to the minus 5 or five e minus five and then it decays down at a constant rate until it hits zero right at the very end of training so why do they call it polynomial and not linear well if you tweak the options you can get a higher order a truly polynomial decay schedule but there's no need to do that right now by default you get a linear schedule and if you were aware that a linear function is a special case of a polynomial you can feel proud so that aside how do we actually use the scheduler so this is easily we just pass it to adam you'll notice the first time when we compiled the model we just passed the string atom keras recognizes the names of common optimizers and loss functions if you pass them as strings so it saves time and it avoids imports to do it that way if you only want the default settings but we're professional machine learners now with our very own learning rate schedule so we have to do things properly so the first thing we do is we import the optimizer then we initialize it with our scheduler in the learning rate argument and then we compile the model using our new optimizer and whatever loss function you want we'll leave that unchanged this will be sparse categorical cross-entropy if you're following along from the fine-tuning video but it can be anything else that you're using yourself so now we have a high performance model ready to go all that remains is to fit the model just like we did before and remember because we've compiled the model with the new optimizer and the new learning rate we actually don't need to change anything about the fit.call at all we just call fit here exactly the same command we used before or if you've seen in other videos but now we get a beautiful training with a nice smooth a good initial learning rate and a solid learning rate decay and you will get much better performance as a result