Spiking readout (off-chip ridge, AltAI-like integer inference) and robustness check

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

readout/robustness.py

@@ -0,0 +1,51 @@
+"""Robustness of the spiking readout: distort raw validation pixels, compute stroke features, run the readout.
+
+usage: python robustness.py <raw.bin> <raw.target> <feat_train.bin> <feat_train.target> --train N [--preset coarse --scale 0.3842]
+"""
+import argparse
+import sys
+
+import numpy as np
+
+sys.path.insert(0, "/mnt/c/knp/features")
+sys.path.insert(0, "/mnt/c/knp/augment")
+import stroke_features as sf  # noqa: E402
+import spiking_readout as sr  # noqa: E402
+from tta_sim import grow, shift, thin  # noqa: E402
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("raw")
+    ap.add_argument("raw_labels")
+    ap.add_argument("feat")
+    ap.add_argument("feat_labels")
+    ap.add_argument("--train", type=int, required=True)
+    ap.add_argument("--eval-first", type=int, default=30000)
+    ap.add_argument("--eval", type=int, default=10000)
+    ap.add_argument("--preset", default="coarse")
+    ap.add_argument("--scale", type=float, default=0.3842)
+    ap.add_argument("--thr", type=float, default=16)
+    ap.add_argument("--lam", type=float, default=10.0)
+    ap.add_argument("--save")
+    a = ap.parse_args()
+    sf.use_preset(a.preset)
+    sf.P["scale"] = a.scale
+    width = 784
+    xf = np.fromfile(a.feat, dtype=np.uint8).reshape(-1, width)[:a.train]
+    yf = np.array([int(v) for v in open(a.feat_labels).read().split()])[:a.train]
+    w = sr.train_ridge(xf, yf, 10, a.lam)
+    wq, scale = sr.quantize(w)
+    raw = np.fromfile(a.raw, dtype=np.uint8).reshape(-1, 28, 28)[a.eval_first:a.eval_first + a.eval].astype(np.int32)
+    y = np.array([int(v) for v in open(a.raw_labels).read().split()])[a.eval_first:a.eval_first + a.eval]
+    row = []
+    for name, v in [("as is", raw), ("shift 2px", shift(raw, 0, 2)), ("thicker 1px", grow(raw)), ("thinner 1px", thin(raw))]:
+        p = sr.spiking_predict(sf.features(v.clip(0, 255).astype(np.uint8)), wq, a.thr)
+        row.append(f"{name} {100 * (p == y).mean():5.1f}%")
+    print(" | ".join(row))
+    if a.save:
+        np.savez(a.save, w=wq, thr=a.thr, scale=scale, preset=a.preset, feat_scale=a.scale)
+
+
+if __name__ == "__main__":
+    main()

readout/spiking_readout.py

