Weight Initialization Strategies

Welcome to the fourth and final lesson in our course on the MLP architecture: activations & initialization! We've come a long way together — we've built a flexible MLP architecture, implemented powerful activation functions like ReLU, and added specialized output activations for classification and regression tasks.

Now, we're going to tackle one of the most critical but often overlooked aspects of neural network design: weight initialization. How we initialize the weights in our network might seem like a minor detail, but it can dramatically impact how (or even whether) our network learns. In particular, in this lesson, we'll explore why proper weight initialization is crucial, examine common issues caused by poor initialization, implement several powerful initialization strategies, and enhance our DenseLayer class to support these strategies.

By the end of this lesson, you'll have a solid understanding of weight initialization and the ability to implement various initialization strategies in your neural networks. This knowledge will significantly improve your models' training speed and overall performance.

Imagine you're starting a journey through a complex, hilly landscape with the goal of finding the lowest valley. The point where you begin this journey greatly affects how quickly (or if) you'll reach your destination. Similarly, the initial values of your neural network weights determine your starting point in the loss landscape and influence the entire training process that allows networks to learn.

Poor weight initialization can lead to several problems:

1. Symmetry Issues: If all weights start with the same value, all neurons in a layer will compute the same output and receive the same gradient updates. This "symmetry" prevents the network from learning diverse features.

2. Vanishing Gradients: If weights are too small, the signals flowing through the network will diminish with each layer, causing gradients to approach zero during backpropagation. This makes learning extremely slow, especially in deeper layers.

3. Exploding Gradients: Conversely, if weights are too large, the signals can grow exponentially through the network, leading to unstable training and numerical overflow.

Let's visualize this with a simple example. Imagine a 10-layer network where each layer either halves or doubles the signal:

• With weights that are too small: 1 → 0.5 → 0.25 → 0.125 → ... → 0.001 (signal vanishes).
• With weights that are too large: 1 → 2 → 4 → 8 → ... → 1024 (signal explodes).

Both scenarios make it difficult for the network to learn efficiently. Proper initialization balances these concerns, allowing signals to flow smoothly through the network without vanishing or exploding.

The simplest approach to weight initialization is to use small random values. Random initialization helps break the symmetry between neurons, allowing them to learn different features. However, the scale of these random values is crucial.

Let's implement a basic random scaled initialization strategy:

#include <random>
#include <vector>

std::vector<std::vector<double>> random_scaled_init(int n_inputs, int n_neurons, double scale = 0.01) {
    // Initialize random number generator
    std::random_device rd;
    std::mt19937 gen(rd());
    std::normal_distribution<double> dist(0.0, 1.0);
    
    // Create weight matrix
    std::vector<std::vector<double>> weights(n_inputs, std::vector<double>(n_neurons));
    
    // Fill with random values scaled by the factor
    for (int i = 0; i < n_inputs; i++) {
        for (int j = 0; j < n_neurons; j++) {
            weights[i][j] = dist(gen) * scale;
        }
    }
    
    return weights;
}

In this approach:

• We draw weights from a normal distribution with mean 0 and standard deviation 1 using std::normal_distribution.
• We multiply by a small scale factor (default 0.01) to control the magnitude.
• The scale parameter lets us adjust how large the initial weights should be.

This method is simple and has been widely used, but the optimal scale factor depends on the network architecture and can be hard to determine. If the scale is too small, we risk vanishing gradients; if too large, exploding gradients.

For years, practitioners used rules of thumb like setting the scale between 0.001 and 0.1, but this approach has largely been superseded by more principled methods that we'll explore next.

Choosing the right initialization strategy for your activation function is crucial for effective training. Here's a quick reference guide:

General Rule: Use Xavier/Glorot initialization for activation functions that are roughly symmetric around zero (sigmoid, tanh, linear), and use He initialization for ReLU-like activations that zero out negative values.

When building networks with mixed activation functions, apply the appropriate initialization strategy to each layer based on its activation function. This targeted approach ensures optimal signal flow throughout your entire network.

Now that we understand different initialization strategies, let's enhance our DenseLayer class to support them. We'll add a parameter that allows us to specify which initialization strategy to use:

class DenseLayer {
private:
    int n_inputs;
    int n_neurons;
    std::string activation_fn_name;
    std::string weight_init_strategy;
    double weight_init_scale;
    std::function<std::vector<std::vector<double>>(const std::vector<std::vector<double>>&)> activation_fn;

public:
    std::vector<std::vector<double>> weights;
    std::vector<std::vector<double>> biases;
    std::vector<std::vector<double>> output;
    
