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
erfFunction 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()
References