""" 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()