Neuron Model Comparison¶

Train the same three-layer architecture on SHD while swapping the hidden neuron model - plain LIF, integrate-and-fire IF, adaptive ALIF, recurrent RLIF, and a mixed ALIF+RLIF stack - to see how the neuron dynamics affect accuracy. The surrogate gradient (arctan) is held fixed so the neuron model is the only changing factor.

Required extras:

pip install "spyx[loaders]"

Expected result: rendered with a reduced 15-epoch budget (the original comparison used up to 500 epochs) to keep five back-to-back trainings tractable. At 15 epochs the models land in the rough 35-60% test-accuracy band and are still climbing; the adaptive/recurrent variants (ALIF, RLIF, and the mixed stack) generally pull ahead of plain LIF/IF given enough epochs because their extra state helps with SHD's long temporal structure. Bump EPOCHS for a clean separation.

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¶

In [2]:

BATCH = 256
SAMPLE_T = 128
CHANNELS = 256

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,))
N_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 = 256 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,)) N_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)

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

Five SNN factories¶

Each factory builds a Sequential with the same wiring (Linear -> spiking -> Linear -> spiking -> Linear -> LI readout) but swaps in a different hidden-layer neuron. All variants share the arctan surrogate gradient.

In [3]:

HIDDEN = 128
ACT = spyx.axn.arctan()


def _wrap(layer_factory, *, rngs):
    return snn.Sequential(
        nnx.Linear(CHANNELS, HIDDEN, use_bias=False, rngs=rngs),
        layer_factory(0, rngs=rngs),
        nnx.Linear(HIDDEN, HIDDEN, use_bias=False, rngs=rngs),
        layer_factory(1, rngs=rngs),
        nnx.Linear(HIDDEN, N_CLASSES, use_bias=False, rngs=rngs),
        snn.LI((N_CLASSES,), rngs=rngs),
    )


def make_lif(seed=0):
    rngs = nnx.Rngs(seed)
    return _wrap(lambda _i, *, rngs: snn.LIF((HIDDEN,), activation=ACT, rngs=rngs), rngs=rngs)


def make_if(seed=0):
    rngs = nnx.Rngs(seed)
    return _wrap(lambda _i, *, rngs: snn.IF((HIDDEN,), activation=ACT), rngs=rngs)


def make_alif(seed=0):
    rngs = nnx.Rngs(seed)
    return _wrap(lambda _i, *, rngs: snn.ALIF((HIDDEN,), activation=ACT, rngs=rngs), rngs=rngs)


def make_rlif(seed=0):
    rngs = nnx.Rngs(seed)
    return _wrap(lambda _i, *, rngs: snn.RLIF((HIDDEN,), activation=ACT, rngs=rngs), rngs=rngs)


def make_mixed(seed=0):
    rngs = nnx.Rngs(seed)
    factories = [
        lambda *, rngs: snn.ALIF((HIDDEN,), activation=ACT, rngs=rngs),
        lambda *, rngs: snn.RLIF((HIDDEN,), activation=ACT, rngs=rngs),
    ]
    return _wrap(lambda i, *, rngs: factories[i](rngs=rngs), rngs=rngs)

HIDDEN = 128 ACT = spyx.axn.arctan() def _wrap(layer_factory, *, rngs): return snn.Sequential( nnx.Linear(CHANNELS, HIDDEN, use_bias=False, rngs=rngs), layer_factory(0, rngs=rngs), nnx.Linear(HIDDEN, HIDDEN, use_bias=False, rngs=rngs), layer_factory(1, rngs=rngs), nnx.Linear(HIDDEN, N_CLASSES, use_bias=False, rngs=rngs), snn.LI((N_CLASSES,), rngs=rngs), ) def make_lif(seed=0): rngs = nnx.Rngs(seed) return _wrap(lambda _i, *, rngs: snn.LIF((HIDDEN,), activation=ACT, rngs=rngs), rngs=rngs) def make_if(seed=0): rngs = nnx.Rngs(seed) return _wrap(lambda _i, *, rngs: snn.IF((HIDDEN,), activation=ACT), rngs=rngs) def make_alif(seed=0): rngs = nnx.Rngs(seed) return _wrap(lambda _i, *, rngs: snn.ALIF((HIDDEN,), activation=ACT, rngs=rngs), rngs=rngs) def make_rlif(seed=0): rngs = nnx.Rngs(seed) return _wrap(lambda _i, *, rngs: snn.RLIF((HIDDEN,), activation=ACT, rngs=rngs), rngs=rngs) def make_mixed(seed=0): rngs = nnx.Rngs(seed) factories = [ lambda *, rngs: snn.ALIF((HIDDEN,), activation=ACT, rngs=rngs), lambda *, rngs: snn.RLIF((HIDDEN,), activation=ACT, rngs=rngs), ] return _wrap(lambda i, *, rngs: factories[i](rngs=rngs), rngs=rngs)

