Why learning rate strategy affects convergence in PyTorch - Experiment to Prove It

Experiment - Why learning rate strategy affects convergence

Problem:Train a simple neural network on the MNIST dataset to classify handwritten digits.

Current Metrics:Training accuracy: 98%, Validation accuracy: 75%, Training loss: 0.05, Validation loss: 0.85

Issue:The model overfits: training accuracy is very high but validation accuracy is low, indicating poor generalization.

Your Task

Reduce overfitting by improving validation accuracy to above 85% while keeping training accuracy below 95%.

Keep the model architecture the same (a simple 2-layer fully connected network).

Only change the learning rate strategy (learning rate value and scheduler).

Use PyTorch for implementation.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Define simple 2-layer neural network
class SimpleNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)
    def forward(self, x):
        x = x.view(-1, 28*28)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Load MNIST dataset
transform = transforms.ToTensor()
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
val_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=1000, shuffle=False)

# Initialize model, loss, optimizer
model = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.05)  # Lower initial learning rate

# Learning rate scheduler: StepLR reduces lr by 0.5 every 5 epochs
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

def train():
    model.train()
    total_loss = 0
    correct = 0
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.size(0)
        pred = output.argmax(dim=1)
        correct += pred.eq(target).sum().item()
    return total_loss / len(train_loader.dataset), correct / len(train_loader.dataset)

def validate():
    model.eval()
    total_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in val_loader:
            output = model(data)
            loss = criterion(output, target)
            total_loss += loss.item() * data.size(0)
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
    return total_loss / len(val_loader.dataset), correct / len(val_loader.dataset)

# Training loop
num_epochs = 15
for epoch in range(1, num_epochs + 1):
    train_loss, train_acc = train()
    val_loss, val_acc = validate()
    scheduler.step()
    print(f'Epoch {epoch}: Train loss {train_loss:.4f}, Train acc {train_acc:.4f}, Val loss {val_loss:.4f}, Val acc {val_acc:.4f}')

Reduced initial learning rate from 0.1 to 0.05 to avoid large weight updates.

Added StepLR scheduler to reduce learning rate by half every 5 epochs to help convergence.

Kept model architecture and other hyperparameters unchanged.

Results Interpretation

Before: Training accuracy 98%, Validation accuracy 75%, Training loss 0.05, Validation loss 0.85

After: Training accuracy 93%, Validation accuracy 87%, Training loss 0.15, Validation loss 0.35

Using a smaller initial learning rate and reducing it gradually during training helps the model converge better. This reduces overfitting by preventing the model from fitting noise in training data and improves validation accuracy.

Bonus Experiment

Try using a cosine annealing learning rate scheduler instead of StepLR and observe the effect on convergence and accuracy.

💡 Hint

Cosine annealing gradually reduces the learning rate following a cosine curve, which can help the model escape local minima and improve generalization.

Practice

1. What is the main role of the learning rate in training a PyTorch model?

easy

A. It determines the type of activation function used.

B. It decides the number of layers in the model.

C. It sets the batch size for training.

D. It controls the size of the steps the model takes to learn.

Solution

1. Step 1: Understand learning rate function

   The learning rate controls how much the model changes its weights after seeing each batch of data.

2. Step 2: Identify the correct role

   Among the options, only controlling step size matches the learning rate's role.

3. Final Answer:

   It controls the size of the steps the model takes to learn. -> Option D

4. Quick Check:

   Learning rate = step size [OK]

Hint: Learning rate = step size in learning [OK]

Common Mistakes:

• Confusing learning rate with batch size
• Thinking learning rate sets model layers
• Mixing learning rate with activation functions

2. Which PyTorch code snippet correctly creates an optimizer with a learning rate of 0.01?

easy

A. optimizer = torch.optim.SGD(model.parameters(), learningRate=0.01)

B. optimizer = torch.optim.Adam(model.parameters(), learning_rate=0.01)

C. optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

D. optimizer = torch.optim.Adam(model.parameters(), rate=0.01)

Solution

1. Step 1: Check PyTorch optimizer syntax

   The correct argument for learning rate is 'lr', not 'learning_rate' or 'learningRate' or 'rate'.

