borrow.py

"""
BORROWING A MEMBER FROM A BASE THAT DOES NOT KNOW YOU.

The account architecture in v5 §4 works and cannot yet be shown to
maintain itself, and this is the first of the three reasons. A member is
fitted to ONE base. Whether it is worth anything to a DIFFERENT base
decides whether members are portable objects or private ones, and with them
whether anything in the possibilities document beyond a single organisation
is reachable.

The scenario is deliberately unhelpful. Two people train two bases from
their own data and their own initialisations. One of them fits a member.
The other wants it. NEITHER KNEW THE OTHER EXISTED at training time, so the
member carries no knowledge of the borrowing base and the borrower pays for
whatever adaptation is needed.

One measurement already bounds the naive case: independently trained models
align at 0.002, so a member fitted on one is arithmetically noise against
another. What is untested is whether SHARING A PARTITION rescues it. A
partition is a coordinate agreement, and two bases built on the same one
agree about what position j MEANS even if they disagree about what they
learned. That is a weaker kind of agreement than sharing weights and it may
or may not be enough.

FIVE ARMS, at each of two partition conditions.

  AT HOME          A's member on A. What it is worth to its owner, and the
                   number every other arm is measured against.
  BORROWED RAW     A's member on B, applied directly. The naive case, and
                   the one the 0.002 alignment predicts will fail.
  BORROWED VIA A   a small map from B's features into A's, fitted once per
  SHIM             base PAIR on a calibration set both parties can run.
                   Cost is quadratic in bases and constant in members,
                   which is workable for a handful of providers and not for
                   an open network.
  BORROWED BY      B fits ITS OWN member to reproduce the EFFECT of A's on
  DISTILLATION     the same public inputs. This needs no feature
                   correspondence and no shared partition at all: only that
                   both parties can run the same data. Cost is per MEMBER
                   rather than per pair, so it is linear in members and
                   constant in bases, which is the opposite trade and the
                   one that scales to a network.
  B'S OWN MEMBER   fitted natively on B. The ceiling for B, and what
                   borrowing has to justify itself against.
  B ALONE          the floor.

and the partition conditions are SHARED, where both bases were built on
the same tying pattern, against INDEPENDENT, where each drew its own. If
sharing a partition is what makes borrowing possible, the two conditions
separate and the partition is worth standardising. If they do not, a member
is a private object and the ecosystem readings of this framework are
finished rather than pending.

The calibration set is swept, because its size is the price of entry. A
translation needing fifty examples is a formality; one needing fifty
thousand is a second training run.
"""

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, tag=0):
    """A 3x3 convolution as a tying pattern. tag selects WHICH pattern: the
    same tag is the same partition, which is the coordinate agreement two
    bases either share or do not."""
    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)
    if tag:
        # A GENUINELY DIFFERENT PARTITION, at matched storage. A first
        # version merely PERMUTED which value sat in which slot, which is
        # not a different partition at all: channel labels are
        # interchangeable by construction, and two seeds were already
        # measured to differ by a permutation of 14.7 of 16 channels. That
        # control controlled nothing, and finding no difference between the
        # conditions said nothing. This ties the same number of values
        # ARBITRARILY, so the two bases agree about how much they store and
        # about nothing else.
        rg = np.random.default_rng(9000 + tag)
        idx = rg.integers(0, K, idx.size).reshape(idx.shape)
    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 ridge(A, R, lam):
    """Least squares with the ridge penalty CHOSEN rather than fixed.

    A fixed lambda means the effective regularisation falls as data
    arrives, because the data term grows with n and the penalty does not.
    That is fine while n is far from the feature count and catastrophic
    near it: a fit on 4,000 examples of 3,136 features sits at n/d = 1.28,
    the interpolation peak, where the variance of a weakly regularised
    solution explodes. Measured, a member fitted on 4,000 examples scored
    WORSE than one fitted on 1,000, which is a property of the solver and
    not of members.

    So lambda is selected on a held-out fifth, from a geometric ladder,
    per fit. It costs a handful of solves and removes an artefact that
    reached the paper."""
    n, d = A.shape
    if n < 8:
        return _ridge_at(A, R, lam)
    k = max(2, n//5)
    tr, va = slice(k, None), slice(0, k)
    best, bw = None, None
    for f in (0.01, 0.1, 1.0, 10.0, 100.0, 1000.0, 1e4):
        W, b = _ridge_at(A[tr], R[tr], lam*f)
        err = float(to_host(((A[va] @ W + b - R[va])**2).mean()))
        if best is None or err < best:
            best, bw = err, lam*f
    return _ridge_at(A, R, bw)


def _ridge_at(A, R, lam):
    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:
        G = A1 @ A1.T + lam*xp.eye(n, dtype=DT)
        W = A1.T @ xp.linalg.solve(G, R)
    return W[:-1], W[-1]


def train_base(Xtr, Ytr, cfg, seed, tag):
    """One person's base: their own data, their own initialisation, and a
    partition that is either shared with the other party or not."""
    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, tag)
        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


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))
    return ((X - X[p[:20000]].mean())/(X[p[:20000]].std()+1e-8)
            ).reshape(len(X), -1), y, p