Reusable training loop¶

run_experiment takes a model factory and a learning rate, trains for EPOCHS, and returns a (train_loss, val_acc, val_loss) history plus the trained model.

In [4]:

Loss = spyx.fn.integral_crossentropy()
Acc = spyx.fn.integral_accuracy()


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


def _forward(model, x_BTC):
    x_TBC = jnp.transpose(x_BTC, (1, 0, 2))
    traces, _ = snn.run(model, x_TBC)
    return jnp.transpose(traces, (1, 0, 2))


def run_experiment(make_model, name, *, epochs=30, lr=3e-4, seed=0):
    model = make_model(seed)
    optimizer = nnx.Optimizer(model, optax.chain(optax.centralize(), optax.lion(lr)), wrt=nnx.Param)
    rng = jax.random.PRNGKey(seed)
    n_train = train_obs.shape[0]
    n_test = test_obs.shape[0]

    @nnx.jit
    def train_step(model, optimizer, events, targets):
        def loss_fn(m):
            traces = _forward(m, events)
            return Loss(traces, targets)
        loss, grads = nnx.value_and_grad(loss_fn)(model)
        optimizer.update(model, grads)
        return loss

    @nnx.jit
    def eval_step(model, events, targets):
        traces = _forward(model, events)
        acc, _preds = Acc(traces, targets)
        loss = Loss(traces, targets)
        return acc, loss

    history = []
    for _ in trange(epochs, desc=name):
        rng, perm_key = jax.random.split(rng)
        order = jax.random.permutation(perm_key, n_train)
        train_losses = []
        for bi in order:
            rng, k = jax.random.split(rng)
            events = augment(_unpack(train_obs[bi]), k)
            targets = train_labels[bi]
            train_losses.append(train_step(model, optimizer, events, targets))

        accs, losses = [], []
        for bi in range(n_test):
            events = _unpack(test_obs[bi])
            targets = test_labels[bi]
            a, l = eval_step(model, events, targets)
            accs.append(a)
            losses.append(l)
        history.append((
            float(jnp.mean(jnp.stack(train_losses))),
            float(jnp.mean(jnp.stack(accs))),
            float(jnp.mean(jnp.stack(losses))),
        ))
    return jnp.array(history), model

Loss = spyx.fn.integral_crossentropy() Acc = spyx.fn.integral_accuracy() def _unpack(batch_obs): obs = jnp.asarray(batch_obs) return jnp.unpackbits(obs, axis=1)[:, :SAMPLE_T, :].astype(jnp.float32) def _forward(model, x_BTC): x_TBC = jnp.transpose(x_BTC, (1, 0, 2)) traces, _ = snn.run(model, x_TBC) return jnp.transpose(traces, (1, 0, 2)) def run_experiment(make_model, name, *, epochs=30, lr=3e-4, seed=0): model = make_model(seed) optimizer = nnx.Optimizer(model, optax.chain(optax.centralize(), optax.lion(lr)), wrt=nnx.Param) rng = jax.random.PRNGKey(seed) n_train = train_obs.shape[0] n_test = test_obs.shape[0] @nnx.jit def train_step(model, optimizer, events, targets): def loss_fn(m): traces = _forward(m, events) return Loss(traces, targets) loss, grads = nnx.value_and_grad(loss_fn)(model) optimizer.update(model, grads) return loss @nnx.jit def eval_step(model, events, targets): traces = _forward(model, events) acc, _preds = Acc(traces, targets) loss = Loss(traces, targets) return acc, loss history = [] for _ in trange(epochs, desc=name): rng, perm_key = jax.random.split(rng) order = jax.random.permutation(perm_key, n_train) train_losses = [] for bi in order: rng, k = jax.random.split(rng) events = augment(_unpack(train_obs[bi]), k) targets = train_labels[bi] train_losses.append(train_step(model, optimizer, events, targets)) accs, losses = [], [] for bi in range(n_test): events = _unpack(test_obs[bi]) targets = test_labels[bi] a, l = eval_step(model, events, targets) accs.append(a) losses.append(l) history.append(( float(jnp.mean(jnp.stack(train_losses))), float(jnp.mean(jnp.stack(accs))), float(jnp.mean(jnp.stack(losses))), )) return jnp.array(history), model

