NonzeroSignLog

class braintools.surrogate.NonzeroSignLog(alpha=1.0)

Judge spiking state with a nonzero sign log function.

The forward function:

\[\begin{split}g(x) = \begin{cases} 1, & x \geq 0 \\ 0, & x < 0 \\ \end{cases}\end{split}\]

The original function:

\[g(x) = \mathrm{NonzeroSign}(x) \log (|\alpha x| + 1)\]

where

\[\begin{split}\begin{split}\mathrm{NonzeroSign}(x) = \begin{cases} 1, & x \geq 0 \\ -1, & x < 0 \\ \end{cases}\end{split}\end{split}\]

Backward function:

\[g'(x) = \frac{\alpha}{1 + |\alpha x|} = \frac{1}{\frac{1}{\alpha} + |x|}\]

This surrogate function has the advantage of low computation cost during the backward.

Parameters:

alpha (float, optional) – Parameter controlling the steepness of the surrogate gradient. Higher values make the transition sharper. Default is 1.0.

See also

nonzero_sign_log

Function version of this class.

Examples

>>> import brainstate
>>> import jax.numpy as jnp
>>>
>>> # Create a nonzero sign log surrogate
>>> nsl_fn = braintools.surrogate.NonzeroSignLog(alpha=1.0)
>>>
>>> # Apply to input
>>> x = jnp.array([-1.0, 0.0, 1.0])
>>> spikes = nsl_fn(x)

>>> 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.]:
>>>   nsl_fn = braintools.surrogate.NonzeroSignLog(alpha=alpha)
>>>   grads = brainstate.augment.vector_grad(nsl_fn)(xs)
>>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
>>> plt.legend()
>>> plt.show()

(Source code)

surrogate_fun(x)

The surrogate function.

surrogate_grad(x)

The gradient function of the surrogate function.