# Spyx: A Library for Just-In-Time Compiled Optimization of Spiking Neural Networks

arXiv: https://arxiv.org/abs/2402.18994 · HTML: https://mlsanity.com/paper/2402.18994/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 29 references that are in the map.
- First posted on arXiv 29 Feb 2024.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Converting Static Image Datasets to Spiking Neuromorphic Datasets Using Saccades](https://mlsanity.com/paper/1507.07629/) (2015)
- [Training Spiking Neural Networks Using Lessons From Deep Learning](https://mlsanity.com/paper/2109.12894/) (2021)
- [SuperSpike: Supervised learning in multi-layer spiking neural networks](https://mlsanity.com/paper/1705.11146/) (2017)
- [SpikingJelly: An open-source machine learning infrastructure platform for spike-based intelligence](https://mlsanity.com/paper/2310.16620/) (2023)
- [The Hardware Lottery](https://mlsanity.com/paper/2009.06489/) (2020)
- [Event-Based Backpropagation can compute Exact Gradients for Spiking Neural Networks](https://mlsanity.com/paper/2009.08378/) (2020)
- [Fast and energy-efficient neuromorphic deep learning with first-spike times](https://mlsanity.com/paper/1912.11443/) (2019)
- [Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning](https://mlsanity.com/paper/2303.17503/) (2023)
- [The BrainScaleS-2 accelerated neuromorphic system with hybrid plasticity](https://mlsanity.com/paper/2201.11063/) (2022)
- [evosax: JAX-based Evolution Strategies](https://mlsanity.com/paper/2212.04180/) (2022)
- [Neuromorphic Intermediate Representation: A Unified Instruction Set for Interoperable Brain-Inspired Computing](https://mlsanity.com/paper/2311.14641/) (2023)
- [Learning Long Sequences in Spiking Neural Networks](https://mlsanity.com/paper/2401.00955/) (2023)
- [jaxsnn: Event-driven Gradient Estimation for Analog Neuromorphic Hardware](https://mlsanity.com/paper/2401.16841/) (2024)
- [Deep Phasor Networks: Connecting Conventional and Spiking Neural Networks](https://mlsanity.com/paper/2106.11908/) (2021)
- [Accelerating SNN Training with Stochastic Parallelizable Spiking Neurons](https://mlsanity.com/paper/2306.12666/) (2023)
- [Harnessing Manycore Processors with Distributed Memory for Accelerated Training of Sparse and Recurrent Models](https://mlsanity.com/paper/2311.04386/) (2023)

## Built on by (most cited first, 2 of 2)

- [SNNAX -- Spiking Neural Networks in JAX](https://mlsanity.com/paper/2409.02842/) (2024)
- [Edge Intelligence with Spiking Neural Networks](https://mlsanity.com/paper/2507.14069/) (2025)

## Abstract (arXiv)

As the role of artificial intelligence becomes increasingly pivotal in modern society, the efficient training and deployment of deep neural networks have emerged as critical areas of focus. Recent advancements in attention-based large neural architectures have spurred the development of AI accelerators, facilitating the training of extensive, multi-billion parameter models. Despite their effectiveness, these powerful networks often incur high execution costs in production environments. Neuromorphic computing, inspired by biological neural processes, offers a promising alternative. By utilizing temporally-sparse computations, Spiking Neural Networks (SNNs) offer to enhance energy efficiency through a reduced and low-power hardware footprint. However, the training of SNNs can be challenging due to their recurrent nature which cannot as easily leverage the massive parallelism of modern AI accelerators. To facilitate the investigation of SNN architectures and dynamics researchers have sought to bridge Python-based deep learning frameworks such as PyTorch or TensorFlow with custom-implemented compute kernels. This paper introduces Spyx, a new and lightweight SNN simulation and optimization library designed in JAX. By pre-staging data in the expansive vRAM of contemporary accelerators and employing extensive JIT compilation, Spyx allows for SNN optimization to be executed as a unified, low-level program on NVIDIA GPUs or Google TPUs. This approach achieves optimal hardware utilization, surpassing the performance of many existing SNN training frameworks while maintaining considerable flexibility.