Xavier/He Weight Initialization - Neural Networks

Xavier/He Weight Initialization

What you'll build

A function that draws a weight matrix with Xavier (Glorot) normal or He (Kaiming) normal initialization and returns the empirical standard deviation of the sampled weights, so you can check it lands on the theoretical value.

Background

If weights are too small, activations shrink layer after layer and gradients vanish; too large, and they explode. Good initialization picks the variance so that the signal keeps roughly the same scale through the network:

• Xavier (Glorot & Bengio, 2010) balances the forward and backward pass for symmetric activations like tanh: Var(W) = 2 / (n_in + n_out).
• He (He et al., 2015) accounts for ReLU zeroing half of its inputs, which halves the variance, by doubling it: Var(W) = 2 / n_in.

Formula

Xavier normal:  W = randn(n_in, n_out) * sqrt(2 / (n_in + n_out))
He normal:      W = randn(n_in, n_out) * sqrt(2 / n_in)
return round(std(W), 4)

Inputs

• n_in: number of input units (fan-in), positive int
• n_out: number of output units (fan-out), positive int
• method: 'xavier' or 'he'

Expected return format

A float: np.std(W) of the sampled (n_in, n_out) matrix, rounded to 4 decimals. It is close to, but not exactly, the theoretical std (sampling noise).

Edge cases and conventions

• Call np.random.seed(42) at the start of the function (the starter code does this), then draw exactly one np.random.randn(n_in, n_out) and scale it. The expected values depend on this exact sequence.
• Any other method may raise a ValueError (not tested).

Constraints

NumPy only.

Follow-up

• The uniform variants draw from U(-a, a) with a = sqrt(3) * std. Why sqrt(3)? (The variance of U(-a, a) is a² / 3.)
• PyTorch's nn.Linear default is Kaiming-uniform with a = sqrt(5), which gives U(-1/sqrt(n_in), 1/sqrt(n_in)). Why do transformers often use a small fixed std like 0.02 instead?

Input: initialize_weights(784, 256, "he")
Output: 0.0505

Input: initialize_weights(100, 50, "xavier")
Output: 0.1151

import numpy as np

def initialize_weights(n_in, n_out, method='xavier'):
    """
    Sample an (n_in, n_out) weight matrix with Xavier or He normal initialization.

    Args:
        n_in: Number of input units (fan-in)
        n_out: Number of output units (fan-out)
        method: 'xavier' (std = sqrt(2 / (n_in + n_out))) or 'he' (std = sqrt(2 / n_in))

    Returns:
        std: Empirical standard deviation of the sampled weights, rounded to 4 decimals
    """
    np.random.seed(42)  # For reproducibility: draw one np.random.randn(n_in, n_out) after this
    # Your code here
    pass