From 75e934a397da418656715fef6c3a20c1d78fa80a Mon Sep 17 00:00:00 2001
From: Thomas Hart <thomashartdev@gmail.com>
Date: Fri, 24 Jul 2026 21:27:58 +0000
Subject: [PATCH] feat: Add activation functions and Xavier/He weight init

---
 README.md                 |  26 ++++
 src/activations.py        | 250 ++++++++++++++++++++++++++++++++++++++
 tests/test_activations.py | 214 ++++++++++++++++++++++++++++++++
 3 files changed, 490 insertions(+)
 create mode 100644 src/activations.py
 create mode 100644 tests/test_activations.py

diff --git a/README.md b/README.md
index 011b29a..3f7e97b 100644
--- a/README.md
+++ b/README.md
@@ -30,6 +30,11 @@ Every deep learning framework is, at its core, a graph that records operations a
 - The `softmax - onehot` output-delta shortcut and `delta_prev = (delta @ Wᵀ) ⊙ act'(z)` for hidden layers
 - Xavier vs He weight initialization chosen from the activation to keep signal variance stable across depth
 - Learning a non-linearly-separable target (XOR, interleaved spirals) that a single linear model cannot fit
+- Local activation derivatives for the chain rule (`act'(z)`): tanh, sigmoid, ReLU, leaky ReLU, identity
+- Glorot/Xavier initialization: variance `2/(fan_in+fan_out)` balancing forward and backward signal
+- He/Kaiming initialization: variance `2/fan_in` restoring the half of the signal ReLU drops
+- Forward variance profiling across depth: naive `N(0,1)` weights explode; correct scales stay `O(1)`
+- Vanishing activations under mismatched init (Xavier on a deep ReLU stack) measured, not just described
 
 ## What's implemented
 
@@ -37,6 +42,7 @@ Every deep learning framework is, at its core, a graph that records operations a
 - **Linear regression, two ways**: `fit_normal_equation` solves ordinary least squares exactly by setting the gradient to zero and solving `(AᵀA + λR)θ = Aᵀy`, with optional ridge and a least-norm fallback when `AᵀA` is singular. `fit_sgd` fits the same model with minibatch gradient descent on standardized features, then maps the weights back to raw feature space. Both are checked to recover the true coefficients and to agree with each other, so the exact-vs-iterative tradeoff is measurable.
 - **Logistic regression + cross-entropy, decision boundary demo**: `src/logreg.py` trains a binary classifier by minibatch SGD, using the fact that the gradient of the mean cross-entropy with respect to the logits is exactly `sigmoid(z) - y`, the same residual form linear regression has. The sigmoid branches on the sign of the logit so `exp` never overflows, and the loss is computed as `softplus(z) - y·z` through `numpy.logaddexp` so a confidently-wrong prediction gives a large finite loss instead of `inf`. `decision_boundary` returns the straight line where the model sits at 50 percent for a two-feature problem, the level set `w·x + b = 0`.
 - **Multilayer perceptron with hand-derived backprop**: `src/mlp.py` stacks affine layers with `tanh` or `relu` and a softmax head, and trains a multiclass classifier by minibatch SGD. The backward pass is written out by hand as one recursion on the per-layer delta rather than delegated to an autodiff engine: the output delta is the `softmax - onehot` residual, each hidden delta is `(delta_next @ W_nextᵀ) ⊙ act'(z)`, and the parameter gradients are `dW = a_prevᵀ @ delta` and `db = Σ delta`. Weights use He init for `relu` and Xavier for `tanh` so the signal variance holds across depth. The gradients are verified against central finite differences to a tight tolerance, and the model learns XOR and a three-arm spiral, targets a single hyperplane provably cannot separate. Ships with `make_xor` and `make_spiral` toy generators.
