MLP Foundations: XOR Data Preparation and Parameter Initialization

Introduction

Welcome, space explorer, to JAX in Action: Neural Networks from Scratch! This lesson marks the exciting start of our third course together. We're about to take all those fantastic JAX skills you've been honing and apply them to building neural networks, piece by piece. Let's quickly chart our journey so far. In our first course, JAX Fundamentals: NumPy Power-Up, you became proficient with JAX's core elements, like immutable arrays, the beauty of pure functions, the magic of automatic differentiation (jax.grad), and the speed boosts from just-in-time (JIT) compilation (jax.jit). Then, in Advanced JAX: Transformations for Speed & Scale, we explored powerful functional transformations. You learned about JAX's explicit random number generation, mastered batch processing with jax.vmap, got a glimpse of multi-device parallelism with jax.shard_map, learned to handle complex data structures with PyTrees, and picked up essential profiling and debugging techniques. Now, in this course, we'll use this powerful toolkit to construct neural networks . We'll begin with the fundamentals, like implementing Multi-Layer Perceptrons (MLPs), and progressively build more complex models, eventually leveraging JAX ecosystem libraries like Flax and Optax. In this first lesson, we'll tackle the classic XOR problem . It's a perfect way to understand non-linear classification. Our focus today will be on the crucial first steps: preparing our data and initializing the network's parameters.

Understanding the XOR Problem

The XOR (exclusive OR) problem is a cornerstone in the study of neural networks. It's a simple yet profound example that highlights why we often need networks with multiple layers. The XOR function takes two binary inputs (0 or 1) and produces a binary output. The rule is: the output is 1 if the inputs are different and 0 if they are the same. Let's visualize this via a truth table: Input x₁ Input x₂ Output (x₁ XOR x₂) 000011101110 If you try to plot these four points on a 2D graph, you'll notice something interesting: you can't draw a single straight line to separate the points that result in an output of 0 from those that result in an output of 1. This property is called linear inseparability . A single-layer perceptron, which can only learn linear decision boundaries, would be stumped by XOR! This is why XOR is such a great example to show the power of multi-layer networks, which can learn the necessary non-linear patterns.

Input x₁	Input x₂	Output (x₁ XOR x₂)
0	0	0
0	1	1
1	0	1
1	1	0

Join the 1M+ learners on CodeSignal

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

Implementing Parameter Initialization

Now, let's write a Python function to initialize the parameters (weights and biases) for our MLP. We'll organize these parameters as a list of dictionaries. Each dictionary in the list will represent a layer and will contain its w (weights) and b (biases). This structure is a PyTree, which JAX handles very efficiently!

Python

def initialize_mlp_params(layer_sizes, parent_key):
    """
    Initializes parameters (weights and biases) for a Multi-Layer Perceptron (MLP).
    
    Args:
        layer_sizes: A list of integers, where each integer is the number of 
                     neurons in a layer. E.g., [input_dim, hidden_dim, output_dim].
        parent_key: The main JAX PRNG key for random number generation.
    
    Returns:
        A list of dictionaries. Each dictionary represents a layer and 
        contains 'w' (weights matrix) and 'b' (bias vector) as JAX arrays.
    """
    params_list = []  # List to hold parameter dictionaries for each layer (weights and biases)
    current_key = parent_key # Start with the parent key
    
    # Iterate through pairs of layers to define weights and biases between them
    for i in range(len(layer_sizes) - 1):
        input_dim, output_dim = layer_sizes[i], layer_sizes[i+1]
        
        # Split the key for this layer's weights and biases
        current_key, w_key, b_key = jax.random.split(current_key, 3)
        
        # Xavier/Glorot initialization for weights
        limit = jnp.sqrt(6.0 / (input_dim + output_dim))
        weights = jax.random.uniform(w_key, 
                                     shape=(input_dim, output_dim), 
                                     minval=-limit, 
                                     maxval=limit)
        
        # Initialize biases to zero
        biases = jnp.zeros((output_dim,)) 
        
        params_list.append({'w': weights, 'b': biases})
        
    return params_list

def initialize_mlp_params(layer_sizes, parent_key):
    """
    Initializes parameters (weights and biases) for a Multi-Layer Perceptron (MLP).
    
    Args:
        layer_sizes: A list of integers, where each integer is the number of 
                     neurons in a layer. E.g., [input_dim, hidden_dim, output_dim].
        parent_key: The main JAX PRNG key for random number generation.
    
    Returns:
        A list of dictionaries. Each dictionary represents a layer and 
        contains 'w' (weights matrix) and 'b' (bias vector) as JAX arrays.
    """
    params_list = []  # List to hold parameter dictionaries for each layer (weights and biases)
    current_key = parent_key # Start with the parent key
    
    # Iterate through pairs of layers to define weights and biases between them
    for i in range(len(layer_sizes) - 1):
        input_dim, output_dim = layer_sizes[i], layer_sizes[i+1]
        
        # Split the key for this layer's weights and biases
        current_key, w_key, b_key = jax.random.split(current_key, 3)
        
        # Xavier/Glorot initialization for weights
        limit = jnp.sqrt(6.0 / (input_dim + output_dim))
        weights = jax.random.uniform(w_key, 
                                     shape=(input_dim, output_dim), 
                                     minval=-limit, 
                                     maxval=limit)
        
        # Initialize biases to zero
        biases = jnp.zeros((output_dim,)) 
        
        params_list.append({'w': weights, 'b': biases})
        
    return params_list

