Training Neural Networks with SGD

Introduction

Welcome back to our course, "Training Neural Networks: the Backpropagation Algorithm"! This is our fifth and final lesson in this course, and we've come a long way. So far, we've explored loss functions, gradient descent, and implemented backpropagation for both individual layers and complete networks. In our previous lesson, we implemented the backward pass for a multi-layer perceptron, allowing us to calculate the gradients of all weights and biases with respect to the loss. However, calculating gradients is only half the story. These gradients tell us which direction to move, but we still need to actually move in that direction! Today, we'll complete our neural network training journey by implementing stochastic gradient descent (SGD), which will use the gradients to update our network's parameters. We'll also create a complete training loop and apply our knowledge to train a model on a real-world regression task. This represents the culmination of everything we've learned so far. After this lesson, you'll be ready to move on to our fourth and final course, "Building and Applying Your Neural Network Library"!

Stochastic Gradient Descent: The Workhorse of Neural Network Training

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Creating a Training Loop

With our SGD optimizer and dataset ready, let's implement a basic training loop to tie everything together. The training loop is the heart of neural network learning, orchestrating the forward pass, loss calculation, backward pass, and parameter updates.

Here's the structure of a typical training loop using mini-batch SGD:

mlp <- create_mlp()
mlp <- add_layer(mlp, create_dense_layer(n_features, 10, 'relu', weight_init_scale = 0.1))
mlp <- add_layer(mlp, create_dense_layer(10, 1, 'linear', weight_init_scale = 0.1))

optimizer <- create_sgd_optimizer(learning_rate = 0.002)  # Small learning rate
epochs <- 100  # Number of epochs
batch_size <- 32  # Define batch size for SGD
num_batches <- ceiling(n_samples / batch_size)  # Total number of batches

cat(sprintf("Training MLP for Diabetes regression for %d epochs with LR=%.3f\n", 
            epochs, optimizer$learning_rate))

