LR not decaying for pytorch AdamW even after hundreds of epochs

Question

I have the following code using AdamW optimizer from Pytorch:

optimizer = AdamW(params=self.model.parameters(), lr=0.00005)

I tried to log in using wandb as follows:

lrs = {f'lr_group_{i}': param_group['lr']
       for i, param_group in enumerate(self.optimizer.param_groups)}
wandb.log({"train_loss": avg_train_loss, "val_loss": val_loss, **lrs})

Note that default value for weight_decay parameter is 0.01 for AdamW.

When I checked wandb dashboard, it showed the same LR for AdamW even after 200 epochs and it was not decaying at all. I tried this in several runs.

Why the LR decay is not happening?

Also, it was showing LR for only one param group. Why is it so? It seems that I miss something basic here. Can someone point out?

Answer 1

As already hinted at in @xdurch0's comment, I think you have fallen trap to a misconception here. Weight decay and decaying the learning rate are two different things – at least conceptually¹:

• Weight decay is used for regularizing the weights in your network by penalizing large weight values. It is used as a means to prevent overfitting and thus to help generalization.
• By decaying the learning rate, one hopes to converge to a good solution: it follows the idea that a larger initial learning rate helps "exploring the loss landscape", while decreasing it in later stages of training helps "homing in" on the final solution.

Crucially, altering the weight decay value does not alter the learning rate value!

Use a learning rate scheduler

To achieve a decaying learning rate in PyTorch, one usually makes use of a learning rate scheduler, typically in the following way:

1. One creates a scheduler instance that receives, as one of its inputs, the optimizer instance.
2. At the end of each training epoch, one calls step() on the scheduler (just like calling step() on the optimizer after processing a training batch). This will trigger the scheduler to calculate the new learning rate value for the next epoch, which it will then use for adjusting the optimizer accordingly.

The following is an excerpt of the PyTorch documentation, demonstrating this approach (comments added by me):

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# 1. Create scheduler, provide it with optimizer
scheduler = ExponentialLR(optimizer, gamma=0.9)

for epoch in range(20):
    for input_, target in dataset:
        optimizer.zero_grad()
        output = model(input_)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
    # 2. Trigger calculation and setting of next learning rate value
    scheduler.step()

Here, an ExponentialLR scheduler is used, which decays the learning rate exponentially, according to a fixed schedule (i.e. only depending on its initial parameters; in particular, the decay factor, but independent of the current network performance). There are also adaptive schedulers, which adjust the learning rate depending on the actual network performance; for example, ReduceLROnPlateau reduces the learning rate by a given factor as soon as a loss plateau of a given length has been detected.

_{¹) Mathematically, for some optimizers, learning rate and weight decay are implicitly coupled, which is one of the reasons why AdamW was derived from the Adam optimizer in the first place. For more information, one might want to read the AdamW paper. But this is not the problem described in the question.}