PiecewiseQuadratic

class braintools.surrogate.PiecewiseQuadratic(alpha=1.0)

Judge spiking state with a piecewise quadratic function.

This class implements a surrogate gradient method using a piecewise quadratic function for training spiking neural networks. It provides smooth gradients within a defined range around zero.

Parameters:

alpha (float) – A parameter controlling the width and steepness of the surrogate gradient. Higher values result in a narrower gradient window. Default is 1.0.

See also

piecewise_quadratic

Function version of this class.

Examples

>>> import brainstate
>>> import jax.numpy as jnp
>>>
>>> # Create a piecewise quadratic surrogate gradient function
>>> pq_fn = braintools.surrogate.PiecewiseQuadratic(alpha=1.0)
>>>
>>> # Apply to membrane potentials
>>> x = jnp.array([-2.0, -0.5, 0.0, 0.5, 2.0])
>>> spikes = pq_fn(x)
>>> print(spikes)  # Binary spike output: [0., 0., 1., 1., 1.]
>>>
>>> # Use in a spiking neural network
>>> import brainstate.nn as nn
>>>
>>> class SpikingNeuron(nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.spike_fn = braintools.surrogate.PiecewiseQuadratic(alpha=2.0)
...         self.membrane = 0.0
...
...     def forward(self, input_current):
...         self.membrane += input_current
...         spike = self.spike_fn(self.membrane)
...         self.membrane = self.membrane * (1 - spike)  # Reset on spike
...         return spike

>>> import jax
>>> import brainstate as brainstate
>>> import matplotlib.pyplot as plt
>>> xs = jax.numpy.linspace(-3, 3, 1000)
>>> for alpha in [0.5, 1., 2., 4.]:
>>>   pq_fn = braintools.surrogate.PiecewiseQuadratic(alpha=alpha)
>>>   grads = brainstate.augment.vector_grad(pq_fn)(xs)
>>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
>>> plt.legend()
>>> plt.show()

(Source code)

Notes

The forward pass uses a Heaviside step function (1 for x >= 0, 0 for x < 0), while the backward pass uses a piecewise quadratic surrogate gradient.

The surrogate gradient is non-zero only within the range \([-1/\\alpha, 1/\\alpha]\), providing localized gradient flow during backpropagation. This helps prevent gradient explosion and vanishing gradients in deep spiking networks.

The surrogate gradient is defined as:

\[\begin{split}g'(x) = \\begin{cases} 0, & |x| > \\frac{1}{\\alpha} \\\\ -\\alpha^2|x| + \\alpha, & |x| \\leq \\frac{1}{\\alpha} \\end{cases}\end{split}\]

References

[1]

Esser S K, Merolla P A, Arthur J V, et al. Convolutional networks for fast, energy-efficient neuromorphic computing[J]. Proceedings of the national academy of sciences, 2016, 113(41): 11441-11446.

[2]

Wu Y, Deng L, Li G, et al. Spatio-temporal backpropagation for training high-performance spiking neural networks[J]. Frontiers in neuroscience, 2018, 12: 331.

[3]

Bellec G, Salaj D, Subramoney A, et al. Long short-term memory and learning-to-learn in networks of spiking neurons[C]//Proceedings of the 32nd International Conference on Neural Information Processing Systems. 2018: 795-805.

[4]

Neftci E O, Mostafa H, Zenke F. Surrogate gradient learning in spiking neural networks: Bringing the power of gradient-based optimization to spiking neural networks[J]. IEEE Signal Processing Magazine, 2019, 36(6): 51-63.

[5]

Panda P, Aketi S A, Roy K. Toward scalable, efficient, and accurate deep spiking neural networks with backward residual connections, stochastic softmax, and hybridization[J]. Frontiers in Neuroscience, 2020, 14.

surrogate_fun(x)

Compute the piecewise quadratic surrogate function.

Parameters:

x (jax.Array) – Input tensor.

Returns:

Output of the surrogate function.

Return type:

jax.Array

surrogate_grad(x)

Compute the gradient of the piecewise quadratic function.

Parameters:

x (jax.Array) – Input tensor.

Returns:

Gradient of the surrogate function.

Return type:

jax.Array