strain.py

"""
DOES A FOLD KNOW WHERE IT HURTS?

A fold forces every weight in a group to hold one value. If those weights
all want to move the same way, the tie costs nothing. If they pull apart,
that group is being asked to do two jobs at once — and the disagreement is
a purely LOCAL quantity, computable from the gradients already flowing,
with nothing deciding anything and nothing learned.

That is the cheapest form of self-adaptation available here: split the
groups under strain, merge the slack ones. No controller, no gate, no
reasoning about it.

But before building a mechanism, the question is whether there is anything
to aim at. If a few groups carry most of the strain, targeting means
something. If strain is spread evenly, there is nothing to target and the
idea dies for the cost of this script.

THE QUANTITY. For group k, with the gradients of its member weights:

    disagreement = var / (mean^2 + var)

which is zero when every weight in the group pulls the same way and one
when they cancel exactly. It is bounded, scale-free, and falls out of the
scatter that already runs every step — one extra accumulation of squares.

THREE THINGS ARE READ FROM IT.

  IS IT CONCENTRATED?   the share of total strain held by the top tenth of
                        groups. Near a tenth means flat and untargetable.

  IS IT STABLE?         rank correlation between epochs. If the strained
                        groups keep changing, targeting them is chasing
                        noise rather than structure.

  DOES IT TELL GOOD PARTITIONS FROM BAD?  a convolution ties weights that
                        share a role; an arbitrary tying at the same
                        storage ties unrelated ones. The arbitrary
                        partition should strain MORE, and if it does not,
                        the measure is not seeing what it claims to.

A caution carried from elsewhere in this programme: growing storage by
splitting groups already FAILED, at -0.0224 against cold training, and the
diagnosis was that the model sat in the small model's basin rather than
that the splits were in the wrong places. A better trigger does not
obviously fix a basin problem, so a positive result here licenses an
experiment, not a mechanism.
"""

import numpy as np
import time
import json

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


def arbitrary(ni, no, K, seed=0):
    """An arbitrary tying at matched storage: the same number of groups,
    assigned without regard to what any weight does."""
    rg = np.random.default_rng(seed)
    return rg.integers(0, K, ni*no).astype(np.int32), K, no


class Grouped:
    """Fixed-order group sums, and group sums of squares beside them.

    The squares are the only addition: everything else already runs every
    step, so the strain measurement costs one more accumulation."""

    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.counts = K, to_dev(counts)
        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 _reduce(self, gs):
        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 sums(self, g, squares=False):
        gs = g.reshape(-1)[self.order]
        s = self._reduce(gs)
        return (s, self._reduce(gs*gs)) if squares else s


def disagreement(s, sq, counts, eps=1e-20):
    """var / (mean^2 + var): zero when a group's weights agree, one when
    they cancel."""
    c = xp.maximum(counts, 1)
    mean = s/c
    var = xp.maximum(sq/c - mean*mean, 0)
    return var/(mean*mean + var + eps)


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)
    if cfg["grid"] != 28:
        s = 28//cfg["grid"]
        X = X.reshape(-1, cfg["grid"], s, cfg["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"]+5000]
    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(grid=14, c_in=1, chan=16, n_train=20000, batch=128, lr=1e-3,
           epochs=40, report=(1, 2, 5, 10, 20, 40), seed=0)


def run(Xtr, Ytr, Xte, yte, idx, K, hid, cfg, seed, label):
    D = Xtr.shape[1]
    rg = np.random.default_rng(seed)
    I = to_dev(idx, np.int32) if _GPU else idx
    grp = Grouped(I, K)
    taps = max(1, int(to_host(grp.counts)[:-1].mean()))
    v = rg.normal(0, np.sqrt(2.0/9), K).astype(np.float32); v[-1] = 0.0
    P = [to_dev(v), to_dev(rg.normal(0, np.sqrt(2.0/hid), (hid, 10))),
         xp.zeros(hid, DT), xp.zeros(10, DT)]
    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)
    snaps = {}
    for ep in range(cfg["epochs"]):
        perm = ag.permutation(n)
        acc, nb = xp.zeros(K, DT), 0
        for st in range(0, n, cfg["batch"]):
            b = perm[st:st+cfg["batch"]]
            x = Xtr[b]; y = Ytr[b]
            W = P[0][I].reshape(D, hid)
            z = x @ W + P[2]; h = xp.maximum(z, 0)
            lg = h @ P[1] + P[3]
            e = xp.exp(lg - lg.max(1, keepdims=True))
            d = (e/e.sum(1, keepdims=True) - y)/len(b)
            d0 = (d @ P[1].T)*(z > 0)
            gW = x.T @ d0
            s, sq = grp.sums(gW, squares=True)
            acc += disagreement(s, sq, grp.counts); nb += 1
            G = [s, h.T @ d, d0.sum(0), d.sum(0)]
            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)
        if ep+1 in cfg["report"]:
            # the padding group holds most of the entries and is not a
            # feature, so it is excluded from the distribution
            snaps[ep+1] = to_host(acc/nb)[:-1]
    out = []
    for s in range(0, Xte.shape[0], 4096):
        hh = xp.maximum(Xte[s:s+4096] @ P[0][I].reshape(D, hid) + P[2], 0)
        out.append(to_host(hh @ P[1] + P[3]))
    acc_ = float((np.concatenate(out).argmax(1) == yte).mean())
    return snaps, acc_, taps


