[!WARNING] this is a work in progress.
SNNs and Surrogate gradient descent
Spiking neural networks (SNNs) are biologically inspired models that compute via discrete, sparse spikes, rather than continuous activations (non differentiable non linearity). This event driven framework not only captures rich temporal patterns (such as inter spike intervals and cross neuron synchrony) but also powers energy efficient neuromorphic hardware.
Surrogate gradient descent (SuGD) answer the most important challenging question: “how to convert precise spike timing into effective learning signals ?”
This problem is solved by Zenke et al 2018 through replacing the non differentiable spike with a smooth surrogate, thereby allowing backpropagation through time (BPTT).
SuGD can learn to make use of information that is not only encoded in the rate of spikes but also timing. SGD training can extract interspike intervals, spatio-temporal spike patterns or polychrony, and coincidence codes (Ziqiao Yu et al 2026)[1].
Example
Heavy side function
$$\Theta(x) = \begin{cases} 0 & x < 0 \\ 1 & x \geq 0 \end{cases} \tag{1}$$
D(Heavy side function)
$$\frac{d\Theta(x)}{dx} = \begin{cases} 0 & x < 0 \\ \text{undefined} & x \geq 0 \end{cases} \tag{2}$$
Surrogate approximation
using a sigmoid or fast sigmoid function (kind of cheat:)
sigmoid or logistic function
$$\sigma(x) = \frac{1}{1 + e^{-\beta x}} \tag{3}$$
D(sigmoid or logistic function)
$$\frac{d\sigma(x)}{dx} = \beta\sigma(x)(1 - \sigma(x))\tag{4}$$
fast sigmoid function
$$\sigma(x) = \frac{x}{1 + |x|} \tag{5}$$
D(fast sigmoid function)
$$\frac{d\sigma(x)}{dx} = \frac{1}{(1 + |x|)^2} \tag{6}$$
Code
$$\text{heaviside}(x) \approx \sigma(x)$$
class SurrogateHeaviside(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return torch.heaviside(input, 0)
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
beta = 5
s = torch.sigmoid(beta*input)
grad = grad_output*beta*s*(1-s)
return grad
surrogate_heaviside = SurrogateHeaviside.apply
$$\text{heaviside}(x) \approx fast\_sigmoid(x)$$
class SurrogateHeaviside(torch.autograd.Function):
scale = 100.0 # controls steepness of surrogate gradient
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input) ## storing context or input to use during backprop
out = torch.zeros_like(input) ## first step make everything zero (i.e. all negative inputs are zero)
out[input > 0] = 1.0 ## mask the inputs > 0 and make them 1
return out
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors ## fetching previously saved input
grad_input = grad_output.clone() ## safely copying the gradient output so we don't accidentally break PyTorch graph
grad = grad_input/(SurrogateHeaviside.scale*torch.abs(input)+1.0)**2 ## computing fast sigmoid instead of using the e^x which takes longer to run
return grad
surrogate_heaviside = SurrogateHeaviside.apply
credits: Tutoral by Dan F M Goodman, 2026 FENS Chen Summer School on Learning with spikes [2]
Tutorial: SuGD for training SNNs
This section is inspired snnTorch tutorial-5 which is in turn inspired by Friedemann Zenke’s extensive work on SNNs. Check out his repo on surrogate gradients spytorch
Now we will implement a basic supervised learing algorithm with SuGD for training spiking neurons to perform image classfication on Static MNIST.
The Recurrent representation of SNNs
If you think of SNNs like a Machine learning researcher you can tell it is type of RNNs with implict recurrence with nonlinear activation function.
Further Infromation on RNNs are found here
LIF Neuron discrete recursive form
where if the membrane potential exceeds the threshold, a spike is emitted:
$$S[t] = \begin{cases} 1, &\text{if}~U[t] > U_{\rm thr} \\ 0, &\text{otherwise}\end{cases}{\tag{8}}$$
Like in RNNs, we have dependence from previous state (here state variable $U[t]$), Also its not even a property of network just the neuron itself.This is illustrated using an implicit recurrent connection for the decay of the membrane potential.
Vanilla-RNN
hidden state: $$h_t = f\left(W_{ih} * x_t + W_{hh} * h_{t-1} + b_h)\right)$$
output: $$ y_t = W_{ho} * h_t + b_o $$
This is almost perfectly poised to take advantage of the developments in training recurrent neural networks (RNNs) and sequence-based models.
there is the special aspect of recurrence called explict recurrence in SNNs where the output spike $S_{\rm out}$ is fed back to the input, i.e. neurons in same layer are connected to each other or even to itself simlar to RLeaky in snnTorch and these connection ($W_{recurr}$ where diagonals represent self connections and off diagonal elements represent recurrent connection between different neurons) fully trainable.
In the figure below, the connection weighted by $-U_{\rm thr}$ represents the reset mechanism $R[t]$.
The benefit of an unrolled graph is that it provides an explicit description of how computations are performed. The process of unfolding illustrates the flow of information forward in time (from left to right) to compute outputs and losses, and backward in time to compute gradients. The more time steps that are simulated, the deeper the graph becomes.
Conventional RNNs treat $\beta$ as a learnable parameter. This is also possible for SNNs, though by default, they are treated as hyperparameters. This replaces the vanishing and exploding gradient problems with a hyperparameter search.
The non differentiablility of spikes
In previous section we have discussed an example of heavy side function and how its derivative is always zero or undefined. This exactly how LIF neurons behave after thresholding.
An alternative way to represent the relationship between $S$ and $U$ in $(1)$ or $(8)$ is:
$$S[t] = \Theta(U[t] - U_{\rm thr}) \tag{9}$$
where $\Theta(\cdot)$ is the Heaviside step function:
Training a network in this form poses some serious challenges. Consider a single, isolated time step of the computational graph from the previous figure titled “Recurrent representation of spiking neurons”, as shown in the forward pass below:
The goal is to train the network using the gradient of the loss with respect to the weights, such that the weights are updated to minimize the loss. The backpropagation algorithm achieves this using the chain rule:
From $(7)$, $\partial I/\partial W=X$, and $\partial U/\partial I=1$. While we have not yet defined a loss function, we can assume $\partial \mathcal{L}/\partial S$ has an analytical solution, in a similar form to the cross-entropy or mean-square error loss (more on that shortly).
However, the term that we are going to grapple with is $\partial S/\partial U$. The derivative of the Heaviside step function from $(9)$ is the Dirac Delta function from $(2)$, which evaluates to 0 everywhere, except at the threshold $U_{\rm thr} = \theta$, where it tends to infinity. This means the gradient will almost always be nulled to zero (or saturated if $U$ sits precisely at the threshold), and no learning can take place. This is known as the dead neuron problem.
Overcoming the Dead Neuron Problem
As we have seen in first section we keep the Heaviside function as it is during the forward pass, but swap the derivative term $\partial S/\partial U$ for something that does not kill the learning process during the backward pass, which will be denoted $\partial \tilde{S}/\partial U$.
This might sound weird, but it turns out that neural networks are quite robust to such approximations. This is commonly known as the surrogate gradient approach.
A variety of options exist to using surrogate gradients, The default method in snnTorch (as of v0.6.0) is to smooth the Heaviside function with the arctangent function.
The backward-pass derivative used is:
$$ \frac{\partial \tilde{S}}{\partial U} \leftarrow \frac{1}{\pi}\frac{1}{(1+[U\pi]^2)} \tag{11}$$
Backpropagation Through Time (BPTT)
Equation $(10)$ only calculates the gradient for one single time step (referred to as the immediate influence in the figure below), but the backpropagation through time (BPTT) algorithm calculates the gradient from the loss to all descendants and sums them together.
Look at the path and we can see that the current output has depedence of previous $U$ which in turn has dependence on previous $W$ and so on.
The weight $W$ is applied at every time step, and so imagine a loss is also calculated at every time step. The influence of the weight on present and historical losses must be summed together to define the global gradient:
credit: Claude 5 Sonnet
So why we have multiple $W[s]$ if all we have is one weight just across the time ?
The total gradient is a sum of per-timestep gradient contributions ($\frac{\partial \mathcal{L[t]}}{\partial W}$). Each of those is itself a sum over every earlier time step $s\leq t$, because $W$ was reused at each of those steps ($W[s]$), and each reuse contributed its own path of influence on the loss at t.
If $W$ appeared only once in the whole computation, we could just chain straight through: $\mathcal{L}[t] \leftarrow U[t] \leftarrow U[t-1] \leftarrow \cdots \leftarrow W$, and $\frac{\partial \mathcal{L}[t]}{\partial W}$ would be one single product of derivatives, no sum needed. like we saw in previous figure and inital explaination in this segment.
But $W$ has multiple separate paths of influence on $\mathcal{L}[t]$: one path through $h[1]$, another through $U[2]$, another through $U[3]$, etc. Changing $W$ nudges all these usages simultaneously, and each nudge sends its own ripple forward to $\mathcal{L}[t]$.
This is just the multivariable chain rule
We’ve probably seen this rule before: if $z = f(x, y)$ and both $x = x(t)$ and $y = y(t)$ depend on the same $t$, then
$$\frac{dz}{dt} = \frac{\partial z}{\partial x}\frac{dx}{dt} + \frac{\partial z}{\partial y}\frac{dy}{dt}$$
We don’t get to pick just one path — we sum over every path by which $t$ reaches $z$.
The point of $(12)$ is to ensure causality: by constraining $s\leq t$, we only account for the contribution of immediate and prior influences of $W$ on the loss. A recurrent system constrains the weight to be shared across all steps: $W[0]=W[1] =~… ~ = W$. Therefore, a change in $W[s]$ will have the same effect on all $W$, which implies that $\partial W[s]/\partial W=1$:
$$\frac{\partial \mathcal{L}}{\partial W}= \sum_t \sum_{s\leq t} \frac{\partial\mathcal{L}[t]}{\partial W[s]} \tag{13} $$
As an example, isolate the prior influence due to $s = t-1$ only; this means the backward pass must track back in time by one step. The influence of $W[t-1]$ on the loss can be written as:
We have already dealt with all of these terms from $(11)$, except for $\partial U[t]/\partial U[t-1]$. From $(7)$, this temporal derivative term simply evaluates to $\beta$. So if we really wanted to, we now know enough to painstakingly calculate the derivative of every weight at every time step by hand, and it’d look something like this for a single neuron:
But thankfully, PyTorch’s autodiff takes care of that in the background for us.
Note: The reset mechanism has been omitted from the above figure. In snnTorch, reset is included in the forward-pass, but detached from the backward pass.
Setting up the loss and output decoding
In a spiking neural net, there are several options to interpreting the output spikes. The most common approaches are:
- Rate coding: Take the neuron with the highest firing rate (or spike count) as the predicted class
- Latency coding: Take the neuron that fires first as the predicted class
This part is similar to encoding. here, we are interpreting (decoding) the output spikes, rather than encoding/converting raw input data into spikes.
Rate code:. When input data is passed to the network, we want the correct neuron class to emit the most spikes over the course of the simulation run. This then corresponds to the highest average firing frequency.
One way to achieve this is to increase the membrane potential of the correct class to $U>U_{\rm thr}$, and that of incorrect classes to $U<U_{\rm thr}$. Applying the target to $U$ serves as a proxy for modulating spiking behavior from $S$.
This can be implemented by taking the softmax of the membrane potential for output neurons, where $C$ is the number of output classes:
The cross-entropy between $p_i$ and the target $y_i \in {0,1}^C$, which is a one-hot target vector, is obtained using:
The practical effect is that the membrane potential of the correct class is encouraged to increase while those of incorrect classes are reduced.
In effect, this means the correct class is encouraged to fire at all time steps, while incorrect classes are suppressed at all steps. This may not be the most efficient implementation of an SNN, but it is among the simplest.
This target is applied at every time step of the simulation, thus also generating a loss at every step. These losses are then summed together at the end of the simulation:
This is just one of many possible ways to apply a loss function to a spiking neural network. A few variety of approaches are available to use in snnTorch.functional.
code
#========================== Training SNN with static MNIST dataset ==========================
# imports
import snntorch as snn
from snntorch import spikeplot as splt
from snntorch import spikegen
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
import itertools
#-------------------------- Surrogate LIF neuron -------------------------
# Leaky neuron model, overriding the backward pass with a custom function
class LeakySurrogate(nn.Module):
def __init__(self, beta, threshold=1.0):
super(LeakySurrogate, self).__init__()
# initialize decay rate beta and threshold
self.beta = beta
self.threshold = threshold
self.spike_gradient = self.ATan.apply
# the forward function is called each time we call Leaky
def forward(self, input_, mem):
spk = self.spike_gradient((mem-self.threshold)) # call the Heaviside function
reset = (self.beta * spk * self.threshold).detach() # remove reset from computational graph as the surrogate gradient should only be applied to $\partial S/\partial U$, and not $\partial R/\partial U$
mem = self.beta * mem + input_ - reset # Eq (1)
return spk, mem
# Forward pass: Heaviside function
# Backward pass: Override Dirac Delta with the ArcTan function
@staticmethod
class ATan(torch.autograd.Function):
@staticmethod
def forward(ctx, mem):
spk = (mem > 0).float() # Heaviside on the forward pass: Eq(2)
ctx.save_for_backward(mem) # store the membrane for use in the backward pass
return spk
@staticmethod
def backward(ctx, grad_output):
(mem,) = ctx.saved_tensors # retrieve the membrane potential
grad = 1 / (1 + (np.pi * mem).pow_(2)) * grad_output # Eqn 5
return grad
lif1 = LeakySurrogate(beta=0.9) # similar to snn.Leaky(beta=0.9) as all snnTorch neurons apply surrogate gradients (Atan) by default
#-------------------------- Load MNIST dataset -------------------------
# dataloader arguments
batch_size = 128
data_path='/tmp/data/mnist'
dtype = torch.float
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
# Define a transform
transform = transforms.Compose([
transforms.Resize((28, 28)),
transforms.Grayscale(),
transforms.ToTensor(),
transforms.Normalize((0,), (1,))])
mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)
mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)
# Create DataLoaders
train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)
test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)
#-------------------------- Define Network model -------------------------
#............. Network Parameters .............
# Network Architecture
num_inputs = 28*28
num_hidden = 1000
num_outputs = 10
# Temporal Dynamics
num_steps = 25
beta = 0.95
#...............Network Class...................
# Define Network
class Net(nn.Module):
def __init__(self):
super().__init__()
# Initialize layers
self.fc1 = nn.Linear(num_inputs, num_hidden)
self.lif1 = snn.Leaky(beta=beta)
self.fc2 = nn.Linear(num_hidden, num_outputs)
self.lif2 = snn.Leaky(beta=beta)
def forward(self, x):
# Initialize hidden states at t=0
mem1 = self.lif1.init_leaky()
mem2 = self.lif2.init_leaky()
# Record the final layer
spk2_rec = []
mem2_rec = []
for step in range(num_steps):
cur1 = self.fc1(x)
spk1, mem1 = self.lif1(cur1, mem1)
cur2 = self.fc2(spk1)
spk2, mem2 = self.lif2(cur2, mem2)
spk2_rec.append(spk2)
mem2_rec.append(mem2)
return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)
# Load the network onto CUDA if available
net = Net().to(device)
#-------------------------- Training and Testing SNN -------------------------
#............. Accuracy metric.............
# pass data into the network, sum the spikes over time
# and compare the neuron with the highest number of spikes
# with the target
def print_batch_accuracy(data, targets, train=False):
output, _ = net(data.view(batch_size, -1))
_, idx = output.sum(dim=0).max(1)
acc = np.mean((targets == idx).detach().cpu().numpy())
if train:
print(f"Train set accuracy for a single minibatch: {acc*100:.2f}%")
else:
print(f"Test set accuracy for a single minibatch: {acc*100:.2f}%")
def train_printer(
data, targets, epoch,
counter, iter_counter,
loss_hist, test_loss_hist, test_data, test_targets):
print(f"Epoch {epoch}, Iteration {iter_counter}")
print(f"Train Set Loss: {loss_hist[counter]:.2f}")
print(f"Test Set Loss: {test_loss_hist[counter]:.2f}")
print_batch_accuracy(data, targets, train=True)
print_batch_accuracy(test_data, test_targets, train=False)
print("\n")
#.............Loss and optimiser .............
loss = nn.CrossEntropyLoss()
# Adaptive Moment Estimation (Adam) optimizer best for training RNNs
optimizer = torch.optim.Adam(net.parameters(), lr=5e-4, betas=(0.9, 0.999))
#............. Training Loop ............
num_epochs = 1
loss_hist = []
test_loss_hist = []
counter = 0
# Outer training loop
for epoch in range(num_epochs):
iter_counter = 0
train_batch = iter(train_loader)
# Minibatch training loop
for data, targets in train_batch:
data = data.to(device)
targets = targets.to(device)
# forward pass
net.train()
spk_rec, mem_rec = net(data.view(batch_size, -1))
# initialize the loss & sum over time
loss_val = torch.zeros((1), dtype=dtype, device=device)
for step in range(num_steps):
loss_val += loss(mem_rec[step], targets)
# Gradient calculation + weight update
optimizer.zero_grad()
loss_val.backward()
optimizer.step()
# Store loss history for future plotting
loss_hist.append(loss_val.item())
# Test set
with torch.no_grad():
net.eval()
test_data, test_targets = next(iter(test_loader))
test_data = test_data.to(device)
test_targets = test_targets.to(device)
# Test set forward pass
test_spk, test_mem = net(test_data.view(batch_size, -1))
# Test set loss
test_loss = torch.zeros((1), dtype=dtype, device=device)
for step in range(num_steps):
test_loss += loss(test_mem[step], test_targets)
test_loss_hist.append(test_loss.item())
# Print train/test loss/accuracy
if counter % 50 == 0:
train_printer(
data, targets, epoch,
counter, iter_counter,
loss_hist, test_loss_hist,
test_data, test_targets)
counter += 1
iter_counter +=1.
#............... Plotting Loss History ............
# Plot Loss
fig = plt.figure(facecolor="w", figsize=(10, 5))
plt.plot(loss_hist)
plt.plot(test_loss_hist)
plt.title("Loss Curves")
plt.legend(["Train Loss", "Test Loss"])
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.show()
#................ Plotting test accuracy............
total = 0
correct = 0
# drop_last switched to False to keep all samples
test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=False)
with torch.no_grad():
net.eval()
for data, targets in test_loader:
data = data.to(device)
targets = targets.to(device)
# forward pass
test_spk, _ = net(data.view(data.size(0), -1))
# calculate total accuracy
_, predicted = test_spk.sum(dim=0).max(1)
total += targets.size(0)
correct += (predicted == targets).sum().item()
print(f"Total correctly classified test set images: {correct}/{total}")
print(f"Test Set Accuracy: {100 * correct / total:.2f}%")
Additional reading on SuGD
Foundational details
- Neftci, Mostafa & Zenke (2019)
- Eshraghian, Ward, Neftci et al. (2021/2023)
reset mechanisms
- Zenke & Vogels (2021)
synaptic/membrane dynamics
- Wu et al. (2018),
Initializations
- Rossbroich, Gygax & Zenke (2022), “Fluctuation-Driven Initialization for Spiking Neural Network Training,” Neuromorphic Computing and Engineering
- Micheli et al. (2024/2025), “Deep Activity Propagation via Weight Initialization in Spiking Neural Networks”
miscellaneous
- Direct Training High-Performance Deep Spiking Neural Networks: A Review of Theories and Methods
- Fractional-order Spiking Neural Network