Measuring Error: Loss Functions and Gradients in JAX

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

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

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.

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.

Python

def 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()

def 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:

text

Initial 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,)

Initial 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!

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