confidence_member.py

"""
A MEMBER THAT CHANGES ONLY THE CONFIDENCE.

Two results boxed the correction loop in. A member that changes the ANSWER
has a damage channel, and the arithmetic is brutal: right answers outnumber
wrong ones 6.6 to 1, so a point of damage costs 6.6 times what a point of
correction earns, and no gate we built came out ahead. And the base cannot
say where it needs help — its confidence separates what it can do from what
it has never seen at AUC 0.4245, BELOW chance.

Lee's move: manipulate the confidence only. The base is confident that its
ignorance is correct, so change that rather than the answer.

A member applied as a positive SCALE on the logits cannot alter the argmax
— accuracy is invariant by construction, not by measurement — so it has no
damage channel at all. What it can alter is how peaked the distribution is,
which is exactly what abstention reads.

    logits -> logits / exp(s(a))

with s a linear read of the frozen base's features, fitted by the same
closed-form solve as any other member. Large s flattens the distribution
and the model reports doubt; small s sharpens it.

THE QUESTION UNDERNEATH is whether the information exists at all. The
output layer demonstrably cannot express unfamiliarity — it was trained to
name one of seven classes and it always does. But the FEATURES might still
carry it, and a probe would find it if so. If a ridge read of the features
separates seen from unseen where the softmax cannot, then ignorance is
visible one layer down and the failsafe has something to run on. If it
cannot, the representation genuinely does not encode unfamiliarity and no
member can supply it.

Measured three ways:

  DOES IT SEPARATE?   AUC of the new confidence against the base's 0.4245
  IS ACCURACY SAFE?   it must be EXACTLY unchanged; anything else is a bug
  IS IT USEFUL?       risk-coverage — accuracy on the most confident X%,
                      before and after, which is what abstention buys
"""

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]




# fit_on: "classes" fits a specialist on a class subset (measured: an
# oracle gate is worth only +0.0070, because the base was already good
# there). "errors" fits on the base's own mistakes, which is the correction
# loop and has 13% of the test set to work with rather than a 1.8-point
# margin on 40% of it.




CFG = dict(grid=14, c_in=1, chan=16, depth=3, n_train=20000, batch=128,
           lr=1e-3, epochs=30, seed=0,
           base_classes=(0, 1, 2, 3, 4, 5, 6),
           new_classes=(7, 8, 9),
           n_holdout=4000, N=200, lam=1.0, draws=5,
           # THE BASE SEED IS SWEPT TOO: the AUCs below are properties of a
           # trained representation, and a first version measured them on
           # one training run of it.
           base_seeds=(0, 1, 2),
           coverage=(0.2, 0.4, 0.6, 0.8, 1.0))


def auc_of(score, pos):
    """Probability a positive ranks above a negative."""
    r = np.argsort(np.argsort(score))
    n1, n0 = pos.sum(), (~pos).sum()
    if n1 == 0 or n0 == 0:
        return float("nan")
    return float((r[pos].sum() - n1*(n1-1)/2)/(n1*n0))


def conf_of(lg):
    e = np.exp(lg - lg.max(1, keepdims=True))
    return (e/e.sum(1, keepdims=True)).max(1)


