Surrogate

class braintools.surrogate.Surrogate

The base surrogate gradient function.

This abstract base class defines the interface for surrogate gradient functions used in training spiking neural networks. Surrogate gradients replace the non-differentiable spike function with smooth approximations during backpropagation.

To customize a surrogate gradient function, inherit this class and implement the surrogate_fun and surrogate_grad methods.

Examples

>>> import braintools
>>> import brainstate
>>> import jax.numpy as jnp
>>>
>>> # Create a custom surrogate gradient function
>>> class MySurrogate(braintools.surrogate.Surrogate):
...     def __init__(self, alpha=1.):
...         super().__init__()
...         self.alpha = alpha
...
...     def surrogate_fun(self, x):
...         # Define the smooth approximation function
...         return jnp.tanh(x * self.alpha) * 0.5 + 0.5
...
...     def surrogate_grad(self, x):
...         # Define its gradient for backpropagation
...         return self.alpha * 0.5 * (1 - jnp.tanh(x * self.alpha) ** 2)
>>>
>>> # Use the custom surrogate
>>> my_surrogate = MySurrogate(alpha=2.0)
>>> x = jnp.array([-1.0, 0.0, 1.0])
>>> spikes = my_surrogate(x)  # Forward: step function
>>> print(spikes)  # [0., 1., 1.]
[0. 1. 1.]

Notes

The forward pass always returns a Heaviside step function (0 or 1), while the backward pass uses the custom surrogate gradient defined in surrogate_grad method. This straight-through estimator approach enables gradient-based training of spiking neural networks.

Implementing surrogate_fun is optional – it is used only for visualization/analysis. Gradient-only surrogates (e.g. ReluGrad, GaussianGrad, MultiGaussianGrad, InvSquareGrad, SlayerGrad) do not define it, so calling surrogate_fun on them raises NotImplementedError. When both are defined, surrogate_grad is the exact derivative of surrogate_fun.

The input x is expected to be a dimensionless array (typically the membrane potential minus the threshold, in matching units). Passing a unitful brainunit.Quantity is not supported by the underlying primitive.

Differentiating the output with respect to the input x yields the surrogate gradient surrogate_grad(x), as intended. The output does not depend on surrogate parameters such as alpha, so its derivative with respect to them is 0. Surrogate parameters still influence training through the gradients they shape; to tune them, compare validation results or differentiate a loss after unrolled training steps with respect to them (second-order derivatives through surrogate_grad are supported).

surrogate_fun(x)

The smooth surrogate function, used for visualization and analysis.

Return type:

Array

surrogate_grad(x)

The gradient of the surrogate function, used during backpropagation.

Return type:

Array