| """ |
| 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) |
| |
| norms = xp.stack([xp.linalg.norm(r, axis=1)**2 for r in R]) |
| chi = float(to_host(norms.mean())) |
| |
| |
| 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() |
|
|