def risk_coverage(conf, correct, cover):
    """Accuracy on the most confident fraction — what abstention buys."""
    order = np.argsort(-conf)
    out = []
    for c in cover:
        k = max(1, int(c*len(conf)))
        out.append(float(correct[order[:k]].mean()))
    return out


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("A MEMBER THAT CHANGES ONLY THE CONFIDENCE")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:14s} = {v}")
    print(f"\n  the base is trained only on {CFG['base_classes']}; a positive")
    print(f"  scale on its logits cannot change the argmax, so accuracy is")
    print(f"  invariant BY CONSTRUCTION and only confidence can move")
    print("=" * 78, flush=True)

    X, Y, y, Xte, yte = load(CFG)
    old_m = np.isin(y, CFG["base_classes"])
    tr_i = np.where(old_m)[0]
    # the held-out split must contain BOTH kinds, or a probe has nothing to
    # contrast: familiar examples the base was not trained on, and unseen
    # ones it has never met
    ho_new = np.where(~old_m)[0][:CFG["n_holdout"]//2]
    ho_old = tr_i[-(CFG["n_holdout"]//2):]
    tr_i = tr_i[:-(CFG["n_holdout"]//2)]
    ho_i = np.concatenate([ho_old, ho_new])
    Xtr, Ytr = to_dev(X[tr_i]), to_dev(Y[tr_i])
    Xho = to_dev(X[ho_i]); yho = y[ho_i]
    unseen_ho = np.isin(yho, CFG["new_classes"])
    Xte_d = to_dev(Xte)

    A_auc, A_err, A_acc, A_rc = [], [], [], []
    B_auc_fam, B_rc_fam, B_auc_all = [], [], []
    base_seen_unseen, base_fam_rw, rc0_all, rc0_fam_all = [], [], [], []
    for bs in CFG["base_seeds"]:
        P, fwd, HEAD, OB, width = train_base(Xtr, Ytr, CFG, bs)
        fho, _ = fwd(Xho); fte, _ = fwd(Xte_d)
        base_te = to_host(fte @ P[HEAD] + P[OB])
        base_ho = to_host(fho @ P[HEAD] + P[OB])
        unseen = np.isin(yte, CFG["new_classes"])
        pred = base_te.argmax(1); correct = pred == yte
        c0 = conf_of(base_te)
        fam_ho = ~unseen_ho
        wrong_ho = base_ho.argmax(1) != yho
        fam_te = ~unseen
        base_seen_unseen.append(auc_of(-c0, unseen))
        base_fam_rw.append(auc_of(-c0[fam_te], ~correct[fam_te]))
        rc0_all.append(risk_coverage(c0, correct, CFG["coverage"]))
        rc0_fam_all.append(risk_coverage(c0[fam_te], correct[fam_te],
                                         CFG["coverage"]))
        print(f"  base seed {bs}: {float(correct[~unseen].mean()):.4f} on its "
              f"seven, confidence separates seen from unseen at AUC "
              f"{base_seen_unseen[-1]:.4f}   [{time.time()-t0:.0f}s]",
              flush=True)
        for dr in range(CFG["draws"]):
            rg = np.random.default_rng(4000 + 83*bs + dr)
            sub = rg.choice(len(ho_i), min(CFG["N"], len(ho_i)), replace=False)
            sd_ = to_dev(sub, np.int64) if _GPU else sub
            tgt = to_dev(np.where(unseen_ho[sub], 1.0, -1.0)
                         .astype(np.float32)[:, None])
            W, b = ridge(fho[sd_], tgt, CFG["lam"])
            sc = to_host(fte @ W + b).ravel()
            scaled = base_te/np.exp(np.clip(sc, -4, 4))[:, None]
            c1 = conf_of(scaled)
            A_auc.append(auc_of(-c1, unseen))
            A_err.append(auc_of(-c1, ~correct))
            A_acc.append(float((scaled.argmax(1) == yte).mean())
                         - float(correct.mean()))
            A_rc.append(risk_coverage(c1, correct, CFG["coverage"]))
            fi = np.where(fam_ho)[0]
            rgb = np.random.default_rng(6000 + 83*bs + dr)
            subb = rgb.choice(fi, min(CFG["N"], len(fi)), replace=False)
            sb = to_dev(subb, np.int64) if _GPU else subb
            tb = to_dev(np.where(wrong_ho[subb], 1.0, -1.0)
                        .astype(np.float32)[:, None])
            Wb, bb = ridge(fho[sb], tb, CFG["lam"])
            s2 = to_host(fte @ Wb + bb).ravel()
            c2 = conf_of(base_te/np.exp(np.clip(s2, -4, 4))[:, None])
            B_auc_fam.append(auc_of(-c2[fam_te], ~correct[fam_te]))
            B_auc_all.append(auc_of(-c2, unseen))
            B_rc_fam.append(risk_coverage(c2[fam_te], correct[fam_te],
                                          CFG["coverage"]))
    a0 = float(np.mean(base_seen_unseen))
    fb = float(np.mean(base_fam_rw))
    aucs, aucs_err, accs, rc_new = A_auc, A_err, A_acc, A_rc
    print(f"\n  {len(CFG['base_seeds'])} bases; their confidence separates")
    print(f"  seen from unseen at {min(base_seen_unseen):.4f} to "
          f"{max(base_seen_unseen):.4f}, and right from wrong on FAMILIAR")
    print(f"  input at {min(base_fam_rw):.4f} to {max(base_fam_rw):.4f}")

    rc0 = np.mean(rc0_all, 0)
    rc1 = np.mean(A_rc, 0)
    print("\n" + "=" * 78)
    print("  DOES A MEMBER SUPPLY WHAT THE OUTPUT LAYER CANNOT?")
    print("=" * 78)
    print(f"  {'':>28s} {'AUC':>8s} {'sd':>7s}")
    print(f"  {'base confidence, seen/unseen':>28s} {a0:8.4f} "
          f"{np.std(base_seen_unseen):7.4f}")
    print(f"  {'scaled confidence, seen/unseen':>28s} {np.mean(aucs):8.4f} "
          f"{np.std(aucs):7.4f}")
    print(f"  {'scaled confidence, right/wrong':>28s} "
          f"{np.mean(aucs_err):8.4f} {np.std(aucs_err):7.4f}")
    print(f"\n  BUT THAT right/wrong FIGURE IS INFLATED: 40% of the test set")
    print(f"  is unseen and therefore wrong by construction. The honest")
    print(f"  question is errors on FAMILIAR input, where there is no")
    print(f"  category boundary to give them away:\n")
    fb = auc_of(-c0[fam_te], ~correct[fam_te])
    print(f"  {'':>34s} {'AUC':>8s} {'sd':>7s}")
    print(f"  {'base confidence, familiar right/wrong':>34s} {fb:8.4f}")
    print(f"  {'probe B on familiar right/wrong':>34s} "
          f"{np.mean(B_auc_fam):8.4f} {np.std(B_auc_fam):7.4f}")
    print(f"  {'probe B also spotting the unseen':>34s} "
          f"{np.mean(B_auc_all):8.4f} {np.std(B_auc_all):7.4f}")
    print(f"\n  accuracy change from the scale, across every base and "
          f"draw: {np.abs(accs).max():.3e}")
    print(f"  (a positive scale cannot move an argmax; this is a bug check)")

    print("\n" + "=" * 78)
    print("  WHAT ABSTENTION BUYS")
    print("=" * 78)
    print(f"  {'coverage':>9s} {'base':>9s} {'with member':>12s} "
          f"{'gain':>8s}")
    for i, c in enumerate(CFG["coverage"]):
        print(f"  {c:8.0%} {rc0[i]:9.4f} {rc1[i]:12.4f} "
              f"{rc1[i]-rc0[i]:+8.4f}")

    rcb = np.mean(B_rc_fam, 0)
    rc0f = np.mean(rc0_fam_all, 0)
    print(f"\n  and abstention on FAMILIAR input alone, which is what a")
    print(f"  correction loop would actually run on:")
    print(f"  {'coverage':>9s} {'base':>9s} {'probe B':>12s} {'gain':>8s}")
    for i, c in enumerate(CFG["coverage"]):
        print(f"  {c:8.0%} {rc0f[i]:9.4f} {rcb[i]:12.4f} "
              f"{rcb[i]-rc0f[i]:+8.4f}")

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    fam_gain = np.mean(B_auc_fam) - fb
    print(f"  THERE ARE TWO BLIND SPOTS AND ONLY ONE IS BLIND.\n")
    print(f"  On FAMILIAR input the base already knows when it is unsure —")
    print(f"  AUC {fb:.4f} — and a {CFG['N']}-example probe on the same")
    print(f"  features reaches {np.mean(B_auc_fam):.4f} ({fam_gain:+.4f}). A")
    print(f"  member adds nothing there.")
    print(f"\n  On UNFAMILIAR input it has no idea: {a0:.4f}, below chance,")
    print(f"  and a member supplies {np.mean(aucs):.4f}.")
    print(f"\n  AND THE TWO DETECTORS DO NOT TRANSFER: probe B, fitted on")
    print(f"  familiar errors, spots the unseen at {np.mean(B_auc_all):.4f}.")
    print(f"  Knowing you are unsure and knowing you are out of your depth")
    print(f"  are different quantities living in different places.")
    print(f"\n  WHICH MEANS THE GATE NEVER FAILED FOR WANT OF A SIGNAL. The")
    print(f"  0.5481 that made the base's confidence look useless was an")
    print(f"  aggregate over a test set 40% of which it had never seen. What")
    print(f"  killed the correction loop was the 6.6-to-1 arithmetic and a")
    print(f"  member that could only reach 17% of irreducible errors.")
    print()
    gain = np.mean(aucs) - a0
    sd = max(np.std(aucs), 1e-9)
    if np.mean(aucs) > 0.7:
        print(f"  IGNORANCE IS VISIBLE ONE LAYER DOWN. The output layer")
        print(f"  separates seen from unseen at {a0:.4f} — below chance —")
        print(f"  and a {CFG['N']}-example linear read of the SAME frozen")
        print(f"  features reaches {np.mean(aucs):.4f}. The representation")
        print(f"  encodes unfamiliarity; the head was trained to name one of")
        print(f"  seven classes and never to admit it could not.")
        print(f"\n  And it costs nothing to try: a positive scale cannot")
        print(f"  change the argmax, so the failure mode of a bad confidence")
        print(f"  member is a worse confidence, never a worse answer.")
    elif np.mean(aucs) > 0.6:
        print(f"  PARTLY VISIBLE: {np.mean(aucs):.4f} against the output")
        print(f"  layer's {a0:.4f}. Real signal, not enough to abstain on")
        print(f"  confidently.")
    else:
        print(f"  NOT VISIBLE EITHER. A probe on the features reaches")
        print(f"  {np.mean(aucs):.4f}, so the representation does not encode")
        print(f"  unfamiliarity and no member can supply it. Ignorance is")
        print(f"  invisible from the inside at every level, which is a")
        print(f"  stronger and more useful claim than the output-layer")
        print(f"  result alone.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote confidence_member.json")
    json.dump(dict(base_auc=a0, member_auc=float(np.mean(aucs)),
                   member_auc_sd=float(np.std(aucs)),
                   err_auc=float(np.mean(aucs_err)),
                   rc_base=rc0, rc_member=list(map(float, rc1))),
              open("confidence_member.json", "w"), indent=2)


if __name__ == "__main__":
    main()