member_weight.py

"""
TURNING A MEMBER UP AND DOWN AT INFERENCE.

A member is a delta on a permanent base head:

    output = base(a) + w * delta(a)

and w is FREE AT INFERENCE. Nothing retrains, nothing is stored, and it can
differ per prompt. At w = 0 the base answers; at w = 1 the member answers as
fitted; between and beyond is unexplored territory that costs nothing to
visit.

Two questions, and the second is the interesting one.

IS IT A DIAL OR A SWITCH? If w = 0.5 gives something coherently halfway,
the weight is a real continuous control. If accuracy falls off a cliff
somewhere, it is an on/off switch with a misleading knob attached. Measured
by sweeping w finely and looking at the shape rather than the endpoints.

WHERE IS THE CONFIDENTLY WRONG BAND? Random output is not confusion, it is
noise — and a model that knows it is lost is not confused either, it is
abstaining. The interesting regime is where ACCURACY FALLS WHILE CONFIDENCE
HOLDS: wrong and committed. That is a band in w, and it can be located.

The member here is a SPECIALIST: a head fitted on four of the ten classes
against a frozen body that saw all ten. Amplifying it should make the model
increasingly insist on its own four, so the sweep separates three things —
accuracy on the member's classes, accuracy on everything else, and how
confident the model is about either.

A caution from the programme's own record: a WRONG member costs far more
than NO member — sixteen points worse in §7F — which is why the failsafe
emits the base when a member is unidentifiable. Amplifying a member
deliberately walks into exactly that failure mode. That is fine if it is
what is wanted; it is worth knowing it is the same mechanism the design
otherwise guards against.
"""

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):
    """The body and its permanent base head, on the whole task."""
    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, 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

    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, keep=True)
            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, layers, fwd, HEAD, OB