CFG = dict(grid=14, c_in=1, chan=16, depth=3, batch=128, lr=1e-3, epochs=25,
           base_classes=(0, 1, 2, 3, 4, 5, 6), new_classes=(7, 8, 9),
           n_each=9000, n_member=400, lam=1.0,
           calib=(50, 200, 1000, 5000), pairs=(0, 1, 2))


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("BORROWING A MEMBER FROM A BASE THAT DOES NOT KNOW YOU")
    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  two people, two bases, their own data and their own")
    print(f"  initialisations. One fits a member for classes "
          f"{CFG['new_classes']};")
    print(f"  the other wants it, and pays for whatever adaptation is needed.")
    print("=" * 78, flush=True)

    X, y, perm = load(CFG)
    old = np.isin(y, CFG["base_classes"])
    oi = perm[np.isin(perm, np.where(old)[0])]
    ni = perm[np.isin(perm, np.where(~old)[0])]
    n = CFG["n_each"]
    # disjoint data: neither party sees the other's training set
    Adata, Bdata = oi[:n], oi[n:2*n]
    memb_pool = ni[:CFG["n_member"]*4]
    calib_pool = oi[2*n:2*n+6000]        # public, both parties can run it
    test = np.concatenate([oi[2*n+6000:2*n+9000], ni[CFG["n_member"]*4:][:2000]])
    yte = y[test]
    new_te = np.isin(yte, CFG["new_classes"])
    oh = lambda i: to_dev(np.eye(10, dtype=np.float32)[y[i]])
    Xd = lambda i: to_dev(X[i])
    print(f"\n  A trains on {len(Adata):,}, B on {len(Bdata):,}, disjoint.")
    print(f"  calibration pool {len(calib_pool):,} (public), test "
          f"{len(test):,} of which {new_te.mean():.0%} is the new classes",
          flush=True)

    res = {}
    for cond, tagB in (("shared partition", 0), ("different partition", 1)):
        for pair in CFG["pairs"]:
            A_P, A_f, A_H, A_O = train_base(Xd(Adata), oh(Adata), CFG,
                                            100+pair, 0)
            B_P, B_f, B_H, B_O = train_base(Xd(Bdata), oh(Bdata), CFG,
                                            500+pair, tagB)
            fA_te, _ = A_f(Xd(test)); fB_te, _ = B_f(Xd(test))
            fA_ca, _ = A_f(Xd(calib_pool)); fB_ca, _ = B_f(Xd(calib_pool))
            fA_mb, _ = A_f(Xd(memb_pool)); fB_mb, _ = B_f(Xd(memb_pool))
            # THE MEMBER HAS ITS OWN OUTPUT. A first version applied it
            # into the base's shared ten-way head at a fixed weight, which
            # violates the design rule v5 §4.1 establishes on measurement:
            # in a shared space raising three classes lowers seven. It
            # crushed the ceiling before borrowing was tested, leaving a
            # native member worth 0.0283. Here the member is a THREE-WAY
            # head over the new classes alone, scored on examples that are
            # one of those three, so nothing competes and the member is the
            # whole of the capability being borrowed.
            nm = np.array(CFG["new_classes"])
            oh3 = lambda i: to_dev(
                (y[i][:, None] == nm[None, :]).astype(np.float32))
            te3 = np.where(new_te)[0]
            y3 = np.searchsorted(nm, yte[te3])
            sub = np.random.default_rng(pair).choice(
                len(memb_pool), CFG["n_member"], replace=False)
            sd_ = to_dev(sub, np.int64) if _GPU else sub
            WA, bA = ridge(fA_mb[sd_], oh3(memb_pool[sub]), CFG["lam"])
            WB, bB = ridge(fB_mb[sd_], oh3(memb_pool[sub]), CFG["lam"])
            sc = lambda F, W, b: float(
                (to_host(F[te3] @ W + b).argmax(1) == y3).mean())
            r = dict(chance=1.0/len(nm),
                     A_home=sc(fA_te, WA, bA),
                     B_own=sc(fB_te, WB, bB),
                     B_raw=sc(fB_te, WA, bA))
            effect_ca = fA_ca @ WA + bA
            for nc in CFG["calib"]:
                cs = to_dev(np.arange(min(nc, len(calib_pool))), np.int64) \
                    if _GPU else np.arange(min(nc, len(calib_pool)))
                T, tb = ridge(fB_ca[cs], fA_ca[cs], CFG["lam"])
                r[f"B_trans_{nc}"] = float(
                    (to_host((fB_te[te3] @ T + tb) @ WA + bA).argmax(1)
                     == y3).mean())
                WD, bD = ridge(fB_ca[cs], effect_ca[cs], CFG["lam"])
                r[f"B_distil_{nc}"] = sc(fB_te, WD, bD)
            for k, v in r.items():
                res.setdefault((cond, k), []).append(v)
            print(f"  {cond:18s} pair {pair}: A at home {r['A_home']:.4f}, "
                  f"B's own {r['B_own']:.4f}, borrowed raw {r['B_raw']:.4f}"
                  f"   [{time.time()-t0:.0f}s]", flush=True)
            _FIXED.clear()
            if _GPU:
                _cp.get_default_memory_pool().free_all_blocks()

    m = {k: (float(np.mean(v)), float(np.std(v))) for k, v in res.items()}
    json.dump({f"{a}/{b}": list(v) for (a, b), v in m.items()},
              open("borrow.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  WHAT IS A BORROWED MEMBER WORTH?")
    print("=" * 78)
    rows = (["chance", "B_raw"]
            + [f"B_trans_{n}" for n in CFG["calib"]]
            + [f"B_distil_{n}" for n in CFG["calib"]]
            + ["B_own", "A_home"])
    lab = {"chance": "chance (3 classes)", "B_raw": "A's member, raw",
           "B_own": "B's own member (ceiling)", "A_home": "A's member at home"}
    for n in CFG["calib"]:
        lab[f"B_trans_{n}"] = f"shim, {n:,} calib"
        lab[f"B_distil_{n}"] = f"distilled, {n:,} calib"
    print(f"  {'':>28s} {'shared partition':>18s} {'different':>18s}")
    for k in rows:
        a = m[("shared partition", k)]; b = m[("different partition", k)]
        print(f"  {lab[k]:>28s} {a[0]:11.4f} +-{a[1]:.4f} "
              f"{b[0]:11.4f} +-{b[1]:.4f}")

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    sh = lambda k: m[("shared partition", k)][0]
    ind = lambda k: m[("different partition", k)][0]
    sd = max(v[1] for v in m.values())
    floor, ceil = sh("chance"), sh("B_own")
    raw = sh("B_raw")
    best_t = max(CFG["calib"], key=lambda n: sh(f"B_trans_{n}"))
    tval = sh(f"B_trans_{best_t}")
    best_d = max(CFG["calib"], key=lambda n: sh(f"B_distil_{n}"))
    dval = sh(f"B_distil_{best_d}")
    print(f"  chance is {floor:.4f}; B's own member {ceil:.4f}; so borrowing")
    print(f"  has {ceil-floor:+.4f} to recover.\n")
    if raw > floor + 2*sd:
        print(f"  A MEMBER TRANSFERS RAW, at {raw:.4f} ({raw-floor:+.4f}),")
        print(f"  which the 0.002 alignment figure says should not happen and")
        print(f"  is the more interesting outcome.")
    else:
        print(f"  RAW BORROWING FAILS, at {raw:.4f} against a floor of")
        print(f"  {floor:.4f}. Two bases that never met do not share a")
        print(f"  coordinate system in any usable sense, which is what the")
        print(f"  0.002 alignment predicted.")
    print()
    if tval > floor + 2*sd:
        print(f"  A TRANSLATION RESCUES IT: {tval:.4f} at {best_t:,}")
        print(f"  calibration examples, recovering "
              f"{(tval-floor)/max(ceil-floor, 1e-9):.0%} of what a native")
        print(f"  member is worth. The shim is fitted once per base PAIR and")
        print(f"  serves every member between them, so its cost is quadratic")
        print(f"  in bases and constant in members.")
        need = next((n for n in CFG["calib"]
                     if sh(f"B_trans_{n}") > floor + 2*sd), None)
        if need:
            print(f"  It first works at {need:,} examples, which is the price")
            print(f"  of entry to someone else's ecosystem.")
    else:
        print(f"  A TRANSLATION DOES NOT RESCUE IT either: best "
              f"{tval:.4f} at")
        print(f"  {best_t:,} examples. A member is a PRIVATE object, tied to")
        print(f"  the base it was fitted on, and the ecosystem readings of")
        print(f"  this framework are finished rather than pending.")
    print()
    print(f"  AND THE OTHER ROUTE, distilling what A's member DOES rather")
    print(f"  than mapping A's features: {dval:.4f} at {best_d:,} examples,")
    print(f"  against the shim's {tval:.4f}.")
    gap = abs(dval - tval)
    if gap < 0.005:
        print(f"  THEY AGREE TO {gap:.4f}, AND THAT IS AN IDENTITY RATHER")
        print(f"  THAN A RESULT. Ridge is linear in its target, so fitting")
        print(f"  B's features to A's and then applying A's member is the")
        print(f"  same computation as fitting B's features directly to what")
        print(f"  A's member outputs. Two names, one operation.")
        print(f"\n  Which settles the cost question by default: the")
        print(f"  distillation form is fitted per MEMBER rather than per base")
        print(f"  PAIR, so it is linear in members and constant in bases")
        print(f"  where the shim is the reverse. Same arithmetic, and the")
        print(f"  cheaper bookkeeping for anything wider than a handful of")
        print(f"  providers.")
    elif dval > tval:
        print(f"  DISTILLATION LEADS BY {dval-tval:+.4f}, which should not")
        print(f"  happen for a linear member and is worth understanding")
        print(f"  before it is used: the two are the same computation up to")
        print(f"  where the regularisation is applied.")
    else:
        print(f"  THE SHIM LEADS BY {tval-dval:+.4f}, which for a linear")
        print(f"  member means the regularisation is landing differently in")
        print(f"  the two forms rather than that the routes differ.")
    print()
    d = sh("B_raw") - ind("B_raw")
    dt = tval - ind(f"B_trans_{best_t}")
    print(f"  DOES SHARING A PARTITION HELP?")
    print(f"    raw:        {d:+.4f}")
    print(f"    translated: {dt:+.4f}")
    if max(d, dt) > 2*sd:
        print(f"  YES. A partition is an agreement about what position j")
        print(f"  MEANS, and it is cheap to standardise: an index, disclosing")
        print(f"  no training data and no weights. That makes it the")
        print(f"  coordination point the possibilities document proposed.")
    else:
        print(f"  NO. Two bases on the same partition transfer no better than")
        print(f"  two on different ones, so a shared coordinate system is not")
        print(f"  sufficient: they must also agree about what they LEARNED,")
        print(f"  and nothing in a published index delivers that.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote borrow.json")


if __name__ == "__main__":
    main()