Starting Right: Initialization and Normalization

# Chapter 12: how big should the starting weights be, and how normalization keeps signals in range.
# Uses the Tensor class from Chapter 6, saved as tensor.py (see Chapter 7).
import numpy as np
from tensor import Tensor

# === 1. Push a signal through 20 ReLU layers ======================================
n, depth = 64, 20
X = np.random.default_rng(0).normal(size=(256, n))          # 256 random inputs, 64 features each
scales = [("σ = 0.05", 0.05), ("σ = 1/√n (Xavier)", 1 / np.sqrt(n)),
          ("σ = √(2/n) (He)", np.sqrt(2 / n)), ("σ = 1", 1.0)]

print("starting scale        per-layer factor   RMS after layer 1, 5, 10, 20         first-layer gradient")
for name, sigma in scales:
    rng = np.random.default_rng(1)
    Ws = [Tensor(rng.normal(0, sigma, (n, n))) for _ in range(depth)]
    h, rms = Tensor(X), []
    for W in Ws:
        h = (h @ W.T).relu()
        rms.append(float(np.sqrt(np.mean(h.data ** 2))))
    (h.sum() / len(X)).backward()                            # any scalar will do: we only want the gradient's size
    factor = np.sqrt(n * sigma ** 2 / 2)
    print(f"{name:20}  {factor:8.3f}           {'  '.join(f'{rms[i]:9.3g}' for i in (0, 4, 9, 19))}"
          f"    {np.sqrt((Ws[0].grad ** 2).sum()):9.3g}")

# He keeps the signal near 1, but each layer's factor is a little random at finite width
for width in [64, 512]:
    Xw = np.random.default_rng(0).normal(size=(256, width))
    finals = []
    for seed in range(10):
        rng, h = np.random.default_rng(seed), Xw
        for _ in range(depth):
            h = np.maximum(0, h @ rng.normal(0, np.sqrt(2 / width), (width, width)).T)
        finals.append(np.sqrt(np.mean(h ** 2)))
    print(f"He, width {width}: RMS after 20 layers over 10 random draws: {min(finals):.2f} to {max(finals):.2f}")

# === 2. Layer normalization, from Chapter 6's operations =============================
def layer_norm(h, gain, shift, eps=1e-5):
    """Standardize each row (one example's features) to mean 0, variance 1, then rescale and shift."""
    average = Tensor(np.ones((h.shape[1], 1)) / h.shape[1])   # (features × 1): a row mean as a product
    mean = h @ average
    centred = h - mean
    var = (centred * centred) @ average
    return centred * (var + eps) ** -0.5 * gain + shift

h = Tensor(np.random.default_rng(2).normal(3.0, 5.0, size=(4, 8)))
out = layer_norm(h, Tensor(np.ones(8)), Tensor(np.zeros(8)))
print("\nlayer norm: row means", out.data.mean(axis=1).round(6) + 0.0, " row variances", out.data.var(axis=1).round(4))

# Normalization ignores the weights' scale, so their gradient shrinks as they grow
x = Tensor(np.random.default_rng(3).normal(size=(16, 8)))
W0 = np.random.default_rng(4).normal(size=(8, 8))
for c in [1.0, 10.0]:
    W = Tensor(c * W0)
    y = layer_norm(x @ W.T, Tensor(np.ones(8)), Tensor(np.zeros(8)))
    loss = (y * Tensor(np.arange(8.0))).sum()                   # an arbitrary scalar of the output
    loss.backward()
    print(f"weights × {c:4}: loss {float(loss.data):.6f}, gradient size {np.sqrt((W.grad ** 2).sum()):.5f}")

# === 3. Training a deep network: Chapter 9's 150 days, 10 hidden layers of 32 ========
def make_days(n_days, seed):
    r = np.random.default_rng(seed)
    X = np.round(r.random((n_days, 3)), 2)
    labels = np.where(X[:, 1] > 0.5, 2, np.where(X[:, 2] > 0.6, 1, 0))
    noisy = r.random(n_days) < 0.1
    labels[noisy] = r.integers(0, 3, noisy.sum())
    return X, labels

X_days, labels = make_days(150, 9)
X_days = (X_days - X_days.mean(axis=0)) / X_days.std(axis=0)  # standardized inputs (Chapter 7)
Y = np.eye(3)[labels]


def cross_entropy(Z):
    E = (Z + (-Z.data.max(axis=1, keepdims=True))).exp()
    P = E / (E @ Tensor(np.ones((3, 1))))
    return -(Tensor(Y) * P.log()).sum() / len(Y)


def train_deep(sigma, normalize, seed, hidden=10, width=32, lr=0.1, epochs=500):
    """Plain gradient descent, full batch, so that only the starting scale differs between runs."""
    rng = np.random.default_rng(seed)
    sizes = [3] + [width] * hidden + [3]
    layers, params = [], []
    for fan_in, fan_out in zip(sizes, sizes[1:]):
        s = np.sqrt(2 / fan_in) if sigma == "he" else sigma
        W, b = Tensor(rng.normal(0, s, (fan_out, fan_in))), Tensor(np.zeros(fan_out))
        layers.append((W, b))
        params += [W, b]
    norms = [(Tensor(np.ones(width)), Tensor(np.zeros(width))) for _ in range(hidden)] if normalize else []
    for gain, shift in norms:
        params += [gain, shift]

    def forward():
        h = Tensor(X_days)
        for i, (W, b) in enumerate(layers):
            h = h @ W.T + b
            if i < hidden:
                if normalize:
                    h = layer_norm(h, *norms[i])
                h = h.relu()
        return h

    for epoch in range(epochs):
        with np.errstate(all="ignore"):
            loss = cross_entropy(forward())
        if not np.isfinite(loss.data):
            return "NaN at epoch %d" % epoch
        for p in params:
            p.grad = np.zeros_like(p.data)
        loss.backward()
        for p in params:
            p.data = p.data - lr * p.grad
    return f"{float(cross_entropy(forward()).data):.3f}"

print("\n10 hidden layers, plain gradient descent (η = 0.1), training loss after 500 epochs, 3 random starts")
print(f"(predicting the class frequencies scores {-(Y.mean(axis=0) * np.log(Y.mean(axis=0))).sum():.3f})")
for name, sigma, normalize in [("σ = 0.01", 0.01, False), ("σ = 1", 1.0, False), ("He", "he", False),
                               ("σ = 1 with layer norm", 1.0, True), ("He with layer norm", "he", True),
                               ("σ = 0.01 with layer norm", 0.01, True)]:
    print(f"  {name:26}", "   ".join(train_deep(sigma, normalize, seed) for seed in range(3)))