ERF#

class braintools.surrogate.ERF(alpha=1.0)#

Judge spiking state with an error function (erf).

The forward function:

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

The original function:

\[\begin{split}\begin{split} g(x) &= \frac{1}{2}(1-\text{erf}(-\alpha x)) \\ &= \frac{1}{2} \text{erfc}(-\alpha x) \\ &= \frac{1}{\sqrt{\pi}}\int_{-\infty}^{\alpha x}e^{-t^2}dt \end{split}\end{split}\]

Backward function:

\[g'(x) = \frac{\alpha}{\sqrt{\pi}}e^{-\alpha^2x^2}\]
Parameters:

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

See also

erf

Function version of this class.

Examples

>>> import brainstate
>>> import jax.numpy as jnp
>>>
>>> # Create an ERF surrogate
>>> erf_fn = braintools.surrogate.ERF(alpha=1.0)
>>>
>>> # Apply to input
>>> x = jnp.array([-1.0, 0.0, 1.0])
>>> spikes = erf_fn(x)
>>> print(spikes)  # [0., 1., 1.]
>>> 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.]:
>>>   erf_fn = braintools.surrogate.ERF(alpha=alpha)
>>>   grads = brainstate.augment.vector_grad(erf_fn)(xs)
>>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
>>> plt.legend()
>>> plt.show()

(Source code)

References

surrogate_fun(x)[source]#

The surrogate function.

surrogate_grad(x)[source]#

The gradient function of the surrogate function.