How do learning rate schedules and warmup work?

Learning rate schedules and warmup work by changing the step size of a neural network during training instead of keeping it fixed. Warmup starts with tiny steps and slowly increases them, and the schedule then slowly decreases the steps so that the model settles at a good answer.

In this blog, we will learn about how learning rate schedules and warmup work. We will also see why a fixed learning rate is not enough, how popular schedules change the learning rate over time, why we need warmup at the very start, how warmup and cosine decay work together, how to write it in code, and when to use which one.

We will cover the following:

• What is a learning rate?
• The problem with a fixed learning rate
• What is a learning rate schedule?
• Step decay
• Exponential decay
• Linear decay
• Cosine decay
• What is warmup?
• Why do we need warmup?
• Warmup plus cosine decay: the modern recipe
• Seeing it in code
• Comparison of schedules
• When to use which one

I am Amit Shekhar, Founder @ Outcome School, I have taught and mentored many developers, and their efforts landed them high-paying tech jobs, helped many tech companies in solving their unique problems, and created many open-source libraries being used by top companies. I am passionate about sharing knowledge through open-source, blogs, and videos.

I teach AI and Machine Learning at Outcome School.

Let's get started.

What is a learning rate?

Before jumping into schedules and warmup, we must know what a learning rate is.

A neural network is a model made of many numbers called weights. A weight is an importance number that decides how much one input matters. When a network starts, its weights are random, so its answers are wrong.

Training means fixing these weights little by little. At every training step, the network does the following:

• It makes a guess.
• It measures how wrong the guess is. This number is called the loss.
• It finds which direction each weight must move to reduce the loss. This direction is called the gradient.
• It moves each weight a small step in that direction.

The learning rate is the number that decides how big that small step is.

The rule is as below:

new_weight = old_weight - learning_rate x gradient

Here, the gradient gives the direction, and the learning rate gives the size of the step.

Let's take a real-world comparison. Imagine we are walking down a hill in thick fog to reach the lowest point of a valley. We cannot see far, but we can feel the slope under our feet. The slope is the gradient. The length of our stride is the learning rate. The lowest point of the valley is the best set of weights, where the loss is smallest.

To master Neural Networks, Gradient Descent, and Backpropagation, and to build a Neural Network from scratch, check out our AI and Machine Learning Program at Outcome School.

The problem with a fixed learning rate

Suppose we pick one learning rate and keep it the same for the whole training. Let's see what happens.

If the learning rate is big: We take long strides. We move fast at the start, which is great. But near the bottom, our long strides keep jumping over the lowest point. We bounce from one side of the valley to the other and never settle.

If the learning rate is small: We take tiny strides. Near the bottom, we settle nicely. But at the start, we move very slowly. Training takes a very long time.

 Big fixed steps:     o -------> o
                      o <------- o      (keeps jumping over the bottom)

 Small fixed steps:   o.o.o.o.o.o.o.o.o.o.o.o.o ...   (very slow)

 What we want:        o -----> o ---> o -> o > o.o   (fast first, careful later)

Here, we can see that neither a big nor a small fixed step is ideal. We want big steps when we are far from the bottom and small steps when we are close to it.

So, here comes the learning rate schedule to the rescue.

What is a learning rate schedule?

A learning rate schedule is a plan that changes the learning rate during training, usually starting high and going lower as training goes on.

It works like a timetable for the step size. It says, "At step 0, use this learning rate. At step 1000, use that learning rate."

The word decay is used a lot here. Decay means "slowly going down". When we say "learning rate decay", we mean the learning rate goes down over time.

There are many schedules. They differ in how they go down: in sudden drops, in a smooth curve, or in a straight line. Let's learn about the popular ones one by one.

Step decay

Step decay keeps the learning rate fixed for a while, then suddenly cuts it by a fixed factor, again and again.

Let's say we start with a learning rate of 0.1 and cut it in half every 10 rounds over the training data. One round over the full training data is called an epoch.

Epoch  1 to 10:  learning rate = 0.1
Epoch 11 to 20:  learning rate = 0.05
Epoch 21 to 30:  learning rate = 0.025

Let's see it in a picture:

 lr
 0.1   |________
       |        |
 0.05  |        |________
       |                 |
 0.025 |                 |________
       +------------------------------ epochs
        0       10       20       30

Here, we can see the learning rate looks like a staircase going down.

Advantage: Very simple to understand and set up.

Disadvantage: The sudden drops are rough. Also, we must choose when to drop and by how much, which needs trial and error.

The issue with this approach is the sudden jumps. Let's see how the next approach solve this issue.

Exponential decay

Exponential decay multiplies the learning rate by a number slightly less than 1 at every step or every epoch, so it goes down smoothly.

Let's say we multiply by 0.9 every epoch, starting at 0.1:

Epoch 1:  0.1
Epoch 2:  0.1   x 0.9 = 0.09
Epoch 3:  0.09  x 0.9 = 0.081
Epoch 4:  0.081 x 0.9 = 0.0729

