sample_efficiency.py

"""
HOW FEW EXAMPLES DOES A MEMBER NEED?

Every member measured in this programme was fitted on thousands of
examples: 15,000 for the CIFAR extension work, 8,045 for the Fashion
specialist. A human in the loop supplies five, or twenty. Against Lee's own
criteria — useful, small, safe, a defined family, ADDABLE — the first four
are settled and this is the whole of the gap.

It is answerable cheaply because the base is frozen, so its features are a
constant function of the input and can be computed once. And because a
member is LINEAR over those features, it does not need gradient descent at
all: a ridge solve gives the exact minimiser in one step, which is what
makes a per-prompt member possible rather than merely small.

    delta = (A'A + lambda I)^-1 A' (target - base logits)

fitted on the RESIDUAL, so a member that explains nothing contributes
nothing. That is the failsafe stated as arithmetic: as lambda dominates,
the solution shrinks toward zero and the base answers. A starved member
degrades to silence rather than to noise.

TWO FEATURE WIDTHS, and the comparison is the point. The flat read is 3,136
numbers per image; a pooled read is 16. At N = 5 examples a 3,136-wide fit
has 3,136 unknowns per class and no hope; a 16-wide fit has sixteen. If
pooling is what makes small-N members work, that is a design rule rather
than an accident — and it connects to the pooling result, where a pooled
head cost nothing at adequate depth and saved several times the storage.

THREE THINGS ARE MEASURED AGAINST N:

  ACCURACY on the member's own task, against the base's accuracy there and
  against a member fitted on everything

  DELTA NORM, to check the failsafe actually operates — a starved member
  should be small, not wrong

  HARM TO THE REST, because a member fitted on four classes should not
  damage the other six, and at small N it might

Gradient descent is run alongside at each N, because the closed form is
only worth having if it matches.
"""

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


_FIXED = {}


class FixedScatter:
    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):
    key = (id(idx), K)
    if key not in _FIXED:
        _FIXED[key] = FixedScatter(idx, K)
    return _FIXED[key](dW)


def train_base(Xtr, Ytr, cfg, seed):
    D, g, ch = Xtr.shape[1], cfg["grid"], cfg["chan"]
    rg = np.random.default_rng(seed)
    layers, cin = [], cfg["c_in"]
    for l in range(cfg["depth"]):
        idx, K, no = windowed(g, cin, 3, ch)
        layers.append(dict(idx=to_dev(idx, np.int32) if _GPU else idx,
                           K=K, out=no, taps=cin*9,
                           ins=D if l == 0 else layers[-1]["out"]))
        cin = ch
    L = cfg["depth"]
    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))
    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):
        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)
            cache.append((h, W, var, zn, zs))
            h = a
        return h, cache

    for ep in range(cfg["epochs"]):
        perm = ag.permutation(n)
        for st in range(0, n, cfg["batch"]):
            b = perm[st:st+cfg["batch"]]
            x = Xtr[b]; y = Ytr[b]
            h, cache = fwd(x)
            lg = h @ P[HEAD] + P[OB]
            e = xp.exp(lg - lg.max(1, keepdims=True))
            d = (e/e.sum(1, keepdims=True) - 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)
    return P, fwd, HEAD, OB, layers[-1]["out"]


def ridge(A, R, lam):
    """The closed-form member: one solve, no steps.

    Fitted on the RESIDUAL between the target and what the base already
    says, so a member that explains nothing contributes nothing — and as
    lambda dominates the solution shrinks toward zero and the base answers.
    That is the failsafe as arithmetic rather than as a rule."""
    n, d = A.shape
    A1 = xp.concatenate([A, xp.ones((n, 1), DT)], 1)
    if d + 1 <= n:
        G = A1.T @ A1 + lam*xp.eye(d+1, dtype=DT)
        W = xp.linalg.solve(G, A1.T @ R)
    else:
        # the wide case: solve in the sample space instead, which is the
        # only tractable form when there are five examples and 3,136
        # features
        G = A1 @ A1.T + lam*xp.eye(n, dtype=DT)
        W = A1.T @ xp.linalg.solve(G, R)
    return W[:-1], W[-1]


