SNNAX -- Spiking Neural Networks in JAX

• URL: http://arxiv.org/abs/2409.02842v1
• Date: Wed, 4 Sep 2024 16:14:14 GMT
• Title: SNNAX -- Spiking Neural Networks in JAX
• Authors: Jamie Lohoff, Jan Finkbeiner, Emre Neftci,
• Abstract summary: Spiking Neural Networks (SNNs) simulators are essential tools to prototype biologically inspired models and neuromorphic hardware architectures. We present SNNAX, a JAX-based framework for simulating and training such models with PyTorch-like intuitiveness and JAX-like execution speed.
• Score: 0.8321953606016751
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• Abstract: Spiking Neural Networks (SNNs) simulators are essential tools to prototype biologically inspired models and neuromorphic hardware architectures and predict their performance. For such a tool, ease of use and flexibility are critical, but so is simulation speed especially given the complexity inherent to simulating SNN. Here, we present SNNAX, a JAX-based framework for simulating and training such models with PyTorch-like intuitiveness and JAX-like execution speed. SNNAX models are easily extended and customized to fit the desired model specifications and target neuromorphic hardware. Additionally, SNNAX offers key features for optimizing the training and deployment of SNNs such as flexible automatic differentiation and just-in-time compilation. We evaluate and compare SNNAX to other commonly used machine learning (ML) frameworks used for programming SNNs. We provide key performance metrics, best practices, documented examples for simulating SNNs in SNNAX, and implement several benchmarks used in the literature.

Related papers

• Scalable Mechanistic Neural Networks [52.28945097811129]
  We propose an enhanced neural network framework designed for scientific machine learning applications involving long temporal sequences. By reformulating the original Mechanistic Neural Network (MNN) we reduce the computational time and space complexities from cubic and quadratic with respect to the sequence length, respectively, to linear. Extensive experiments demonstrate that S-MNN matches the original MNN in precision while substantially reducing computational resources.
  arXiv  Detail & Related papers  (2024-10-08T14:27:28Z)
• Towards Scalable GPU-Accelerated SNN Training via Temporal Fusion [8.995682796140429]
  Spiking Neural Networks (SNNs) emerge as a transformative development in artificial intelligence. SNNs show promising efficiency on specialized sparse-computational hardware, but their practical training often relies on conventional GPU. We present a novel temporal fusion method, specifically designed to expedite the propagation dynamics of SNNs on GPU platforms.
  arXiv  Detail & Related papers  (2024-08-01T04:41:56Z)
• Spyx: A Library for Just-In-Time Compiled Optimization of Spiking Neural Networks [0.08965418284317034]
  Spiking Neural Networks (SNNs) offer to enhance energy efficiency through a reduced and low-power hardware footprint. This paper introduces Spyx, a new and lightweight SNN simulation and optimization library designed in JAX.
  arXiv  Detail & Related papers  (2024-02-29T09:46:44Z)
• SparseProp: Efficient Event-Based Simulation and Training of Sparse Recurrent Spiking Neural Networks [4.532517021515834]
  Spiking Neural Networks (SNNs) are biologically-inspired models that are capable of processing information in streams of action potentials. We introduce SparseProp, a novel event-based algorithm for simulating and training sparse SNNs.
  arXiv  Detail & Related papers  (2023-12-28T18:48:10Z)
• Intelligence Processing Units Accelerate Neuromorphic Learning [52.952192990802345]
  Spiking neural networks (SNNs) have achieved orders of magnitude improvement in terms of energy consumption and latency. We present an IPU-optimized release of our custom SNN Python package, snnTorch.
  arXiv  Detail & Related papers  (2022-11-19T15:44:08Z)
• Training High-Performance Low-Latency Spiking Neural Networks by Differentiation on Spike Representation [70.75043144299168]
  Spiking Neural Network (SNN) is a promising energy-efficient AI model when implemented on neuromorphic hardware. It is a challenge to efficiently train SNNs due to their non-differentiability. We propose the Differentiation on Spike Representation (DSR) method, which could achieve high performance.
  arXiv  Detail & Related papers  (2022-05-01T12:44:49Z)
• Rethinking Pretraining as a Bridge from ANNs to SNNs [13.984523794353477]
  Spiking neural networks (SNNs) are known as a typical kind of brain-inspired models with their unique features. How to obtain a high-accuracy model has always been the main challenge in the field of SNN.
  arXiv  Detail & Related papers  (2022-03-02T14:59:57Z)
• Sub-bit Neural Networks: Learning to Compress and Accelerate Binary Neural Networks [72.81092567651395]
  Sub-bit Neural Networks (SNNs) are a new type of binary quantization design tailored to compress and accelerate BNNs. SNNs are trained with a kernel-aware optimization framework, which exploits binary quantization in the fine-grained convolutional kernel space. Experiments on visual recognition benchmarks and the hardware deployment on FPGA validate the great potentials of SNNs.
  arXiv  Detail & Related papers  (2021-10-18T11:30:29Z)
• Fully Spiking Variational Autoencoder [66.58310094608002]
  Spiking neural networks (SNNs) can be run on neuromorphic devices with ultra-high speed and ultra-low energy consumption. In this study, we build a variational autoencoder (VAE) with SNN to enable image generation.
  arXiv  Detail & Related papers  (2021-09-26T06:10:14Z)
• Progressive Tandem Learning for Pattern Recognition with Deep Spiking Neural Networks [80.15411508088522]
  Spiking neural networks (SNNs) have shown advantages over traditional artificial neural networks (ANNs) for low latency and high computational efficiency. We propose a novel ANN-to-SNN conversion and layer-wise learning framework for rapid and efficient pattern recognition.
  arXiv  Detail & Related papers  (2020-07-02T15:38:44Z)
• Compiling Spiking Neural Networks to Neuromorphic Hardware [4.273223677453178]
  Spiking Neural Network (SNN) can lower the energy consumption of machine learning applications executed on neuromorphic hardware. We propose an approach to analyze and compile SNNs on a resource-constrained neuromorphic hardware.
  arXiv  Detail & Related papers  (2020-04-07T21:13:27Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.