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

TL;DR

Spyx is a lightweight JAX-based library that enables efficient, just-in-time compiled optimization of spiking neural networks on modern accelerators, improving performance and flexibility over existing frameworks.

Contribution

The paper introduces Spyx, a novel JAX-based library that leverages JIT compilation and data pre-staging to optimize SNN training on GPUs and TPUs, addressing efficiency and flexibility challenges.

Findings

Surpasses performance of existing SNN frameworks

Utilizes JAX's JIT compilation for optimized execution

Achieves high hardware utilization on GPUs and TPUs

Abstract

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.