Here, we can notice that there are no sudden drops anymore. The learning rate slides down smoothly.

Disadvantage: It drops very fast at the beginning and then becomes almost flat. Sometimes, the learning rate becomes tiny too early, and the model stops improving while there is still a lot to learn.

Planning your move into AI?

No matter which tech domain you work in today, with your effort, you can transition to these roles:

• AI Engineer
• Forward Deployed Engineer
• AI Agent Engineer
• Harness Engineer
• Voice AI Engineer
• AI Architect
• LLM Engineer
• LLM Inference Engineer
• Gen AI Engineer
• MLOps Engineer
• ML Engineer
• Deep Learning Engineer

Our AI and Machine Learning Program at Outcome School takes you there.

Linear decay

Linear decay lowers the learning rate in a straight line, from a starting value down to a final value at the end of training.

 lr
 max |\
     |  \
     |    \
     |      \
     |        \
 min |          \
     +------------ steps

Here, we can see that the learning rate goes down by the same amount at every step.

Advantage: Simple and predictable. We need to know the starting value, the final value, and the total number of steps.

Linear decay is popular when fine-tuning language models. Fine-tuning means taking an already trained model and training it a little more on our own data.

Cosine decay

Cosine decay lowers the learning rate along a smooth curve shaped like half of a wave: slow at the start, faster in the middle, and slow again at the end.

The name comes from the cosine function in math, which draws a smooth wave. Here, we need to understand only the shape.

 lr
 max |---__
     |      `-.
     |         `.
     |           `.
     |             `-.
 min |                 `---___
     +---------------------------- steps

Here, we can see three phases:

• At the start: The learning rate stays high for a while, so the model keeps learning fast.
• In the middle: The learning rate goes down steadily.
• At the end: The learning rate stays low for a while, so the model can settle gently into the bottom of the valley.

Cosine decay works very well in practice, and it needs very few settings: the maximum learning rate, the minimum learning rate, and the total number of steps.

Now, it's time to learn about warmup, which raises the learning rate at the start.

What is warmup?

Warmup is a short phase at the very beginning of training where the learning rate starts very small and slowly increases up to its full value.

Let's say our target learning rate is 0.001 and we use 100 warmup steps. Then, the learning rate climbs like below:

Step   0:  0.00001
Step  49:  0.0005
Step  99:  0.001   (warmup is done, full learning rate reached)

Here, we can see that the learning rate grows a little at every step, in a straight line, until it reaches the target value. This is called linear warmup, and it is the most common type.

A quick note for you

No matter which tech domain you work in, get familiar with these topics:

• LLM
• RAG
• MCP
• Agent
• Fine-tuning
• Quantization

We put it all together in one video:

AI Engineering Explained: LLM, RAG, MCP, Agent, Fine-Tuning, and Quantization

No need to stop reading - bookmark it and watch later when you get time. Future you will thank you.

Now, let's get back to the topic.

Why do we need warmup?

Now, the next big question is: why should we start slow, when the beginning is the place where we are farthest from the answer?

The answer is that the very beginning of training is the most unstable moment. There are two main reasons.

Reason 1: The weights are random, so the gradients are wild.

At step 0, the network knows nothing. Its guesses are badly wrong, so the loss is very high. The gradients at this moment can be very large and point in strange directions. If we take a big step using such a wild gradient, the weights jump to a bad place. Sometimes, the loss even explodes to a huge number, and the training crashes. This crash is called divergence. We have a detailed blog on Vanishing and Exploding Gradients that explains why gradients can grow so large.

Reason 2: Optimizers like Adam need time to collect good averages.

An optimizer is the rule that updates the weights. One of the most popular optimizers, called Adam, keeps a running average of the past gradients to decide how big each step must be. At step 0, there is no history. These averages are based on only a few steps, so they are not reliable yet. If the learning rate is already high, these unreliable averages lead to bad, oversized steps.

Warmup gives the optimizer time to collect good averages before we let it take big steps.

Note: Warmup is especially important for Transformers, which are the architecture behind LLMs (Large Language Models). The original Transformer paper, "Attention Is All You Need", used 4000 warmup steps. Without warmup, Transformers often fail to train properly.

Warmup plus cosine decay: the modern recipe

Now, it's time to put the schedules and warmup together.

Many modern LLMs are trained with this recipe:

• Phase 1, Warmup: Increase the learning rate in a straight line from almost 0 to the maximum value.
• Phase 2, Cosine decay: Lower the learning rate along the smooth cosine curve, from the maximum value down to a small minimum value.

 lr
 max |      /---__
     |     /      `-.
     |    /          `.
     |   /             `.
     |  /                `-.
 min | /                     `---___
     +/------------------------------- steps
      |warmup|     cosine decay      |

Here, we can see the full journey of the learning rate. It climbs up carefully, stays high while the model learns the most, and then comes down smoothly so the model can settle.

This is how learning rate schedules and warmup work together.