for (epoch in 1:epochs) {
  # Shuffle the data at the beginning of each epoch
  indices <- sample(n_samples)
  X_shuffled <- X_train[indices, , drop = FALSE]
  y_shuffled <- y_train[indices, , drop = FALSE]
  
  total_loss <- 0
  
  # Process data in mini-batches
  for (i in seq(1, n_samples, by = batch_size)) {
    end_idx <- min(i + batch_size - 1, n_samples)
    X_batch <- X_shuffled[i:end_idx, , drop = FALSE]
    y_batch <- y_shuffled[i:end_idx, , drop = FALSE]

mlp <- create_mlp()
mlp <- add_layer(mlp, create_dense_layer(n_features, 10, 'relu', weight_init_scale = 0.1))
mlp <- add_layer(mlp, create_dense_layer(10, 1, 'linear', weight_init_scale = 0.1))

optimizer <- create_sgd_optimizer(learning_rate = 0.002)  # Small learning rate
epochs <- 100  # Number of epochs
batch_size <- 32  # Define batch size for SGD
num_batches <- ceiling(n_samples / batch_size)  # Total number of batches

cat(sprintf("Training MLP for Diabetes regression for %d epochs with LR=%.3f\n", 
            epochs, optimizer$learning_rate))

for (epoch in 1:epochs) {
  # Shuffle the data at the beginning of each epoch
  indices <- sample(n_samples)
  X_shuffled <- X_train[indices, , drop = FALSE]
  y_shuffled <- y_train[indices, , drop = FALSE]
  
  total_loss <- 0
  
  # Process data in mini-batches
  for (i in seq(1, n_samples, by = batch_size)) {
    end_idx <- min(i + batch_size - 1, n_samples)
    X_batch <- X_shuffled[i:end_idx, , drop = FALSE]
    y_batch <- y_shuffled[i:end_idx, , drop = FALSE]

Let's analyze this initial part of our training loop:

We first create our MLP with two layers: a hidden layer with 10 neurons and ReLU activation, and an output layer with 1 neuron (for regression) and linear activation.
We initialize our SGD optimizer with a learning rate of 0.002, which is relatively small to ensure stable training.
We set up hyperparameters: 100 epochs and a batch size of 32.
For each epoch, we shuffle the data to prevent the model from learning patterns based on data order.
We then process the data in mini-batches, selecting subsets of the training data for each update.

The shuffling step is crucial in stochastic gradient descent. It ensures that each epoch presents the data in a different order, which helps prevent the model from getting stuck in local minima and improves generalization.

Completing the Training Loop

Now let's complete our training loop by implementing the forward pass, backward pass, and parameter updates for each mini-batch:

    # Within the mini-batch loop:
    # 1. Forward pass
    forward_result <- forward_mlp(mlp, X_batch)
    mlp <- forward_result$mlp
    y_pred <- forward_result$output
    # 2. Calculate loss
    loss <- mse_loss(y_batch, y_pred)
    total_loss <- total_loss + loss
    # 3. Calculate derivative of loss
    d_loss_wrt_pred <- mse_loss_derivative(y_batch, y_pred)
    # 4. Backward pass
    mlp <- backward_mlp(mlp, d_loss_wrt_pred)
    # 5. Update weights
    for (j in 1:length(mlp$layers)) {
      mlp$layers[[j]] <- optimizer$update(mlp$layers[[j]])
    }
  }
  
  # Average loss over all batches for reporting
  avg_loss <- total_loss / num_batches
  if ((epoch %% max(1, epochs %/% 10) == 0) || epoch == 1) {
    cat(sprintf("Epoch %4d/%d, Loss: %.2f\n", epoch, epochs, avg_loss))
  }
}

For each mini-batch, we:

Forward pass: Use our MLP to generate predictions. Note that forward_mlp returns both the modified mlp object (with stored layer inputs and outputs) and the predictions, so we capture both and update our mlp reference.
Loss calculation: Compute the MSE loss between predictions and targets.
Loss derivative: Calculate how the loss changes with respect to our predictions.
Backward pass: Propagate this gradient backward through the network.
Parameter updates: Use our SGD optimizer to update each layer's weights and biases.

After processing all mini-batches, we calculate the average loss for the epoch and periodically print it to track training progress. This allows us to monitor whether the model is learning effectively.

The training loop is the orchestrator that brings together all the components we've built throughout this course: the layer structure, forward propagation, loss functions, backpropagation, and now parameter updates through SGD. This elegant cycle of computation and learning is what gives neural networks their power.

Evaluating Training Results

After training is complete, we want to evaluate our model's performance and examine some sample predictions. Let's add this final piece to our implementation:

cat("\n--- Training Over ---\n")
forward_result <- forward_mlp(mlp, X_train)
mlp <- forward_result$mlp
final_preds <- forward_result$output
final_loss <- mse_loss(y_train, final_preds)
cat(sprintf("Final MSE Loss on training data: %.2f\n", final_loss))

cat("\nSample Predictions:\n")
for (i in 1:min(5, n_samples)) {  # Print first 5 samples
  cat(sprintf("  True: %8.2f, Predicted: %8.2f\n", 
              y_train[i, 1], final_preds[i, 1]))
}

cat("\n--- Training Over ---\n")
forward_result <- forward_mlp(mlp, X_train)
mlp <- forward_result$mlp
final_preds <- forward_result$output
final_loss <- mse_loss(y_train, final_preds)
cat(sprintf("Final MSE Loss on training data: %.2f\n", final_loss))

cat("\nSample Predictions:\n")
for (i in 1:min(5, n_samples)) {  # Print first 5 samples
  cat(sprintf("  True: %8.2f, Predicted: %8.2f\n", 
              y_train[i, 1], final_preds[i, 1]))
}

After running this code, we might see output like:

text

Training data shape: X_train: (442, 10), y_train: (442, 1)
Training MLP for Diabetes regression for 100 epochs with LR=0.002
Epoch    1/100, Loss: 22878.82
Epoch   10/100, Loss: 4148.13
Epoch   20/100, Loss: 3505.00
Epoch   30/100, Loss: 3527.20
Epoch   40/100, Loss: 3590.56
Epoch   50/100, Loss: 3069.60
Epoch   60/100, Loss: 2950.78
Epoch   70/100, Loss: 3104.61
Epoch   80/100, Loss: 3110.34
Epoch   90/100, Loss: 3422.17
Epoch  100/100, Loss: 3023.98

--- Training Over ---
Final MSE Loss on training data: 2927.23

Sample Predictions:
  True:   151.00, Predicted:   193.71
  True:    75.00, Predicted:    65.94
  True:   141.00, Predicted:   166.13
  True:   206.00, Predicted:   156.16
  True:   135.00, Predicted:   121.92

Training data shape: X_train: (442, 10), y_train: (442, 1)
Training MLP for Diabetes regression for 100 epochs with LR=0.002
Epoch    1/100, Loss: 22878.82
Epoch   10/100, Loss: 4148.13
Epoch   20/100, Loss: 3505.00
Epoch   30/100, Loss: 3527.20
Epoch   40/100, Loss: 3590.56
Epoch   50/100, Loss: 3069.60
Epoch   60/100, Loss: 2950.78
Epoch   70/100, Loss: 3104.61
Epoch   80/100, Loss: 3110.34
Epoch   90/100, Loss: 3422.17
Epoch  100/100, Loss: 3023.98

--- Training Over ---
Final MSE Loss on training data: 2927.23

Sample Predictions:
  True:   151.00, Predicted:   193.71
  True:    75.00, Predicted:    65.94
  True:   141.00, Predicted:   166.13
  True:   206.00, Predicted:   156.16
  True:   135.00, Predicted:   121.92

The training results show that:

Loss Reduction: The loss drops dramatically in the early epochs (from 22878.82 to around 4148), demonstrating that our model is learning rapidly. The loss continues to decrease more gradually in later epochs, with some fluctuation due to the stochastic nature of mini-batch training.
Final Performance: The final MSE on the training data is approximately 2927.23, which gives us a sense of the model's fit. For the Diabetes dataset, this is a reasonable result for a simple MLP model without any feature engineering or hyperparameter tuning. Note that the "Loss" reported at epoch 100/100 (3023.98) is the average loss per mini-batch during that last epoch, while the "Final MSE Loss" (2927.23) is calculated over the entire training set after training is complete, which is why they differ slightly.
Predictions vs. Targets: Looking at the sample predictions, we can see that our model makes reasonable predictions for most samples, with errors typically ranging from a few units to a few dozen units of the true values. This suggests our implementation is working correctly.

This analysis provides valuable insights into our model's performance and confirms that our implementation of SGD and the training loop is working correctly.

Introduction

Stochastic Gradient Descent: The Workhorse of Neural Network Training

Join the 1M+ learners on CodeSignal

The SGD Algorithm: Pseudocode and Key Steps

Implementing the SGD Optimizer

The Diabetes Dataset

Creating a Training Loop

Completing the Training Loop

Evaluating Training Results

Conclusion