Weight Initialization & Layer Normalization

Deep neural networks consist of stacked non-linear matrix transformations. Without careful parameter initialization, activations and backpropagated gradients either explode exponentially towards numerical infinity or collapse towards zero as depth increases.

Modern deep architecture stability rests on two foundational mathematical developments: Variance-Preserving Weight Initialization (Glorot and He) and Activation Normalization (LayerNorm and BatchNorm).

1. Mathematical Proof: Variance Collapse in Naive Networks#

Consider a linear layer computing pre-activations $y$:

Assume input features $x_i$ and weights $w_i$ are independent random variables with zero mean ($\mathbb{E}[x_i] = 0$, $\mathbb{E}[w_i] = 0$) and identical variances $\text{Var}(x)$ and $\text{Var}(w)$.

To preserve activation scale across the forward pass, we must enforce $n_{\text{in}} \text{Var}(w) = 1$, which requires $\text{Var}(w) = \frac{1}{n_{\text{in}}}$.

2. Xavier/Glorot vs He/Kaiming Derivations#

Xavier/Glorot Initialization (Symmetric Activations: Tanh / Linear)#

Glorot and Bengio (2010) analyzed both the forward pass and backward gradient flow:

• Forward pass stability requires: $\text{Var}(w) = \frac{1}{n_{\text{in}}}$
• Backward pass stability requires: $\text{Var}(w) = \frac{1}{n_{\text{out}}}$

Taking the harmonic mean between forward and backward constraints:

For a uniform distribution $U[-a, a]$, where $\text{Var} = \frac{a^2}{3}$:

He/Kaiming Initialization (Rectified Linear Unit: ReLU)#

Xavier initialization assumes activations are linear or zero-mean symmetric. Under the ReLU activation $f(z) = \max(0, z)$, all negative inputs are set to zero:

Because ReLU halves the active variance of incoming signals, He et al. (2015) proved that weights must be scaled by an additional factor of 2:

For a Gaussian normal distribution:

3. Layer Normalization (LayerNorm)#

While Batch Normalization computes statistics across the mini-batch dimension, it degrades when batch sizes are small and fails on variable-length autoregressive sequential tokens.

Layer Normalization (Ba, Kiros, Hinton, 2016) computes the mean and variance across the feature (channel) dimension for each individual sample independently:

Given feature vector $x \in \mathbb{R}^d$:

Where $\gamma \in \mathbb{R}^d$ and $\beta \in \mathbb{R}^d$ are learnable affine parameters initialized to $\gamma = 1$ and $\beta = 0$.

Because LayerNorm has zero cross-sample dependencies, it performs identically during training and evaluation, serving as the universal normalization layer across all modern Transformer architectures (GPT, LLaMA, BERT).

4. From-Scratch Implementation: Zero Dependencies#

The following implementations construct He/Kaiming Initialization and a complete LayerNorm Forward Pass from scratch in pure Python and pure Rust.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import math
import random
from typing import List

class DeepLearningPrimitives:
    @staticmethod
    def he_normal_init(fan_in: int, fan_out: int) -> List[List[float]]:
        # Variance = 2 / fan_in => StdDev = sqrt(2 / fan_in)
        std = math.sqrt(2.0 / fan_in)
        random.seed(42)
        weights = []
        for _ in range(fan_out):
            row = [random.gauss(0.0, std) for _ in range(fan_in)]
            weights.append(row)
        return weights

    @staticmethod
    def layer_norm(
        x: List[List[float]],
        gamma: List[float],
        beta: List[float],
        eps: float = 1e-5
    ) -> List[List[float]]:
        num_samples = len(x)
        d = len(x[0])
        output = []

        for i in range(num_samples):
            # Compute mean across feature dimension d
            mean = sum(x[i][j] for j in range(d)) / d
            # Compute variance across feature dimension d
            variance = sum((x[i][j] - mean) ** 2 for j in range(d)) / d

            std_inv = 1.0 / math.sqrt(variance + eps)
            # Normalize, scale, and shift
            normed_row = [
                gamma[j] * ((x[i][j] - mean) * std_inv) + beta[j]
                for j in range(d)
            ]
            output.append(normed_row)

        return output

if __name__ == "__main__":
    # Test He initialization
    W = DeepLearningPrimitives.he_normal_init(fan_in=4, fan_out=2)
    print("He Initialized Weight Matrix (2x4):")
    for row in W:
        print(" ", [round(v, 4) for v in row])

    # Test LayerNorm
    tokens = [
        [1.0, 2.0, 3.0, 4.0],
        [10.0, 20.0, 30.0, 40.0]
    ]
    gamma = [1.0] * 4
    beta = [0.0] * 4
    normed = DeepLearningPrimitives.layer_norm(tokens, gamma, beta)
    print("\nLayerNorm Normalized Tokens:")
    for row in normed:
        print(" ", [round(v, 4) for v in row])

5. Active Knowledge Checks#