    DenseLayer(int n_inputs, int n_neurons, 
               const std::string& activation_fn_name = "sigmoid",
               const std::string& weight_init_strategy = "random_scaled", 
               double weight_init_scale = 0.01) 
        : n_inputs(n_inputs), n_neurons(n_neurons), 
          activation_fn_name(activation_fn_name),
          weight_init_strategy(weight_init_strategy),
          weight_init_scale(weight_init_scale) {
        
        // Initialize weights based on the selected strategy
        if (weight_init_strategy == "random_scaled") {
            // Simple scaled random initialization
            std::random_device rd;
            std::mt19937 gen(rd());
            std::normal_distribution<double> dist(0.0, 1.0);
            
            weights = std::vector<std::vector<double>>(n_inputs, std::vector<double>(n_neurons));
            for (int i = 0; i < n_inputs; i++) {
                for (int j = 0; j < n_neurons; j++) {
                    weights[i][j] = dist(gen) * weight_init_scale;
                }
            }
        } 
        else if (weight_init_strategy == "xavier_normal") {
            // Xavier/Glorot normal initialization (good for sigmoid/tanh)
            double stddev = std::sqrt(2.0 / (n_inputs + n_neurons));
            std::random_device rd;
            std::mt19937 gen(rd());
            std::normal_distribution<double> dist(0.0, stddev);
            
            weights = std::vector<std::vector<double>>(n_inputs, std::vector<double>(n_neurons));
            for (int i = 0; i < n_inputs; i++) {
                for (int j = 0; j < n_neurons; j++) {
                    weights[i][j] = dist(gen);
                }
            }
        } 
        else if (weight_init_strategy == "he_uniform") {
            // He uniform initialization (good for ReLU)
            double limit = std::sqrt(6.0 / n_inputs);
            std::random_device rd;
            std::mt19937 gen(rd());
            std::uniform_real_distribution<double> dist(-limit, limit);
            
            weights = std::vector<std::vector<double>>(n_inputs, std::vector<double>(n_neurons));
            for (int i = 0; i < n_inputs; i++) {
                for (int j = 0; j < n_neurons; j++) {
                    weights[i][j] = dist(gen);
                }
            }
        } 
        else {
            throw std::invalid_argument("Unsupported weight initialization strategy: " + weight_init_strategy);
        }

        // Initialize biases to zero (standard practice)
        biases = std::vector<std::vector<double>>(1, std::vector<double>(n_neurons, 0.0));
        
        // Set activation function
        if (activation_fn_name == "sigmoid") {
            activation_fn = sigmoid;
        } else if (activation_fn_name == "relu") {
            activation_fn = relu;
        } else if (activation_fn_name == "linear") {
            activation_fn = linear;
        } else {
            throw std::invalid_argument("Unsupported activation: " + activation_fn_name);
        }
    }
    
    // Forward method remains the same as before
    std::vector<std::vector<double>> forward(const std::vector<std::vector<double>>& inputs) {
        // Matrix multiplication and bias addition
        std::vector<std::vector<double>> z = add_matrices(dot_product(inputs, weights), biases);
        
        // Apply activation function
        output = activation_fn(z);
        return output;
    }
};

Key enhancements in this updated class:

• Added parameters to specify the initialization strategy and scale.
• Implemented conditional logic to select the appropriate initialization method.
• Maintained bias initialization at zero (this is standard practice).
• Kept our existing activation function selection logic.

The beauty of this approach is that we can now easily experiment with different initialization strategies for different layers in our network. For example, we might use He initialization for ReLU layers and Xavier initialization for sigmoid layers.

Before we can verify our initialization strategies, we need some utility functions to calculate statistics:

#include <numeric>
#include <cmath>

double calculate_mean(const std::vector<std::vector<double>>& matrix) {
    double sum = 0.0;
    int count = 0;
    
    for (const auto& row : matrix) {
        for (double val : row) {
            sum += val;
            count++;
        }
    }
    
    return sum / count;
}

double calculate_std(const std::vector<std::vector<double>>& matrix) {
    double mean = calculate_mean(matrix);
    double sum_squared_diff = 0.0;
    int count = 0;
    
    for (const auto& row : matrix) {
        for (double val : row) {
            sum_squared_diff += (val - mean) * (val - mean);
            count++;
        }
    }
    
    return std::sqrt(sum_squared_diff / count);
}

To ensure our initialization strategies are working as expected, let's build a simple neural network and verify the statistical properties of the initialized weights:

#include <iostream>
#include <iomanip>

