Add SPSN

[kmheckel]

Add Stochastic Parallelizable Spiking Neuron model.

Paper: https://arxiv.org/abs/2306.12666#:~:text=In%20this%20paper%2C%20we%20propose,run%20in%20parallel%20over%20time.

Torch implementation: https://github.com/NECOTIS/Stochastic-Parallelizable-Spiking-Neuron-SPSN/blob/main/neurons/spsn.py

• [x] Implement prototype SPSN model
• [ ] Ensure implementation is correct
• [ ] Determine proper way to write training loop for S4 that follows the philosophy of Haiku/Spyx

[kmheckel]

https://arxiv.org/abs/2401.00955

Spiking State Space Machine Paper.

[kmheckel]

A prototype implementation has been implemented in spyx.experimental. Need to verify the dimensions of internal calculation as right now there's broadcasting issues if the output dimension isn't a multiple/fraction of the input dim.

Training accuracy on SHD comparable to regularly trained recurrent SNNs in Spyx. Performance is comparable to recurrent model at small scales tested locally with short sequences (64). Testing notebook is at research/SPSN.

Considering creating a stochastic Axon class for _SigmoidBernoulli and refactoring it.

• [ ] Investigate limiting the sigmoid maximum to <1.0 to prevent extreme firing when oversaturated.
• [ ] Implement CuBaSPSN with second constant.
• [ ] Consider changing name schame to ParaLIF and ParaCuBaLIF

[kmheckel]

Found this:

Parallel Spiking Unit: https://arxiv.org/abs/2402.00449