Train all five variants¶

In [5]:

EPOCHS = 15  # rendered with a reduced budget; see the note above.

lif_hist, lif_model = run_experiment(make_lif, "LIF", epochs=EPOCHS, lr=1e-4)
if_hist, if_model = run_experiment(make_if, "IF", epochs=EPOCHS, lr=1e-4)
alif_hist, alif_model = run_experiment(make_alif, "ALIF", epochs=EPOCHS, lr=2e-4)
rlif_hist, rlif_model = run_experiment(make_rlif, "RLIF", epochs=EPOCHS, lr=3e-4)
mixed_hist, mixed_model = run_experiment(make_mixed, "ALIF+RLIF", epochs=EPOCHS, lr=1e-4)

EPOCHS = 15 # rendered with a reduced budget; see the note above. lif_hist, lif_model = run_experiment(make_lif, "LIF", epochs=EPOCHS, lr=1e-4) if_hist, if_model = run_experiment(make_if, "IF", epochs=EPOCHS, lr=1e-4) alif_hist, alif_model = run_experiment(make_alif, "ALIF", epochs=EPOCHS, lr=2e-4) rlif_hist, rlif_model = run_experiment(make_rlif, "RLIF", epochs=EPOCHS, lr=3e-4) mixed_hist, mixed_model = run_experiment(make_mixed, "ALIF+RLIF", epochs=EPOCHS, lr=1e-4)

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

LIF:   7%|▋         | 1/15 [00:04<00:57,  4.09s/it]

LIF:  13%|█▎        | 2/15 [00:06<00:43,  3.35s/it]

LIF:  20%|██        | 3/15 [00:09<00:36,  3.08s/it]

LIF:  27%|██▋       | 4/15 [00:12<00:32,  2.99s/it]

LIF:  33%|███▎      | 5/15 [00:15<00:29,  2.91s/it]

LIF:  40%|████      | 6/15 [00:18<00:25,  2.88s/it]

LIF:  47%|████▋     | 7/15 [00:20<00:22,  2.84s/it]

LIF:  53%|█████▎    | 8/15 [00:23<00:19,  2.83s/it]

LIF:  60%|██████    | 9/15 [00:26<00:17,  2.84s/it]

LIF:  67%|██████▋   | 10/15 [00:29<00:14,  2.86s/it]

LIF:  73%|███████▎  | 11/15 [00:32<00:11,  2.91s/it]

LIF:  80%|████████  | 12/15 [00:35<00:08,  2.89s/it]

LIF:  87%|████████▋ | 13/15 [00:38<00:05,  2.86s/it]

LIF:  93%|█████████▎| 14/15 [00:41<00:02,  2.90s/it]

LIF: 100%|██████████| 15/15 [00:44<00:00,  2.93s/it]

LIF: 100%|██████████| 15/15 [00:44<00:00,  2.94s/it]

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

IF:   7%|▋         | 1/15 [00:02<00:40,  2.91s/it]

IF:  13%|█▎        | 2/15 [00:05<00:33,  2.59s/it]

IF:  20%|██        | 3/15 [00:07<00:29,  2.42s/it]

IF:  27%|██▋       | 4/15 [00:09<00:25,  2.34s/it]

IF:  33%|███▎      | 5/15 [00:11<00:23,  2.31s/it]

