scramble.py

"""
STARTING WITH THE PICTURE IN PIECES.

The curriculum that worked outside this programme destroyed structure and
then restored it: jumbled words, then jumbled sentences, then jumbled
paragraphs, then the real text. Each stage removed a scale of arrangement
and handed it back later.

The image version is to SHUFFLE BLOCKS, per image and differently for every
image, so no fixed unscrambling can be learned. A convolution reads a 3x3
neighbourhood; shuffle single pixels and every neighbourhood is nonsense,
so nothing about arrangement is available and only per-pixel statistics
remain. Shuffle 2x2 blocks and local texture survives while global layout
does not. Leave the image alone and everything is available.

This is a different lever from the two curricula already tried here.
Growing width and growing storage changed what the model could REPRESENT,
and both failed on models that were not short of capacity. This changes
what the DATA contains, which is the thing the original result actually
manipulated.

Four arms, all trained for the same number of epochs on the same data:

  INTACT              the control, ordinary training throughout
  5 HARD              five epochs on pixel-shuffled images, then intact
  15 HARD             fifteen, to separate severity from duration
  5 MILD              five epochs on 2x2-block shuffle, then intact

and a floor arm trained on shuffled images throughout, which says how much
of this task needs arrangement at all.

The claim to check is not merely the final accuracy. It is whether the
model RECOVERS FASTER after the images are restored than a fresh model
reaches the same point — because that is what "the hard phase taught it
something" would look like. The epoch-by-epoch trace is printed for
exactly this.

A caution worth stating in advance: these models UNDERFIT. Test loss sits
below train loss throughout and accuracy is still climbing at the end of
training, so a model already short of signal may not benefit from being
given less of it. The original result came from the opposite regime, where
a degenerate shortcut was available and the jumbling removed it.
"""

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


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 shuffle_blocks(X, side, block, rg):
    """Permute the blocks of each image INDEPENDENTLY.

    A fixed permutation could be undone by the first layer, which would
    make the task no harder at all — the point is that no unscrambling
    generalises, exactly as every sentence being jumbled differently means
    word order carries nothing.

    block=1 destroys everything but the pixel histogram; block=side leaves
    the image alone."""
    if block >= side:
        return X
    n = X.shape[0]
    nb = side//block
    im = X.reshape(n, nb, block, nb, block)
    im = im.transpose(0, 1, 3, 2, 4).reshape(n, nb*nb, block*block)
    r = rg.random((n, nb*nb))
    r = xp.asarray(r) if _GPU else r
    perm = xp.argsort(r, axis=1)
    im = xp.take_along_axis(im, perm[:, :, None], axis=1)
    im = im.reshape(n, nb, nb, block, block).transpose(0, 1, 3, 2, 4)
    return im.reshape(n, side*side)