@@ -0,0 +1,114 @@
+"""Spiking readout: weights trained off-chip (ridge regression), inference by AltAI-like integer neurons.
+
+Input neurons fire with the same incrementing raster as mnist-learn (value/255 per step, 10 active steps
+of 15). One output neuron per class integrates integer weights of the spikes it receives, fires at
+`thr` and resets to 0; the potential is clipped at 0 like the AltAI negative threshold. A bias input
+fires on every active step. The answer is the class whose neuron fired most (earliest on ties).
+
+usage: python spiking_readout.py <data.bin> <labels> --width 784 --train 30000 --eval-first 30000 --eval 10000
+                                  [--save weights.npz] [--thr 50,100,200]
+"""
+import argparse
+
+import numpy as np
+
+STEPS, ACTIVE = 15, 10
+
+
+def spike_counts(x):
+    """x: [n, width] uint8 → spikes per input over ACTIVE steps (raster from rest)."""
+    st = np.zeros(x.shape, dtype=np.float32)
+    inc = x.astype(np.float32) * np.float32(1 / 255.0)
+    cnt = np.zeros(x.shape, dtype=np.int16)
+    for _ in range(ACTIVE):
+        st += inc
+        f = st >= 1.0
+        cnt += f
+        st[f] -= 1.0
+    return cnt
+
+
+def raster(x):
+    """x: [n, width] uint8 → [ACTIVE, n, width] bool."""
+    st = np.zeros(x.shape, dtype=np.float32)
+    inc = x.astype(np.float32) * np.float32(1 / 255.0)
+    out = np.zeros((ACTIVE,) + x.shape, dtype=bool)
+    for t in range(ACTIVE):
+        st += inc
+        out[t] = st >= 1.0
+        st[out[t]] -= 1.0
+    return out
+
+
+def train_ridge(x, y, classes, lam):
+    f = np.hstack([spike_counts(x).astype(np.float32), np.full((len(x), 1), ACTIVE, np.float32)])
+    t = np.eye(classes, dtype=np.float32)[y] * 2 - 1
+    return np.linalg.solve(f.T @ f + lam * np.eye(f.shape[1], dtype=np.float32), f.T @ t)
+
+
+def quantize(w, wmax=255):
+    scale = wmax / np.abs(w).max()
+    return np.round(w * scale).astype(np.int32), scale
+
+
+def spiking_predict(x, wq, thr, batch=2000):
+    preds = []
+    for i in range(0, len(x), batch):
+        r = raster(x[i:i + batch])
+        n = r.shape[1]
+        pot = np.zeros((n, wq.shape[1]), dtype=np.int64)
+        count = np.zeros_like(pot)
+        first = np.full(pot.shape, STEPS, dtype=np.int64)
+        bias = wq[-1]
+        for t in range(STEPS):
+            if t < ACTIVE:
+                pot += r[t].astype(np.int64) @ wq[:-1] + bias
+            fired = pot >= thr
+            count += fired
+            first = np.where(fired & (first == STEPS), t, first)
+            pot[fired] = 0
+            pot = np.maximum(pot, 0)
+        # most spikes; ties broken by earliest first spike; no spikes → -1
+        score = count * (STEPS + 1) - first
+        p = score.argmax(1)
+        p[count.max(1) == 0] = -1
+        preds.append(p)
+    return np.concatenate(preds)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("data")
+    ap.add_argument("labels")
+    ap.add_argument("--width", type=int, default=784)
+    ap.add_argument("--classes", type=int, default=10)
+    ap.add_argument("--train", type=int, required=True)
+    ap.add_argument("--eval-first", type=int, required=True)
+    ap.add_argument("--eval", type=int, required=True)
+    ap.add_argument("--lam", type=float, default=10.0)
+    ap.add_argument("--thr", default="")
+    ap.add_argument("--save")
+    a = ap.parse_args()
+    x = np.fromfile(a.data, dtype=np.uint8).reshape(-1, a.width)
+    y = np.array([int(v) for v in open(a.labels).read().split()])
+    w = train_ridge(x[:a.train], y[:a.train], a.classes, a.lam)
+    xe, ye = x[a.eval_first:a.eval_first + a.eval], y[a.eval_first:a.eval_first + a.eval]
+    fe = np.hstack([spike_counts(xe).astype(np.float32), np.full((len(xe), 1), ACTIVE, np.float32)])
+    print(f"float linear on spike counts: {100 * ((fe @ w).argmax(1) == ye).mean():.2f}%")
+    wq, scale = quantize(w)
+    ref = np.abs(wq).max()
+    thrs = [float(t) for t in a.thr.split(",")] if a.thr else [ref * k for k in (0.5, 1, 2, 4, 8)]
+    best = None
+    for thr in thrs:
+        p = spiking_predict(xe, wq, thr)
+        acc = 100 * (p == ye).mean()
+        print(f"spiking, thr {thr:g}: correct {acc:.2f}%  silent {100 * (p < 0).mean():.2f}%")
+        if best is None or acc > best[0]:
+            best = (acc, thr)
+    if a.save:
+        np.savez(a.save, w=wq, thr=best[1], scale=scale, width=a.width)
+        print("saved", a.save, "thr", best[1])
+
+
+if __name__ == "__main__":
+    main()