Surrogate-Gradient Template (with regularization)¶

A reusable template for training a Spyx SNN on the Spiking Heidelberg Digits (SHD) dataset using BPTT, surrogate gradients, and per-layer activity regularization. Use this as the starting point for the scaling experiments in research/scaling_experiments/.

Required extras:

pip install "spyx[loaders]"

In [1]:

import os

os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", ".80")

import jax
import jax.numpy as jnp
import optax
from flax import nnx
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
from tqdm import trange
import matplotlib.pyplot as plt

import spyx
import spyx.nn as snn

import os os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", ".80") import jax import jax.numpy as jnp import optax from flax import nnx from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix from tqdm import trange import matplotlib.pyplot as plt import spyx import spyx.nn as snn

/home/kade/Code/spyx/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Data loading and augmentation¶

spyx.data.SHD_loader wraps Tonic's SHD dataset behind a Google Grain pipeline. Spikes are bit-packed along the time axis; we recover dense spikes inline with jnp.unpackbits. spyx.data.shift_augment returns a JIT-compiled function that randomly rolls the input along its channel axis - a cheap form of data augmentation.

In [2]:

BATCH = 256
SAMPLE_T = 128
CHANNELS = 128

shd_dl = spyx.data.SHD_loader(batch_size=BATCH, sample_T=SAMPLE_T, channels=CHANNELS)
augment = spyx.data.shift_augment(max_shift=16, axes=(2,))
print("observation shape:", shd_dl.obs_shape, "num classes:", int(shd_dl.act_shape[0]))

# Prestage the whole dataset onto the accelerator once, then scan over the
# prestaged batches each epoch. This is dramatically faster than re-decoding
# SHD through the streaming grain pipeline every epoch, and the whole dataset
# fits comfortably in memory. `prestage` returns per-batch arrays shaped
# (num_batches, batch, packed_time, channels).
train_obs, train_labels = shd_dl.prestage("train")
test_obs, test_labels = shd_dl.prestage("test")
print("prestaged train:", train_obs.shape, " test:", test_obs.shape)

BATCH = 256 SAMPLE_T = 128 CHANNELS = 128 shd_dl = spyx.data.SHD_loader(batch_size=BATCH, sample_T=SAMPLE_T, channels=CHANNELS) augment = spyx.data.shift_augment(max_shift=16, axes=(2,)) print("observation shape:", shd_dl.obs_shape, "num classes:", int(shd_dl.act_shape[0])) # Prestage the whole dataset onto the accelerator once, then scan over the # prestaged batches each epoch. This is dramatically faster than re-decoding # SHD through the streaming grain pipeline every epoch, and the whole dataset # fits comfortably in memory. `prestage` returns per-batch arrays shaped # (num_batches, batch, packed_time, channels). train_obs, train_labels = shd_dl.prestage("train") test_obs, test_labels = shd_dl.prestage("test") print("prestaged train:", train_obs.shape, " test:", test_obs.shape)

observation shape: (128,) num classes: 20

prestaged train: (31, 256, 16, 128)  test: (8, 256, 16, 128)

SNN with intermediate spike taps¶

The model's __call__ returns both readout traces and the spike trains of every hidden layer. We feed those spike trains directly to silence_reg / sparsity_reg in the loss - this avoids the need to track mutable activity counters inside JIT.

In [3]:

class SHDSNN(nnx.Module):
    def __init__(self, in_dim, hidden, n_classes, *, rngs):
        self.l1 = nnx.Linear(in_dim, hidden, use_bias=False, rngs=rngs)
        self.lif1 = snn.LIF((hidden,), activation=spyx.axn.triangular(), rngs=rngs)
        self.l2 = nnx.Linear(hidden, hidden, use_bias=False, rngs=rngs)
        self.lif2 = snn.LIF((hidden,), activation=spyx.axn.triangular(), rngs=rngs)
        self.l3 = nnx.Linear(hidden, n_classes, use_bias=False, rngs=rngs)
        self.li = snn.LI((n_classes,), rngs=rngs)

    def __call__(self, x_BTC):
        x_TBC = jnp.transpose(x_BTC, (1, 0, 2))
        T, B, _ = x_TBC.shape

        s1 = self.lif1.initial_state(B)
        s2 = self.lif2.initial_state(B)
        s3 = self.li.initial_state(B)

        def step(carry, x_t):
            s1, s2, s3 = carry
            sp1, s1 = self.lif1(self.l1(x_t), s1)
            sp2, s2 = self.lif2(self.l2(sp1), s2)
            v, s3 = self.li(self.l3(sp2), s3)
            return (s1, s2, s3), (sp1, sp2, v)

        _, (sp1_TB, sp2_TB, v_TB) = jax.lax.scan(step, (s1, s2, s3), x_TBC)
        traces = jnp.transpose(v_TB, (1, 0, 2))
        spikes_l1 = jnp.transpose(sp1_TB, (1, 0, 2))
        spikes_l2 = jnp.transpose(sp2_TB, (1, 0, 2))
        return traces, [spikes_l1, spikes_l2]