def fit_member(feats, Y, mask, base_lg, cfg, seed):
    """A DELTA on the base head, fitted only on the member's own examples.

    The base is frozen throughout — this is the injection construction, and
    the delta starts at zero so w = 0 reproduces the base exactly."""
    rg = np.random.default_rng(seed)
    w = feats.shape[1]
    Dh = xp.zeros((w, 10), DT); Db = xp.zeros(10, DT)
    P = [Dh, Db]
    M = [xp.zeros_like(p) for p in P]; V = [xp.zeros_like(p) for p in P]
    n = feats.shape[0]; t = 0
    sel = xp.asarray(mask) if _GPU else mask
    for ep in range(cfg["member_epochs"]):
        perm = rg.permutation(n)
        for st in range(0, n, cfg["batch"]):
            b = perm[st:st+cfg["batch"]]
            a = feats[b]; y = Y[b]
            lg = base_lg[b] + a @ P[0] + P[1]
            e = xp.exp(lg - lg.max(1, keepdims=True))
            d = (e/e.sum(1, keepdims=True) - y)*sel[b][:, None]/len(b)
            G = [a.T @ d, 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)
    return P


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, member_epochs=25, seed=0,
           member_classes=(0, 1, 2, 3),
           weights=(0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0, 3.0, 5.0,
                    8.0, -0.5, -1.0))


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("TURNING A MEMBER UP AND DOWN AT INFERENCE")
    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  output = base(a) + w * delta(a), with w free at inference.")
    print(f"  the member is a SPECIALIST fitted on classes "
          f"{CFG['member_classes']} against a frozen body that saw all ten.")
    print(f"  turning it up should make the model insist on its own four.")
    print("=" * 78, flush=True)

    Xtr, Ytr, ytr, Xte, yte = load(CFG)
    Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
    P, layers, fwd, HEAD, OB = 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]
    print(f"\n  base alone: {float((to_host(base_te).argmax(1) == yte).mean()):.4f}"
          f"   [{time.time()-t0:.0f}s]", flush=True)

    mask = np.isin(ytr, CFG["member_classes"]).astype(np.float32)
    Dh, Db = fit_member(ftr, Ytr, mask, base_tr, CFG, CFG["seed"]+3)
    delta_te = fte @ Dh + Db

    own = np.isin(yte, CFG["member_classes"])
    print(f"  member fitted on {int(mask.sum()):,} examples "
          f"({own.mean()*100:.0f}% of the test set is its own classes)")

    print(f"\n  {'w':>6s} {'overall':>8s} {'its own':>8s} {'the rest':>9s} "
          f"{'confidence':>11s} {'over':>7s} {'claims own':>11s}")
    rows = []
    for w in CFG["weights"]:
        lg = to_host(base_te + w*delta_te)
        e = np.exp(lg - lg.max(1, keepdims=True))
        p = e/e.sum(1, keepdims=True)
        pred = p.argmax(1)
        conf = p.max(1)
        acc = float((pred == yte).mean())
        a_own = float((pred[own] == yte[own]).mean())
        a_rest = float((pred[~own] == yte[~own]).mean())
        claims = float(np.isin(pred, CFG["member_classes"]).mean())
        rows.append(dict(w=w, acc=acc, own=a_own, rest=a_rest,
                         conf=float(conf.mean()), over=float(conf.mean())-acc,
                         claims=claims))
        print(f"  {w:6.2f} {acc:8.4f} {a_own:8.4f} {a_rest:9.4f} "
              f"{conf.mean():11.4f} {conf.mean()-acc:+7.4f} {claims:10.1%}")
    json.dump(rows, open("member_weight.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  DIAL OR SWITCH?")
    print("=" * 78)
    pos = [r for r in rows if 0 <= r["w"] <= 2.0]
    ws = np.array([r["w"] for r in pos]); ac = np.array([r["acc"] for r in pos])
    steps = np.abs(np.diff(ac)/np.diff(ws))
    print(f"  accuracy changes per unit w, over 0 to 2: "
          + "  ".join(f"{s:.3f}" for s in steps))
    if steps.max() < 3*max(steps.min(), 1e-6) and steps.max() < 0.2:
        print(f"  SMOOTH — no step is more than a few times any other, so w")
        print(f"  is a genuine continuous control and half a member means")
        print(f"  something.")
    else:
        print(f"  UNEVEN — the steepest stretch is {steps.max()/max(steps.min(),1e-6):.0f}"
              f" times the flattest, so w behaves")
        print(f"  more like a switch with a knob drawn on it than a dial.")

    print("\n" + "=" * 78)
    print("  WHERE IS WRONG-BUT-COMMITTED?")
    print("=" * 78)
    base_acc = rows[0]["acc"]; base_conf = rows[0]["conf"]
    print(f"  at w = 0 the base is {base_acc:.4f} accurate and "
          f"{base_conf:.4f} confident\n")
    print(f"  {'w':>6s} {'accuracy lost':>14s} {'confidence lost':>16s} "
          f"{'ratio':>8s}")
    band = []
    for r in rows:
        if r["w"] <= 0:
            continue
        da = base_acc - r["acc"]; dc = base_conf - r["conf"]
        ratio = da/max(dc, 1e-6) if dc > 0 else float("inf")
        band.append((ratio, r))
        print(f"  {r['w']:6.2f} {da:14.4f} {dc:16.4f} "
              + (f"{ratio:8.1f}" if np.isfinite(ratio) else "     inf"))
    lost = [(r["acc"], r) for _, r in band if base_acc - r["acc"] > 0.05]
    print()
    if lost:
        worst = min(lost)[1]
        print(f"  the model gives up the most accuracy at w = {worst['w']}:")
        print(f"    accuracy {worst['acc']:.4f} (from {base_acc:.4f})")
        print(f"    confidence {worst['conf']:.4f} (from {base_conf:.4f})")
        print(f"    and it names one of its own four classes "
              f"{worst['claims']:.0%} of the time")
        if worst["conf"] > base_conf - 0.05:
            print(f"\n  WRONG AND COMMITTED. Accuracy falls a long way while")
            print(f"  confidence barely moves, so this is not the model")
            print(f"  becoming unsure — it is the model becoming sure of")
            print(f"  something else. That is the band worth having.")
        else:
            print(f"\n  WRONG AND KNOWS IT. Confidence falls with accuracy, so")
            print(f"  amplifying the member produces hesitancy rather than")
            print(f"  misplaced conviction — closer to noise than to")
            print(f"  confusion.")
    else:
        print(f"  no weight in this sweep costs more than five points of")
        print(f"  accuracy, so the member is too weak to push the model")
        print(f"  anywhere interesting. Fit it harder or choose a member")
        print(f"  that disagrees with the base more.")
    neg = [r for r in rows if r["w"] < 0]
    if neg:
        print(f"\n  and NEGATIVE weights, which invert the member rather than")
        print(f"  removing it:")
        for r in neg:
            print(f"    w = {r['w']:5.2f}: overall {r['acc']:.4f}, its own "
                  f"{r['own']:.4f}, claims own {r['claims']:.1%}")
    print(f"\n  total {time.time()-t0:.0f}s; wrote member_weight.json")


if __name__ == "__main__":
    main()