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!

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.

JAX provides powerful utilities for working with PyTrees through the jax.tree_util module; recall that the tree_map function that we encountered previously is particularly useful for applying an operation element-wise across corresponding leaves of one or more PyTrees that share the same structure.

Let's say we have our current parameters (current_params) and their corresponding gradients, both structured as PyTrees. We want to apply the gradient descent update to each parameter. We can define a simple function for the update rule:

import jax.tree_util as jtu

# This function defines the update for a single parameter (leaf)
def apply_gradient_step(param_leaf, grad_leaf):
    """Apply gradient descent update to a single parameter leaf."""
    # 'learning_rate' would be accessible from an outer scope
    return param_leaf - learning_rate * grad_leaf

# To apply this to all parameters:
new_params = jtu.tree_map(apply_gradient_step, current_params, gradients)

Here's what makes this approach powerful:

• apply_gradient_step takes a single parameter value (param_leaf) and its corresponding gradient (grad_leaf). It uses learning_rate (which we'll define in our training setup) to compute the updated parameter.
• jtu.tree_map(apply_gradient_step, current_params, gradients) then intelligently traverses both the current_params PyTree and the gradients PyTree. For every matching pair of parameter leaf and gradient leaf, it calls apply_gradient_step and constructs a new PyTree (new_params) with the updated values. This means we don't need to manually iterate through layers or dictionaries; JAX handles the PyTree structure for us.

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.

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!

To understand how well our training is progressing, we need to monitor and display the loss values. This feedback helps us verify that learning is occurring and provides insights into the training dynamics. We can simply add a print statement inside our loop.

# (Inside the training loop, after parameter update)
# Print progress at regular intervals (e.g., every 200 epochs) and on the first epoch
if (epoch + 1) % 200 == 0 or epoch == 0:
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {loss_value:.6f}")

# After the loop finishes
print("Training finished.")

When we integrate this into our training loop and run it (assuming all previous functions like initialize_mlp_params, compute_loss_for_mlp, and the XOR data are set up), we should see output similar to this:

Starting training for 2000 epochs with learning rate 0.5
Epoch 1/2000, Loss: 0.778152
Epoch 200/2000, Loss: 0.662008
Epoch 400/2000, Loss: 0.575099
Epoch 600/2000, Loss: 0.520910
Epoch 800/2000, Loss: 0.440773
Epoch 1000/2000, Loss: 0.197226
Epoch 1200/2000, Loss: 0.081303
Epoch 1400/2000, Loss: 0.046896
Epoch 1600/2000, Loss: 0.032172
Epoch 1800/2000, Loss: 0.024233
Epoch 2000/2000, Loss: 0.019327
Training finished.

This progress log tells a beautiful story of learning! Notice how the loss starts around 0.778 (quite high for our normalized BCE loss) and steadily decreases to about 0.019 by the end of 2000 epochs. The most dramatic improvement often happens once the network starts to "understand" the underlying pattern, as seen by the significant drop around epoch 1000. This steady decrease in loss is exactly what we want to see during successful training, indicating our MLP is getting better at predicting the XOR outputs.

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.

# (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:

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.

Congratulations on this momentous milestone! You've successfully completed JAX in Action: Neural Networks from Scratch, building a neural network training system entirely from scratch. You've mastered forward propagation, loss functions, automatic differentiation, and now, the complete gradient descent training loop using pure JAX.

The foundational principles you've learned here are pivotal for modern deep learning. Get ready for your next adventure, Beyond Pure JAX: Flax & Optax for Elegant ML! In this upcoming course, you'll leverage powerful JAX ecosystem libraries like Flax and Optax to construct sophisticated models with greater ease and efficiency.