model = SHDSNN(
    in_dim=CHANNELS,
    hidden=64,
    n_classes=int(shd_dl.act_shape[0]),
    rngs=nnx.Rngs(0),
)

class SHDSNN(nnx.Module): def __init__(self, in_dim, hidden, n_classes, *, rngs): self.l1 = nnx.Linear(in_dim, hidden, use_bias=False, rngs=rngs) self.lif1 = snn.LIF((hidden,), activation=spyx.axn.triangular(), rngs=rngs) self.l2 = nnx.Linear(hidden, hidden, use_bias=False, rngs=rngs) self.lif2 = snn.LIF((hidden,), activation=spyx.axn.triangular(), rngs=rngs) self.l3 = nnx.Linear(hidden, n_classes, use_bias=False, rngs=rngs) self.li = snn.LI((n_classes,), rngs=rngs) def __call__(self, x_BTC): x_TBC = jnp.transpose(x_BTC, (1, 0, 2)) T, B, _ = x_TBC.shape s1 = self.lif1.initial_state(B) s2 = self.lif2.initial_state(B) s3 = self.li.initial_state(B) def step(carry, x_t): s1, s2, s3 = carry sp1, s1 = self.lif1(self.l1(x_t), s1) sp2, s2 = self.lif2(self.l2(sp1), s2) v, s3 = self.li(self.l3(sp2), s3) return (s1, s2, s3), (sp1, sp2, v) _, (sp1_TB, sp2_TB, v_TB) = jax.lax.scan(step, (s1, s2, s3), x_TBC) traces = jnp.transpose(v_TB, (1, 0, 2)) spikes_l1 = jnp.transpose(sp1_TB, (1, 0, 2)) spikes_l2 = jnp.transpose(sp2_TB, (1, 0, 2)) return traces, [spikes_l1, spikes_l2] model = SHDSNN( in_dim=CHANNELS, hidden=64, n_classes=int(shd_dl.act_shape[0]), rngs=nnx.Rngs(0), )

Training loop with combined loss¶

We combine integral_crossentropy on the readout traces with two activity regularizers - silence_reg to encourage every neuron to fire at least a little, and sparsity_reg to discourage layers from firing too much. The regularizers are scaled by a small REG_WEIGHT: silence_reg is on the order of 1e5 at initialisation, so a weight around 1e-6 turns it into a gentle nudge rather than a term that swamps the classification loss.

Expected result: with the default EPOCHS = 30 this 64-unit regularized network reaches roughly 40-45% test accuracy on SHD (20 classes; chance is 5%), while the regularizers hold the hidden layers to a modest per-neuron firing rate. This is a reduced budget so the template stays quick to render; SHD keeps improving toward ~80% with 50-100 epochs, a wider hidden layer, and a tuned REG_WEIGHT.

In [4]:

Loss = spyx.fn.integral_crossentropy()
Acc = spyx.fn.integral_accuracy()
Silence = spyx.fn.silence_reg(min_spikes=4.0)
Sparsity = spyx.fn.sparsity_reg(max_spikes=16.0)
REG_WEIGHT = 1e-6  # silence_reg is O(1e5) at init; keep the nudge gentle

optimizer = nnx.Optimizer(model, optax.chain(
    optax.centralize(),
    optax.lion(learning_rate=3e-4),
), wrt=nnx.Param)


def _unpack(batch_obs):
    obs = jnp.asarray(batch_obs)
    return jnp.unpackbits(obs, axis=1)[:, :SAMPLE_T, :].astype(jnp.float32)


@nnx.jit
def train_step(model, optimizer, events, targets):
    def loss_fn(m):
        traces, spikes = m(events)
        return Loss(traces, targets) + REG_WEIGHT * (Silence(spikes) + Sparsity(spikes))

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss


