| name | surrogate-gradient-snn-training |
| description | Surrogate Gradient Learning for Spiking Neural Networks - comprehensive training framework using differentiable surrogate functions to overcome the non-differentiability of spike functions. Includes multiple surrogate gradient types (fast-sigmoid, exponential, arctan, erf), temporal batch normalization, neuron normalization, and advanced training strategies for deep SNNs. Activation: surrogate gradient SNN, differentiable spike, spiking neural network training, SNN backpropagation, time surrogate gradient. |
| tags | ["spiking-neural-networks","surrogate-gradient","SNN-training","backpropagation-through-time","differentiable-spike","temporal-credit-assignment","neuromorphic-deep-learning"] |
Surrogate Gradient Learning for Spiking Neural Networks
Overview
Surrogate gradient learning is the dominant training method for Spiking Neural Networks (SNNs), overcoming the non-differentiability of the spike function by using smooth approximations during backpropagation while maintaining discrete spikes during forward passes.
The Problem: Non-Differentiable Spike Functions
Spike Function (Heaviside Step):
s[t] = Θ(v[t] - v_th)
where Θ(x) = {1 if x >= 0, 0 if x < 0}
Derivative:
dΘ/dx = 0 (almost everywhere) or undefined (at x=0)
This prevents gradient flow during backpropagation!
The Solution: Surrogate Gradients
Forward Pass: Use discrete spikes (non-differentiable)
Backward Pass: Use smooth surrogate gradient (differentiable)
Surrogate Function σ(x) ≈ Θ(x) but with well-defined derivatives
Surrogate Gradient Functions
1. Fast Sigmoid (Most Common)
import torch
import torch.nn as nn
def fast_sigmoid_surrogate(x, alpha=1.0):
"""
Fast sigmoid surrogate gradient.
σ(x) = x / (1 + |x|)
dσ/dx = 1 / (1 + |x|)²
Args:
x: Membrane potential - threshold (v - v_th)
alpha: Steepness parameter (higher = steeper)
"""
return x / (1.0 + torch.abs(x * alpha))
class FastSigmoidSurrogate(torch.autograd.Function):
"""
Fast sigmoid surrogate with custom forward/backward.
"""
@staticmethod
def forward(ctx, x, alpha=1.0):
ctx.save_for_backward(x)
ctx.alpha = alpha
return (x >= 0).float()
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
alpha = ctx.alpha
grad_input = grad_output * alpha / (1.0 + torch.abs(x * alpha))**2
return grad_input, None
2. Exponential Surrogate
class ExponentialSurrogate(torch.autograd.Function):
"""
Exponential surrogate gradient.
σ(x) = exp(-|x|)
dσ/dx = -sign(x) * exp(-|x|)
Properties:
- Maximum gradient at threshold
- Smooth everywhere
- Bounded derivative
"""
@staticmethod
def forward(ctx, x, alpha=1.0):
ctx.save_for_backward(x)
ctx.alpha = alpha
return (x >= 0).float()
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
alpha = ctx.alpha
abs_x = torch.abs(x * alpha)
grad_input = grad_output * alpha * torch.exp(-abs_x)
return grad_input, None
3. Arctangent Surrogate
class ArctanSurrogate(torch.autograd.Function):
"""
Arctangent surrogate gradient.
σ(x) = arctan(αx) / π + 0.5
dσ/dx = α / (π(1 + (αx)²))
Properties:
- Normalized to [0, 1]
- Smooth transitions
- Used in some implementations
"""
@staticmethod
def forward(ctx, x, alpha=1.0):
ctx.save_for_backward(x)
ctx.alpha = alpha
return (x >= 0).float()
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
alpha = ctx.alpha
grad_input = grad_output * alpha / (torch.pi * (1 + (x * alpha)**2))
return grad_input, None
4. Super (SuperSpike) Surrogate
class SuperSurrogate(torch.autograd.Function):
"""
SuperSpike surrogate with adaptive width.
σ(x) = 1 / (1 + |x|)²
Based on: Zenke & Vogels (2021) "The Remarkable Robustness of Surrogate Gradient Learning"
"""
@staticmethod
def forward(ctx, x, beta=0.3):
ctx.save_for_backward(x)
ctx.beta = beta
return (x >= 0).float()
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
beta = ctx.beta
abs_x = torch.abs(x)
grad_input = grad_output / (beta * (1.0 + abs_x / beta)**2)
return grad_input, None
5. Sigmoid Surrogate
class SigmoidSurrogate(torch.autograd.Function):
"""
Standard sigmoid surrogate.
σ(x) = 1 / (1 + exp(-αx))
dσ/dx = α * σ(x) * (1 - σ(x))
"""
@staticmethod
def forward(ctx, x, alpha=10.0):
ctx.save_for_backward(x)
ctx.alpha = alpha
return (x >= 0).float()
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
alpha = ctx.alpha
sigmoid = torch.sigmoid(x * alpha)
grad_input = grad_output * alpha * sigmoid * (1 - sigmoid)
return grad_input, None
6. Gaussian (Spike-Response Model) Surrogate
class GaussianSurrogate(torch.autograd.Function):
"""
Gaussian/error function surrogate.
σ(x) = exp(-x² / 2σ²) / √(2πσ²)
Matches the derivative of the error function,
often used with probabilistic neuron models.
"""
@staticmethod
def forward(ctx, x, sigma=0.5):
ctx.save_for_backward(x)
ctx.sigma = sigma
return (x >= 0).float()
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
sigma = ctx.sigma
grad_input = grad_output * torch.exp(-x**2 / (2 * sigma**2)) / (sigma * torch.sqrt(torch.tensor(2 * torch.pi)))
return grad_input, None
Complete Surrogate Neuron Implementation
import torch
import torch.nn as nn
class SurrogateGradientNeuron(nn.Module):
"""
Spiking neuron with configurable surrogate gradient.
Supports multiple neuron models:
- LIF: Leaky Integrate-and-Fire
- ALIF: Adaptive LIF
- PLIF: Parametric LIF
"""
SURROGATE_FUNCTIONS = {
'fast_sigmoid': FastSigmoidSurrogate,
'exponential': ExponentialSurrogate,
'arctan': ArctanSurrogate,
'super': SuperSurrogate,
'sigmoid': SigmoidSurrogate,
'gaussian': GaussianSurrogate,
}
def __init__(
self,
neuron_type='LIF',
surrogate_type='fast_sigmoid',
surrogate_params=None,
tau_mem=20.0,
tau_adapt=None,
v_th=1.0,
v_reset=0.0,
spike_fn=None
):
super().__init__()
self.neuron_type = neuron_type
self.surrogate_type = surrogate_type
self.surrogate_params = surrogate_params or {}
self.tau_mem = tau_mem
self.v_th = v_th
self.v_reset = v_reset
if neuron_type == 'ALIF':
self.tau_adapt = tau_adapt or 100.0
self.beta_adapt =
spike_fn :
.spike_fn = .SURROGATE_FUNCTIONS[surrogate_type]
:
.spike_fn = spike_fn
():
batch_size = x.shape[]
state :
state = .reset_state(batch_size, x.device)
v = state[]
.neuron_type == :
v_th = state.get(, .v_th)
:
v_th = .v_th
alpha = torch.exp(- / .tau_mem)
v = alpha * v + ( - alpha) * x
spike = .spike_fn.apply(v - v_th, **.surrogate_params)
v = v * ( - spike) + .v_reset * spike
.neuron_type == :
alpha_adapt = torch.exp(- / .tau_adapt)
v_th = alpha_adapt * v_th + spike * .beta_adapt
new_state = {: v, : v_th}
:
new_state = {: v}
spike, new_state
():
state = {: torch.zeros(batch_size, device=device)}
.neuron_type == :
state[] = torch.full((batch_size,), .v_th, device=device)
state
SNN Layer with Temporal Processing
class SpikingLayer(nn.Module):
"""
Spiking neural network layer with temporal dynamics.
Processes input over time steps and maintains state.
"""
def __init__(
self,
in_features,
out_features,
neuron_type='LIF',
surrogate_type='fast_sigmoid',
surrogate_alpha=1.0,
tau_mem=20.0,
time_steps=100,
dropout=0.0,
use_recurrent=False,
recurrent_tau=5.0
):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.time_steps = time_steps
self.use_recurrent = use_recurrent
self.linear = nn.Linear(in_features, out_features)
if use_recurrent:
self.recurrent = nn.Linear(out_features, out_features, bias=False)
self.recurrent_tau = recurrent_tau
self.neuron = SurrogateGradientNeuron(
neuron_type=neuron_type,
surrogate_type=surrogate_type,
surrogate_params={'alpha': surrogate_alpha},
tau_mem=tau_mem,
v_th=1.0
)
self.dropout = nn.Dropout(dropout) if dropout > 0 else None
():
batch_size = x.shape[]
x.dim() == :
x = x.view(batch_size // .time_steps, .time_steps, -)
states :
states = .reset_states(batch_size, x.device)
spike_list = []
membrane_list = []
t (.time_steps):
x_t = x[:, t, :]
current = .linear(x_t)
.use_recurrent:
rec_current = .recurrent(states[])
alpha_rec = torch.exp(- / .recurrent_tau)
states[] = alpha_rec * states[] + ( - alpha_rec) * rec_current
current = current + states[]
.dropout :
current = .dropout(current)
spike, new_neuron_state = .neuron(current, states[])
states[] = new_neuron_state
.use_recurrent:
states[] = spike
spike_list.append(spike)
membrane_list.append(new_neuron_state[])
spikes = torch.stack(spike_list, dim=)
membrane_trace = torch.stack(membrane_list, dim=)
spikes, membrane_trace, states
():
states = {
: .neuron.reset_state(batch_size, device),
}
.use_recurrent:
states[] = torch.zeros(batch_size, .out_features, device=device)
states[] = torch.zeros(batch_size, .out_features, device=device)
states
Temporal Batch Normalization
class TemporalBatchNorm(nn.Module):
"""
Batch normalization adapted for temporal sequences.
Normalizes across batch and time dimensions while
maintaining temporal statistics separately.
"""
def __init__(self, num_features, eps=1e-5, momentum=0.1):
super().__init__()
self.num_features = num_features
self.eps = eps
self.momentum = momentum
self.weight = nn.Parameter(torch.ones(num_features))
self.bias = nn.Parameter(torch.zeros(num_features))
self.register_buffer('running_mean', torch.zeros(num_features))
self.register_buffer('running_var', torch.ones(num_features))
def forward(self, x):
"""
Args:
x: (batch, time, features) or (batch, features)
"""
if self.training:
if x.dim() == 3:
mean = x.mean(dim=[0, 1])
var = x.var(dim=[0, 1], unbiased=False)
else:
mean = x.mean(dim=0)
var = x.var(dim=0, unbiased=False)
.running_mean = ( - .momentum) * .running_mean + .momentum * mean
.running_var = ( - .momentum) * .running_var + .momentum * var
:
mean = .running_mean
var = .running_var
x_normalized = (x - mean) / torch.sqrt(var + .eps)
x_normalized * .weight + .bias
(nn.Module):
():
().__init__()
.eps = eps
.scale = nn.Parameter(torch.ones(num_neurons))
():
mean = x.mean(dim=-, keepdim=)
std = x.std(dim=-, keepdim=)
.scale * (x - mean) / (std + .eps)
Deep SNN Architecture
class DeepSNN(nn.Module):
"""
Deep Spiking Neural Network with surrogate gradient training.
Architecture: Input -> [Conv/ReLU] -> SNN Layers -> Output
"""
def __init__(
self,
input_size,
hidden_sizes=[512, 256],
output_size=10,
time_steps=100,
neuron_type='LIF',
surrogate_type='fast_sigmoid',
surrogate_alpha=1.0,
tau_mem=20.0,
use_readout='mean',
dropout=0.2
):
super().__init__()
self.input_size = input_size
self.output_size = output_size
self.time_steps = time_steps
self.use_readout = use_readout
layers = []
prev_size = input_size
for i, hidden_size in enumerate(hidden_sizes):
layers.append(
SpikingLayer(
in_features=prev_size,
out_features=hidden_size,
neuron_type=neuron_type,
surrogate_type=surrogate_type,
surrogate_alpha=surrogate_alpha,
tau_mem=tau_mem,
time_steps=time_steps,
dropout=dropout if i < len(hidden_sizes) - 1 else 0.0,
use_recurrent=(i == 0)
)
)
layers.append(TemporalBatchNorm(hidden_size))
prev_size = hidden_size
self.snn_layers = nn.ModuleList(layers)
.readout = nn.Linear(prev_size, output_size)
():
x.dim() == :
x = x.unsqueeze().repeat(, .time_steps, )
batch_size = x.shape[]
current = x
spike_counts = []
states =
i, layer (.snn_layers):
(layer, SpikingLayer):
current, membrane, states = layer(current, states)
spike_counts.append(current.())
:
original_shape = current.shape
current = layer(current.view(-, current.shape[-]))
current = current.view(original_shape)
.use_readout == :
readout_input = current.mean(dim=)
.use_readout == :
readout_input = current[:, -, :]
.use_readout == :
readout_input = current.(dim=)
.use_readout == :
readout_input = current.(dim=)[]
output = .readout(readout_input)
output, spike_counts
():
layer .snn_layers:
(layer, SpikingLayer):
layer.reset_states()
Training Pipeline
class SNNTrainer:
"""
Training pipeline for Spiking Neural Networks.
Handles:
- Surrogate gradient backpropagation
- Temporal credit assignment
- Activity regularization
- Spike rate monitoring
"""
def __init__(
self,
model,
optimizer,
criterion,
device='cuda',
reg_lambda=1e-5,
target_spike_rate=0.1
):
self.model = model
self.optimizer = optimizer
self.criterion = criterion
self.device = device
self.reg_lambda = reg_lambda
self.target_spike_rate = target_spike_rate
def train_step(self, x, y):
"""Single training step."""
self.model.train()
x = x.to(self.device)
y = y.to(self.device)
output, spike_counts = self.model(x)
loss = self.criterion(output, y)
total_spikes = sum(spike_counts)
avg_spike_rate = total_spikes / (x.shape[0] * self.model.time_steps)
rate_loss = self.reg_lambda * (avg_spike_rate - self.target_spike_rate)**2
total_loss = loss + rate_loss
.optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(.model.parameters(), max_norm=)
.optimizer.step()
{
: loss.item(),
: rate_loss.item(),
: total_loss.item(),
: (output.argmax(dim=) == y).().mean().item(),
: avg_spike_rate.item()
}
():
.model.()
total_loss =
total_acc =
total_spikes =
torch.no_grad():
x, y dataloader:
x = x.to(.device)
y = y.to(.device)
output, spike_counts = .model(x)
loss = .criterion(output, y)
total_loss += loss.item() * x.shape[]
total_acc += (output.argmax(dim=) == y).().item()
total_spikes += (s.item() s spike_counts)
n_samples = (dataloader.dataset)
{
: total_loss / n_samples,
: total_acc / n_samples,
: total_spikes / (n_samples * .model.time_steps)
}
():
train_metrics = []
x, y train_loader:
metrics = .train_step(x, y)
train_metrics.append(metrics)
avg_train = {
k: (m[k] m train_metrics) / (train_metrics)
k train_metrics[].keys()
}
val_loader :
val_metrics = .validate(val_loader)
avg_train, val_metrics
avg_train,
():
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)
criterion = nn.CrossEntropyLoss()
trainer = SNNTrainer(
model=model,
optimizer=optimizer,
criterion=criterion,
device=device
)
best_acc =
history = []
epoch (epochs):
train_metrics, val_metrics = trainer.train_epoch(train_loader, val_loader)
scheduler.step()
history.append({
: epoch,
: train_metrics,
: val_metrics
})
val_metrics val_metrics[] > best_acc:
best_acc = val_metrics[]
torch.save(model.state_dict(), )
epoch % == :
(
)
history
Advanced Techniques
1. Adaptive Surrogate Gradients
class AdaptiveSurrogate(nn.Module):
"""
Adaptive surrogate gradient that adjusts steepness during training.
"""
def __init__(self, initial_alpha=1.0, min_alpha=0.5, max_alpha=10.0):
super().__init__()
self.alpha = nn.Parameter(torch.tensor(initial_alpha))
self.min_alpha = min_alpha
self.max_alpha = max_alpha
def forward(self, x):
alpha = torch.clamp(self.alpha, self.min_alpha, self.max_alpha)
spike = (x >= 0).float()
if self.training:
spike.register_hook(
lambda grad: grad * alpha / (1.0 + torch.abs(x * alpha))**2
)
return spike
2. Population Coding Loss
class PopulationCodeLoss(nn.Module):
"""
Loss function for population-coded outputs in SNNs.
Uses spike timing or rates to encode continuous values.
"""
def __init__(self, coding_scheme='rate', tau_readout=20.0):
super().__init__()
self.coding_scheme = coding_scheme
self.tau_readout = tau_readout
def forward(self, spike_trains, target):
"""
Args:
spike_trains: (batch, time, neurons) spike trains
target: (batch, output_dim) target values
"""
if self.coding_scheme == 'rate':
rates = spike_trains.sum(dim=1)
decoded = self._decode_population(rates)
loss = F.mse_loss(decoded, target)
elif self.coding_scheme == 'time_to_first_spike':
spike_times = self._get_first_spike_times(spike_trains)
decoded = self._decode_temporal(spike_times)
loss = F.mse_loss(decoded, target)
elif self.coding_scheme == 'rank_order':
spike_order = self._get_spike_order(spike_trains)
loss = self._rank_order_loss(spike_order, target)
return loss
():
rates / (rates.(dim=-, keepdim=) + )
():
spike_times = (spike_trains > ).().argmax(dim=)
no_spike = (spike_trains.(dim=) == )
spike_times = spike_times.() + no_spike.() *
spike_times
3. Multi-Timescale Learning
class MultiTimescaleSNN(nn.Module):
"""
SNN with multiple timescales for different neurons.
Mimics biological heterogeneity in membrane time constants.
"""
def __init__(self, sizes, tau_range=(10, 100)):
super().__init__()
self.tau_mem = nn.ParameterList()
for size in sizes:
log_tau_min, log_tau_max = np.log(tau_range[0]), np.log(tau_range[1])
log_taus = torch.linspace(log_tau_min, log_tau_max, size)
self.tau_mem.append(nn.Parameter(torch.exp(log_taus)))
def get_decay_constants(self, layer_idx):
"""Get decay constants for a layer."""
return torch.exp(-1.0 / self.tau_mem[layer_idx])
Performance Comparison
def compare_surrogate_gradients():
"""
Compare different surrogate gradient functions on benchmark task.
"""
surrogates = [
('fast_sigmoid', {'alpha': 1.0}),
('exponential', {'alpha': 1.0}),
('super', {'beta': 0.3}),
('arctan', {'alpha': 2.0}),
]
results = {}
for name, params in surrogates:
model = DeepSNN(
input_size=784,
hidden_sizes=[512, 256],
output_size=10,
surrogate_type=name,
surrogate_alpha=params.get('alpha', 1.0)
)
history = train_snn(model, train_loader, val_loader, epochs=50)
results[name] = history
return results
References
- Neftci, E. O., Mostafa, H., & Zenke, F. (2019). Surrogate gradient learning in spiking neural networks. IEEE Signal Processing Magazine.
Advanced: Quantum-Assisted SNN Training (June 2026)
- QDS-SNN uses TSA-LIF neurons with adaptive thresholds for gradient flow with ≤6 time steps
- QACM (Quantum-Assisted Classifier Module) provides quantum-deep supervision
- See
quantum-neuromorphic-computing skill for cross-domain methodology
- Zenke, F., & Vogels, T. P. (2021). The remarkable robustness of surrogate gradient learning for instilling complex function in spiking neural networks. Neural Computation.
- Kaiser, J., Mostafa, H., & Neftci, E. (2020). Synaptic plasticity dynamics for deep continuous local learning. Frontiers in Neuroscience.
- Fang, W., et al. (2021). Incorporating learnable membrane time constant to enhance learning of spiking neural networks. IEEE/CVF ICCV.
- Zenke, F., & Ganguli, S. (2018). Superspike: Supervised learning in multi-layer spiking neural networks. Neural Computation.
Activation Keywords
- surrogate gradient SNN
- differentiable spike
- spiking neural network training
- SNN backpropagation
- surrogate gradient function
- temporal credit assignment
- surrogate gradient descent
- spike function derivative
- SNN optimization
- neuromorphic deep learning
- spiking neuron backpropagation