def concentration(v):
    """Share of the total held by the top tenth of groups. A tenth means
    perfectly flat."""
    s = np.sort(v)[::-1]
    k = max(1, len(s)//10)
    return float(s[:k].sum()/max(s.sum(), 1e-12))


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("DOES A FOLD KNOW WHERE IT HURTS?")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:8s} = {v}")
    print(f"\n  disagreement = var / (mean^2 + var) per group: 0 when every")
    print(f"  weight in the group pulls the same way, 1 when they cancel")
    print("=" * 78, flush=True)

    Xtr, Ytr, Xte, yte = load(CFG)
    Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
    g, ci, ch = CFG["grid"], CFG["c_in"], CFG["chan"]
    ci_, K, hid = windowed(g, ci, 3, ch)
    arms = {"convolution": (ci_, K, hid),
            "arbitrary (same K)": arbitrary(ci*g*g, hid, K, CFG["seed"])}
    res = {}
    for nm, (idx, KK, hh) in arms.items():
        snaps, acc, taps = run(Xtr, Ytr, Xte, yte, idx, KK, hh, CFG,
                               CFG["seed"], nm)
        res[nm] = dict(acc=acc, taps=taps,
                       snaps={k: v.tolist() for k, v in snaps.items()})
        print(f"\n  {nm}  (accuracy {acc:.4f}, {taps} weights a group)")
        print(f"  {'epoch':>6s} {'mean':>8s} {'median':>8s} {'max':>8s} "
              f"{'top-10% share':>14s}")
        for ep, v in snaps.items():
            print(f"  {ep:6d} {v.mean():8.4f} {np.median(v):8.4f} "
                  f"{v.max():8.4f} {concentration(v):13.1%}")
        eps = sorted(snaps)
        if len(eps) > 1:
            a, b = snaps[eps[0]], snaps[eps[-1]]
            ra = np.argsort(np.argsort(a)); rb = np.argsort(np.argsort(b))
            rho = float(np.corrcoef(ra, rb)[0, 1])
            print(f"  rank correlation between epoch {eps[0]} and "
                  f"{eps[-1]}: {rho:+.3f}")
            res[nm]["stability"] = rho
        json.dump({k: {kk: vv for kk, vv in v.items() if kk != "snaps"}
                   for k, v in res.items()},
                  open("strain.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    c = res["convolution"]; a = res["arbitrary (same K)"]
    last = max(CFG["report"])
    cv = np.array(c["snaps"][last]); av = np.array(a["snaps"][last])
    print(f"  {'':>20s} {'accuracy':>9s} {'mean strain':>12s} "
          f"{'top-10%':>9s} {'stability':>10s}")
    for nm, r, v in (("convolution", c, cv), ("arbitrary", a, av)):
        print(f"  {nm:>20s} {r['acc']:9.4f} {v.mean():12.4f} "
              f"{concentration(v):8.1%} {r.get('stability', float('nan')):+10.3f}")
    print()
    conc, stab = concentration(cv), c.get("stability", 0.0)
    if conc > 0.25 and stab > 0.5:
        print(f"  THERE IS SOMETHING TO AIM AT. The top tenth of groups holds")
        print(f"  {conc:.0%} of the strain and the ranking is stable across")
        print(f"  training ({stab:+.2f}), so a split targeted by disagreement")
        print(f"  would land in the same places consistently. That licenses")
        print(f"  the experiment — remembering that uniform splitting already")
        print(f"  failed at -0.0224, and for a reason (basins) that a better")
        print(f"  trigger may not fix.")
    elif conc > 0.25:
        print(f"  CONCENTRATED BUT UNSTABLE. The top tenth holds {conc:.0%} of")
        print(f"  the strain, but which groups those are keeps changing")
        print(f"  ({stab:+.2f}), so targeting them is chasing the gradient's")
        print(f"  own noise rather than structure.")
    else:
        print(f"  NOTHING TO AIM AT. The top tenth holds {conc:.0%} against")
        print(f"  the 10% a flat distribution would give, so strain is spread")
        print(f"  evenly and there is no subset of groups worth splitting.")
        print(f"  Adaptive splitting has no signal here, and the idea costs")
        print(f"  this script rather than a day.")
    print()
    if av.mean() > cv.mean() + 0.02:
        print(f"  AND THE MEASURE DISCRIMINATES: an arbitrary tying strains")
        print(f"  {av.mean():.3f} against the convolution's {cv.mean():.3f}. Tying")
        print(f"  weights that share a role really does cost less than tying")
        print(f"  unrelated ones, which is the blindness argument showing up")
        print(f"  in the gradients rather than in the accuracy.")
    elif abs(av.mean() - cv.mean()) < 0.02:
        print(f"  BUT THE MEASURE DOES NOT DISCRIMINATE: arbitrary strains")
        print(f"  {av.mean():.3f} against the convolution's {cv.mean():.3f}, so it")
        print(f"  cannot tell a partition that shares a role from one that")
        print(f"  does not — and it is not measuring what it claims.")
    else:
        print(f"  UNEXPECTED: the convolution strains MORE than an arbitrary")
        print(f"  tying ({cv.mean():.3f} against {av.mean():.3f}). Worth")
        print(f"  understanding before anything is built on this quantity.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote strain.json")


if __name__ == "__main__":
    main()