Ctrl+K

braincore.surrogate.s2nn

Contents

braincore.surrogate.s2nn#

braincore.surrogate.s2nn(x, alpha=4.0, beta=1.0, epsilon=1e-08)[source]#

Judge spiking state with the S2NN surrogate spiking function [1].

If origin=False, computes the forward function:

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

If origin=True, computes the original function:

\[\begin{split} \begin{split}g(x) = \begin{cases} \mathrm{sigmoid} (\alpha x), x < 0 \\ \beta \ln(|x + 1|) + 0.5, x \ge 0 \end{cases}\end{split}\end{split}\]

Backward function:

\[\begin{split} \begin{split}g'(x) = \begin{cases} \alpha * (1 - \mathrm{sigmoid} (\alpha x)) \mathrm{sigmoid} (\alpha x), x < 0 \\ \frac{\beta}{(x + 1)}, x \ge 0 \end{cases}\end{split}\end{split}\]
>>> import braincore as bc
>>> import matplotlib.pyplot as plt
>>> xs = jax.numpy.linspace(-3, 3, 1000)
>>> grads = bc.vector_grad(bc.surrogate.s2nn)(xs, 4., 1.)
>>> plt.plot(xs, grads, label=r'$\alpha=4, \beta=1$')
>>> grads = bc.vector_grad(bc.surrogate.s2nn)(xs, 8., 2.)
>>> plt.plot(xs, grads, label=r'$\alpha=8, \beta=2$')
>>> plt.legend()
>>> plt.show()

(Source code)

Parameters:

  • x (jax.Array, Array) – The input data.

  • alpha (float) – The param that controls the gradient when x < 0.

  • beta (float) – The param that controls the gradient when x >= 0

  • epsilon (float) – Avoid nan

Returns:

out – The spiking state.

Return type:

jax.Array

References

Contents