+- **Activation functions + weight initialization (Xavier/He) and why they matter**: `src/activations.py` is the dedicated treatment of the nonlinearity and the initial scale. Each activation (`linear`, `tanh`, `sigmoid`, `relu`, `leaky_relu`) exposes `forward` and a local `backward(z, grad_out)` that multiplies by `act'(z)`, so a hand-written backprop step can drop it in. Xavier/Glorot draws `N(0, 2/(fan_in+fan_out))` (or the matching uniform bound) to keep both forward and backward variance stable for symmetric activations; He/Kaiming draws `N(0, 2/fan_in)` so a ReLU stack does not quietly die after a few layers. `forward_variance_profile` stacks affine+activation layers from unit-variance noise and returns the per-layer activation variance: naive `N(0,1)` weights explode, He keeps a ReLU stack `O(1)`, and Xavier on the same ReLU stack fades, which is the usual silent failure mode when the scheme and the nonlinearity disagree.
 
 ## Usage
 
@@ -98,6 +104,26 @@ print(history[0], history[-1])   # cross-entropy falls over training
 print(model.predict_proba(X[:3]))  # per-class probabilities that sum to 1
 ```
 
+Compare init schemes by watching activation variance with depth:
+
+```python
+from src.activations import (
+    get_activation,
+    he_normal,
+    recommended_init,
+    forward_variance_profile,
+)
+
+act = get_activation("relu")
+print(recommended_init("relu"))   # "he"
+print(he_normal(64, 64).std())    # ~sqrt(2/64)
+
+# unit-scale noise, ten ReLU layers: He stays O(1), Xavier fades, naive explodes
+print(forward_variance_profile(10, 64, "relu", "he")[-1])
+print(forward_variance_profile(10, 64, "relu", "xavier")[-1])
+print(forward_variance_profile(6, 64, "linear", "naive")[-1])
+```
+
 Differentiate an arbitrary scalar expression:
 
 ```python
