grow_depth.py

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