int main() {
    // Create sample input
    std::vector<std::vector<double>> X_sample = {{0.2, -0.4, 1.5}};
    std::cout << "Input X (shape " << X_sample.size() << "x" << X_sample[0].size() << "):\n";
    for (const auto& row : X_sample) {
        std::cout << "[";
        for (size_t i = 0; i < row.size(); i++) {
            std::cout << std::setw(6) << std::fixed << std::setprecision(1) << row[i];
            if (i < row.size() - 1) std::cout << " ";
        }
        std::cout << "]\n";
    }
    std::cout << "\n";

    // Build layers with different initialization strategies
    // Layer 1: ReLU with random scaled (custom scale)
    DenseLayer l1(3, 64, "relu", "random_scaled", 0.1);

    // Layer 2: Sigmoid with Xavier normal
    DenseLayer l2(64, 32, "sigmoid", "xavier_normal");

    // Layer 3: ReLU with He uniform
    DenseLayer l3(32, 2, "relu", "he_uniform");

    // Calculate and compare expected vs. actual standard deviations
    std::cout << "Standard deviations for initialized weights:\n";

    // Random scaled: should be close to the scale parameter
    double expected_std_l1 = 0.1;  // weight_init_scale
    double actual_std_l1 = calculate_std(l1.weights);
    std::cout << "  L1 (random_scaled): Expected ~" << std::fixed << std::setprecision(4) 
              << expected_std_l1 << ", Actual " << actual_std_l1 << "\n";

    // Xavier normal: stddev = sqrt(2/(fan_in + fan_out))
    double expected_std_l2 = std::sqrt(2.0 / (64 + 32));
    double actual_std_l2 = calculate_std(l2.weights);
    std::cout << "  L2 (xavier_normal): Expected ~" << expected_std_l2 
              << ", Actual " << actual_std_l2 << "\n";

    // He uniform: for U(-limit,limit), stddev = sqrt(2/fan_in)
    double expected_std_l3 = std::sqrt(2.0 / 32);
    double actual_std_l3 = calculate_std(l3.weights);
    std::cout << "  L3 (he_uniform): Expected ~" << expected_std_l3 
              << ", Actual " << actual_std_l3 << "\n";

    return 0;
}

This code:

• Creates a sample input and three layers with different initialization strategies.
• Calculates the expected standard deviation for each initialization method.
• Compares it with the actual standard deviation of the initialized weights.
• Helps us confirm that our implementation matches the theoretical expectations.

When we run this code, we can see how our different initialization strategies produce weight distributions with the expected statistical properties:

Input X (shape 1x3):
[   0.2  -0.4   1.5]

Standard deviations for initialized weights:
  L1 (random_scaled): Expected ~0.1000, Actual 0.0987
  L2 (xavier_normal): Expected ~0.1443, Actual 0.1421
  L3 (he_uniform): Expected ~0.2500, Actual 0.2234

Looking at the results, we can see that the actual standard deviations closely match our expected values, with minor variations due to random sampling. This confirms that our implementations are working correctly:

• For the random scaled initialization (L1), the standard deviation is very close to our specified scale of 0.1.
• The Xavier normal initialization (L2) produces weights with a standard deviation near the theoretical value based on fan-in and fan-out.
• The He uniform initialization (L3) generates weights with a standard deviation that approximates our expected value for ReLU layers.

This verification step is crucial because it confirms that our implementations are working correctly. Proper initialization ensures that signals can flow through the network without vanishing or exploding, setting the stage for effective training.

While proper initialization sets your network up for success, it's also valuable to monitor how weight distributions evolve during training. As your network learns, the weights will naturally shift from their initial distribution. Tracking statistics like mean, standard deviation, and the percentage of weights near zero can help you identify potential training issues such as vanishing or exploding gradients. Many deep learning frameworks provide built-in tools for visualizing weight distributions through histograms or summary statistics, which can be invaluable for debugging and optimizing your models.

Congratulations! You've now mastered weight initialization strategies, a critical component in building effective neural networks. We've explored why initialization matters, implemented powerful strategies like Xavier/Glorot and He initialization, and enhanced our DenseLayer class to support different initialization methods based on the specific needs of each layer. You've learned how to choose the right strategy for different activation functions and how to verify that your initialization is working as expected.

In the upcoming practice section, you'll have the opportunity to experiment with these initialization strategies and observe how they impact network behavior. After completing this course, you'll be ready to move on to the third course in our series, titled "Training Neural Networks: the Backpropagation Algorithm", where we'll learn how to efficiently train our networks using gradient-based optimization, building on the solid foundation of network architecture and initialization we've established.