Measuring Error: Loss Functions and Gradients in JAX

Measuring Error: Loss Functions and Gradients in JAX | CodeSignal Learn

Introduction

Welcome back to another exciting chapter in our JAX in Action: Neural Networks from Scratch journey! You've made excellent progress so far. In our first lesson, we tackled the XOR challenge and established our MLP architecture with proper parameter initialization. Then, in our second lesson, we brought our network to life by implementing forward propagation, watching as data flowed through layers to produce predictions. Now we're ready for a crucial next step: measuring how wrong our predictions are and figuring out how to improve them. Today, we'll dive into loss functions and JAX's automatic differentiation capabilities. We'll implement a Binary Cross Entropy (BCE) loss function to quantify the difference between our network's predictions and the true XOR targets. More importantly, we'll harness the power of jax.value_and_grad to compute both the loss value and the gradients we need for training. By the end of this lesson, you'll have the essential tools to measure your network's performance and understand exactly how to adjust each parameter to make it better!

Understanding Loss Functions: The Heart of Learning

\text{BCE}(y, \hat{y}) = -\frac{1}{N}\sum_{i=1}^{N} [y_i \log(\hat{y}_i) + (1-y_i) \log(1-\hat{y}_i)]

Join the 1M+ learners on CodeSignal

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

Implementing Binary Cross Entropy Loss

Let's implement our bce_loss function. This function will be the cornerstone of our training process, providing the error signal that drives learning. Python import jax.numpy as jnp import jax # (mlp_forward_pass and initialize_mlp_params are defined as in previous lessons) def bce_loss(predictions, targets): """ Computes Binary Cross Entropy loss between predictions and targets. Args: predictions: Model predictions (values between 0 and 1). targets: True binary labels (0 or 1). Returns: Scalar loss value. """ # Add a small epsilon to prevent log(0) which would cause NaN (Not a Number) eps = 1e-8 predictions = jnp.clip(predictions, eps, 1 - eps) # Compute BCE loss # targets * jnp.log(predictions): loss for positive class # (1 - targets) * jnp.log(1 - predictions): loss for negative class loss = -jnp.mean(targets * jnp.log(predictions) + (1 - targets) * jnp.log(1 - predictions)) return loss Let's break this down: The bce_loss function takes predictions from our model and the true targets . A crucial numerical stability trick is eps = 1e-8 and jnp.clip(predictions, eps, 1 - eps) . This prevents predictions from being exactly 0 or 1, which would lead to jnp.log(0) (negative infinity) or jnp.log(1-1) (also negative infinity), causing issues. Clipping keeps predictions within a safe range like [1e-8, 1-1e-8] . The core BCE computation targets * jnp.log(predictions) + (1 - targets) * jnp.log(1 - predictions) calculates the loss for each sample. jnp.mean(...) averages the loss across all samples in our batch, giving us a single scalar value representing the overall error. The negative sign ensures that we minimize this value.

Connecting Loss to Model Predictions

Now we need to create a bridge between our MLP's forward pass (which you built in the previous lesson!) and our new loss calculation. We'll build a function that takes our network parameters and a batch of data, runs the forward pass to get predictions, and then computes the loss based on those predictions and the true labels. Python def compute_loss_for_mlp(params_list, x_batch, y_batch): """ Computes the loss for the MLP given parameters and data. Args: params_list: PyTree of MLP parameters (list of layer dicts). x_batch: Batch of input data. y_batch: Batch of true labels. Returns: Scalar loss value. """ # Get predictions using our forward pass function from the previous lesson predictions = mlp_forward_pass(params_list, x_batch) # Compute BCE loss between predictions and true targets loss = bce_loss(predictions, y_batch) return loss This function, compute_loss_for_mlp, is straightforward but very important. It serves as the complete evaluation pipeline for our network's performance on a given batch of data: It accepts the current params_list (our PyTree of weights and biases), a batch of inputs x_batch, and the corresponding true labels y_batch . It calls mlp_forward_pass (which you developed earlier) to get the network's predictions for x_batch . It then uses our bce_loss function to calculate the scalar loss between these predictions and y_batch . This single scalar loss value tells us how well our MLP is doing with its current parameters.

JAX's Automatic Differentiation Magic: value_and_grad

