Weight Initialization Strategies for MLPs in JavaScript

Introduction

Welcome to the fourth and final lesson in our course on The MLP Architecture: Activations & Initialization! So far, you'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 this lesson, you'll learn why proper weight initialization is crucial, see 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.

Why Weight Initialization Matters

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. Poor weight initialization can lead to several problems: 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. Vanishing Gradients : If weights are too small, the signals flowing through the network will diminish with each layer, causing gradients to approach zero during training. This makes learning extremely slow, especially in deeper layers. Exploding Gradients : 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 \to 0.5 \to 0.25 \to 0.125 \to ... \to 0.001 (signal vanishes) With weights that are too large: 1 \to 2 \to 4 \to 8 \to ... \to 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.

Join the 1M+ learners on CodeSignal

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

Implementing Different Strategies in Our Layer

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

JavaScript

// --- Enhanced DenseLayer Class ---

class DenseLayer {
    /**
     * @param {number} nInputs
     * @param {number} nNeurons
     * @param {string} activationFnName
     * @param {string} weightInitStrategy
     * @param {number} weightInitScale
     */
    constructor(
        nInputs,
        nNeurons,
        activationFnName = 'sigmoid',
        weightInitStrategy = 'random_scaled',
        weightInitScale = 0.01
    ) {
        this.nInputs = nInputs;
        this.nNeurons = nNeurons;
        this.activationFnName = activationFnName;
        this.weightInitStrategy = weightInitStrategy;
        this.weightInitScale = weightInitScale;

        // Initialize weights based on the selected strategy
        if (weightInitStrategy === 'random_scaled') {
            this.weights = randomScaledInit(nInputs, nNeurons, weightInitScale);
        } else if (weightInitStrategy === 'xavier_normal') {
            this.weights = xavierNormalInit(nInputs, nNeurons);
        } else if (weightInitStrategy === 'he_uniform') {
            this.weights = heUniformInit(nInputs, nNeurons);
        } else {
            throw new Error(`Unsupported weight initialization strategy: ${weightInitStrategy}`);
        }

        // Initialize biases to zero (standard practice)
        this.biases = math.zeros(1, nNeurons);
        this.output = null;

        // Set activation function
        if (activationFnName === 'sigmoid') {
            this.activationFn = sigmoid;
        } else if (activationFnName === 'relu') {
            this.activationFn = relu;
        } else if (activationFnName === 'softmax') {
            this.activationFn = softmax;
        } else if (activationFnName === 'linear') {
            this.activationFn = linear;
        } else {
            throw new Error(`Unsupported activation function: ${activationFnName}`);
        }
    }

    forward(inputs) {
        const weightedSum = math.multiply(inputs, this.weights);
        const outputBeforeActivation = math.add(weightedSum, this.biases);
        this.output = this.activationFn(outputBeforeActivation);
        return this.output;
    }
}

// --- Enhanced DenseLayer Class ---

class DenseLayer {
    /**
     * @param {number} nInputs
     * @param {number} nNeurons
     * @param {string} activationFnName
     * @param {string} weightInitStrategy
     * @param {number} weightInitScale
     */
    constructor(
        nInputs,
        nNeurons,
        activationFnName = 'sigmoid',
        weightInitStrategy = 'random_scaled',
        weightInitScale = 0.01
    ) {
        this.nInputs = nInputs;
        this.nNeurons = nNeurons;
        this.activationFnName = activationFnName;
        this.weightInitStrategy = weightInitStrategy;
        this.weightInitScale = weightInitScale;

        // Initialize weights based on the selected strategy
        if (weightInitStrategy === 'random_scaled') {
            this.weights = randomScaledInit(nInputs, nNeurons, weightInitScale);
        } else if (weightInitStrategy === 'xavier_normal') {
            this.weights = xavierNormalInit(nInputs, nNeurons);
        } else if (weightInitStrategy === 'he_uniform') {
            this.weights = heUniformInit(nInputs, nNeurons);
        } else {
            throw new Error(`Unsupported weight initialization strategy: ${weightInitStrategy}`);
        }

        // Initialize biases to zero (standard practice)
        this.biases = math.zeros(1, nNeurons);
        this.output = null;

        // Set activation function
        if (activationFnName === 'sigmoid') {
            this.activationFn = sigmoid;
        } else if (activationFnName === 'relu') {
            this.activationFn = relu;
        } else if (activationFnName === 'softmax') {
            this.activationFn = softmax;
        } else if (activationFnName === 'linear') {
            this.activationFn = linear;
        } else {
            throw new Error(`Unsupported activation function: ${activationFnName}`);
        }
    }

