criticality.py

"""
ORDERED, CRITICAL, OR CHAOTIC — WHERE DO THESE STACKS ACTUALLY SIT?

A network has a critical initialisation scale where a signal neither dies
nor explodes as it passes through layers, and networks near it train far
deeper than networks that are not. Below it a stack is ORDERED: each layer
contracts, and by the fourth or fifth there is nothing left to work with.
Above it the stack is CHAOTIC: nearby inputs diverge and nothing
generalises.

This programme has been there once without naming it. The fold-aware
initialisation fix was exactly a criticality repair — getting the scale
wrong by the square root of the tying ratio made every weight 85 times too
large, deep in the chaotic regime, and training collapsed to 0.168 where
the corrected scale gave 0.390. It was recorded as a bug. It is the
criticality condition, found the hard way.

And there is a live question it might answer. Depth stops paying early
here: three layers beat five on Fashion, and five layers monotonically LOST
accuracy at 7x7. If these stacks are sub-critical, the fourth and fifth
layers are receiving a signal that has already decayed, and that is a
different diagnosis from "the model has enough capacity" — it would mean
depth is being wasted rather than unneeded.

TWO QUANTITIES, measured per layer on real activations.

  CHI, the expansion factor. Perturb a layer's input by a small random
  vector and see how much the output moves: chi = E[||J v||^2 / ||v||^2].
  Below one the layer contracts, above one it expands, near one it
  preserves. This is the mean squared singular value of the layer's
  Jacobian, estimated by probing rather than by forming it.

  PARTICIPATION RATIO, how many directions survive. A layer can preserve
  the average scale while collapsing everything onto a few directions, and
  that is invisible to chi. Estimated from the Gram matrix of the probe
  responses: near one means the layer projects onto a line, near the probe
  count means it preserves the space.

Both are measured AT INITIALISATION and AFTER TRAINING, because a stack
that starts critical and trains itself into contraction is a different
story from one that was never critical at all.

Layer normalisation pins the forward norm, so measuring how big the
activations get would say nothing. The Jacobian is what carries the answer.
"""

import numpy as np
import time
import json

try:
    import cupy as _cp
    _GPU = _cp.cuda.runtime.getDeviceCount() > 0
except Exception:
    _GPU = False
xp = _cp if _GPU else np
DT = np.float32


def to_dev(a, dtype=DT):
    a = np.asarray(a, dtype=dtype)
    return xp.asarray(a) if _GPU else a


def to_host(a):
    return _cp.asnumpy(a) if _GPU and isinstance(a, _cp.ndarray) else np.asarray(a)


