From b8ea9dc6ce68d3a2e9d0aabd726742ccc6c4aada Mon Sep 17 00:00:00 2001 From: int Date: Wed, 14 Oct 2026 20:58:39 +0300 Subject: [PATCH] Spiking readout (off-chip ridge, AltAI-like integer inference) and robustness check Co-Authored-By: Claude Opus 5.5 --- readout/robustness.py | 51 +++++++++++++++++ readout/spiking_readout.py | 114 +++++++++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+) create mode 100644 readout/robustness.py create mode 100644 readout/spiking_readout.py diff --git a/readout/robustness.py b/readout/robustness.py new file mode 100644 index 0000000..65044f9 --- /dev/null +++ b/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 --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() diff --git a/readout/spiking_readout.py b/readout/spiking_readout.py new file mode 100644 index 0000000..5d63ce9 --- /dev/null +++ b/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 --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()