""" GROWING DEPTH DURING A RUN — A PROTOTYPE ON FASHION. Two curricula have already failed in this programme and one succeeded, and the difference between them is what makes this worth testing. GROWING WIDTH failed: two 8-channel parents joined into a 16-channel child ended 0.0156 BELOW the same width trained from scratch, and joining early beat joining late by 0.0003, which is nothing. GROWING STORAGE failed: splitting value groups 64 -> 128 -> 256 -> 512 ended 0.0224 below cold training at the final budget, and the rungs gained +0.0076, -0.0007, +0.0255. SHRINKING STORAGE works: §6's annealing reaches a good solution at a small budget where a cold start does not, and restarts helped the cold arm while doing nothing to the annealed one — which points at the schedule avoiding bad basins rather than spending compute. Both failures added CAPACITY to a model that was not capacity-bound. Growing DEPTH is a different proposition: if four and five layers are suboptimal from a cold start because they are hard to OPTIMISE rather than too small, a curriculum addresses exactly that, and it is the same thing §6's annealing was measured doing. The design separates two effects that the obvious experiment confounds. A regime that grows depth AND slows the learning rate cannot say which did the work, so this is a 2x2: constant rate slowed at each growth grown 1->3->5 [ ] [ ] 5 from scratch [ ] [ ] with 1-layer and 3-layer references, since three layers is the current best and the thing to beat. A new layer is initialised as a DELTA FILTER — centre tap one where the input channel matches the output channel, zero elsewhere — which the fold represents exactly, because the value index is literally (in-channel, out-channel, tap). Layer normalisation then re-standardises, so the pass-through is close but NOT exact, unlike the width and storage experiments where it was exact to machine precision. The accuracy either side of each growth is printed, and how large that step is matters: a big drop means the curriculum is paying a real cost for the freedom it buys. """ 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): """A 3x3 convolution as a tying pattern. The value index is (input channel, output channel, tap), which is what lets a new layer be initialised as an exact delta filter.""" 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. Rectangle for groups near the median; any oversized group summed over its own contiguous slice. A convolution index is wildly skewed — every out-of-window position dumps into one padding group — so a single rectangle would be enormous.""" 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, self.n = K, len(h) 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 delta_values(K, c_in, c_out, k): """A layer that passes its input through: centre tap one where the input channel matches the output channel. Exactly representable, because the value index IS (in-channel, out-channel, tap).""" v = np.zeros(K, np.float32) centre = (k//2)*k + k//2 for c in range(min(c_in, c_out)): v[(c*c_out + c)*k*k + centre] = 1.0 return v class Stack: """A stack of folded convolutions that can grow a layer mid-run.""" def __init__(self, g, c_in, k, chan, nc, seed): self.g, self.c_in, self.k, self.chan, self.nc = g, c_in, k, chan, nc self.rg = np.random.default_rng(seed) self.layers = [] self.D = c_in*g*g self.add_layer(first=True) self.head = to_dev(self.rg.normal( 0, np.sqrt(2.0/self.width), (self.width, nc))) self.ob = xp.zeros(nc, DT) self.t = 0 self._reset_moments() def add_layer(self, first=False, identity=True): cin = self.c_in if first else self.chan idx, K, no, macs = windowed(self.g, cin, self.k, self.chan) v = (self.rg.normal(0, np.sqrt(2.0/(cin*self.k*self.k)), K) .astype(np.float32) if first or not identity else delta_values(K, cin, self.chan, self.k)) if first: v[-1] = 0.0 self.layers.append(dict( idx=to_dev(idx, np.int32) if _GPU else idx, K=K, out=no, macs=macs, v=to_dev(v), gam=xp.ones(no, DT), bet=xp.zeros(no, DT), ins=self.D if first else self.layers[-1]["out"])) self.width = no def _reset_moments(self): self.P = self._params() self.M = [xp.zeros_like(p) for p in self.P] self.V = [xp.zeros_like(p) for p in self.P] def _params(self): p = [] for l in self.layers: p += [l["v"], l["gam"], l["bet"]] return p + [self.head, self.ob] def _store(self, P): i = 0 for l in self.layers: l["v"], l["gam"], l["bet"] = P[i], P[i+1], P[i+2] i += 3 self.head, self.ob = P[i], P[i+1] def grow(self, sample=None): """Add a layer on top, initialised to pass its input through. A delta filter reproduces the input exactly, but the layer normalisation that follows re-standardises it, so the new layer is NOT transparent by default — measured, the logits moved by 1.28 and a fifth of predictions changed. If a sample batch is given, the normalisation's own scale and shift are set to undo it on average: gamma to the typical spread across features, beta to the typical centre. That cannot be exact, because both vary per example, but it removes the systematic part. The head reads the same width, so it survives untouched. Adam's moments for existing parameters are KEPT — only the new layer starts cold.""" old = [(m, v) for m, v in zip(self.M, self.V)] stats = None if sample is not None: h, _ = self.fwd(sample) stats = (float(to_host(h.mean(1).mean())), float(to_host(h.std(1).mean()))) self.add_layer() if stats is not None: mu, sd = stats self.layers[-1]["gam"] = xp.full(self.width, DT(max(sd, 1e-6))) self.layers[-1]["bet"] = xp.full(self.width, DT(mu)) newP = self._params() self.M, self.V = [], [] # three new entries appear before the head's two n_old_layers = len(self.layers) - 1 for i in range(len(newP)): if i < 3*n_old_layers: self.M.append(old[i][0]); self.V.append(old[i][1]) elif i < 3*len(self.layers): self.M.append(xp.zeros_like(newP[i])) self.V.append(xp.zeros_like(newP[i])) else: j = i - 3 self.M.append(old[j][0]); self.V.append(old[j][1]) self.P = newP def fwd(self, x, keep=False): cache = []; h = x for l in self.layers: W = l["v"][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*l["gam"] + l["bet"] a = xp.maximum(zs, 0) if keep: cache.append((h, W, var, zn, zs)) h = a return h, cache def acc(self, Xte, yte): out = [] for s in range(0, Xte.shape[0], 4096): h, _ = self.fwd(Xte[s:s+4096]) out.append(to_host(h @ self.head + self.ob)) return float((np.concatenate(out).argmax(1) == yte).mean()) def epoch(self, Xtr, Ytr, cfg, lr, rg, side): n = Xtr.shape[0] L = len(self.layers) P = self._params() tot, nb = 0.0, 0 perm = rg.permutation(n) for st in range(0, n, cfg["batch"]): b = perm[st:st+cfg["batch"]] x = Xtr[b]; y = Ytr[b] h, cache = self.fwd(x, keep=True) lg = h @ P[3*L] + P[3*L+1] mx = lg.max(1, keepdims=True) e = xp.exp(lg - mx); se = e.sum(1, keepdims=True) tot += float(to_host((-(lg-mx-xp.log(se))*y).sum(1).mean())) nb += 1 d = (e/se - y)/len(b) G = [xp.zeros_like(p) for p in P] G[3*L] = h.T @ d; G[3*L+1] = d.sum(0) dh = d @ P[3*L].T for li in range(L-1, -1, -1): hin, W, var, zn, zs = cache[li] l = self.layers[li] dzs = dh*(zs > 0) G[3*li+1] = (dzs*zn).sum(0) G[3*li+2] = dzs.sum(0) dzn = dzs*l["gam"] dz = (dzn - dzn.mean(1, keepdims=True) - zn*(dzn*zn).mean(1, keepdims=True))/xp.sqrt(var) G[3*li] = scatter(hin.T @ dz, l["idx"], l["K"]) if li > 0: dh = dz @ W.T self.t += 1 for i, (p_, gr) in enumerate(zip(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 P[i] = p_ - lr*(self.M[i]/(1-0.9**self.t)) \ / (xp.sqrt(self.V[i]/(1-0.999**self.t))+1e-8) self._store(P) return tot/max(nb, 1) def macs(self): return sum(l["macs"] for l in self.layers) def stored(self): return sum(l["K"] for l in self.layers) 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) side = 28 if cfg["grid"] != side: s = side//cfg["grid"] X = X.reshape(-1, cfg["grid"], s, cfg["grid"], s).mean(axis=(2, 4)) side = cfg["grid"] 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], side CFG = dict(grid=14, c_in=1, k=3, chan=16, n_train=20000, batch=128, lr=1e-3, stage_epochs=30, stages=3, slow=0.5, seeds=(0, 1, 2)) def run(Xtr, Ytr, Xte, yte, side, cfg, seed, grown, slowdown, target=5): """grown: start at one layer and add. Otherwise start at target. slowdown: multiply the rate at each growth point, whether or not the model actually grows — so the two arms share a schedule.""" rg = np.random.default_rng(seed + 77) m = Stack(cfg["grid"], cfg["c_in"], cfg["k"], cfg["chan"], 10, seed) if not grown: while len(m.layers) < target: m.add_layer(identity=False) m.head = to_dev(np.random.default_rng(seed).normal( 0, np.sqrt(2.0/m.width), (m.width, 10))) m._reset_moments() plan = [1, 3, target][:cfg["stages"]] lr = cfg["lr"] hist, events = [], [] for si in range(cfg["stages"]): if si and grown: before = m.acc(Xte, yte) while len(m.layers) < plan[si]: m.grow(sample=Xte[:512]) after = m.acc(Xte, yte) events.append(dict(to=len(m.layers), before=before, after=after)) if si and slowdown: lr *= cfg["slow"] for ep in range(cfg["stage_epochs"]): tl = m.epoch(Xtr, Ytr, cfg, lr, rg, side) hist.append(dict(stage=si, layers=len(m.layers), lr=lr, train_loss=tl, acc=m.acc(Xte, yte))) return m.acc(Xte, yte), hist, events, m.stored(), m.macs() def main(**over): CFG.update(over) t0 = time.time() print("=" * 78) print("GROWING DEPTH DURING A RUN — A PROTOTYPE ON FASHION") print("=" * 78) print(f" backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}") for k, v in CFG.items(): print(f" {k:13s} = {v}") E = CFG["stage_epochs"]*CFG["stages"] print(f"\n every arm trains for {E} epochs. The grown arms spend the") print(f" first {CFG['stage_epochs']} at one layer and the next at three,") print(f" so they are CHEAPER as well as different — the table reports") print(f" wall clock so that is visible.") print("=" * 78, flush=True) Xtr, Ytr, Xte, yte, side = load(CFG) Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte) arms = [("grown 1>3>5, constant", True, False, 5), ("grown 1>3>5, slowed", True, True, 5), ("5 layers, constant", False, False, 5), ("5 layers, slowed", False, True, 5), ("3 layers, constant", False, False, 3), ("1 layer, constant", False, False, 1)] res = {} for nm, grown, slow, tgt in arms: accs, hs, evs, K, mc, secs = [], None, None, 0, 0, time.time() for s in CFG["seeds"]: a, h, e, K, mc = run(Xtr, Ytr, Xte, yte, side, CFG, s, grown, slow, tgt) accs.append(a) if hs is None: hs, evs = h, e res[nm] = dict(acc=float(np.mean(accs)), sd=float(np.std(accs)), stored=K, macs=mc, hist=hs, events=evs, secs=time.time()-secs) print(f" {nm:>22s} {np.mean(accs):.4f} sd {np.std(accs):.4f}" f" {K:,} values [{time.time()-t0:.0f}s]", flush=True) json.dump({k: {kk: vv for kk, vv in v.items() if kk != "hist"} for k, v in res.items()}, open("grow_depth.json", "w"), indent=2) print("\n" + "=" * 78) print(" WHAT DOES ADDING A LAYER COST AT THE MOMENT IT HAPPENS?") print("=" * 78) for nm in ("grown 1>3>5, constant", "grown 1>3>5, slowed"): print(f" {nm}") for e in res[nm]["events"]: print(f" to {e['to']} layers: {e['before']:.4f} -> " f"{e['after']:.4f} ({e['after']-e['before']:+.4f})") print(f"\n a delta filter passes its input through exactly, but layer") print(f" normalisation re-standardises afterwards, so these steps are") print(f" small rather than zero — unlike the width and storage") print(f" experiments, where the join was exact to machine precision.") print("\n" + "=" * 78) print(" READOUT") print("=" * 78) print(f" {'arm':>22s} {'values':>8s} {'multiplies':>11s} {'time':>7s} " f"{'accuracy':>9s} {'sd':>7s}") for nm, *_ in arms: r = res[nm] print(f" {nm:>22s} {r['stored']:8,} {r['macs']:11,} " f"{r['secs']:6.0f}s {r['acc']:9.4f} {r['sd']:7.4f}") gc = res["grown 1>3>5, constant"]["acc"] gs = res["grown 1>3>5, slowed"]["acc"] sc = res["5 layers, constant"]["acc"] ss = res["5 layers, slowed"]["acc"] l3 = res["3 layers, constant"]["acc"] sd = max(r["sd"] for r in res.values()) print(f"\n the 2x2, and the interaction is the interesting cell:\n") print(f" {'':>16s} {'constant':>10s} {'slowed':>10s} {'effect of slowing':>19s}") print(f" {'grown':>16s} {gc:10.4f} {gs:10.4f} {gs-gc:+19.4f}") print(f" {'from scratch':>16s} {sc:10.4f} {ss:10.4f} {ss-sc:+19.4f}") print(f" {'effect of growing':>16s} {gc-sc:+10.4f} {gs-ss:+10.4f} " f"{(gs-gc)-(ss-sc):+19.4f}") print(f"\n seed spread (worst) {sd:.4f}") print(f" three layers, the current best, scores {l3:.4f}") print() best = max((gc, "grown constant"), (gs, "grown slowed"), (sc, "5 from scratch"), (ss, "5 slowed"), (l3, "3 layers")) if max(gc, gs) > max(sc, ss) + 2*sd and max(gc, gs) > l3 + 2*sd: print(" GROWING DEPTH WORKS. It beats the same depth trained cold") print(" AND the three-layer model, so five layers are hard to") print(" optimise rather than too large, and a curriculum reaches") print(" what a cold start cannot. Worth repeating on CIFAR.") elif max(gc, gs) > max(sc, ss) + 2*sd: print(" GROWING BEATS COLD AT THE SAME DEPTH but not the shallower") print(" model, so the curriculum helps with the optimisation and") print(" five layers are still the wrong size for this problem.") elif abs(max(gc, gs) - max(sc, ss)) < 2*sd: print(" GROWING CHANGES NOTHING at the same depth — a third") print(" curriculum that adds capacity to a model which was not") print(" short of it. The pattern across width, storage and now") print(" depth is consistent, and consistent enough to stop.") else: print(" GROWING LOSES. Read the table; the transitions above say") print(" whether the cost is paid at the moment of growth or") print(" afterwards.") print(f"\n best arm: {best[1]} at {best[0]:.4f}") print(f" total {time.time()-t0:.0f}s; wrote grow_depth.json") if __name__ == "__main__": main()