Deep Neural Representations & Normalization Mathematics

💡 1. Core Intuition & Concepts

Deep neural networks transform non-linear input manifolds through successive affine and activation layers. Normalization techniques enforce scale-invariance and variance bounds, preventing internal covariate shift and vanishing/exploding gradients across hundreds of residual layers.

📐 2. Mathematical Formulations & Derivations

Universal Approximation Theorem (Hornik / Cybenko)

\forall f \in C(K), \; \forall \epsilon > 0, \; \exists N \in \mathbb{N}: \; \left| f(x) - \sum_{i=1}^N v_i \sigma(w_i^T x + b_i) \right| < \epsilon

A single hidden layer with non-linear activation σ can approximate any continuous function on compact domain K.

Layer Normalization Formulation

\text{LN}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \gamma + \beta \qquad \mu = \frac{1}{D}\sum_{i=1}^D x_i, \quad \sigma^2 = \frac{1}{D}\sum_{i=1}^D (x_i - \mu)^2

Normalizes across feature dimension D with learned scale γ and shift β parameters.

Root Mean Square Normalization (RMSNorm)

\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \odot \gamma \qquad \text{RMS}(x) = \sqrt{\frac{1}{D}\sum_{i=1}^D x_i^2 + \epsilon}

Omits mean re-centering, reducing memory bandwidth by 7-10% while preserving scale invariance.

Scale-Invariance & Gradient Damping Property

\frac{\partial \text{RMSNorm}(\alpha x)}{\partial (\alpha x)} = \frac{1}{\alpha} \frac{\partial \text{RMSNorm}(x)}{\partial x} \implies \text{Gradients scale inversely with activation magnitude}

Automatically suppresses exploding activation surges during deep transformer training.

• Let y_i = RMSNorm(x)_i = (x_i / RMS(x)) * gamma_i, where RMS(x) = ((1/D) * sum_{k=1}^D x_k^2 + eps)^(1/2).

• Compute partial derivative ∂RMS(x)/∂x_i = (1 / (2 RMS(x))) * (2 x_i / D) = x_i / (D * RMS(x)).

• By the multivariable chain rule: ∂L/∂x_i = sum_{j=1}^D (∂L/∂y_j) * (∂y_j/∂x_i).

• Differentiating y_j: ∂y_j/∂x_i = (gamma_j / RMS(x)) * delta_{ij} - (x_j gamma_j / RMS(x)^2) * (∂RMS(x)/∂x_i) = (gamma_i delta_{ij} / RMS(x)) - (x_j gamma_j x_i / (D RMS(x)^3)).

• Substitute and factor: ∂L/∂x_i = (gamma_i / RMS(x)) * (∂L/∂y_i) - (x_i / (D RMS(x)^3)) * sum_{j=1}^D (∂L/∂y_j * gamma_j x_j). Q.E.D.

⚙️ 3. Step-by-Step Computational Mechanism

1

Coordinate Standardization

Normalization projects activations onto an invariant hypersphere with unit variance.

2

Residual Stream Stabilization

In Pre-LN/Pre-RMSNorm architectures, signal variance across L layers grows as O(sqrt(L)) rather than exponentially O(2^L).

3

GPU Memory Bandwidth Optimization

RMSNorm replaces two reduction passes (mean + variance) with a single sum-of-squares pass, accelerating fused CUDA kernels.

💻 4. Code from Scratch (python)

math_neural_norm.py

import numpy as np

# Exact Vectorized RMSNorm Forward & Backward Implementation
class RMSNormScratch:
    def __init__(self, dim: int, eps: float = 1e-6):
        self.eps = eps
        self.gamma = np.ones(dim)

    def forward(self, x: np.ndarray) -> np.ndarray:
        self.x = x
        self.rms = np.sqrt(np.mean(x**2, axis=-1, keepdims=True) + self.eps)
        self.x_norm = x / self.rms
        return self.x_norm * self.gamma

    def backward(self, dout: np.ndarray) -> np.ndarray:
        D = self.x.shape[-1]
        # dL/dgamma
        self.dgamma = np.sum(dout * self.x_norm, axis=0)
        # Analytical dL/dx
        sum_term = np.sum(dout * self.gamma * self.x, axis=-1, keepdims=True)
        dx = (dout * self.gamma / self.rms) - (self.x / (D * self.rms**3)) * sum_term
        return dx

if __name__ == "__main__":
    layer = RMSNormScratch(dim=4)
    x = np.array([[1.0, 2.0, 3.0, 4.0]])
    out = layer.forward(x)
    dx = layer.backward(np.ones_like(out))
    print("Forward output: ", np.round(out, 4))
    print("Backward grad dx:", np.round(dx, 4))

🧠 5. Comprehension Checkpoint

Answer all 1 questions correctly to complete the chapter · 0 / 1 done

Q1/1 Why is RMSNorm preferred over standard LayerNorm in modern LLMs (Llama 3, DeepSeek-V3, Mistral)?

Learn ML from Scratch: Beginner to Advanced