def descent(A, R, lam, steps, lr):
    """The same fit by gradient descent, for comparison. The closed form is
    only worth having if it matches."""
    n, d = A.shape
    W = xp.zeros((d, R.shape[1]), DT); b = xp.zeros(R.shape[1], DT)
    M = [xp.zeros_like(W), xp.zeros_like(b)]
    V = [xp.zeros_like(W), xp.zeros_like(b)]
    for t in range(1, steps+1):
        E = A @ W + b - R
        G = [A.T @ E/n + lam*W/n, E.mean(0)]
        for i, (p_, gr) in enumerate(zip([W, b], G)):
            M[i] = 0.9*M[i] + 0.1*gr
            V[i] = 0.999*V[i] + 0.001*gr*gr
            upd = p_ - lr*(M[i]/(1-0.9**t))/(xp.sqrt(V[i]/(1-0.999**t))+1e-8)
            if i == 0:
                W = upd
            else:
                b = upd
    return W, b


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"]+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, y[tr], f(X[te]), y[te]


CFG = dict(grid=14, c_in=1, chan=16, depth=3, n_train=20000, batch=128,
           lr=1e-3, epochs=30, seed=0, member_classes=(0, 1, 2, 3),
           Ns=(2, 5, 10, 20, 50, 100, 200, 500, 1000, 3000),
           draws=5, lam=1.0, gd_steps=300, gd_lr=0.05)


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("HOW FEW EXAMPLES DOES A MEMBER NEED?")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:15s} = {v}")
    print(f"\n  the member is a closed-form ridge solve on the RESIDUAL — one")
    print(f"  step, no gradient descent — fitted for classes "
          f"{CFG['member_classes']}")
    print("=" * 78, flush=True)

    Xtr, Ytr, ytr, Xte, yte = load(CFG)
    Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
    P, fwd, HEAD, OB, width = train_base(Xtr, Ytr, CFG, CFG["seed"])
    ftr, _ = fwd(Xtr); fte, _ = fwd(Xte)
    base_tr = ftr @ P[HEAD] + P[OB]
    base_te = fte @ P[HEAD] + P[OB]
    own_te = np.isin(yte, CFG["member_classes"])
    b_all = float((to_host(base_te).argmax(1) == yte).mean())
    b_own = float((to_host(base_te).argmax(1)[own_te] == yte[own_te]).mean())
    print(f"\n  base: {b_all:.4f} overall, {b_own:.4f} on the member's four"
          f"   [{time.time()-t0:.0f}s]", flush=True)

    ch, g = CFG["chan"], CFG["grid"]
    views = {"flat (3,136 wide)": (ftr, fte),
             "pooled (16 wide)": (ftr.reshape(-1, ch, g*g).mean(2),
                                  fte.reshape(-1, ch, g*g).mean(2))}
    # the base's own logit scale, so w = 1 is a sensible default rather
    # than an arbitrary one
    margin = float(to_host(base_tr.std()))
    print(f"  base logits have a spread of {margin:.2f}, so a member's target")
    print(f"  is to lift the true class by that much")
    pool_idx = np.where(np.isin(ytr, CFG["member_classes"]))[0]
    res = {}
    for vn, (A_tr, A_te) in views.items():
        print(f"\n  {vn}   member is {A_tr.shape[1]}x10 = "
              f"{A_tr.shape[1]*10+10:,} values")
        print(f"  {'N':>6s} {'own (ridge)':>12s} {'sd':>7s} "
              f"{'own (descent)':>14s} {'the rest':>9s} {'|delta|':>9s} "
              f"{'solve ms':>9s}")
        for N in CFG["Ns"]:
            accs, gds, rests, norms, ms = [], [], [], [], []
            for dr in range(CFG["draws"]):
                rg = np.random.default_rng(1000*dr + N)
                sub = rg.choice(pool_idx, min(N, len(pool_idx)),
                                replace=False)
                s = to_dev(sub, np.int64) if _GPU else sub
                # THE TARGET: lift the true class by the base logits' own
                # scale, leaving the others alone. The first version used
                # (one_hot - base_logits), which asked the member to drag
                # every logit down to a 0/1 vector — a mis-specified
                # regression, and it made accuracy FALL as N rose because
                # the curve measured how completely a bad target got
                # fitted. Gradient descent scored HIGHER there precisely
                # because it did not converge.
                A = A_tr[s]; R = margin*Ytr[s]
                t1 = time.time()
                W, b = ridge(A, R, CFG["lam"])
                ms.append((time.time()-t1)*1000)
                lg = to_host(base_te + A_te @ W + b)
                pr = lg.argmax(1)
                accs.append(float((pr[own_te] == yte[own_te]).mean()))
                rests.append(float((pr[~own_te] == yte[~own_te]).mean()))
                norms.append(float(to_host(xp.linalg.norm(W))))
                if dr == 0:
                    Wg, bg = descent(A, R, CFG["lam"], CFG["gd_steps"],
                                     CFG["gd_lr"])
                    lgg = to_host(base_te + A_te @ Wg + bg).argmax(1)
                    gds.append(float((lgg[own_te] == yte[own_te]).mean()))
            res[f"{vn}/{N}"] = dict(own=float(np.mean(accs)),
                                    sd=float(np.std(accs)),
                                    gd=float(np.mean(gds)),
                                    rest=float(np.mean(rests)),
                                    norm=float(np.mean(norms)))
            print(f"  {N:6d} {np.mean(accs):12.4f} {np.std(accs):7.4f} "
                  f"{np.mean(gds):14.4f} {np.mean(rests):9.4f} "
                  f"{np.mean(norms):9.3f} {np.mean(ms):9.2f}", flush=True)
            json.dump(res, open("sample_efficiency.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    print(f"  the base already scores {b_own:.4f} on these four classes, so a")
    print(f"  member has to beat that to be worth adding\n")
    for vn in views:
        ns = [N for N in CFG["Ns"] if f"{vn}/{N}" in res]
        best = max(res[f"{vn}/{N}"]["own"] for N in ns)
        beats = next((N for N in ns if res[f"{vn}/{N}"]["own"] > b_own), None)
        near = next((N for N in ns
                     if res[f"{vn}/{N}"]["own"] > best - 0.01), None)
        print(f"  {vn}")
        print(f"    beats the base from N = {beats if beats else 'never'}")
        print(f"    within a point of its own best from N = "
              f"{near if near else 'never'}")
        print(f"    best {best:.4f} at N = "
              f"{max(ns, key=lambda N: res[f'{vn}/{N}']['own'])}")
    fl = "flat (3,136 wide)"; po = "pooled (16 wide)"
    small = [N for N in CFG["Ns"] if N <= 20]
    df = np.mean([res[f"{po}/{N}"]["own"] - res[f"{fl}/{N}"]["own"]
                  for N in small if f"{po}/{N}" in res])
    print(f"\n  at N <= 20 the pooled member is {df:+.4f} against the flat one")
    if df > 0.02:
        print(f"  POOLING IS WHAT MAKES SMALL N WORK. Sixteen unknowns a class")
        print(f"  can be fitted from a handful of examples where 3,136 cannot,")
        print(f"  so a human-in-the-loop member should read a pooled view —")
        print(f"  which the pooling result already said costs nothing at")
        print(f"  adequate depth.")
    elif df < -0.02:
        print(f"  THE FLAT VIEW WINS EVEN AT SMALL N, which is surprising and")
        print(f"  means ridge is handling the wide case better than the")
        print(f"  parameter count suggests.")
    else:
        print(f"  THE TWO VIEWS ARE CLOSE AT SMALL N, so the width is not the")
        print(f"  binding constraint and ridge is doing the work.")
    ns = CFG["Ns"]
    nf = [res[f"{po}/{N}"]["norm"] for N in ns if f"{po}/{N}" in res]
    print(f"\n  the failsafe: |delta| against N, pooled — "
          + "  ".join(f"{x:.2f}" for x in nf))
    if nf[0] < nf[-1]/2:
        print(f"  IT OPERATES. A starved member is SMALL rather than wrong, so")
        print(f"  the failure mode at low N is the base answering rather than")
        print(f"  noise being added.")
    else:
        print(f"  IT DOES NOT OPERATE as expected — a member fitted on two")
        print(f"  examples is as large as one fitted on thousands, so lambda")
        print(f"  is too small to protect the low-N case.")
    hr = [res[f"{po}/{N}"]["rest"] for N in ns if f"{po}/{N}" in res]
    print(f"\n  harm to the other six classes, pooled — "
          + "  ".join(f"{x:.3f}" for x in hr))
    print(f"\n  total {time.time()-t0:.0f}s; wrote sample_efficiency.json")


if __name__ == "__main__":
    main()