@nnx.jit
def eval_step(model, events, targets):
    traces, spikes = model(events)
    acc, preds = Acc(traces, targets)
    loss = Loss(traces, targets)
    layer_rates = jnp.stack([s.mean(axis=(0, 1)) for s in spikes])
    return acc, loss, preds, layer_rates

Loss = spyx.fn.integral_crossentropy() Acc = spyx.fn.integral_accuracy() Silence = spyx.fn.silence_reg(min_spikes=4.0) Sparsity = spyx.fn.sparsity_reg(max_spikes=16.0) REG_WEIGHT = 1e-6 # silence_reg is O(1e5) at init; keep the nudge gentle optimizer = nnx.Optimizer(model, optax.chain( optax.centralize(), optax.lion(learning_rate=3e-4), ), wrt=nnx.Param) def _unpack(batch_obs): obs = jnp.asarray(batch_obs) return jnp.unpackbits(obs, axis=1)[:, :SAMPLE_T, :].astype(jnp.float32) @nnx.jit def train_step(model, optimizer, events, targets): def loss_fn(m): traces, spikes = m(events) return Loss(traces, targets) + REG_WEIGHT * (Silence(spikes) + Sparsity(spikes)) loss, grads = nnx.value_and_grad(loss_fn)(model) optimizer.update(model, grads) return loss @nnx.jit def eval_step(model, events, targets): traces, spikes = model(events) acc, preds = Acc(traces, targets) loss = Loss(traces, targets) layer_rates = jnp.stack([s.mean(axis=(0, 1)) for s in spikes]) return acc, loss, preds, layer_rates

In [5]:

EPOCHS = 30  # bump this up for a real run; SHD typically wants 50-100 epochs.
rng = jax.random.PRNGKey(0)
N_TRAIN = train_obs.shape[0]
N_TEST = test_obs.shape[0]
history = []

for epoch in trange(EPOCHS):
    rng, perm_key = jax.random.split(rng)
    order = jax.random.permutation(perm_key, N_TRAIN)
    train_losses = []
    for bi in order:
        rng, aug_key = jax.random.split(rng)
        events = augment(_unpack(train_obs[bi]), aug_key)
        targets = train_labels[bi]
        train_losses.append(train_step(model, optimizer, events, targets))

    val_acc, val_loss = [], []
    for bi in range(N_TEST):
        events = _unpack(test_obs[bi])
        targets = test_labels[bi]
        a, l, _, _ = eval_step(model, events, targets)
        val_acc.append(a)
        val_loss.append(l)

    history.append((
        float(jnp.mean(jnp.stack(train_losses))),
        float(jnp.mean(jnp.stack(val_acc))),
        float(jnp.mean(jnp.stack(val_loss))),
    ))

print("final:", history[-1])

EPOCHS = 30 # bump this up for a real run; SHD typically wants 50-100 epochs. rng = jax.random.PRNGKey(0) N_TRAIN = train_obs.shape[0] N_TEST = test_obs.shape[0] history = [] for epoch in trange(EPOCHS): rng, perm_key = jax.random.split(rng) order = jax.random.permutation(perm_key, N_TRAIN) train_losses = [] for bi in order: rng, aug_key = jax.random.split(rng) events = augment(_unpack(train_obs[bi]), aug_key) targets = train_labels[bi] train_losses.append(train_step(model, optimizer, events, targets)) val_acc, val_loss = [], [] for bi in range(N_TEST): events = _unpack(test_obs[bi]) targets = test_labels[bi] a, l, _, _ = eval_step(model, events, targets) val_acc.append(a) val_loss.append(l) history.append(( float(jnp.mean(jnp.stack(train_losses))), float(jnp.mean(jnp.stack(val_acc))), float(jnp.mean(jnp.stack(val_loss))), )) print("final:", history[-1])

  0%|          | 0/30 [00:00<?, ?it/s]

  3%|▎         | 1/30 [00:02<01:14,  2.59s/it]

  7%|▋         | 2/30 [00:03<00:49,  1.78s/it]

 10%|█         | 3/30 [00:05<00:41,  1.54s/it]

 13%|█▎        | 4/30 [00:06<00:36,  1.41s/it]

 17%|█▋        | 5/30 [00:07<00:33,  1.35s/it]

 20%|██        | 6/30 [00:08<00:31,  1.31s/it]

 23%|██▎       | 7/30 [00:09<00:28,  1.26s/it]

 27%|██▋       | 8/30 [00:11<00:28,  1.28s/it]

 30%|███       | 9/30 [00:12<00:26,  1.26s/it]

 33%|███▎      | 10/30 [00:13<00:25,  1.25s/it]

 37%|███▋      | 11/30 [00:14<00:23,  1.26s/it]

 40%|████      | 12/30 [00:16<00:22,  1.26s/it]

 43%|████▎     | 13/30 [00:17<00:21,  1.27s/it]

 47%|████▋     | 14/30 [00:18<00:20,  1.27s/it]

 50%|█████     | 15/30 [00:20<00:19,  1.27s/it]

 53%|█████▎    | 16/30 [00:21<00:18,  1.29s/it]

 57%|█████▋    | 17/30 [00:22<00:16,  1.28s/it]

 60%|██████    | 18/30 [00:23<00:15,  1.26s/it]

 63%|██████▎   | 19/30 [00:25<00:13,  1.24s/it]

 67%|██████▋   | 20/30 [00:26<00:12,  1.23s/it]

 70%|███████   | 21/30 [00:27<00:11,  1.26s/it]

 73%|███████▎  | 22/30 [00:28<00:10,  1.27s/it]

 77%|███████▋  | 23/30 [00:29<00:08,  1.22s/it]

 80%|████████  | 24/30 [00:31<00:07,  1.19s/it]

 83%|████████▎ | 25/30 [00:32<00:05,  1.18s/it]

 87%|████████▋ | 26/30 [00:33<00:04,  1.16s/it]

 90%|█████████ | 27/30 [00:34<00:03,  1.18s/it]

 93%|█████████▎| 28/30 [00:35<00:02,  1.22s/it]

 97%|█████████▋| 29/30 [00:37<00:01,  1.26s/it]

