Feed-Forward Net — Formation Energy

What you will learn

• Build a configurable MLP in PyTorch using nn.Module, nn.Linear, and nn.ReLU
• Implement a full training loop: forward pass, MSE loss, backward pass, Adam step
• Apply early stopping on a validation set to prevent overfitting
• Use a learning-rate scheduler (ReduceLROnPlateau) and plot the loss and LR curves

Instructions

Getting started

Open the Colab notebook using the button above. Run each cell in order; cells marked Exercise require you to fill in code.

1. Set up the environment. Install torch, numpy, and matplotlib. Confirm torch.cuda.is_available() and choose the appropriate device.

2. Load and preprocess the dataset. Download the OQMD formation energy dataset subset. Standardise input features (composition-derived descriptors) using the training set statistics only.

3. Build the MLP. Complete the FormationEnergyMLP class: add hidden layers, ReLU activations, and dropout. Expose layer widths and dropout rate as constructor arguments.

4. Implement the training loop. Fill in the train_one_epoch() and evaluate() functions. Log train and validation MAE each epoch. Implement early stopping with patience=20.

5. Tune and analyse. Experiment with depth (2 vs. 4 layers) and width (128 vs. 512 neurons). Plot train/val MAE curves. Identify the epoch where the best validation MAE occurs.

Questions

Answer these questions as you work through the notebook. Discuss with your neighbour — some have no single right answer.

Warm-up
What is the MAE of a baseline model that always predicts the training set mean? How does your MLP compare?
Easy
Core
How does adding dropout (p=0.3) affect training vs. validation MAE at convergence? Does it help or hurt, and why?
Medium
Core
Plot the gradient norm of the first layer over training. Do you observe vanishing or exploding gradients? How would you fix either?
Medium
Challenge
Replace the Adam optimiser with SGD + momentum (lr=0.01, momentum=0.9) with a cosine annealing schedule. Compare convergence speed and final MAE to Adam.
Challenge

Resources