This function, initialize_mlp_params, iterates through the layer_sizes list. For each connection between an input layer and an output layer (e.g., input layer to hidden layer, then hidden layer to output layer), it:

Splits the current_key to get new, unique keys (w_key, b_key) for initializing weights and biases. This is crucial for JAX's functional approach to randomness.
Calculates the limit for Xavier initialization based on the input_dim (number of neurons in the current layer) and output_dim (number of neurons in the next layer).
Initializes the weights matrix with shape (input_dim, output_dim). This shape is standard for a forward pass computation like output = jnp.dot(input_activations, weights) + biases.
Initializes the biases vector with shape (output_dim,) to all zeros.
Appends a dictionary containing these weights and biases to our params_list.

Verifying Our Setup: Data and Parameters

Let's put all the pieces together and inspect what we've created. We'll print our XOR data and the structure of our initialized MLP parameters. This helps confirm that the shapes and data types are as we expect.

Python

def main():
    # Display our XOR dataset
    print("XOR Input Features (X):")
    print(xor_X)
    print(f"Shape of X: {xor_X.shape}, dtype: {xor_X.dtype}")

    print("\nXOR Target Labels (y):")
    print(xor_y)
    print(f"Shape of y: {xor_y.shape}, dtype: {xor_y.dtype}")

    # Display the defined MLP layer sizes
    print(f"\nMLP Layer Sizes: {mlp_layer_sizes}")

    # Initialize parameters using a PRNG key
    # JAX requires an explicit key for all random operations
    prng_key = jax.random.key(0) 
    mlp_params = initialize_mlp_params(mlp_layer_sizes, prng_key)

    # Display the structure and shapes of the initialized parameters
    print("\nInitialized MLP Parameters (PyTree structure):")
    for i, layer_params in enumerate(mlp_params):
        print(f"  Layer {i+1}:") # Layer numbering starts from 1 for human readability
        print(f"    Weights 'w' shape: {layer_params['w'].shape}")
        print(f"    Biases 'b' shape: {layer_params['b'].shape}")

if __name__ == "__main__":
    main()

def main():
    # Display our XOR dataset
    print("XOR Input Features (X):")
    print(xor_X)
    print(f"Shape of X: {xor_X.shape}, dtype: {xor_X.dtype}")

    print("\nXOR Target Labels (y):")
    print(xor_y)
    print(f"Shape of y: {xor_y.shape}, dtype: {xor_y.dtype}")

    # Display the defined MLP layer sizes
    print(f"\nMLP Layer Sizes: {mlp_layer_sizes}")

    # Initialize parameters using a PRNG key
    # JAX requires an explicit key for all random operations
    prng_key = jax.random.key(0) 
    mlp_params = initialize_mlp_params(mlp_layer_sizes, prng_key)

    # Display the structure and shapes of the initialized parameters
    print("\nInitialized MLP Parameters (PyTree structure):")
    for i, layer_params in enumerate(mlp_params):
        print(f"  Layer {i+1}:") # Layer numbering starts from 1 for human readability
        print(f"    Weights 'w' shape: {layer_params['w'].shape}")
        print(f"    Biases 'b' shape: {layer_params['b'].shape}")

if __name__ == "__main__":
    main()

When you run this code (for example, by placing the xor_X, xor_y, mlp_layer_sizes definitions and the initialize_mlp_params function before main, and then calling main()), you'll see output like this:

text

XOR Input Features (X):
[[0. 0.]
 [0. 1.]
 [1. 0.]
 [1. 1.]]
Shape of X: (4, 2), dtype: float32

XOR Target Labels (y):
[[0.]
 [1.]
 [1.]
 [0.]]
Shape of y: (4, 1), dtype: float32

MLP Layer Sizes: [2, 3, 1]

Initialized MLP Parameters (PyTree structure):
  Layer 1:
    Weights 'w' shape: (2, 3)
    Biases 'b' shape: (3,)
  Layer 2:
    Weights 'w' shape: (3, 1)
    Biases 'b' shape: (1,)

XOR Input Features (X):
[[0. 0.]
 [0. 1.]
 [1. 0.]
 [1. 1.]]
Shape of X: (4, 2), dtype: float32

XOR Target Labels (y):
[[0.]
 [1.]
 [1.]
 [0.]]
Shape of y: (4, 1), dtype: float32

MLP Layer Sizes: [2, 3, 1]

Initialized MLP Parameters (PyTree structure):
  Layer 1:
    Weights 'w' shape: (2, 3)
    Biases 'b' shape: (3,)
  Layer 2:
    Weights 'w' shape: (3, 1)
    Biases 'b' shape: (1,)

Look at that! Our input features X are (4, 2) and labels y are (4, 1), just as planned. The MLP parameters also have the correct shapes:

Layer 1 (connecting the 2 input neurons to the 3 hidden neurons): Weights are (2, 3) and biases are (3,).
Layer 2 (connecting the 3 hidden neurons to the 1 output neuron): Weights are (3, 1) and biases are (1,).

This confirms our setup is correct and ready for the next steps! The PyTree structure (list of dictionaries) will make it easy to access these parameters when we build the forward pass of our network.

Introduction

Understanding the XOR Problem

Join the 1M+ learners on CodeSignal

Setting Up XOR Data with JAX Arrays

Designing the MLP Architecture

Parameter Initialization Strategy

Implementing Parameter Initialization

Verifying Our Setup: Data and Parameters

Conclusion and Next Steps