braincore.surrogate.q_pseudo_spike

Contents

braincore.surrogate.q_pseudo_spike#

braincore.surrogate.q_pseudo_spike(x, alpha=2.0)[source]#

Judge spiking state with the q-PseudoSpike surrogate 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} \frac{1}{2}(1-\frac{2x}{\alpha-1})^{1-\alpha}, & x < 0 \\ 1 - \frac{1}{2}(1+\frac{2x}{\alpha-1})^{1-\alpha}, & x \geq 0. \end{cases}\end{split}\end{split}\]

Backward function:

\[g'(x) = (1+\frac{2|x|}{\alpha-1})^{-\alpha}\]
>>> import braincore as bc
>>> import matplotlib.pyplot as plt
>>> xs = jax.numpy.linspace(-3, 3, 1000)
>>> for alpha in [0.5, 1., 2., 4.]:
>>>   grads = bc.vector_grad(bc.surrogate.q_pseudo_spike)(xs, alpha)
>>>   plt.plot(xs, grads, label=r'$\alpha=$' + str(alpha))
>>> plt.legend()
>>> plt.show()

(Source code)

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

  • alpha (float) – The parameter to control tail fatness of gradient.

Returns:

out – The spiking state.

Return type:

jax.Array

References