diff --git a/src/activations.py b/src/activations.py
new file mode 100644
index 0000000..65ecaaa
--- /dev/null
+++ b/src/activations.py
@@ -0,0 +1,250 @@
+"""Activation functions and variance-preserving weight initialization.
+
+A deep stack of affine layers multiplies variances: if each weight matrix has
+entries of order one, the pre-activation variance grows or shrinks exponentially
+with depth, and gradients follow. Activations change the story further (ReLU
+zeros half its inputs; tanh saturates), so the right scale depends on both fan
+sizes and the nonlinearity. Glorot/Xavier balances fan-in and fan-out for
+symmetric activations; He/Kaiming uses only fan-in and a factor of two for ReLU.
+This module implements the common activations with their local derivatives and
+those two init schemes, plus a forward variance profile that makes the blow-up
+under naive N(0,1) weights measurable.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+from numpy.typing import NDArray
+
+Array = NDArray[np.float64]
+
+InitScheme = Literal["xavier", "he", "naive"]
+ActivationName = Literal["linear", "tanh", "sigmoid", "relu", "leaky_relu"]
+
+
+@dataclass(frozen=True)
+class Activation:
+    """Named nonlinearity with a local backward for the chain rule.
+
+    `backward(z, grad_out)` multiplies the upstream gradient by act'(z), the
+    form a hand-written backprop step needs. `z` is the pre-activation.
+    """
+
+    name: str
+    forward: Callable[[Array], Array]
+    backward: Callable[[Array, Array], Array]
+
+
+def _linear_forward(z: Array) -> Array:
+    return np.asarray(z, dtype=np.float64)
+
+
+def _linear_backward(z: Array, grad_out: Array) -> Array:
+    return np.asarray(grad_out, dtype=np.float64) * np.ones_like(z, dtype=np.float64)
+
+
+def _tanh_forward(z: Array) -> Array:
+    return np.tanh(z)
+
+
+def _tanh_backward(z: Array, grad_out: Array) -> Array:
+    t = np.tanh(z)
+    return grad_out * (1.0 - t * t)
+
+
+def _sigmoid_forward(z: Array) -> Array:
+    # sign branch keeps exp argument non-positive so it never overflows
+    out = np.empty_like(z, dtype=np.float64)
+    pos = z >= 0.0
+    out[pos] = 1.0 / (1.0 + np.exp(-z[pos]))
+    ez = np.exp(z[~pos])
+    out[~pos] = ez / (1.0 + ez)
+    return out
+
+
+def _sigmoid_backward(z: Array, grad_out: Array) -> Array:
+    s = _sigmoid_forward(z)
+    return grad_out * s * (1.0 - s)
+
+
+def _relu_forward(z: Array) -> Array:
+    return np.maximum(z, 0.0)
+
+
+def _relu_backward(z: Array, grad_out: Array) -> Array:
+    # subgradient at 0 is taken as 0 (dead unit stays dead)
+    return grad_out * (z > 0.0)
+
+
+def _leaky_relu_forward(z: Array, alpha: float = 0.01) -> Array:
+    return np.where(z > 0.0, z, alpha * z)
+
+
+def _leaky_relu_backward(z: Array, grad_out: Array, alpha: float = 0.01) -> Array:
+    return grad_out * np.where(z > 0.0, 1.0, alpha)
+
+
+ACTIVATIONS: dict[str, Activation] = {
+    "linear": Activation("linear", _linear_forward, _linear_backward),
+    "tanh": Activation("tanh", _tanh_forward, _tanh_backward),
+    "sigmoid": Activation("sigmoid", _sigmoid_forward, _sigmoid_backward),
+    "relu": Activation("relu", _relu_forward, _relu_backward),
+    "leaky_relu": Activation("leaky_relu", _leaky_relu_forward, _leaky_relu_backward),
+}
+
+
+def get_activation(name: str) -> Activation:
+    if name not in ACTIVATIONS:
+        known = ", ".join(sorted(ACTIVATIONS))
+        raise ValueError(f"unknown activation {name!r}; choose one of: {known}")
+    return ACTIVATIONS[name]
+
+
+def recommended_init(activation: str) -> InitScheme:
+    """Pick Xavier for symmetric/saturating acts, He for ReLU-family."""
+    act = get_activation(activation)
+    if act.name in ("relu", "leaky_relu"):
+        return "he"
+    return "xavier"
+
+
+def _validate_fans(fan_in: int, fan_out: int) -> None:
+    if fan_in <= 0 or fan_out <= 0:
+        raise ValueError(
+            f"fan_in and fan_out must be positive, got {fan_in}, {fan_out}"
+        )
+
+
+def xavier_normal(
+    fan_in: int,
+    fan_out: int,
+    rng: np.random.Generator | None = None,
+    gain: float = 1.0,
+) -> Array:
+    """Glorot normal: N(0, gain^2 * 2 / (fan_in + fan_out)).
+
+    Balances variance of the forward pass (fan_in) and the backward pass
+    (fan_out). Default gain is 1; tanh is usually left at 1, linear too.
+    """
+    _validate_fans(fan_in, fan_out)
+    if gain <= 0.0:
+        raise ValueError("gain must be positive")
+    rng = rng if rng is not None else np.random.default_rng()
+    std = gain * np.sqrt(2.0 / (fan_in + fan_out))
+    return rng.standard_normal((fan_in, fan_out)) * std
+
+
+def xavier_uniform(
+    fan_in: int,
+    fan_out: int,
+    rng: np.random.Generator | None = None,
+    gain: float = 1.0,
+) -> Array:
+    """Glorot uniform: U(-a, a) with a = gain * sqrt(6 / (fan_in + fan_out))."""
+    _validate_fans(fan_in, fan_out)
+    if gain <= 0.0:
+        raise ValueError("gain must be positive")
+    rng = rng if rng is not None else np.random.default_rng()
+    bound = gain * np.sqrt(6.0 / (fan_in + fan_out))
+    return rng.uniform(-bound, bound, size=(fan_in, fan_out))
+
+
+def he_normal(
+    fan_in: int,
+    fan_out: int,
+    rng: np.random.Generator | None = None,
+    gain: float = 1.0,
+) -> Array:
+    """Kaiming normal: N(0, gain^2 * 2 / fan_in).
+
+    ReLU zeros roughly half the units, so the 2 restores unit variance after
+    the rectification. Only fan_in appears because the analysis is for the
+    forward signal; gain multiplies the std (leave at 1 for plain ReLU).
+    """
+    _validate_fans(fan_in, fan_out)
+    if gain <= 0.0:
+        raise ValueError("gain must be positive")
+    rng = rng if rng is not None else np.random.default_rng()
+    std = gain * np.sqrt(2.0 / fan_in)
+    return rng.standard_normal((fan_in, fan_out)) * std
+
+
+def he_uniform(
+    fan_in: int,
+    fan_out: int,
+    rng: np.random.Generator | None = None,
+    gain: float = 1.0,
+) -> Array:
+    """Kaiming uniform: U(-a, a) with a = gain * sqrt(6 / fan_in)."""
+    _validate_fans(fan_in, fan_out)
+    if gain <= 0.0:
+        raise ValueError("gain must be positive")
+    rng = rng if rng is not None else np.random.default_rng()
+    bound = gain * np.sqrt(6.0 / fan_in)
+    return rng.uniform(-bound, bound, size=(fan_in, fan_out))
+
+
+def init_weights(
+    fan_in: int,
+    fan_out: int,
+    scheme: InitScheme | str,
+    rng: np.random.Generator | None = None,
+    *,
+    distribution: Literal["normal", "uniform"] = "normal",
+    gain: float = 1.0,
+) -> Array:
+    """Dispatch to Xavier, He, or naive N(0,1) (the last is the control case)."""
+    rng = rng if rng is not None else np.random.default_rng()
+    if scheme == "naive":
+        _validate_fans(fan_in, fan_out)
+        if distribution == "normal":
+            return rng.standard_normal((fan_in, fan_out))
+        return rng.uniform(-1.0, 1.0, size=(fan_in, fan_out))
+    if scheme == "xavier":
+        if distribution == "normal":
+            return xavier_normal(fan_in, fan_out, rng, gain=gain)
+        return xavier_uniform(fan_in, fan_out, rng, gain=gain)
+    if scheme == "he":
+        if distribution == "normal":
+            return he_normal(fan_in, fan_out, rng, gain=gain)
+        return he_uniform(fan_in, fan_out, rng, gain=gain)
+    raise ValueError(f"unknown init scheme {scheme!r}; use xavier, he, or naive")
+
+
+def forward_variance_profile(
+    n_layers: int,
+    width: int,
+    activation: str,
+    scheme: InitScheme | str,
+    *,
+    n_samples: int = 2000,
+    seed: int = 0,
+    distribution: Literal["normal", "uniform"] = "normal",
+) -> list[float]:
+    """Mean activation variance after each of `n_layers` affine+act layers.
+
+    Starts from unit-variance isotropic inputs and applies W_l with the given
+    init, then the activation. A good scheme keeps the returned variances near
+    order 1; naive N(0,1) weights explode (or vanish after saturating acts).
+    """
+    if n_layers <= 0:
+        raise ValueError("n_layers must be positive")
+    if width <= 0:
+        raise ValueError("width must be positive")
+    if n_samples <= 0:
+        raise ValueError("n_samples must be positive")
+
+    act = get_activation(activation)
+    rng = np.random.default_rng(seed)
+    x = rng.standard_normal((n_samples, width))
+    variances: list[float] = []
+    for _ in range(n_layers):
+        w = init_weights(width, width, scheme, rng, distribution=distribution)
+        z = x @ w
+        x = act.forward(z)
+        variances.append(float(np.var(x)))
+    return variances
diff --git a/tests/test_activations.py b/tests/test_activations.py
new file mode 100644
index 0000000..a474a40
--- /dev/null
+++ b/tests/test_activations.py
@@ -0,0 +1,214 @@
+import numpy as np
+import pytest
+
+from src.activations import (
+    ACTIVATIONS,
+    forward_variance_profile,
+    get_activation,
+    he_normal,
+    he_uniform,
+    init_weights,
+    recommended_init,
+    xavier_normal,
+    xavier_uniform,
+)
+
+
+def _numeric_deriv(fn, z0: float, eps: float = 1e-6) -> float:
+    return (fn(z0 + eps) - fn(z0 - eps)) / (2.0 * eps)
+
+
+@pytest.mark.parametrize("name", sorted(ACTIVATIONS))
+def test_forward_backward_shapes_and_finite(name):
+    act = get_activation(name)
+    z = np.linspace(-3.0, 3.0, 21).reshape(3, 7)
+    a = act.forward(z)
+    g = act.backward(z, np.ones_like(z))
+    assert a.shape == z.shape
+    assert g.shape == z.shape
+    assert np.all(np.isfinite(a))
+    assert np.all(np.isfinite(g))
+
+
+@pytest.mark.parametrize(
+    "name,points",
+    [
+        ("tanh", (-1.5, -0.3, 0.0, 0.7, 2.0)),
+        ("sigmoid", (-2.0, -0.5, 0.0, 0.5, 2.0)),
+        ("relu", (-1.0, 0.5, 1.5)),
+        ("leaky_relu", (-1.0, 0.5, 1.5)),
+        ("linear", (-2.0, 0.0, 3.0)),
+    ],
+)
+def test_derivative_matches_finite_differences(name, points):
+    act = get_activation(name)
+    for z0 in points:
+        z = np.array([[z0]], dtype=np.float64)
+        analytic = float(act.backward(z, np.ones_like(z))[0, 0])
+
+        def _f(t: float) -> float:
+            return float(act.forward(np.array([[t]]))[0, 0])
+
+        numeric = _numeric_deriv(_f, z0)
+        assert analytic == pytest.approx(numeric, abs=1e-5, rel=1e-4)
+
+
+def test_relu_dead_at_and_below_zero():
+    act = get_activation("relu")
+    z = np.array([[-2.0, 0.0, 3.0]])
+    assert np.allclose(act.forward(z), [[0.0, 0.0, 3.0]])
+    assert np.allclose(act.backward(z, np.ones_like(z)), [[0.0, 0.0, 1.0]])
+
+
+def test_leaky_relu_leaks_on_negative():
+    act = get_activation("leaky_relu")
+    z = np.array([[-2.0, 4.0]])
+    a = act.forward(z)
+    assert a[0, 0] == pytest.approx(-0.02)
+    assert a[0, 1] == pytest.approx(4.0)
+    g = act.backward(z, np.ones_like(z))
+    assert g[0, 0] == pytest.approx(0.01)
+    assert g[0, 1] == pytest.approx(1.0)
+
+
+def test_sigmoid_stays_in_unit_interval_on_huge_logits():
+    act = get_activation("sigmoid")
+    z = np.array([[1e3, -1e3, 0.0]])
+    s = act.forward(z)
+    assert s[0, 0] == pytest.approx(1.0)
+    assert s[0, 1] == pytest.approx(0.0)
+    assert s[0, 2] == pytest.approx(0.5)
+    assert np.all(np.isfinite(s))
+
+
+def test_tanh_range_and_oddness():
+    act = get_activation("tanh")
+    z = np.array([[-2.0, 0.0, 2.0]])
+    t = act.forward(z)
+    assert t[0, 1] == pytest.approx(0.0)
+    assert t[0, 0] == pytest.approx(-t[0, 2])
+    assert np.all(np.abs(t) <= 1.0 + 1e-12)
+
+
+def test_unknown_activation_lists_choices():
+    with pytest.raises(ValueError, match="unknown activation"):
+        get_activation("swish")
+
+
+@pytest.mark.parametrize(
+    "fn,expected_var",
+    [
+        (xavier_normal, lambda fi, fo: 2.0 / (fi + fo)),
+        (he_normal, lambda fi, fo: 2.0 / fi),
+    ],
+)
+def test_normal_init_sample_variance_matches_formula(fn, expected_var):
+    fan_in, fan_out = 64, 128
+    rng = np.random.default_rng(0)
+    # draw many matrices so sample var converges to the theoretical scale
+    samples = np.concatenate(
+        [fn(fan_in, fan_out, rng).ravel() for _ in range(40)]
+    )
+    assert float(np.var(samples)) == pytest.approx(
+        expected_var(fan_in, fan_out), rel=0.08
+    )
+
+
+def test_xavier_uniform_stays_inside_bound():
+    fan_in, fan_out = 50, 30
+    bound = np.sqrt(6.0 / (fan_in + fan_out))
+    w = xavier_uniform(fan_in, fan_out, np.random.default_rng(1))
+    assert w.shape == (fan_in, fan_out)
+    assert np.all(w >= -bound - 1e-12)
+    assert np.all(w <= bound + 1e-12)
+
+
+def test_he_uniform_stays_inside_bound():
+    fan_in, fan_out = 40, 20
+    bound = np.sqrt(6.0 / fan_in)
+    w = he_uniform(fan_in, fan_out, np.random.default_rng(2))
+    assert w.shape == (fan_in, fan_out)
+    assert np.all(np.abs(w) <= bound + 1e-12)
+
+
+def test_gain_scales_std():
+    rng = np.random.default_rng(3)
+    base = xavier_normal(80, 80, rng, gain=1.0)
+    rng = np.random.default_rng(3)
+    scaled = xavier_normal(80, 80, rng, gain=2.0)
+    assert float(np.std(scaled)) == pytest.approx(2.0 * float(np.std(base)), rel=0.05)
+
+
+def test_init_weights_dispatch_and_recommended():
+    assert recommended_init("relu") == "he"
+    assert recommended_init("leaky_relu") == "he"
+    assert recommended_init("tanh") == "xavier"
+    assert recommended_init("sigmoid") == "xavier"
+    assert recommended_init("linear") == "xavier"
+
+    w = init_weights(10, 5, "he", np.random.default_rng(0))
+    assert w.shape == (10, 5)
+    w2 = init_weights(10, 5, "xavier", np.random.default_rng(0), distribution="uniform")
+    assert w2.shape == (10, 5)
+    w3 = init_weights(10, 5, "naive", np.random.default_rng(0))
+    assert w3.shape == (10, 5)
+
+
+def test_bad_fans_and_schemes_raise():
+    with pytest.raises(ValueError, match="fan_in"):
+        xavier_normal(0, 4)
+    with pytest.raises(ValueError, match="fan_in"):
+        he_normal(3, 0)
+    with pytest.raises(ValueError, match="gain"):
+        he_normal(4, 4, gain=0.0)
+    with pytest.raises(ValueError, match="unknown init"):
+        init_weights(4, 4, "orthogonal")  # type: ignore[arg-type]
+
+
+def test_forward_profile_rejects_bad_args():
+    with pytest.raises(ValueError, match="n_layers"):
+        forward_variance_profile(0, 16, "relu", "he")
+    with pytest.raises(ValueError, match="width"):
+        forward_variance_profile(4, 0, "relu", "he")
+    with pytest.raises(ValueError, match="n_samples"):
+        forward_variance_profile(4, 16, "relu", "he", n_samples=0)
+
+
+def test_naive_init_explodes_variance_with_depth():
+    """Unit-scale weights without the 1/sqrt(fan) factor blow up signal var."""
+    vars_naive = forward_variance_profile(
+        n_layers=6, width=64, activation="linear", scheme="naive", seed=0
+    )
+    # each layer multiplies var by ~width under N(0,1) weights
+    assert vars_naive[-1] > 1e6
+    assert vars_naive[-1] > vars_naive[0] * 100
+
+
+def test_xavier_keeps_tanh_variance_stable():
+    vars_x = forward_variance_profile(
+        n_layers=8, width=64, activation="tanh", scheme="xavier", seed=1
+    )
+    # tanh is contractive, but should not collapse to machine zero
+    assert all(0.01 < v < 2.0 for v in vars_x)
+    assert max(vars_x) / min(vars_x) < 20.0
+
+
+def test_he_keeps_relu_variance_stable_while_xavier_fades():
+    vars_he = forward_variance_profile(
+        n_layers=10, width=64, activation="relu", scheme="he", seed=2
+    )
+    vars_x = forward_variance_profile(
+        n_layers=10, width=64, activation="relu", scheme="xavier", seed=2
+    )
+    # He restores the half of the signal ReLU drops; variance stays O(1)
+    assert all(0.05 < v < 5.0 for v in vars_he)
+    # Xavier under-scales ReLU stacks, so later layers quietly die
+    assert vars_x[-1] < vars_he[-1] * 0.5
+    assert vars_x[-1] < 0.5
+
+
+def test_profile_length_matches_depth():
+    profile = forward_variance_profile(
+        n_layers=5, width=32, activation="relu", scheme="he", seed=0
+    )
+    assert len(profile) == 5