IF:  40%|████      | 6/15 [00:14<00:20,  2.30s/it]

IF:  47%|████▋     | 7/15 [00:16<00:18,  2.30s/it]

IF:  53%|█████▎    | 8/15 [00:18<00:16,  2.30s/it]

IF:  60%|██████    | 9/15 [00:21<00:13,  2.30s/it]

IF:  67%|██████▋   | 10/15 [00:23<00:11,  2.30s/it]

IF:  73%|███████▎  | 11/15 [00:25<00:09,  2.30s/it]

IF:  80%|████████  | 12/15 [00:28<00:06,  2.31s/it]

IF:  87%|████████▋ | 13/15 [00:30<00:04,  2.33s/it]

IF:  93%|█████████▎| 14/15 [00:32<00:02,  2.35s/it]

IF: 100%|██████████| 15/15 [00:35<00:00,  2.35s/it]

IF: 100%|██████████| 15/15 [00:35<00:00,  2.35s/it]

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

ALIF:   7%|▋         | 1/15 [00:04<00:56,  4.03s/it]

ALIF:  13%|█▎        | 2/15 [00:06<00:43,  3.37s/it]

ALIF:  20%|██        | 3/15 [00:10<00:41,  3.42s/it]

ALIF:  27%|██▋       | 4/15 [00:13<00:38,  3.46s/it]

ALIF:  33%|███▎      | 5/15 [00:17<00:35,  3.51s/it]

ALIF:  40%|████      | 6/15 [00:21<00:31,  3.52s/it]

ALIF:  47%|████▋     | 7/15 [00:24<00:28,  3.60s/it]

ALIF:  53%|█████▎    | 8/15 [00:28<00:24,  3.50s/it]

ALIF:  60%|██████    | 9/15 [00:31<00:20,  3.50s/it]

ALIF:  67%|██████▋   | 10/15 [00:35<00:17,  3.49s/it]

ALIF:  73%|███████▎  | 11/15 [00:38<00:13,  3.39s/it]

ALIF:  80%|████████  | 12/15 [00:41<00:10,  3.41s/it]

ALIF:  87%|████████▋ | 13/15 [00:45<00:06,  3.44s/it]

ALIF:  93%|█████████▎| 14/15 [00:49<00:03,  3.55s/it]

ALIF: 100%|██████████| 15/15 [00:53<00:00,  3.68s/it]

ALIF: 100%|██████████| 15/15 [00:53<00:00,  3.53s/it]

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

RLIF:   7%|▋         | 1/15 [00:03<00:50,  3.58s/it]

RLIF:  13%|█▎        | 2/15 [00:06<00:43,  3.31s/it]

RLIF:  20%|██        | 3/15 [00:09<00:38,  3.24s/it]

RLIF:  27%|██▋       | 4/15 [00:13<00:35,  3.24s/it]

RLIF:  33%|███▎      | 5/15 [00:16<00:32,  3.21s/it]

RLIF:  40%|████      | 6/15 [00:19<00:28,  3.20s/it]

RLIF:  47%|████▋     | 7/15 [00:22<00:25,  3.24s/it]

RLIF:  53%|█████▎    | 8/15 [00:25<00:22,  3.22s/it]

RLIF:  60%|██████    | 9/15 [00:29<00:19,  3.19s/it]

RLIF:  67%|██████▋   | 10/15 [00:32<00:15,  3.16s/it]

RLIF:  73%|███████▎  | 11/15 [00:35<00:12,  3.18s/it]

RLIF:  80%|████████  | 12/15 [00:38<00:09,  3.17s/it]

RLIF:  87%|████████▋ | 13/15 [00:41<00:06,  3.14s/it]

RLIF:  93%|█████████▎| 14/15 [00:44<00:03,  3.11s/it]

RLIF: 100%|██████████| 15/15 [00:47<00:00,  3.11s/it]

RLIF: 100%|██████████| 15/15 [00:47<00:00,  3.18s/it]

ALIF+RLIF:   0%|          | 0/15 [00:00<?, ?it/s]

ALIF+RLIF:   7%|▋         | 1/15 [00:04<00:58,  4.16s/it]

