| """ |
| 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 = [], [] |
| |
| 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() |
|
|