    forward(inputs) {
        const weightedSum = math.multiply(inputs, this.weights);
        const outputBeforeActivation = math.add(weightedSum, this.biases);
        this.output = this.activationFn(outputBeforeActivation);
        return this.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 using math.zeros(1, nNeurons) for consistency.
Kept our existing activation function selection logic.
All weights are now mathjs matrices for consistent matrix operations.

Now you can easily experiment with different initialization strategies for different layers in your network. For example, you might use He initialization for ReLU layers and Xavier for sigmoid layers.

Verifying Our Initialization Strategies

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.

JavaScript

const math = require('mathjs');

// Create sample input as mathjs matrix
const X_sample = math.matrix([[0.2, -0.4, 1.5]]);
console.log(`Input X (shape [${math.size(X_sample)._data.join(', ')}]):`);
console.log(X_sample._data);

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

// Layer 2: Sigmoid with Xavier normal
const l2 = new DenseLayer(64, 32, 'sigmoid', 'xavier_normal');

// Layer 3: ReLU with He uniform
const l3 = new DenseLayer(32, 2, 'relu', 'he_uniform');

// Calculate and compare expected vs. actual standard deviations
console.log('\nStandard deviations for initialized weights:');

// Random scaled: should be close to the scale parameter
const expectedStdL1 = l1.weightInitScale;
const actualStdL1 = math.std(l1.weights._data.flat());
console.log(`  L1 (random_scaled): Expected ~${expectedStdL1.toFixed(4)}, Actual ${actualStdL1.toFixed(4)}`);

// Xavier normal: stddev = sqrt(2/(fan_in + fan_out))
const expectedStdL2 = Math.sqrt(2.0 / (l2.nInputs + l2.nNeurons));
const actualStdL2 = math.std(l2.weights._data.flat());
console.log(`  L2 (xavier_normal): Expected ~${expectedStdL2.toFixed(4)}, Actual ${actualStdL2.toFixed(4)}`);

// He uniform: for U(-limit,limit), stddev = sqrt(2/fan_in)
const expectedStdL3 = Math.sqrt(2.0 / l3.nInputs);
const actualStdL3 = math.std(l3.weights._data.flat());
console.log(`  L3 (he_uniform): Expected ~${expectedStdL3.toFixed(4)}, Actual ${actualStdL3.toFixed(4)}`);

// Test forward pass to ensure everything works
console.log('\nTesting forward pass:');
const out1 = l1.forward(X_sample);
const out2 = l2.forward(out1);
const out3 = l3.forward(out2);
console.log(`Final output shape: [${math.size(out3)._data.join(', ')}]`);
console.log('Final output:', out3._data);

This code:

Creates a sample input as a mathjs matrix 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.
Tests a forward pass to ensure all components work together.
Helps us confirm that our implementation matches the theoretical expectations.

/** * Initialize weights using Xavier/Glorot normal initialization. * @param {number} nInputs - Number of input features. * @param {number} nNeurons - Number of neurons in the layer. * @returns {math.Matrix} Initialized weight matrix. */ function xavierNormalInit(nInputs, nNeurons) { // Calculate standard deviation based on fan-in and fan-out const stddev = Math.sqrt(2.0 / (nInputs + nNeurons)); const weights = []; for (let i = 0; i < nInputs; i++) { const row = []; for (let j = 0; j < nNeurons; j++) { // Draw from normal distribution with calculated stddev row.push(randomNormal({ mean: 0, dev: stddev })); } weights.push(row); } return math.matrix(weights); }

/** * Initialize weights using He uniform initialization. * @param {number} nInputs - Number of input features. * @param {number} nNeurons - Number of neurons in the layer. * @returns {math.Matrix} Initialized weight matrix. */ function heUniformInit(nInputs, nNeurons) { // Calculate limit based on fan-in (input connections) const limit = Math.sqrt(6.0 / nInputs); const weights = []; for (let i = 0; i < nInputs; i++) { const row = []; for (let j = 0; j < nNeurons; j++) { // Draw from uniform distribution within calculated bounds row.push(math.random(-limit, limit)); } weights.push(row); } return math.matrix(weights); }

Introduction

Why Weight Initialization Matters

Join the 1M+ learners on CodeSignal

Random Scaled Initialization

Xavier/Glorot Initialization

He Initialization

Implementing Different Strategies in Our Layer

Verifying Our Initialization Strategies

Output Discussion

Conclusion and Next Steps