ALIF+RLIF:  13%|█▎        | 2/15 [00:07<00:45,  3.50s/it]

ALIF+RLIF:  20%|██        | 3/15 [00:10<00:39,  3.25s/it]

ALIF+RLIF:  27%|██▋       | 4/15 [00:13<00:35,  3.23s/it]

ALIF+RLIF:  33%|███▎      | 5/15 [00:16<00:31,  3.17s/it]

ALIF+RLIF:  40%|████      | 6/15 [00:19<00:28,  3.15s/it]

ALIF+RLIF:  47%|████▋     | 7/15 [00:22<00:25,  3.18s/it]

ALIF+RLIF:  53%|█████▎    | 8/15 [00:25<00:22,  3.14s/it]

ALIF+RLIF:  60%|██████    | 9/15 [00:28<00:18,  3.14s/it]

ALIF+RLIF:  67%|██████▋   | 10/15 [00:32<00:15,  3.17s/it]

ALIF+RLIF:  73%|███████▎  | 11/15 [00:35<00:12,  3.14s/it]

ALIF+RLIF:  80%|████████  | 12/15 [00:38<00:09,  3.13s/it]

ALIF+RLIF:  87%|████████▋ | 13/15 [00:41<00:06,  3.17s/it]

ALIF+RLIF:  93%|█████████▎| 14/15 [00:44<00:03,  3.15s/it]

ALIF+RLIF: 100%|██████████| 15/15 [00:47<00:00,  3.14s/it]

ALIF+RLIF: 100%|██████████| 15/15 [00:47<00:00,  3.19s/it]

In [6]:

for name, hist in [
    ("LIF", lif_hist), ("IF", if_hist), ("ALIF", alif_hist),
    ("RLIF", rlif_hist), ("ALIF+RLIF", mixed_hist),
]:
    plt.plot(hist[:, 1], label=f"{name} val acc")
plt.title("Neuron model validation accuracy comparison")
plt.xlabel("epoch")
plt.legend()
plt.show()

for name, hist in [ ("LIF", lif_hist), ("IF", if_hist), ("ALIF", alif_hist), ("RLIF", rlif_hist), ("ALIF+RLIF", mixed_hist), ]: plt.plot(hist[:, 1], label=f"{name} val acc") plt.title("Neuron model validation accuracy comparison") plt.xlabel("epoch") plt.legend() plt.show()

Final test evaluation¶

Pick the strongest model and walk the full test set, then plot a confusion matrix. Swap model below to inspect any variant.

In [7]:

@nnx.jit
def test_step(model, events, targets):
    traces = _forward(model, events)
    acc, preds = Acc(traces, targets)
    loss = Loss(traces, targets)
    return acc, loss, preds


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

preds = jnp.concatenate(all_preds)
tgts = jnp.concatenate(all_tgts)
print("Accuracy:", float(jnp.mean(jnp.stack(all_acc))), "Loss:", float(jnp.mean(jnp.stack(all_loss))))

cm = confusion_matrix(tgts, preds)
ConfusionMatrixDisplay(cm).plot()
plt.title("ALIF model on SHD test set")
plt.show()

@nnx.jit def test_step(model, events, targets): traces = _forward(model, events) acc, preds = Acc(traces, targets) loss = Loss(traces, targets) return acc, loss, preds model = alif_model all_preds, all_tgts, all_acc, all_loss = [], [], [], [] for bi in range(test_obs.shape[0]): events = _unpack(test_obs[bi]) targets = test_labels[bi] a, l, preds = test_step(model, events, targets) all_acc.append(a) all_loss.append(l) all_preds.append(preds) all_tgts.append(targets) preds = jnp.concatenate(all_preds) tgts = jnp.concatenate(all_tgts) print("Accuracy:", float(jnp.mean(jnp.stack(all_acc))), "Loss:", float(jnp.mean(jnp.stack(all_loss)))) cm = confusion_matrix(tgts, preds) ConfusionMatrixDisplay(cm).plot() plt.title("ALIF model on SHD test set") plt.show()

Accuracy: 0.49267578125 Loss: 2.4504354000091553