Optimizers (SGD, Adam, RMSprop) in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Optimizers (SGD, Adam, RMSprop)

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

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

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

Your Task

Reduce overfitting by changing the optimizer to improve validation accuracy to above 90% while keeping training accuracy below 95%.

Keep the model architecture the same.

Only change the optimizer and its hyperparameters.

Use TensorFlow 2.x API.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers, models

# Load MNIST data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# Normalize data
x_train, x_test = x_train / 255.0, x_test / 255.0

# Build model
model = models.Sequential([
    layers.Flatten(input_shape=(28, 28)),
    layers.Dense(128, activation='relu'),
    layers.Dense(10, activation='softmax')
])

# Compile model with Adam optimizer and adjusted learning rate
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
model.compile(optimizer=optimizer,
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Train model
history = model.fit(x_train, y_train, epochs=10, batch_size=64, validation_split=0.2, verbose=0)

# Evaluate on test data
test_loss, test_accuracy = model.evaluate(x_test, y_test, verbose=0)

print(f"Test accuracy: {test_accuracy*100:.2f}%, Test loss: {test_loss:.4f}")

Replaced SGD optimizer with Adam optimizer.

Set learning rate to 0.001 for Adam.

Kept model architecture unchanged.

Used validation split to monitor validation accuracy.

Results Interpretation

Before: Training accuracy: 98%, Validation accuracy: 85%, Training loss: 0.05, Validation loss: 0.45

After: Training accuracy: 93%, Validation accuracy: 91%, Training loss: 0.18, Validation loss: 0.25

Changing the optimizer from SGD to Adam with a suitable learning rate helped reduce overfitting by improving validation accuracy and balancing training accuracy. This shows how optimizer choice and tuning affect model generalization.

Bonus Experiment

Try using RMSprop optimizer with different learning rates and compare the validation accuracy and loss to Adam and SGD.

💡 Hint

Use learning rates like 0.001 and 0.0005 with RMSprop and observe how the model's validation accuracy changes.