Stay updated: Subscribe to our newsletter to get our latest AI and Machine Learning blogs straight to your inbox.

Seeing it in code

The best way to learn this is by taking an example. Let's write a small Python function that gives the learning rate for any training step, using warmup plus cosine decay.

import math

def get_lr(step, max_lr=0.001, min_lr=0.0001, warmup_steps=100, total_steps=1000):
    # Phase 1: linear warmup
    if step < warmup_steps:
        return max_lr * (step + 1) / warmup_steps

    # Phase 2: cosine decay
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    cosine = 0.5 * (1 + math.cos(math.pi * progress))
    return min_lr + (max_lr - min_lr) * cosine

for step in [0, 49, 99, 325, 550, 775, 999]:
    print(f"step {step:4d}  lr = {get_lr(step):.6f}")

Here, we have:

• max_lr is the highest learning rate, which we reach at the end of warmup.
• min_lr is the lowest learning rate, which we reach at the very end of training.
• In the warmup phase, the learning rate grows in a straight line with the step number. We use step + 1 so that the very first step does not have a learning rate of exactly 0.
• In the cosine phase, progress goes from 0 to 1 as training moves from the end of warmup to the last step.
• cosine goes smoothly from 1 down to 0 as progress goes from 0 to 1. So, the learning rate goes smoothly from max_lr down to min_lr.

This will print the following:

step    0  lr = 0.000010
step   49  lr = 0.000500
step   99  lr = 0.001000
step  325  lr = 0.000868
step  550  lr = 0.000550
step  775  lr = 0.000232
step  999  lr = 0.000100

Here, we can notice that the learning rate starts tiny at step 0, reaches its maximum of 0.001 at step 99, stays high for a while, reaches the halfway value of 0.00055 in the middle at step 550, and finally ends at the minimum of 0.0001.

In a real training loop, we call this function at every step and give the result to the optimizer as below:

for step in range(total_steps):
    lr = get_lr(step)
    for group in optimizer.param_groups:
        group["lr"] = lr  # set the new learning rate

    # your code for the forward pass, loss, and backward pass
    optimizer.step()

Here, we are using PyTorch. The optimizer is one we have already created, for example AdamW, and total_steps is 1000 like before. PyTorch keeps the settings of an optimizer, including the learning rate, inside param_groups. Before every update, we set the learning rate in each group to the value from our schedule. Then, optimizer.step() updates the weights using that learning rate.

Note: PyTorch also has ready-made schedulers inside torch.optim.lr_scheduler, such as StepLR, ExponentialLR, and CosineAnnealingLR. Writing our own function, like above, makes it clear what is happening.

If we want to go deep into PyTorch, Optimizers, and Model Development and Training, we have a complete program on it. Check out our AI and Machine Learning Program at Outcome School to learn it end to end.

Comparison of schedules

Let me tabulate the differences between these schedules for your better understanding.

When to use which one

Now, let's see the practical guidance.

• Training a Transformer or an LLM from scratch: Use linear warmup followed by cosine decay. This is the most common recipe.
• Fine-tuning a pre-trained language model: Use a short warmup followed by linear decay. This is the default in many popular fine-tuning tools.
• Training a small model for learning or quick experiments: A fixed learning rate or step decay is fine.
• Training crashes or the loss explodes in the first few steps: Add warmup or make the warmup longer.
• The loss stops improving at the end of training: Lower the minimum learning rate or let the decay go on longer.

A common choice is to keep the warmup short, often a few percent of the total training steps. Warmup plus cosine decay is a good starting point. Later, depending on our use case, we can experiment.

Now, we must have understood how learning rate schedules and warmup work.

Frequently Asked Questions

How many warmup steps should we use?

A common choice is to keep the warmup short, often a few percent of the total training steps. The original Transformer paper, "Attention Is All You Need", used 4000 warmup steps. If training crashes or the loss explodes in the first few steps, we can make the warmup longer.

Does PyTorch have built-in learning rate schedulers?

Yes. PyTorch has ready-made schedulers inside torch.optim.lr_scheduler, such as StepLR, ExponentialLR, and CosineAnnealingLR. We can also write our own function that returns the learning rate for each step and set it in the optimizer's param_groups before every update, which makes it clear what is happening.

What happens if we train a Transformer without warmup?

Without warmup, Transformers often fail to train properly. At the start, the weights are random and the gradients can be very large, so big steps can push the weights to a bad place. The loss can even explode to a huge number, and the training crashes. This crash is called divergence.

What settings does cosine decay need?

Cosine decay needs very few settings: the maximum learning rate, the minimum learning rate, and the total number of steps. Because it needs the total number of steps, we must know how long training will run before we start. With warmup, we also choose the number of warmup steps.

That's it for now.

Thanks

Amit Shekhar
Founder @ Outcome School

You can connect with me on:

• X
• LinkedIn
• GitHub

Follow Outcome School on:

• X
• YouTube
• LinkedIn
• GitHub

Read all of our high-quality blogs here.