Introduction

Using different learning rates for different parameter groups is standard in PyTorch, especially for transfer learning and models with separate backbones and heads. Learning-rate schedulers work with those groups, but you need to understand whether the scheduler applies the same multiplier to every group or whether you must define separate behavior yourself.

Define Parameter Groups Explicitly

The optimizer is where group-specific base learning rates are declared.

1import torch
2import torch.nn as nn
3import torch.optim as optim
4
5
6class Net(nn.Module):
7    def __init__(self):
8        super().__init__()
9        self.backbone = nn.Sequential(
10            nn.Linear(20, 64),
11            nn.ReLU(),
12            nn.Linear(64, 32),
13            nn.ReLU(),
14        )
15        self.head = nn.Linear(32, 3)
16
17    def forward(self, x):
18        return self.head(self.backbone(x))
19
20
21model = Net()
22optimizer = optim.Adam(
23    [
24        {"params": model.backbone.parameters(), "lr": 1e-4},
25        {"params": model.head.parameters(), "lr": 1e-3},
26    ],
27    weight_decay=1e-5,
28)

Here the head learns ten times faster than the backbone, which is a common fine-tuning setup.

Standard Schedulers Scale All Groups from Their Own Base Rates

Most built-in schedulers do not erase the fact that the groups started at different rates. They apply a schedule to each group while preserving the relative difference.

scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

In a training loop:

1loss_fn = nn.CrossEntropyLoss()
2
3for epoch in range(12):
4    optimizer.zero_grad()
5    x = torch.randn(32, 20)
6    y = torch.randint(0, 3, (32,))
7
8    logits = model(x)
9    loss = loss_fn(logits, y)
10    loss.backward()
11    optimizer.step()
12    scheduler.step()
13
14    lrs = [group["lr"] for group in optimizer.param_groups]
15    print(f"epoch={epoch + 1}, lrs={lrs}")

Both groups decay on the same schedule, but the backbone remains lower than the head.

Use LambdaLR for Different Group Schedules

If the groups need different schedules rather than just different starting rates, use LambdaLR with one lambda function per parameter group.

1scheduler = optim.lr_scheduler.LambdaLR(
2    optimizer,
3    lr_lambda=[
4        lambda epoch: 1.0 if epoch < 5 else 0.5,
5        lambda epoch: 1.0 if epoch < 3 else 0.1,
6    ],
7)

This example keeps the backbone steady longer while decaying the head earlier and more aggressively. That pattern is useful when the task head converges quickly but the pretrained backbone needs smaller, more stable updates.

Log Effective Learning Rates During Training

Do not trust your mental model alone. Print the effective learning rates and confirm they follow the intended schedule.

1def current_lrs(opt):
2    return [pg["lr"] for pg in opt.param_groups]
3
4print("initial", current_lrs(optimizer))
5for epoch in range(5):
6    optimizer.step()
7    scheduler.step()
8    print("after epoch", epoch + 1, current_lrs(optimizer))

This is a simple check, but it catches many configuration mistakes before they affect model quality.

Save Scheduler State with the Optimizer

If you pause and resume training, save the scheduler state along with the model and optimizer. Otherwise resumed training can continue with the wrong rates.

1torch.save(
2    {
3        "model": model.state_dict(),
4        "optimizer": optimizer.state_dict(),
5        "scheduler": scheduler.state_dict(),
6    },
7    "checkpoint.pt",
8)

When loading, restore all three components so the schedule continues from the correct point.

Know When Step Timing Matters

One common source of confusion is when to call scheduler.step(). For many epoch-based schedulers, the typical pattern is to call it after optimizer.step(). Some schedulers are intended for batch-level updates instead, so always check the specific scheduler contract you are using.

The important part is consistency. If you designed the schedule per epoch, step it per epoch. If the scheduler is tied to batches, step it per batch.

Common Pitfalls

A common mistake is expecting a standard scheduler to produce different decay shapes for different groups automatically. Most of them apply the same rule to every group unless you configure separate behavior.

Another problem is forgetting to inspect the actual learning rates after a checkpoint restore. If the scheduler state was not restored, the optimizer may continue at the wrong stage of the schedule.

Developers also sometimes call scheduler.step() in the wrong place or at the wrong frequency, which shifts the schedule even though the code still runs.

Summary

• Use optimizer parameter groups to assign different base learning rates.
• Standard schedulers usually scale each group while preserving relative differences.
• Use LambdaLR when groups need genuinely different schedule shapes.
• Log learning rates during training to verify that the schedule matches your intent.
• Save and restore scheduler state together with model and optimizer state.