Learning by Example: The Manual Training Loop

Introduction

Welcome back to the final lesson of JAX in Action: Neural Networks from Scratch! What an incredible journey we've taken together through the fundamentals of building neural networks from the ground up. We started by tackling the classic XOR problem and establishing our MLP architecture with proper parameter initialization. We then implemented forward propagation to see our network generate predictions, and in our previous lesson, we developed the crucial ability to measure prediction quality using Binary Cross Entropy loss and compute gradients with JAX's automatic differentiation. Today marks the culmination of our neural network construction project. We'll implement a complete training loop that iteratively improves our MLP's performance using gradient descent. You'll learn to orchestrate the training process by repeatedly computing predictions, calculating loss and gradients, and updating parameters to minimize error. By the end of this lesson, you'll witness your network actually learning to solve the XOR problem through experience — a truly magical moment in machine learning!

The Core of Machine Learning: Iterative Improvement

At its heart, machine learning is about iterative improvement . Think of learning to play a musical instrument: you don't master it overnight, but through countless repetitions, gradually adjusting your technique based on feedback. Neural network training follows this same principle. Our training process operates in cycles called epochs . During each epoch, we present our network with the training data and perform these key steps: compute predictions using forward propagation, calculate how wrong these predictions are using our loss function, determine the direction to adjust parameters using gradients, and finally update parameters to (hopefully) improve performance. This cycle repeats hundreds or thousands of times until our network performs well. The beauty of this approach lies in its simplicity — we don't need to manually figure out what each parameter should be. Instead, we let the mathematical framework of gradient descent guide us toward better solutions automatically. Each iteration brings us closer to a network that can accurately solve our XOR problem.

Gradient Descent: The Update Rule

W_{new} = W_{old} - \alpha \cdot \frac{\partial \mathcal{L}}{\partial W}

Join the 1M+ learners on CodeSignal

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

Building the Complete Training Loop

Now we'll construct the main training loop that orchestrates the entire learning process. This loop will cycle through epochs, and in each epoch, it will compute loss and gradients, then apply updates to our network's parameters.

First, let's set up our training configuration and recall our loss_and_grad_fn from the previous lesson, which gives us both the loss and the gradients.

Python

import jax
import jax.numpy as jnp
# Training configuration
learning_rate = 0.5
num_epochs = 2000

# Create loss and gradient function (JIT-compiled for performance)
loss_and_grad_fn = jax.jit(jax.value_and_grad(compute_loss_for_mlp))

print(f"Starting training for {num_epochs} epochs with learning rate {learning_rate}")

# Main training loop
for epoch in range(num_epochs):
    # Compute loss and gradients for current parameters using the whole dataset
    loss_value, gradients = loss_and_grad_fn(current_params, xor_X, xor_y)
    
    # Define the gradient descent update function (captures current learning_rate)
    def apply_gradient_step_for_loop(param_leaf, grad_leaf):
        return param_leaf - learning_rate * grad_leaf
    
    # Apply updates to all parameters using tree_map
    current_params = jtu.tree_map(apply_gradient_step_for_loop, current_params, gradients)
    
    # We'll add progress printing in the next step

Let's break down this crucial piece of code:

We set a learning_rate and the num_epochs (how many times we'll iterate through the data).
loss_and_grad_fn is our JIT-compiled function from the last lesson that efficiently calculates both the loss and the gradients for current_params given the xor_X data and xor_y targets.
The for loop iterates for num_epochs.
Inside the loop, we first call loss_and_grad_fn to get the loss_value and the gradients PyTree.
We then define apply_gradient_step_for_loop. Notice it's defined inside the loop. This is a common Python pattern that allows the function to "capture" the learning_rate from the loop's current scope.
Finally, current_params = jtu.tree_map(...) updates all parameters in our current_params PyTree by applying apply_gradient_step_for_loop to each parameter and its corresponding gradient. The result is the new set of parameters for the next iteration.

This loop is the engine of learning for our MLP!

Evaluating the Trained Model

After training is complete, we need to evaluate how well our network has learned to solve the XOR problem. This involves running the forward pass with the trained parameters, calculating the final loss, and, importantly, checking the actual predictions against the targets to compute accuracy.

Python

# (After the training loop has finished)
# current_params now holds the trained parameters

# Evaluate final performance
final_predictions = mlp_forward_pass(current_params, xor_X) # Use the trained params
final_loss = bce_loss(final_predictions, xor_y) # Calculate loss on trained model

print(f"\nFinal Loss after training: {final_loss:.6f}")
print("Final Predictions vs Targets:")

for i in range(xor_X.shape[0]):
    input_data = xor_X[i]
    target_label = xor_y[i, 0]
    predicted_value = final_predictions[i, 0]
    # Convert prediction to binary (0 or 1) based on a 0.5 threshold
    binary_prediction = 1 if predicted_value > 0.5 else 0 
    print(f"  Input: {input_data}, Target: {target_label}, Prediction: {predicted_value:.4f} (Binary: {binary_prediction})")
    
# Calculate overall accuracy
# Convert all predictions to binary
binary_predictions_all = (final_predictions > 0.5).astype(jnp.float32)
# Compare binary predictions to true targets and average
accuracy = jnp.mean(binary_predictions_all == xor_y)
print(f"Final Accuracy: {accuracy*100:.2f}%")

Running this evaluation code after our training loop should produce results like this:

text

Final Loss after training: 0.019308
Final Predictions vs Targets:
  Input: [0. 0.], Target: 0.0, Prediction: 0.0131 (Binary: 0)
  Input: [0. 1.], Target: 1.0, Prediction: 0.9829 (Binary: 1)
  Input: [1. 0.], Target: 1.0, Prediction: 0.9828 (Binary: 1)
  Input: [1. 1.], Target: 0.0, Prediction: 0.0291 (Binary: 0)
Final Accuracy: 100.00%

These results are remarkable! Our network, trained from scratch using JAX, has achieved 100% accuracy on the XOR problem. Notice how the raw prediction values are very close to the target binary values (e.g., 0.0131 for a target of 0.0, and 0.9829 for a target of 1.0). This confirms that our MLP has successfully learned the non-linear XOR logical function through iterative gradient descent optimization.

Introduction

The Core of Machine Learning: Iterative Improvement

Gradient Descent: The Update Rule

Join the 1M+ learners on CodeSignal

Implementing Parameter Updates with tree_map

Building the Complete Training Loop

Monitoring Training Progress

Evaluating the Trained Model

Conclusion and What's Next