""" THE HEAD HAS BEEN MOST OF THE MODEL ALL ALONG. Every script in this programme ends with a head that reads EVERY hidden unit: 3,136 of them here, so 31,370 parameters. Measured against it, the folded body is 1.9% of a five-layer model at 28x28 and 82.5% at 7x7 — and the depth sweep that produced those numbers could not separate "deeper suits a larger grid" from "layers are cheap at a larger grid", because scaling channels inversely with area made a layer cost 36,865 values at 7x7 and 145 at 28x28. A conventional convolutional network does not do this. It POOLS over spatial positions before the head, so the head reads one number per channel rather than one per channel per position. Here that would take the head from 31,370 parameters to CHANNELS x 10 — 650 at 7x7, 170 at 14x14, 50 at 28x28. The body would become the model instead of a rounding error. Nothing in this programme has ever used pooling. That is not a small omission: it is most of why these models are head-dominated, most of why the storage ratios were inflated, and possibly why depth has looked so weak. So the same depth sweep is run twice, with and without pooling before the head, and three things are read from it. DOES POOLING COST ACCURACY? it discards where a feature fired and keeps only how much, which on a centred dataset may be free and may not be DOES IT CHANGE THE DEPTH STORY? if depth only ever looked useful because layers were cheap relative to a huge head, pooling should change which depth wins WHAT DOES IT COST TO STORE? the head shrinks by a factor of the grid area, which is 49 to 784 times The honest expectation is that pooling costs some accuracy on a task this small, and that the question is whether it costs less than the storage it saves. """ import numpy as np import time import json import os 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, no*(c_in*k*k) DETERMINISTIC = True _FIXED = {} class FixedScatter: """Fixed-order accumulation, so a rerun reproduces exactly.""" 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): if not DETERMINISTIC: g = xp.zeros(K, DT) if _GPU: import cupyx cupyx.scatter_add(g, idx, dW.reshape(-1)) else: np.add.at(g, idx, dW.reshape(-1)) return g key = (id(idx), K) if key not in _FIXED: _FIXED[key] = FixedScatter(idx, K) return _FIXED[key](dW) def train(Xtr, Ytr, Xte, yte, g, chan, depth, cfg, seed, pool=False): D = Xtr.shape[1] rg = np.random.default_rng(seed) layers, cin = [], 1 for l in range(depth): idx, K, no, macs = windowed(g, cin, 3, chan) layers.append(dict(idx=to_dev(idx, np.int32) if _GPU else idx, K=K, out=no, macs=macs, ins=D if l == 0 else layers[-1]["out"], taps=cin*9)) cin = chan L = depth P = [] for l in layers: v = rg.normal(0, np.sqrt(2.0/l["taps"]), l["K"]).astype(np.float32) v[-1] = 0.0 P.append(to_dev(v)) P += [xp.ones(l["out"], DT) for l in layers] P += [xp.zeros(l["out"], DT) for l in layers] hw = chan if pool else layers[-1]["out"] P += [to_dev(rg.normal(0, np.sqrt(2.0/hw), (hw, 10))), xp.zeros(10, DT)] HEAD, OB = 3*L, 3*L+1 M = [xp.zeros_like(p) for p in P]; V = [xp.zeros_like(p) for p in P] n = Xtr.shape[0]; t = 0 ag = np.random.default_rng(seed + 991) npos = g*g def fwd(x, keep=False): cache = []; h = x for li, l in enumerate(layers): 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) zs = zn*P[L+li] + P[2*L+li] a = xp.maximum(zs, 0) if keep: cache.append((h, W, var, zn, zs)) h = a if pool: # average over spatial positions: the head then reads one # number per channel instead of one per channel per position h = h.reshape(h.shape[0], chan, npos).mean(2) return h, cache for ep in range(cfg["epochs"]): perm = ag.permutation(n) for st in range(0, n, cfg["batch"]): b = perm[st:st+cfg["batch"]] x = Xtr[b]; y = Ytr[b] h, cache = fwd(x, keep=True) lg = h @ P[HEAD] + P[OB] 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 P] G[HEAD] = h.T @ d; G[OB] = d.sum(0) dh = d @ P[HEAD].T if pool: dh = xp.broadcast_to(dh[:, :, None]/npos, (dh.shape[0], chan, npos) ).reshape(dh.shape[0], chan*npos) 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*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, layers[li]["idx"], layers[li]["K"]) if li > 0: dh = dz @ W.T t += 1 for i, (p_, gr) in enumerate(zip(P, G)): M[i] = 0.9*M[i] + 0.1*gr V[i] = 0.999*V[i] + 0.001*gr*gr P[i] = p_ - cfg["lr"]*(M[i]/(1-0.9**t)) \ / (xp.sqrt(V[i]/(1-0.999**t))+1e-8) out = [] for s in range(0, Xte.shape[0], 4096): h, _ = fwd(Xte[s:s+4096]) out.append(to_host(h @ P[HEAD] + P[OB])) acc = float((np.concatenate(out).argmax(1) == yte).mean()) # COUNT THE HEAD. It is dense and unfolded, and on Fashion at 14x14 it # is 99.5% of a one-layer model's parameters — reporting the folded # body alone made a 15% storage difference look like 65x. head = (chan if pool else layers[-1]["out"])*10 + 10 return (acc, sum(l["K"] for l in layers) + head, sum(l["macs"] for l in layers) + head - 10, head) def load(grid, 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 grid != 28: s = 28//grid X = X.reshape(-1, grid, s, 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"]+10000] 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(n_train=20000, batch=128, lr=1e-3, epochs=40, seeds=(0, 1, 2), hidden_target=3136, grids=(7, 14, 28), depths=(1, 2, 3, 4, 5)) def main(**over): CFG.update(over) t0 = time.time() print("=" * 78) print("THE HEAD HAS BEEN MOST OF THE MODEL ALL ALONG") print("=" * 78) print(f" backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}") for k, v in CFG.items(): print(f" {k:14s} = {v}") print(f"\n channels scale inversely with area to hold the hidden width") print(f" near {CFG['hidden_target']}, so no arm is favoured by having a") print(f" bigger matrix product:") for g in CFG["grids"]: c = max(1, CFG["hidden_target"]//(g*g)) print(f" {g:2d}x{g:<2d} -> {c:3d} channels, hidden {c*g*g:,}") print(f"\n reach says L >= (n-1)/2 to cover the grid: " + ", ".join(f"{g}x{g} needs {int(np.ceil((g-1)/2))}" for g in CFG["grids"])) print("=" * 78, flush=True) res = {} for g in CFG["grids"]: chan = max(1, CFG["hidden_target"]//(g*g)) Xtr, Ytr, Xte, yte = load(g, CFG) Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte) print(f"\n {g}x{g}, {chan} channels") print(f" (the head alone is " f"{chan*g*g*10 + 10:,} parameters, counted in every row)") print(f" {'head':>7s} {'L':>4s} {'TOTAL vals':>11s} " f"{'multiplies':>12s} {'accuracy':>9s} {'sd':>7s} {'reach':>7s}") for pool, L in [(p, L) for p in (False, True) for L in CFG["depths"]]: # the fixed-order tables are keyed by the index array's identity # and each arm builds new ones, so the cache would grow across # fifteen arms until the device ran out _FIXED.clear() if _GPU: _cp.get_default_memory_pool().free_all_blocks() out = [train(Xtr, Ytr, Xte, yte, g, chan, L, CFG, s, pool=pool) for s in CFG["seeds"]] acc = [o[0] for o in out] K, mc, hd = out[0][1], out[0][2], out[0][3] reach = min(g, 2*L+1) res[f"{g}/{L}/{int(pool)}"] = dict(acc=float(np.mean(acc)), sd=float(np.std(acc)), K=K, macs=mc, reach=reach, chan=chan, head=hd) print(f" {'pooled' if pool else 'flat':>7s} {L:4d} " f"{K:11,} {mc:12,} {np.mean(acc):9.4f} " f"{np.std(acc):7.4f} {reach:5d}x{reach:<2d}" f" [{time.time()-t0:.0f}s]", flush=True) json.dump(res, open("pooling.json", "w"), indent=2) print("\n" + "=" * 78) print(" WHAT DOES POOLING COST, AND WHAT DOES IT SAVE?") print("=" * 78) print(f" {'grid':>5s} {'L':>3s} {'flat':>8s} {'pooled':>8s} " f"{'accuracy':>9s} {'head flat':>10s} {'head pooled':>12s} " f"{'saved':>7s}") for g in CFG["grids"]: for L in CFG["depths"]: a = res.get(f"{g}/{L}/0"); b = res.get(f"{g}/{L}/1") if not (a and b): continue print(f" {g:5d} {L:3d} {a['acc']:8.4f} {b['acc']:8.4f} " f"{b['acc']-a['acc']:+9.4f} {a['head']:10,} " f"{b['head']:12,} {a['head']/b['head']:6.0f}x") print("\n" + "=" * 78) print(" READOUT") print("=" * 78) sd = max(r["sd"] for r in res.values()) d = [res[f"{g}/{L}/1"]["acc"] - res[f"{g}/{L}/0"]["acc"] for g in CFG["grids"] for L in CFG["depths"] if f"{g}/{L}/1" in res and f"{g}/{L}/0" in res] d = np.array(d) print(f" pooling changes accuracy by {d.mean():+.4f} on average, " f"range {d.min():+.4f} to {d.max():+.4f}") print(f" seed spread (worst) {sd:.4f}\n") print(f" best depth, flat head vs pooled:") for g in CFG["grids"]: bf = max(CFG["depths"], key=lambda L: res[f"{g}/{L}/0"]["acc"]) bp = max(CFG["depths"], key=lambda L: res[f"{g}/{L}/1"]["acc"]) print(f" {g:2d}x{g:<2d} flat {bf} pooled {bp}" + (" <- the depth story CHANGES" if bf != bp else "")) print() if d.mean() > -2*sd: print(" POOLING IS FREE. The head shrinks by the grid area and") print(" accuracy does not move, so every model in this programme") print(" has been carrying a head one to two orders of magnitude") print(" larger than it needed — which is most of why the storage") print(" ratios were inflated and the body looked negligible.") elif d.mean() > -0.02: print(f" POOLING COSTS {abs(d.mean()):.4f} AND SAVES THE HEAD. Whether") print(" that is worth it depends on the budget, but it should be a") print(" stated choice rather than an omission — and the ratios in") print(" the paper were computed without it.") else: print(f" POOLING COSTS TOO MUCH ({d.mean():+.4f}). Discarding WHERE a") print(" feature fired matters on this task, so the large head is") print(" earning its storage and the programme's design was right by") print(" accident rather than wrong.") print(f"\n total {time.time()-t0:.0f}s; wrote pooling.json") if __name__ == "__main__": main()