100%|██████████| 30/30 [00:38<00:00,  1.28s/it]

100%|██████████| 30/30 [00:38<00:00,  1.29s/it]

final: (2.8177902698516846, 0.421875, 2.500636577606201)

In [6]:

plt.plot(history, label=["train loss", "val acc", "val loss"])
plt.title("SHD surrogate gradient + dual regularization")
plt.legend()
plt.show()

plt.plot(history, label=["train loss", "val acc", "val loss"]) plt.title("SHD surrogate gradient + dual regularization") plt.legend() plt.show()

Evaluation and per-neuron firing rates¶

Use the spike traces returned by the SNN to inspect which hidden neurons are doing work.

In [7]:

all_preds, all_tgts, all_acc, all_loss, all_rates = [], [], [], [], []
for bi in range(test_obs.shape[0]):
    events = _unpack(test_obs[bi])
    targets = test_labels[bi]
    a, l, preds, rates = eval_step(model, events, targets)
    all_acc.append(a)
    all_loss.append(l)
    all_preds.append(preds)
    all_tgts.append(targets)
    all_rates.append(rates)

preds = jnp.concatenate(all_preds)
tgts = jnp.concatenate(all_tgts)
mean_rates = jnp.mean(jnp.stack(all_rates), axis=0)
print("Test accuracy:", float(jnp.mean(jnp.stack(all_acc))))
print("Layer-1 mean per-neuron rate:", float(jnp.mean(mean_rates[0])))
print("Layer-2 mean per-neuron rate:", float(jnp.mean(mean_rates[1])))

cm = confusion_matrix(tgts, preds)
ConfusionMatrixDisplay(cm).plot()
plt.show()

all_preds, all_tgts, all_acc, all_loss, all_rates = [], [], [], [], [] for bi in range(test_obs.shape[0]): events = _unpack(test_obs[bi]) targets = test_labels[bi] a, l, preds, rates = eval_step(model, events, targets) all_acc.append(a) all_loss.append(l) all_preds.append(preds) all_tgts.append(targets) all_rates.append(rates) preds = jnp.concatenate(all_preds) tgts = jnp.concatenate(all_tgts) mean_rates = jnp.mean(jnp.stack(all_rates), axis=0) print("Test accuracy:", float(jnp.mean(jnp.stack(all_acc)))) print("Layer-1 mean per-neuron rate:", float(jnp.mean(mean_rates[0]))) print("Layer-2 mean per-neuron rate:", float(jnp.mean(mean_rates[1]))) cm = confusion_matrix(tgts, preds) ConfusionMatrixDisplay(cm).plot() plt.show()

Test accuracy: 0.421875
Layer-1 mean per-neuron rate: 0.15934425592422485
Layer-2 mean per-neuron rate: 0.04123491048812866