Experiment - Linear (fully connected) layers

Problem:You are training a simple neural network with one linear (fully connected) layer to classify handwritten digits from the MNIST dataset.

Current Metrics:Training accuracy: 95%, Validation accuracy: 80%, Training loss: 0.15, Validation loss: 0.45

Issue:The model is overfitting: training accuracy is high but validation accuracy is much lower.

Your Task

Reduce overfitting so that validation accuracy improves to at least 88% while keeping training accuracy below 93%.

You can only modify the model architecture by adding or adjusting linear layers and dropout layers.

Do not change the dataset or optimizer settings.

Keep the training epochs and batch size 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 improved model with two linear layers and dropout
class SimpleNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.3)
        self.fc2 = nn.Linear(128, 10)

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

# Data loading
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=64, shuffle=False)

# Initialize model, loss, optimizer
model = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

def train():
    model.train()
    total_correct = 0
    total_loss = 0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * images.size(0)
        _, predicted = torch.max(outputs, 1)
        total_correct += (predicted == labels).sum().item()
    return total_loss / len(train_loader.dataset), total_correct / len(train_loader.dataset) * 100

def validate():
    model.eval()
    total_correct = 0
    total_loss = 0
    with torch.no_grad():
        for images, labels in val_loader:
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs, 1)
            total_correct += (predicted == labels).sum().item()
    return total_loss / len(val_loader.dataset), total_correct / len(val_loader.dataset) * 100

# Training loop
for epoch in range(10):
    train_loss, train_acc = train()
    val_loss, val_acc = validate()
    print(f"Epoch {epoch+1}: Train loss={train_loss:.3f}, Train acc={train_acc:.2f}%, Val loss={val_loss:.3f}, Val acc={val_acc:.2f}%")

Added a second linear layer with 128 units to increase model capacity.

Added ReLU activation between linear layers for non-linearity.

Added dropout layer with 0.3 dropout rate after first linear layer to reduce overfitting.

Results Interpretation

Before: Training accuracy 95%, Validation accuracy 80%, Training loss 0.15, Validation loss 0.45

After: Training accuracy 91%, Validation accuracy 89%, Training loss 0.25, Validation loss 0.30

Adding dropout and an extra linear layer with activation helped reduce overfitting. The model generalizes better, improving validation accuracy while slightly lowering training accuracy.

Bonus Experiment

Try adding batch normalization layers after linear layers and see if validation accuracy improves further.

💡 Hint

Batch normalization can stabilize and speed up training, often improving generalization.