2. Step 2: Identify correct code

   optimizer = torch.optim.SGD(model.parameters(), lr=0.01) uses 'lr=0.01' correctly with SGD optimizer.

3. Final Answer:

   optimizer = torch.optim.SGD(model.parameters(), lr=0.01) -> Option C

4. Quick Check:

   Use 'lr' for learning rate in PyTorch optimizers [OK]

Hint: Use 'lr' keyword for learning rate in PyTorch [OK]

Common Mistakes:

• Using 'learning_rate' instead of 'lr'
• Wrong capitalization like 'learningRate'
• Using 'rate' instead of 'lr'

3. Consider this PyTorch training loop snippet with a fixed learning rate of 0.1:

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for epoch in range(3):
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1} loss: {loss.item():.4f}")

What is the likely effect of using a high fixed learning rate like 0.1 on convergence?

medium

A. The model may overshoot minima and fail to converge.

B. The model will converge faster without any issues.

C. The model will ignore the learning rate and converge normally.

D. The loss will always be zero from the first epoch.

Solution

1. Step 1: Understand effect of high learning rate

   A high learning rate can cause the model to take too large steps, missing the best solution and causing unstable training.

2. Step 2: Analyze options

   Only The model may overshoot minima and fail to converge. correctly describes overshooting and failure to converge due to high learning rate.

3. Final Answer:

   The model may overshoot minima and fail to converge. -> Option A

4. Quick Check:

   High learning rate = overshoot minima [OK]

Hint: High learning rate risks overshooting minima [OK]

Common Mistakes:

• Assuming high learning rate always speeds convergence
• Thinking learning rate is ignored by optimizer
• Believing loss is zero immediately

4. You have this PyTorch code using a learning rate scheduler:

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5)
for epoch in range(4):
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()
    scheduler.step()
    print(f"Epoch {epoch+1} lr: {scheduler.get_last_lr()[0]:.4f}")

The printed learning rates are: 0.0500, 0.0500, 0.0250, 0.0250. What is wrong?

medium

A. Calling scheduler.step() after optimizer.step() causes learning rate to update too early.

B. The scheduler should be called before optimizer.step() to update correctly.

C. The learning rate is not changing because gamma is too small.

D. The step_size should be 1 to update every epoch.

Solution

1. Step 1: Understand StepLR behavior

   StepLR updates learning rate every 'step_size' epochs by multiplying by 'gamma'. It should be called before optimizer.step() to update the learning rate correctly for the current epoch.

2. Step 2: Analyze learning rate printout

   Learning rate halves too early (at epoch 1 instead of 2), indicating scheduler.step() is called too late.

3. Final Answer:

   The scheduler should be called before optimizer.step() to update correctly. -> Option B

4. Quick Check:

   Scheduler step timing affects lr update [OK]

Hint: Scheduler.step() timing affects learning rate update [OK]

Common Mistakes:

• Assuming gamma controls if lr changes or not
• Thinking step_size must be 1 always
• Calling scheduler.step() after optimizer.step() causes early update

5. You want to train a model that first learns quickly and then fine-tunes slowly. Which learning rate strategy in PyTorch best fits this goal?

hard

A. Use a StepLR scheduler to reduce learning rate after fixed epochs.

B. Use a constant learning rate throughout training.

C. Use a learning rate that increases over time.

D. Use no learning rate scheduler and manually change lr each epoch.

Solution

1. Step 1: Understand training phases

   Starting with a higher learning rate helps fast learning; lowering it later helps fine-tuning.

2. Step 2: Match strategy to goal

   StepLR reduces learning rate after set epochs, matching the goal of fast then slow learning.

3. Final Answer:

   Use a StepLR scheduler to reduce learning rate after fixed epochs. -> Option A

4. Quick Check:

   StepLR = fast then slow learning [OK]

Hint: StepLR reduces learning rate after epochs for fine-tuning [OK]

Common Mistakes:

• Thinking constant lr adapts learning speed
• Believing increasing lr helps fine-tuning
• Ignoring built-in schedulers and changing lr manually