MLP Foundations: XOR Data Preparation and Parameter Initialization

MLP Foundations: XOR Data Preparation and Parameter Initialization | CodeSignal Learn

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

Setting Up XOR Data with JAX Arrays

First things first, let's prepare our dataset for the XOR problem using JAX arrays. It's important to structure our data correctly. The input features will be a 2D array, where each row is a distinct data point (an input pair like [0,0]), and each column represents a feature. The target labels, which are the desired outputs, will be a 2D column vector. Python import jax import jax.numpy as jnp # XOR input features: each row is a data point ([x1, x2]) xor_X = jnp.array([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]], dtype=jnp.float32) # XOR target labels: each row is the corresponding output y # We shape it as (4, 1) to be a column vector, a common convention for targets xor_y = jnp.array([[0.0], [1.0], [1.0], [0.0]], dtype=jnp.float32) import jax import jax.numpy as jnp # XOR input features: each row is a data point ([x1, x2]) xor_X = jnp.array([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]], dtype=jnp.float32) # XOR target labels: each row is the corresponding output y # We shape it as (4, 1) to be a column vector, a common convention for targets xor_y = jnp.array([[0.0], [1.0], [1.0], [0.0]], dtype=jnp.float32) In this snippet, we define xor_X to hold our four input pairs and xor_y for their corresponding outputs. We've explicitly set dtype=jnp.float32 . Neural networks typically perform calculations using 32-bit floating-point numbers for a good balance of precision and computational efficiency. Our input xor_X has a shape of (4, 2), meaning 4 data samples, each with 2 features. The target xor_y has a shape of (4, 1), making it a column vector. Maintaining consistent shapes like this is vital for smooth operations later on, especially during the forward and backward passes of our network.

Designing the MLP Architecture

To solve the XOR problem, we need a Multi-Layer Perceptron (MLP) with at least one hidden layer. This hidden layer is what allows the network to learn non-linear relationships. We'll design a simple yet effective architecture: An input layer that matches the number of features in our data. A hidden layer with a few neurons to capture the XOR logic. An output layer that produces our single binary prediction. Let's define the sizes for these layers: Python # Define MLP layer sizes for XOR input_size = xor_X.shape[1] # Number of features from data (2 for XOR) hidden_size = 3 # Number of neurons in the hidden layer output_size = xor_y.shape[1] # Number of output neurons (1 for XOR) mlp_layer_sizes = [input_size, hidden_size, output_size] # Define MLP layer sizes for XOR input_size = xor_X.shape[1] # Number of features from data (2 for XOR) hidden_size = 3 # Number of neurons in the hidden layer output_size = xor_y.shape[1] # Number of output neurons (1 for XOR) mlp_layer_sizes = [input_size, hidden_size, output_size] This code sets up mlp_layer_sizes as [2, 3, 1] . This means our network will have: 2 neurons in the input layer (one for each input feature of XOR). 3 neurons in the hidden layer. While 2 hidden neurons are theoretically enough for XOR, using 3 gives a bit more capacity and makes our example slightly more general. 1 neuron in the output layer (to predict the single 0 or 1 output). This list, mlp_layer_sizes, will be very handy when we initialize the weights and biases for our network, as it dictates their dimensions.

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! Pythondef 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_listdef 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. Pythondef 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: textXOR 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.

Initializing the weights and biases of a neural network correctly is super important for effective training. If all weights start the same, all neurons in a layer will learn the same thing. Random initialization breaks this symmetry. The scale of these initial random values also matters; if they're too large or too small, it can lead to problems like vanishing or exploding gradients during training, making it hard for the network to learn.

We'll use a popular and effective method called Xavier/Glorot initialization for our weights. This technique helps keep the signal and its gradients flowing smoothly through the network by scaling the initial weights based on the number of input and output neurons for a layer. The formula for the limit used in uniform Xavier initialization is:

\text{limit} = \sqrt{\frac{6}{\text{input\_dim} + \text{output\_dim}}}

Weights are then sampled from a uniform distribution between -limit and +limit. For biases, we'll take a simpler approach and initialize them all to zero. This is a common practice, allowing the network to learn the appropriate bias values during training.