As you recall from our earlier work with JAX, one of its most powerful features is automatic differentiation through functions like jax.grad; here, we'll apply these familiar tools to our neural network training pipeline. Since we need both the loss value (to monitor training progress) and the gradients (to update parameters), jax.value_and_grad is perfect for our needs. As you remember, it transforms a scalar-valued function to return both the original output and the gradients with respect to the first argument: Python # Transform our loss function to also compute gradients # jax.value_and_grad computes gradients with respect to the first argument (params_list) loss_and_gradients_fn = jax.value_and_grad(compute_loss_for_mlp) # We can also JIT-compile this new function for performance loss_and_gradients_fn = jax.jit(loss_and_gradients_fn) Notice how we're stacking transformations here — a common and powerful pattern in JAX . We first apply jax.value_and_grad to get both loss and gradients, then wrap the result with jax.jit for optimized performance. This gives us a highly efficient function that, when called like loss_and_gradients_fn(mlp_params, xor_X, xor_y), returns both the loss_value and a gradients_pytree with the exact same structure as our parameters. The gradients will tell us exactly how to adjust each weight and bias to reduce the loss — the foundation of gradient-based training that we'll fully implement in our next lesson!

Testing Our Complete Implementation

Let's put all these pieces together: our XOR data, initialized parameters and the forward pass from the previous lessons, as well as our new BCE loss, the combined loss computation function, and the value_and_grad transformation. We'll see the initial loss of our untrained network and inspect the structure of the gradients. Pythondef main(): # Setup XOR data (as we've used before) xor_X = jnp.array([[0., 0.], [0., 1.], [1., 0.], [1., 1.]], dtype=jnp.float32) xor_y = jnp.array([[0.], [1.], [1.], [0.]], dtype=jnp.float32) # Initialize MLP parameters (using function from a previous lesson) layer_sizes = [2, 3, 1] # 2 inputs, 3 hidden neurons, 1 output prng_key = jax.random.key(0) # For reproducible random numbers mlp_params = initialize_mlp_params(layer_sizes, prng_key) # Test the loss computation directly first initial_loss = compute_loss_for_mlp(mlp_params, xor_X, xor_y) print(f"Initial Loss (BCE) with random parameters: {initial_loss:.4f}") # Enable automatic gradient computation and JIT-compile loss_and_gradients_fn = jax.value_and_grad(compute_loss_for_mlp) loss_and_gradients_fn = jax.jit(loss_and_gradients_fn) loss_value, gradients_pytree = loss_and_gradients_fn(mlp_params, xor_X, xor_y) print(f"\nLoss value (from value_and_grad): {loss_value:.4f}") print("Gradients PyTree Structure:") for i, layer_grads in enumerate(gradients_pytree): print(f" Layer {i+1} Gradients:") print(f" Weights 'w' gradient shape: {layer_grads['w'].shape}") print(f" Biases 'b' gradient shape: {layer_grads['b'].shape}") if __name__ == "__main__": main() When you run this code, you should see output similar to this: textInitial Loss (BCE) with random parameters: 0.7782 Loss value (from value_and_grad): 0.7782 Gradients PyTree Structure: Layer 1 Gradients: Weights 'w' gradient shape: (2, 3) Biases 'b' gradient shape: (3,) Layer 2 Gradients: Weights 'w' gradient shape: (3, 1) Biases 'b' gradient shape: (1,) Let's interpret this output: The "Initial Loss (BCE)" is around 0.7782. Since our network parameters are random, this loss value reflects how far off its initial guesses are from the true XOR outputs. The "Loss value (from value_and_grad)" is identical, confirming that jax.value_and_grad correctly computes the loss value alongside the gradients. The "Gradients PyTree Structure" shows the shapes of the gradients. Notice how they perfectly match the shapes of our parameters: Layer 1 (input 2, output 3): weight gradients are (2, 3), bias gradients are (3,). Layer 2 (input 3, output 1): weight gradients are (3, 1), bias gradients are (1,). This confirms that JAX has calculated a gradient for every single parameter in our network!

Measuring Error: Loss Functions and Gradients in JAX

Introduction

Understanding Loss Functions: The Heart of Learning

Join the 1M+ learners on CodeSignal

Implementing Binary Cross Entropy Loss

Connecting Loss to Model Predictions

JAX's Automatic Differentiation Magic: value_and_grad

Testing Our Complete Implementation

Conclusion and Next Steps