I’m working on a project to learn the PyTorch library hands-on. I’m currently tackling a regression task with a dataset containing 14 continuous features. My model’s training loss is plateauing and remains far from the validation loss, so I suspect an issue with the architecture.

I’m using a feedforward neural network in PyTorch, tuned with Optuna for learning rate, hidden layers, hidden sizes, and batch size. The resulting architecture has two hidden layers, with 2 neurons in the first layer and 44 in the second, and a batch size of 128. The model includes dropout layers with a 0.2 probability and uses the LeakyReLU activation function.

I’m using MSELoss as the criterion and Adam as the optimizer. Could someone give advice on how I could improve my model or probable cause for the plaeauing?

fnn(
(layers): Sequential(
(0): Linear(in_features=14, out_features=2, bias=True)
(1): BatchNorm1d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): LeakyReLU(negative_slope=0.2)
(3): Dropout(p=0.2, inplace=False)
(4): Linear(in_features=2, out_features=44, bias=True)
(5): BatchNorm1d(44, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(6): LeakyReLU(negative_slope=0.2)
(7): Dropout(p=0.2, inplace=False)
(8): Linear(in_features=44, out_features=1, bias=True)
)
)

Train & Validation loss

• I have tried several different architectures resulting more or less in a similiar train/valid loss.
• I have tried ReduceLROnPlateau scheduler, but did not yield in any better results.