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

JavaScript

const math = require('mathjs');
const randomNormal = require('random-normal');

/**
 * Initialize weights with random values scaled by a factor.
 * @param {number} nInputs - Number of input features.
 * @param {number} nNeurons - Number of neurons in the layer.
 * @param {number} scale - Scaling factor for the weights.
 * @returns {math.Matrix} Initialized weight matrix.
 */
function randomScaledInit(nInputs, nNeurons, scale = 0.01) {
    // Create a matrix of shape [nInputs, nNeurons] with random normal values
    const weights = [];
    for (let i = 0; i < nInputs; i++) {
        const row = [];
        for (let j = 0; j < nNeurons; j++) {
            // Draw from normal distribution with mean 0, std 1, then scale
            row.push(randomNormal({ mean: 0, dev: 1 }) * scale);
        }
        weights.push(row);
    }
    return math.matrix(weights);
}

const math = require('mathjs');
const randomNormal = require('random-normal');

/**
 * Initialize weights with random values scaled by a factor.
 * @param {number} nInputs - Number of input features.
 * @param {number} nNeurons - Number of neurons in the layer.
 * @param {number} scale - Scaling factor for the weights.
 * @returns {math.Matrix} Initialized weight matrix.
 */
function randomScaledInit(nInputs, nNeurons, scale = 0.01) {
    // Create a matrix of shape [nInputs, nNeurons] with random normal values
    const weights = [];
    for (let i = 0; i < nInputs; i++) {
        const row = [];
        for (let j = 0; j < nNeurons; j++) {
            // Draw from normal distribution with mean 0, std 1, then scale
            row.push(randomNormal({ mean: 0, dev: 1 }) * scale);
        }
        weights.push(row);
    }
    return math.matrix(weights);
}

JavaScript

/**
 * 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);
}

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;
    }
}

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);

text

Input X (shape [1, 3]):
[ [ 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.2325

Testing forward pass:
Final output shape: [1, 2]
Final output: [ [ 0.8234, 1.2156 ] ]

Input X (shape [1, 3]):
[ [ 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.2325

Testing forward pass:
Final output shape: [1, 2]
Final output: [ [ 0.8234, 1.2156 ] ]

JavaScript

const math = require('mathjs');
const randomNormal = require('random-normal');

/**
 * Initialize weights with random values scaled by a factor.
 * @param {number} nInputs - Number of input features.
 * @param {number} nNeurons - Number of neurons in the layer.
 * @param {number} scale - Scaling factor for the weights.
 * @returns {math.Matrix} Initialized weight matrix.
 */
function randomScaledInit(nInputs, nNeurons, scale = 0.01) {
    // Create a matrix of shape [nInputs, nNeurons] with random normal values
    const weights = [];
    for (let i = 0; i < nInputs; i++) {
        const row = [];
        for (let j = 0; j < nNeurons; j++) {
            // Draw from normal distribution with mean 0, std 1, then scale
            row.push(randomNormal({ mean: 0, dev: 1 }) * scale);
        }
        weights.push(row);
    }
    return math.matrix(weights);
}

JavaScript

/**
 * 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);
}

JavaScript

/**
 * 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);
}

/**
 * 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);
}

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;
    }
}

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);

text

Input X (shape [1, 3]):
[ [ 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.2325

Testing forward pass:
Final output shape: [1, 2]
Final output: [ [ 0.8234, 1.2156 ] ]

JavaScript

/**
 * 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