Why regularization controls overfitting in PyTorch - Experiment to Prove It

Experiment - Why regularization controls overfitting

Problem:We want to train a neural network to classify handwritten digits from the MNIST dataset. The current model fits the training data very well but performs poorly on new data.

Current Metrics:Training accuracy: 98%, Validation accuracy: 82%, Training loss: 0.05, Validation loss: 0.45

Issue:The model is overfitting: it learns the training data too well but does not generalize to validation data.

Your Task

Reduce overfitting by applying regularization techniques so that validation accuracy improves to at least 90% while keeping training accuracy below 95%.

You can only add L2 weight decay and dropout layers to the model.

Do not change the model architecture or dataset.

Keep the number of training epochs the same.

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 the neural network with dropout
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 256)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# Load data
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 with weight decay (L2 regularization)
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)  # L2 weight decay

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)

# Train for 10 epochs
for epoch in range(10):
    train_loss, train_acc = train()
    val_loss, val_acc = validate()
    print(f'Epoch {epoch+1}: Train loss {train_loss:.4f}, Train acc {train_acc:.4f}, Val loss {val_loss:.4f}, Val acc {val_acc:.4f}')

Added dropout layer with 0.5 probability after the first fully connected layer to reduce neuron co-adaptation.

Added L2 weight decay (1e-4) to the Adam optimizer to penalize large weights and encourage simpler models.

Results Interpretation

Before regularization: Training accuracy was 98%, validation accuracy was 82%. The large gap shows overfitting.

After regularization: Training accuracy dropped to 93%, validation accuracy improved to 91%. The gap narrowed, showing better generalization.

Regularization techniques like dropout and L2 weight decay help prevent the model from memorizing training data. This leads to better performance on new, unseen data by encouraging simpler, more general patterns.

Bonus Experiment

Try using early stopping to stop training when validation loss stops improving to further reduce overfitting.

💡 Hint

Monitor validation loss each epoch and stop training if it does not improve for 3 consecutive epochs.

Practice

1. Why does regularization help prevent overfitting in a PyTorch model?

easy

A. It keeps the model weights small by adding a penalty to the loss.

B. It increases the size of the training dataset automatically.

C. It removes layers from the neural network during training.

D. It speeds up the training process by skipping some data points.

Solution

1. Step 1: Understand what overfitting means

   Overfitting happens when a model learns the training data too well, including noise, causing poor performance on new data.

2. Step 2: Explain how regularization affects model weights

   Regularization adds a penalty to large weights, encouraging smaller weights that generalize better to new data.

3. Final Answer:

   It keeps the model weights small by adding a penalty to the loss. -> Option A

4. Quick Check:

   Regularization = penalty on weights = less overfitting [OK]

Hint: Regularization adds penalty to weights to reduce overfitting [OK]

Common Mistakes:

• Thinking regularization increases data size
• Believing regularization removes layers
• Assuming regularization speeds training

2. Which PyTorch code snippet correctly applies L2 regularization (weight decay) during optimizer setup?

easy

A. optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.1)

B. optimizer = torch.optim.SGD(model.parameters(), lr=0.01, dropout=0.1)

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

D. optimizer = torch.optim.SGD(model.parameters(), lr=0.01, decay=0.1)

Solution

1. Step 1: Identify correct parameter for L2 regularization in PyTorch

   PyTorch uses weight_decay in optimizers to apply L2 regularization.

2. Step 2: Check the code options for correct usage

   Only optimizer = torch.optim.SGD(model.parameters(), lr=0.01, weight_decay=0.1) uses weight_decay=0.1, which is the correct way to add L2 regularization.

3. Final Answer:

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

4. Quick Check:

   weight_decay = L2 regularization in PyTorch [OK]

Hint: Use weight_decay param for L2 regularization in PyTorch optimizers [OK]

Common Mistakes:

• Using dropout parameter in optimizer
• Confusing momentum with regularization
• Using decay instead of weight_decay

3. Consider this PyTorch training loop snippet with L2 regularization applied:

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.01)
for data, target in dataloader:
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

What effect does the weight_decay=0.01 have during training?

medium

A. It adds a penalty to large weights, helping reduce overfitting.

B. It increases the learning rate by 0.01 each step.

C. It drops 1% of neurons randomly during training.

D. It stops training early when loss is below 0.01.

Solution

1. Step 1: Understand weight_decay in optimizer

   The weight_decay parameter adds L2 regularization, penalizing large weights during training.

2. Step 2: Identify the effect on training

   This penalty helps the model avoid overfitting by keeping weights smaller and more generalizable.

3. Final Answer:

   It adds a penalty to large weights, helping reduce overfitting. -> Option A

4. Quick Check:

   weight_decay = L2 penalty = less overfitting [OK]

Hint: weight_decay adds penalty to weights, not learning rate or dropout [OK]

Common Mistakes:

• Confusing weight_decay with learning rate changes
• Thinking weight_decay is dropout
• Assuming weight_decay controls early stopping

4. You have this PyTorch code snippet intended to apply L2 regularization:

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for data, target in dataloader:
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target) + 0.01 * torch.sum(model.parameters())
    loss.backward()
    optimizer.step()

What is wrong with this code regarding regularization?

medium

A. It uses SGD optimizer which does not support regularization.

B. It forgets to call optimizer.zero_grad() before backward.

C. It applies regularization after optimizer.step(), so no effect.

D. It incorrectly sums parameters instead of their squares for L2 penalty.

Solution

1. Step 1: Check how L2 regularization is computed

   L2 regularization requires summing the squares of parameters, not just their values.

2. Step 2: Analyze the code's regularization term

   The code sums parameters directly with torch.sum(model.parameters()), which is incorrect for L2 penalty.

3. Final Answer:

   It incorrectly sums parameters instead of their squares for L2 penalty. -> Option D

4. Quick Check:

   L2 penalty = sum of squares, not sum of values [OK]

Hint: L2 regularization sums squares of weights, not weights themselves [OK]

Common Mistakes:

• Summing parameters instead of squared parameters
• Thinking SGD can't use regularization
• Misplacing optimizer.zero_grad() call

5. You train two PyTorch models on the same dataset: Model A uses no regularization, Model B uses L2 regularization with weight_decay=0.05. After training, Model A has training accuracy 98% but test accuracy 70%, while Model B has training accuracy 90% and test accuracy 85%. What explains this difference?

hard

A. Model A's higher training accuracy means it will always perform better on test data.

B. Model B's regularization reduced overfitting by keeping weights smaller, improving test accuracy.

C. Model B used a larger learning rate, causing better generalization.

D. Model A trained longer, so it has better test accuracy.

Solution

1. Step 1: Compare training and test accuracies

   Model A fits training data very well but performs poorly on test data, indicating overfitting.

2. Step 2: Understand effect of L2 regularization on Model B

   Model B has lower training accuracy but better test accuracy because regularization keeps weights smaller, improving generalization.

3. Final Answer:

   Model B's regularization reduced overfitting by keeping weights smaller, improving test accuracy. -> Option B

4. Quick Check:

   Regularization = smaller weights = better test accuracy [OK]

Hint: Better test accuracy with regularization means less overfitting [OK]

Common Mistakes:

• Assuming higher training accuracy means better test accuracy
• Confusing learning rate with regularization effect
• Ignoring the role of weight size in generalization