def train(Xtr, Ytr, Xte, yte, cfg, seed, schedule, side, label=""):
    """schedule maps an epoch to a block size; side means intact."""
    D = Xtr.shape[1]
    rg = np.random.default_rng(seed)
    layers = []
    cin = cfg["c_in"]
    for l in range(cfg["depth"]):
        idx, K, no = windowed(cfg["grid"], cin, 3, cfg["chan"])
        layers.append(dict(idx=to_dev(idx, np.int32) if _GPU else idx,
                           K=K, out=no, ins=(D if l == 0 else layers[-1]["out"]),
                           taps=cin*9))
        cin = cfg["chan"]
    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))
    L = len(layers)
    P += [xp.ones(l["out"], DT) for l in layers]
    P += [xp.zeros(l["out"], DT) for l in layers]
    P += [to_dev(rg.normal(0, np.sqrt(2.0/layers[-1]["out"]),
                           (layers[-1]["out"], 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)

    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
        return h, cache

    hist = []
    for ep in range(cfg["epochs"]):
        blk = schedule(ep)
        perm = ag.permutation(n)
        tot, nb_ = 0.0, 0
        for st in range(0, n, cfg["batch"]):
            b = perm[st:st+cfg["batch"]]
            x = Xtr[b]
            if blk < side:
                x = shuffle_blocks(x, side, blk, ag)
            y = Ytr[b]
            h, cache = fwd(x, keep=True)
            lg = h @ P[HEAD] + P[OB]
            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[HEAD] = h.T @ d; G[OB] = d.sum(0)
            dh = d @ P[HEAD].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*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 s2 in range(0, Xte.shape[0], 4096):
            h, _ = fwd(Xte[s2:s2+4096])
            out.append(to_host(h @ P[HEAD] + P[OB]))
        acc = float((np.concatenate(out).argmax(1) == yte).mean())
        hist.append(dict(epoch=ep+1, block=blk, train_loss=tot/max(nb_, 1),
                         acc=acc))
        if label and ((ep+1) % cfg["report_every"] == 0 or ep < 6):
            print(f"    {label} ep {ep+1:3d}  block {blk:2d}  "
                  f"train {tot/max(nb_,1):.4f}  acc {acc:.4f}", flush=True)
    return hist


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, chan=16, depth=3, n_train=20000, batch=128,
           lr=1e-3, epochs=60, report_every=10, seeds=(0, 1))


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("STARTING WITH THE PICTURE IN PIECES")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:12s} = {v}")
    print(f"\n  blocks are permuted PER IMAGE, so no fixed unscrambling")
    print(f"  generalises — the analogue of every sentence being jumbled")
    print(f"  differently. block 1 leaves only the pixel histogram.")
    print("=" * 78, flush=True)

    Xtr, Ytr, Xte, yte, side = load(CFG)
    Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
    S = side
    arms = [
        ("intact", lambda ep: S),
        ("5 hard (block 1)", lambda ep: 1 if ep < 5 else S),
        ("15 hard (block 1)", lambda ep: 1 if ep < 15 else S),
        ("5 mild (block 2)", lambda ep: 2 if ep < 5 else S),
        ("shuffled throughout", lambda ep: 1),
    ]
    res = {}
    for nm, sch in arms:
        hs = [train(Xtr, Ytr, Xte, yte, CFG, s, sch, side,
                    label=f"{nm[:19]:>19s} s{s}")
              for s in CFG["seeds"]]
        fin = [h[-1]["acc"] for h in hs]
        res[nm] = dict(acc=float(np.mean(fin)), sd=float(np.std(fin)),
                       hist=hs[0])
        print(f"  {nm:>19s} {np.mean(fin):.4f}  sd {np.std(fin):.4f}"
              f"   [{time.time()-t0:.0f}s]\n", flush=True)
        json.dump({k: dict(acc=v["acc"], sd=v["sd"]) for k, v in res.items()},
                  open("scramble.json", "w"), indent=2)

    print("=" * 78)
    print("  DOES THE HARD PHASE TEACH ANYTHING?")
    print("=" * 78)
    ctrl = res["intact"]["hist"]
    print(f"  accuracy by epoch, first seed:\n")
    print(f"  {'epoch':>6s} " + " ".join(f"{n[:14]:>15s}" for n, _ in arms))
    for e in list(range(0, 8)) + list(range(9, CFG["epochs"], 10)):
        row = []
        for nm, _ in arms:
            h = res[nm]["hist"]
            row.append(f"{h[e]['acc']:.4f}" if e < len(h) else "")
        print(f"  {e+1:6d} " + " ".join(f"{c:>15s}" for c in row))
    print(f"\n  the arms that start hard are BEHIND at epoch 5 by")
    print(f"  construction. What matters is whether they catch up faster")
    print(f"  than the control got there — that is what the hard phase")
    print(f"  teaching something would look like.")
    for nm in ("5 hard (block 1)", "15 hard (block 1)", "5 mild (block 2)"):
        h = res[nm]["hist"]
        start = 5 if "5 " in nm else 15
        a0 = h[start-1]["acc"]
        # how many epochs the control needed to reach where this arm
        # resumes, and how many this arm needs to reach the control's next
        catch = next((i+1 for i in range(start, len(h))
                      if h[i]["acc"] >= ctrl[start-1]["acc"]), None)
        print(f"    {nm:>19s}: resumed at {a0:.4f}, reached the control's "
              f"epoch-{start} score at epoch "
              f"{catch if catch else 'never'}")

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    for nm, _ in arms:
        r = res[nm]
        print(f"  {nm:>19s} {r['acc']:.4f}  sd {r['sd']:.4f}")
    base = res["intact"]["acc"]
    sd = max(r["sd"] for r in res.values())
    best = max((res[nm]["acc"], nm) for nm, _ in arms if nm != "intact"
               and nm != "shuffled throughout")
    print(f"\n  against ordinary training ({base:.4f}):")
    for nm, _ in arms:
        if nm != "intact":
            print(f"    {nm:>19s} {res[nm]['acc']-base:+.4f}")
    print(f"\n  seed spread (worst) {sd:.4f}")
    print(f"  shuffled throughout reaches "
          f"{res['shuffled throughout']['acc']:.4f}, which is how much of")
    print(f"  this task needs no arrangement at all")
    print()
    if best[0] > base + 2*sd:
        print(f"  A HARD START HELPS. {best[1]} beats ordinary training by")
        print(f"  {best[0]-base:+.4f}, so removing arrangement early and")
        print(f"  handing it back leaves the model somewhere better than")
        print(f"  starting with everything available.")
    elif best[0] < base - 2*sd:
        print(f"  A HARD START COSTS. The epochs spent on scrambled images")
        print(f"  are not recovered, so this is a third curriculum that does")
        print(f"  not transfer — after width and storage, the pattern is")
        print(f"  consistent enough to be worth stating.")
    else:
        print(f"  NO DIFFERENCE. The hard phase neither helps nor hurts: the")
        print(f"  model reaches the same place having wasted the epochs.")
        print(f"  Which is itself informative — these models UNDERFIT, and a")
        print(f"  model short of signal has no shortcut for a curriculum to")
        print(f"  remove.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote scramble.json")


if __name__ == "__main__":
    main()