Spiking neural networks

Neuron models

All neuron models share one membrane dynamic per step:

h     = decay * v + x
spike = heaviside(h - threshold)
v_new = h - threshold * spike             (soft reset)
v_new = h * (1 - spike) + v_reset * spike (hard reset)

Each neuron module keeps its membrane potential in the buffer v, supports step_mode="s" (single step) and step_mode="m" (spike trains of shape [T, ...]), and exposes reset() to clear the state. LIFNode and its subclasses accept:

• learnable_threshold=True – registers threshold as a trainable Parameter instead of a plain number.

• return_v_seq=True – in step_mode="m", forward returns (spikes, v_seq) where v_seq is the per-step voltage trajectory of shape [T, ...].

ALIFNode adds the parameters base_threshold and decay_th (the leak of the adaptation variable). Its effective threshold is a + base_threshold, with a = decay_th * a + v after each step, so the threshold rises after strong activity and decays back to the base value. Because the effective threshold changes every step, ALIFNode runs the native single-step operator in a tight loop for multi-step inputs.

Surrogate gradients

The spike step is the Heaviside function, whose derivative is zero almost everywhere. Surrogate-gradient training replaces that derivative with a smooth approximation during backward:

The fused LIF kernels support the single-parameter families (ids 0-9); the multi-parameter surrogates (ids 10-14) are available through snn.spike() and snn.surrogate_grad().

detach_reset=True treats the spike in the voltage-update path as a constant during backward, which is the standard stabilizing choice.

Functional primitives

• snn.heaviside(x) – the spike step.

• snn.spike(x, surrogate) – Heaviside forward, surrogate backward.

• snn.surrogate_grad(x, surrogate, alpha, beta, w, c, n, T_period) – evaluate a surrogate gradient (all 15 surrogate ids).

• snn.lif_step(x, v, decay, threshold, v_reset, surrogate, detach_reset) – one fused neuron update, returns (spike, v_new).

• snn.lif_seq(x, v_init, decay, threshold, v_reset, surrogate, detach_reset, return_v_seq) – a fused multi-step update on [T, B, ...] input, returns (spikes, v_final) (or (spikes, v_seq, v_final) with return_v_seq=True). The time loop runs inside one native kernel, so multi-step training (BPTT) does not pay a per-step launch overhead. The returned v_seq is differentiable: its gradient is injected into the per-step voltage back-propagation.

• snn.spike_rate(spikes, dim) – average firing rate along the time axis.

• snn.spike_rate_loss(spikes, target_rate, reduction) – mean-squared-error between the firing rate and a target rate.

Native implementation

The neuron dynamics are implemented as fused CPU/CUDA kernels. The single- step operators (snn_lif_step / snn_lif_step_backward) process every element in one pass (CPU: OpenMP parallel_for; CUDA: grid-stride threads) and support float32, float64, float16 and bfloat16 inputs, with per-neuron or scalar decay / threshold / v_reset broadcasting. The multi-step operators (snn_lif_seq / snn_lif_seq_backward) iterate the membrane dynamics over T inside the same kernel, so the whole spike-train forward/backward is a single kernel launch per device. The backward recomputes the membrane trajectory from the stored spikes (no full forward state is saved) and back-propagates the voltage gradient across time (BPTT), so decay, threshold, v_init and every input frame receive gradients.

v_reset selects the reset mode: None for the soft reset (subtract the threshold) and a scalar or broadcastable tensor for the hard reset (reset to that voltage). detach_reset controls whether the spike inside the voltage-update path is treated as a constant during backward; the default True is the standard stabilizing choice.

A throughput benchmark is provided in benchmark/snn_benchmark.py; it compares the fused multi-step operator with a Python time loop over snn.lif_step and reports spikes per second on both CPU and CUDA.

Spike encoders

• snn.rate_encode(x, T) – Bernoulli spikes over T steps.

• snn.poisson_encode(x, T) – Poisson-process encoding (same draw).

• snn.temporal_encode(x, T, tau) – time-to-first-spike encoding.

• snn.PeriodicEncoder(encoding) – repeats a fixed periodic spike pattern.

• snn.GaussianTuningCurveEncoder(mu, sigma) – population encoding where each input feature drives a bank of Gaussian-tuned Bernoulli neurons.

• snn.LatencyEncoder(T, enc_function) – one-hot time-to-first-spike encoder with linear or logarithmic firing-time map.

• snn.WeightedPhaseEncoder(T) – binary-decomposition weighted-phase encoding.

Deep spiking layers

• snn.SeqToANNContainer(module) – apply any ANN layer to every time step of a [T, B, ...] input.

• snn.SpikingConv2d(...) – Conv2d (applied per time step) followed by a multi-step LIFNode, for [T, B, C, H, W] spike trains.

• snn.SpikingBatchNorm1d / snn.SpikingBatchNorm2d – batch normalization with the time axis folded into the batch dimension, the standard stabilizer for deep SNN training.

Monitoring

snn.SpikeMonitor(module) attaches a forward hook to any nn.Module, records every spike tensor the module returns, and exposes spike_rate(), firing_rate(dim), total_spikes() and last_spike(). Call reset() to clear the records and detach() to remove the hook.

Training example

A complete convolutional spiking network (SpikingConv2d + SpikingBatchNorm2d + max pooling + rate readout) with a full surrogate-gradient BPTT training loop:

import tensorplay as tp
from tensorplay import snn

class ConvSpikingNetwork(tp.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = snn.SpikingConv2d(1, 32, 3, padding=1)
        self.bn1 = snn.SpikingBatchNorm2d(32)
        self.pool1 = snn.SeqToANNContainer(tp.nn.MaxPool2d(2))
        self.conv2 = snn.SpikingConv2d(32, 64, 3, padding=1)
        self.bn2 = snn.SpikingBatchNorm2d(64)
        self.pool2 = snn.SeqToANNContainer(tp.nn.MaxPool2d(2))
        self.fc = snn.SpikeLinear(64 * 7 * 7, 10)
        self.readout = snn.LIFNode(step_mode="m")

    def forward(self, x):          # x: [T, B, 1, 28, 28] Poisson spikes
        x = self.pool1(self.bn1(self.conv1(x)))   # [T, B, 32, 14, 14]
        x = self.pool2(self.bn2(self.conv2(x)))   # [T, B, 64, 7, 7]
        x = x.reshape(x.shape[0], x.shape[1], -1)  # [T, B, 64*7*7]
        return self.readout(self.fc(x)).mean(0)    # [B, 10] firing rates

    def reset(self):
        for m in self.modules():
            if m is not self and hasattr(m, "reset"):
                m.reset()

# Full training loop (MNIST):
model = ConvSpikingNetwork()
opt = tp.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
T = 8
for epoch in range(epochs):
    for x, y in dataloader:          # x: [B, 1, 28, 28] uint8, y: [B] labels
        spikes = snn.poisson_encode(x.to(tp.float32) / 255.0, T=T)
        out = model(spikes)                            # [B, 10]
        target = tp.nn.functional.one_hot(y, 10).to(tp.float32)
        loss = ((out - target) ** 2).mean()
        opt.zero_grad()
        loss.backward()            # surrogate-gradient BPTT
        opt.step()
        model.reset()              # clear membrane potentials per batch