feat: Add activation functions and Xavier/He weight init

README.md

@@ -30,13 +30,19 @@ 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
 
 - **Gradient descent + autodiff-lite**: a scalar reverse-mode autograd engine (`Value`) with `+`, `*`, `**`, division, and `relu`/`tanh`/`exp`/`log`/`sigmoid`, plus an `SGD` optimizer (with momentum) and a `minimize` training loop. Gradients are checked against finite differences and against numpy for softmax cross-entropy.
 - **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

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