def windowed(g, c_in, k, c_out):
    ni, no = c_in*g*g, c_out*g*g
    ii, jj = np.meshgrid(np.arange(ni), np.arange(no), indexing='ij')
    ci, pi = ii // (g*g), ii % (g*g)
    co, po = jj // (g*g), jj % (g*g)
    dr = pi // g - (po // g - k//2)
    dc = pi % g - (po % g - k//2)
    inside = (dr >= 0) & (dr < k) & (dc >= 0) & (dc < k)
    K = c_in*c_out*k*k + 1
    idx = np.where(inside, (ci*c_out + co)*k*k + dr*k + dc, K-1)
    return idx.ravel().astype(np.int32), K, no


_FIXED = {}


class FixedScatter:
    def __init__(self, idx, K, cap=8192):
        h = to_host(idx).astype(np.int64).reshape(-1)
        order = np.argsort(h, kind="stable")
        counts = np.bincount(h, minlength=K)
        starts = np.cumsum(counts) - counts
        big = np.where(counts > cap)[0]
        small = np.where(counts <= cap)[0]
        self.K = K
        self.order = to_dev(order, np.int64) if _GPU else order
        self.big = [(int(b), int(starts[b]), int(starts[b]+counts[b]))
                    for b in big]
        self.small = to_dev(small, np.int64) if _GPU else small
        self.width = int(counts[small].max()) if len(small) else 0
        if self.width:
            pos = np.concatenate([np.arange(counts[s]) for s in small])
            src = np.concatenate([np.arange(starts[s], starts[s]+counts[s])
                                  for s in small])
            row = np.repeat(np.arange(len(small)), counts[small])
            self.src = to_dev(src, np.int64) if _GPU else src
            sl = row*self.width + pos
            self.slot = to_dev(sl, np.int64) if _GPU else sl
            self.buf = xp.zeros(len(small)*self.width, DT)
        self._keep = idx

    def __call__(self, g):
        gs = g.reshape(-1)[self.order]
        out = xp.zeros(self.K, DT)
        if self.width:
            self.buf[:] = 0
            self.buf[self.slot] = gs[self.src]
            out[self.small] = self.buf.reshape(-1, self.width).sum(1)
        for b, a, z in self.big:
            out[b] = gs[a:z].sum()
        return out


def scatter(dW, idx, K):
    key = (id(idx), K)
    if key not in _FIXED:
        _FIXED[key] = FixedScatter(idx, K)
    return _FIXED[key](dW)


class Stack:
    def __init__(self, g, c_in, chan, depth, nc, seed):
        self.g, self.chan, self.depth = g, chan, depth
        self.D = c_in*g*g
        rg = np.random.default_rng(seed)
        self.layers, cin = [], c_in
        for l in range(depth):
            idx, K, no = windowed(g, cin, 3, chan)
            self.layers.append(dict(
                idx=to_dev(idx, np.int32) if _GPU else idx, K=K, out=no,
                ins=self.D if l == 0 else self.layers[-1]["out"],
                taps=cin*9))
            cin = chan
        L = depth
        self.P = []
        for l in self.layers:
            v = rg.normal(0, np.sqrt(2.0/l["taps"]), l["K"]).astype(np.float32)
            v[-1] = 0.0
            self.P.append(to_dev(v))
        self.P += [xp.ones(l["out"], DT) for l in self.layers]
        self.P += [xp.zeros(l["out"], DT) for l in self.layers]
        self.P += [to_dev(rg.normal(0, np.sqrt(2.0/self.layers[-1]["out"]),
                                    (self.layers[-1]["out"], nc))),
                   xp.zeros(nc, DT)]
        self.L = L
        self.M = [xp.zeros_like(p) for p in self.P]
        self.V = [xp.zeros_like(p) for p in self.P]
        self.t = 0

    def layer(self, li, h):
        """One layer, exactly as training applies it."""
        L, P, l = self.L, self.P, self.layers[li]
        W = P[li][l["idx"]].reshape(l["ins"], l["out"])
        z = h @ W
        var = z.var(1, keepdims=True) + 1e-5
        zn = (z - z.mean(1, keepdims=True))/xp.sqrt(var)
        return xp.maximum(zn*P[L+li] + P[2*L+li], 0)

    def acts(self, x):
        """The activations entering each layer, and leaving the last."""
        out = [x]; h = x
        for li in range(self.L):
            h = self.layer(li, h)
            out.append(h)
        return out

    def acc(self, Xte, yte):
        L = self.L
        preds = []
        for s in range(0, Xte.shape[0], 4096):
            h = self.acts(Xte[s:s+4096])[-1]
            preds.append(to_host(h @ self.P[3*L] + self.P[3*L+1]))
        return float((np.concatenate(preds).argmax(1) == yte).mean())

    def fit(self, Xtr, Ytr, cfg, seed, epochs):
        L = self.L
        rg = np.random.default_rng(seed + 991)
        n = Xtr.shape[0]
        for ep in range(epochs):
            perm = rg.permutation(n)
            for st in range(0, n, cfg["batch"]):
                b = perm[st:st+cfg["batch"]]
                x = Xtr[b]; y = Ytr[b]
                cache = []; h = x
                for li, l in enumerate(self.layers):
                    W = self.P[li][l["idx"]].reshape(l["ins"], l["out"])
                    z = h @ W
                    var = z.var(1, keepdims=True) + 1e-5
                    zn = (z - z.mean(1, keepdims=True))/xp.sqrt(var)
                    zs = zn*self.P[L+li] + self.P[2*L+li]
                    a = xp.maximum(zs, 0)
                    cache.append((h, W, var, zn, zs)); h = a
                lg = h @ self.P[3*L] + self.P[3*L+1]
                e = xp.exp(lg - lg.max(1, keepdims=True))
                d = (e/e.sum(1, keepdims=True) - y)/len(b)
                G = [xp.zeros_like(p) for p in self.P]
                G[3*L] = h.T @ d; G[3*L+1] = d.sum(0)
                dh = d @ self.P[3*L].T
                for li in range(L-1, -1, -1):
                    hin, W, var, zn, zs = cache[li]
                    dzs = dh*(zs > 0)
                    G[L+li] = (dzs*zn).sum(0); G[2*L+li] = dzs.sum(0)
                    dzn = dzs*self.P[L+li]
                    dz = (dzn - dzn.mean(1, keepdims=True)
                          - zn*(dzn*zn).mean(1, keepdims=True))/xp.sqrt(var)
                    G[li] = scatter(hin.T @ dz, self.layers[li]["idx"],
                                    self.layers[li]["K"])
                    if li > 0:
                        dh = dz @ W.T
                self.t += 1
                for i, (p_, gr) in enumerate(zip(self.P, G)):
                    self.M[i] = 0.9*self.M[i] + 0.1*gr
                    self.V[i] = 0.999*self.V[i] + 0.001*gr*gr
                    self.P[i] = p_ - cfg["lr"]*(self.M[i]/(1-0.9**self.t)) \
                        / (xp.sqrt(self.V[i]/(1-0.999**self.t))+1e-8)
        return self


def probe(model, X, n_probe=48, rel=1e-3, seed=0):
    """chi and participation ratio per layer, by finite differences.

    With ReLU the map is piecewise linear, so a small enough step is exact
    as long as no unit changes sign — and small enough is set relative to
    the activation's own scale rather than absolutely, because layer
    normalisation fixes that scale and a fixed epsilon would be wrong at
    the input and right nowhere else."""
    rg = np.random.default_rng(seed)
    A = model.acts(X)
    out = []
    for li in range(model.L):
        h = A[li]
        base = model.layer(li, h)
        hn = float(to_host(xp.linalg.norm(h, axis=1).mean()))
        R = []
        for p in range(n_probe):
            v = to_dev(rg.normal(size=h.shape))
            v = v/xp.linalg.norm(v, axis=1, keepdims=True)
            eps = rel*hn
            r = (model.layer(li, h + eps*v) - base)/eps
            R.append(r)
        # chi: how much a unit perturbation grows
        norms = xp.stack([xp.linalg.norm(r, axis=1)**2 for r in R])
        chi = float(to_host(norms.mean()))
        # participation ratio of the response Gram matrix, averaged over
        # examples: how many directions the layer actually keeps
        prs = []
        for i in range(0, min(16, h.shape[0])):
            M = xp.stack([r[i] for r in R])
            Gm = to_host(M @ M.T)
            w = np.linalg.eigvalsh(Gm)
            w = np.clip(w, 0, None)
            if w.sum() > 0:
                prs.append(float(w.sum()**2/(w**2).sum()))
        out.append(dict(chi=chi, pr=float(np.mean(prs)) if prs else np.nan,
                        n_probe=n_probe))
    return out


def load(cfg):
    from tensorflow import keras
    (a, b), (c, d) = keras.datasets.fashion_mnist.load_data()
    X = np.concatenate([a, c]).astype(np.float32)/255.0
    y = np.concatenate([b, d]).ravel().astype(np.int64)
    if cfg["grid"] != 28:
        s = 28//cfg["grid"]
        X = X.reshape(-1, cfg["grid"], s, cfg["grid"], s).mean(axis=(2, 4))
    rg = np.random.default_rng(0); p = rg.permutation(len(X))
    tr, te = p[:cfg["n_train"]], p[cfg["n_train"]:cfg["n_train"]+5000]
    mu, sd = X[tr].mean(), X[tr].std()+1e-8
    f = lambda Z: ((Z-mu)/sd).reshape(len(Z), -1)
    Y = np.zeros((len(tr), 10), np.float32); Y[np.arange(len(tr)), y[tr]] = 1
    return f(X[tr]), Y, f(X[te]), y[te]


CFG = dict(grid=14, c_in=1, chan=16, n_train=20000, batch=128, lr=1e-3,
           epochs=40, depth=5, n_probe=48, seed=0)


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("ORDERED, CRITICAL, OR CHAOTIC?")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:9s} = {v}")
    print(f"\n  chi is the mean squared singular value of a layer's Jacobian:")
    print(f"    chi < 1  the layer CONTRACTS, and a signal dies with depth")
    print(f"    chi ~ 1  CRITICAL, the regime where deep stacks train")
    print(f"    chi > 1  the layer EXPANDS, and nearby inputs diverge")
    print(f"  the participation ratio says how many of {CFG['n_probe']} probe")
    print(f"  directions survive — a layer can hold the scale and still")
    print(f"  collapse everything onto a line")
    print("=" * 78, flush=True)

    Xtr, Ytr, Xte, yte = load(CFG)
    Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
    probe_x = Xte[:256]

    m = Stack(CFG["grid"], CFG["c_in"], CFG["chan"], CFG["depth"], 10,
              CFG["seed"])
    at_init = probe(m, probe_x, CFG["n_probe"], seed=CFG["seed"])
    acc0 = m.acc(Xte, yte)
    print(f"\n  at initialisation (accuracy {acc0:.4f}):")
    print(f"  {'layer':>6s} {'chi':>8s} {'cumulative':>11s} "
          f"{'participation':>14s}")
    cum = 1.0
    for i, r in enumerate(at_init):
        cum *= r["chi"]
        print(f"  {i+1:6d} {r['chi']:8.4f} {cum:11.3e} {r['pr']:14.1f}")

    m.fit(Xtr, Ytr, CFG, CFG["seed"], CFG["epochs"])
    trained = probe(m, probe_x, CFG["n_probe"], seed=CFG["seed"])
    acc1 = m.acc(Xte, yte)
    print(f"\n  after {CFG['epochs']} epochs (accuracy {acc1:.4f}):")
    print(f"  {'layer':>6s} {'chi':>8s} {'cumulative':>11s} "
          f"{'participation':>14s} {'chi moved':>10s}")
    cum = 1.0
    for i, r in enumerate(trained):
        cum *= r["chi"]
        print(f"  {i+1:6d} {r['chi']:8.4f} {cum:11.3e} {r['pr']:14.1f} "
              f"{r['chi']-at_init[i]['chi']:+10.4f}")

    json.dump(dict(init=at_init, trained=trained, acc_init=acc0,
                   acc_trained=acc1), open("criticality.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    ci = np.array([r["chi"] for r in trained])
    pr = np.array([r["pr"] for r in trained])
    cum = float(np.prod(ci))
    print(f"  chi by layer, trained: "
          + "  ".join(f"{c:.3f}" for c in ci))
    print(f"  cumulative through {CFG['depth']} layers: {cum:.3e}")
    print(f"  participation: " + "  ".join(f"{p:.0f}" for p in pr)
          + f"   (of {CFG['n_probe']} probes)")
    print()
    if cum < 0.1:
        print(f"  SUB-CRITICAL — THE STACK CONTRACTS. A perturbation at the")
        print(f"  input is {1/cum:.0f} times smaller by the last layer, so the")
        print(f"  later layers see a signal that has largely died. That is a")
        print(f"  different diagnosis from 'enough capacity': depth here is")
        print(f"  being WASTED rather than being unneeded, and holding chi")
        print(f"  near one is a specific, testable repair.")
    elif cum > 10:
        print(f"  SUPER-CRITICAL — THE STACK EXPANDS by {cum:.1f} times over")
        print(f"  {CFG['depth']} layers, so nearby inputs diverge as they pass")
        print(f"  through and the later layers are amplifying noise.")
    else:
        print(f"  NEAR CRITICAL. The stack neither dies nor explodes over")
        print(f"  {CFG['depth']} layers, so signal propagation is NOT why depth")
        print(f"  stops paying — the ceiling is somewhere else, and the")
        print(f"  edge-of-chaos angle has nothing to fix here.")
    drop = pr[0]/max(pr[-1], 1e-9)
    print()
    if drop > 3:
        print(f"  AND THE RANK COLLAPSES: {pr[0]:.0f} directions at the first")
        print(f"  layer against {pr[-1]:.0f} at the last, a factor of {drop:.1f}.")
        print(f"  The stack is funnelling everything onto a few directions,")
        print(f"  which chi alone would not have shown.")
    else:
        print(f"  and the rank holds up: {pr[0]:.0f} directions at the first")
        print(f"  layer against {pr[-1]:.0f} at the last, so the layers are not")
        print(f"  collapsing the space.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote criticality.json")


if __name__ == "__main__":
    main()