PaperlayerView exact revision on arXiv

Source-linked AI summary

SpikingJelly: An open-source machine learning infrastructure platform for spike-based intelligence

Wei Fang, Yanqi Chen, Jianhao Ding, Zhaofei Yu, Timothée Masquelier, Ding Chen, Liwei Huang, Huihui Zhou, Guoqi Li, Yonghong Tian

arXiv:2310.16620v1cs.NEcs.LGcs.SE

TL;DR

Spiking deep learning lacks mature frameworks that integrate automatic differentiation, acceleration, and end-to-end SNN workflows. SpikingJelly addresses this gap with an open-source full-stack framework, and the paper reports state-of-the-art simulation speed. Its future scope includes broader support for timing-based SNNs.

  • Problem

    Existing frameworks do not adequately integrate automatic differentiation with the acceleration and full-stack capabilities needed for spiking deep learning.

  • Method

    SpikingJelly is an open-source full-stack framework for building, training, analyzing, simulating, and deploying SNNs.

  • Results

    SpikingJelly achieves state-of-the-art simulation speed, reducing time spent waiting for lengthy deep SNN training processes.

  • Takeaways & Limitations

    SpikingJelly provides an integrated framework intended to support the growing spiking deep learning research community.

  • Takeaways & Limitations

    Timing-based SNN toolkits and completion of the timing_based subpackage remain worthwhile future development directions.

Abstract

from arXiv · show

Spiking neural networks (SNNs) aim to realize brain-inspired intelligence on neuromorphic chips with high energy efficiency by introducing neural dynamics and spike properties. As the emerging spiking deep learning paradigm attracts increasing interest, traditional programming frameworks cannot meet the demands of the automatic differentiation, parallel computation acceleration, and high integration of processing neuromorphic datasets and deployment. In this work, we present the SpikingJelly framework to address the aforementioned dilemma. We contribute a full-stack toolkit for pre-processing neuromorphic datasets, building deep SNNs, optimizing their parameters, and deploying SNNs on neuromorphic chips. Compared to existing methods, the training of deep SNNs can be accelerated $11\times$, and the superior extensibility and flexibility of SpikingJelly enable users to accelerate custom models at low costs through multilevel inheritance and semiautomatic code generation. SpikingJelly paves the way for synthesizing truly energy-efficient SNN-based machine intelligence systems, which will enrich the ecology of neuromorphic computing.

Introduction

SNNs offer biologically inspired, event-driven processing with strong energy-efficiency potential, but their complex dynamics make deep learning difficult. Existing frameworks do not adequately combine automatic differentiation, acceleration, and full-stack support, motivating SpikingJelly.

  • Motivation: SNNs process information through spikes and nonlinear neuronal dynamics, making them suitable for spatiotemporal data.
  • Motivation: 1000× greater power efficiency than ANNs is reported for SNNs on tailored neuromorphic chips.
  • Challenges: Primitive supervised methods are constrained by single-spike neurons or inability to train multilayer SNNs, limiting tasks to roughly MNIST-level classification.
  • Deep Spiking Learning: Surrogate-gradient training reaches complex datasets including CIFAR, DVS Gesture, and ImageNet, while ANN-to-SNN conversion approaches original ANN accuracy on ImageNet.
  • Framework Gap: Existing SNN frameworks either emphasize biologically precise simulation without automatic differentiation or provide limited machine-learning capabilities.
  • SpikingJelly: SpikingJelly presents an open-source deep-learning framework with dataset processing, SNN construction, training, analysis, simulation, and neuromorphic deployment support.

Results

SpikingJelly provides flexible APIs, full-stack support for neuromorphic data and devices, and efficient SNN simulation. Experiments report up to 11× training speedup and up to 2× inference speedup over other frameworks.

  • Convenient and Flexible SNN Construction: SpikingJelly offers convenient, flexible PyTorch-style APIs for constructing and running SNNs with sequential layers and stored hidden states.Its multi_step_forward function iterates over time steps and concatenates the network outputs.
  • High-Performance Simulation: Up to 11× training speedup was achieved on Spiking ResNet-18 at T = 32 compared with Norse and SNNTorch.The evaluation measured execution time for a single training iteration using the surrogate gradient method.
  • Neuromorphic Dataset Support: The framework integrates event-based and frame-based neuromorphic data processing, including decoding, event-to-frame conversion, and dataset workflows.It implements decoding methods for different raw formats and provides event-to-frame downsampling methods.
  • Neuromorphic Device Support: SpikingJelly supports conversion and deployment of native SNN modules on Loihi and Lynxi KA200 neuromorphic chips.Other chips can be supported by developing specific exchange modules.
  • Framework Scope: Its full-stack integration covers neuromorphic datasets and chips, ANN2SNN, surrogate gradients, biologically plausible learning rules, and spike-operation optimizations.The framework is described as supporting a broad range of SNN applications and functional modules.

Discussion

SpikingJelly addresses the lack of mature spiking deep-learning frameworks with flexible APIs, computational acceleration, and a full-stack solution for building, training, and deploying SNNs.

  • SpikingJelly is designed to address the absence of mature frameworks for spiking deep learning.
  • Its convenient APIs and included models and scripts let users build and run varied SNNs with few lines of code.The design reduces the need to understand underlying implementation details and frees researchers from painstaking coding operations.
  • Multilevel inheritance and flexible APIs let researchers modify modules or define new classes by changing or overriding selected functions and properties.Most modules inherit from parent classes, allowing substantial behavioral changes with limited code.
  • SpikingJelly emphasizes computational efficiency through PyTorch acceleration, dimension merging, fused operations, generated CUDA kernels, and JIT compilation.The framework reports state-of-the-art simulation speed and aims to reduce time spent waiting for deep-SNN training.
  • Its full-stack solution extends deep-SNN use from toy classification toward practical classification, network deployment, and event-data processing.

Materials and Methods

This section describes SpikingJelly’s module organization, neuron dynamics, surrogate-gradient training, extensible inheritance design, and acceleration-oriented components.

  • Module organization: SpikingJelly categorizes modules as component modules for defining SNNs and functional modules for simulation, acceleration, recording, or variable modification.
  • Neuron model: LIF neurons fire when membrane potential crosses threshold and can use hard or soft reset mechanisms.Hard reset sets V(t) to Vreset, whereas soft reset subtracts Vth after firing.
  • Neuron model: Its discrete-time neuron model separates charging, firing, and resetting, with charging specific to each neuron and firing and resetting shared.The framework uses equations for these three processes and distinguishes pre-firing membrane potential H[t] from post-reset potential V[t].
  • Extensibility and acceleration: Multilevel inheritance makes it possible to extend neuron modules and their optional CUDA kernels by completing or overriding selected functions.The design is intended to enhance extensibility and reduce the workload of developing new models.
  • Surrogate gradients: Surrogate gradients replace the nondifferentiable spike-trigger derivative during backpropagation with the derivative of a smoother surrogate function.SpikingJelly packages the surrogate function as a component of the spiking neuron.
  • Extensibility and acceleration: In sequential 2D convolution, SpikingJelly’s SJ method is faster than RAW and scales more slowly with increasing time steps than RAW and 3D alternatives.The cited comparison reports quadratic growth for RAW during training, while SJ increases more slowly.

Funding

The work reports support from national, laboratory, and municipal Chinese research funding sources, and declares no competing interests.

  • The work is supported by grants from the National Natural Science Foundation of China.The listed grants are 62027804, 61825101, 62088102, 62236009, and 62176003.
  • Additional support comes from the major key project of Peng Cheng Laboratory and the Beijing Natural Science Foundation.The listed project and grant identifiers are PCL2021A13 and JQ21015.
  • The authors declare no competing interests.

Supplementary Information

SpikingJelly organizes SNN development into modular components, functional interfaces, acceleration tools, and reusable network structures. Its examples cover recurrent SNN construction, MNIST training, and device-compatible workflows.

  • Framework organization: SpikingJelly is organized into Components, Functions, Acceleration, and Networks subpackages.Components provide SNN building blocks, Functions provide training and analysis operations, Acceleration supports faster simulation, and Networks provide reusable architectures.
  • Framework organization: Components include spiking neurons, layers, encoders, and surrogate functions, including Parametric LIF neurons with learnable τm.The PLIF charging function uses the learnable parameter controlling τm(a).
  • Framework organization: Functions support training, simulating, analyzing, converting, quantizing, and deploying SNNs, while functional formulations allow forward propagation to be overridden without changing loaded weights.This supports reuse of trained ANN weights when modifying forward functions for SNNs.
  • Framework organization: SpikingJelly provides classic and large-scale network structures, including PLIF Net, Spiking VGG, Spiking ResNet, and SEW ResNet, with training scripts for reuse.Tutorial structures are intended to provide a low-cost starting process for new users.
  • Examples: Recurrent SNNs combine recurrent and feedforward modules, and SpikingJelly supports user-defined element-wise recurrent functions.The recurrent structure uses a fully connected layer and LIF neuron layer inside an element-wise recurrent container.
  • Examples: 92.9% test accuracy is achieved by a simple SNN trained on MNIST using a Poisson encoder, flatten layer, fully connected layer, and LIF neuron layer.The example repeats static images across T time steps, computes firing-rate loss, and uses surrogate learning for backpropagation.

Design of the Step Module

The StepModule unifies single-step and multistep execution, allowing SNNs to switch propagation patterns while managing stateful modules and hardware acceleration. These choices trade memory, speed, latency, and compatibility across workloads.

  • Step modes: StepModule inheritance lets stateless modules support both step modes, with multistep execution merging time and batch dimensions.When modules share the StepModule base, SNNs can switch between step-by-step and layer-by-layer propagation.
  • Step modes: StepModule uses step_mode='s' for data shaped (N, ...) and step_mode='m' for data shaped (T, N, ...).This setting controls whether modules operate one time step at a time or process complete sequences.
  • State management: SpikingJelly stores hidden states inside stateful modules through MemoryModule, including membrane potential v in the neuron base.The memory dictionary provides module-level storage for hidden states.
  • Complexity and latency: O(N · L) spatial inference complexity makes step-by-step propagation suitable for inference with many time steps.Earlier time-step outputs and hidden states can be discarded after each step.
  • Complexity and latency: O(N · T) spatial inference complexity makes layer-by-layer propagation appropriate for extremely deep SNNs.Layer-local computation allows data from previous layers to be discarded during inference.
  • Complexity and latency: dm[i] ≤ T · ds[i] yields Dm ≤ Ds, while multistep execution reduces kernel-calling overhead and can run faster when T is large.The analysis assumes available parallel processing and does not include input-sequence latency.
  • Compatibility: Some networks, tasks, and hardware are incompatible with both propagation patterns, including time-to-first-spike coding, event-driven backpropagation, and attention SNNs.These cases involve operations over all time steps and therefore can only be simulated with the layer-by-layer pattern.

Examples of Implementing Complex Neurons

SpikingJelly’s inheritance design supports concise implementations of complex neuron models while preserving configurable dynamics. Examples include adaptive-threshold ALIF, Izhikevich, and learnable-dynamics neurons.

  • Inheritance-based neuron design: Multilevel inheritance enables users to define complex neurons from base neurons with little effort.The section illustrates this principle through ALIF and Izhikevich implementations.
  • ALIF neuron: ALIF extends a base neuron with adaptive threshold dynamics using the membrane potential, threshold, and dynamic offset B[t].After firing, the offset updates as B[t + 1] = ρ · B[t] + (1 − ρ) · S[t], while the threshold adapts over time.
  • ALIF neuron: 21 lines of code complete the ALIF implementation, whose threshold increases after firing and then decreases exponentially under constant input.The plotted V[t] and Vth[t] curves demonstrate the adaptive-threshold behavior.
  • Izhikevich neuron: The Izhikevich neuron extends an adaptive base neuron with recovery dynamics, resetting the recovery state after firing.Its implementation is completed in 50 lines of code and models chattering dynamics.
  • Izhikevich neuron: Izhikevich neurons model chattering dynamics resembling pyramidal-neuron recordings from cat visual cortex.Chattering neurons generate high-frequency repetitive bursts in response to injected depolarizing currents.

Examples of Implementing New Learning Rules

SpikingJelly supports implementation of custom learning rules and optimizers through its extensible PyTorch-style interfaces. Examples cover learnable neuronal dynamics, surrogate gradients, Grad R pruning, and local learning rules.

  • Extensible learning: SpikingJelly demonstrates extensibility by implementing new learning rules within the framework.The examples are presented as evidence of the platform’s flexibility.
  • Learnable neuronal dynamics: A general linear neuron can use learnable threshold Vth and learnable charging parameters α and β.The charging function is H[t] = α · (V[t −1] − Vreset) + β · X[t], with α clamped to [0, 1].
  • Learnable neuronal dynamics: Gradient checks show that the implemented learnable-dynamics neuron works normally with learnable parameters.The test applies forward and backward propagation to a random input sequence and inspects parameter gradients.
  • Surrogate gradients: The custom STE surrogate function defines forward and backward operations and can be attached to an IF neuron layer.A one-time-step test checks the resulting input gradients.
  • Grad R pruning: Grad R is a plug-and-play PyTorch optimizer that prunes SNNs during training without a pretraining stage.It extends Adam with target sparsity and sparse-prior hyperparameters while maintaining weight signs.
  • Local learning rules: Local learning rules such as STDP can be implemented with monitors that record input and output spikes in single-step or multistep mode.The recorded firing rates generate weight updates through a broadcast mechanism.

Ablation Study of Acceleration Methods

The ablation study evaluates acceleration methods on Spiking ResNet-18 across training and inference time steps. Combining layer-by-layer propagation with MergeTB and CuPy provides the fastest training in most cases.

  • JIT effects: JIT accelerates inference noticeably in the step-by-step pattern, but its training effect varies and can slow layer-by-layer execution.When layer-by-layer propagation merges time and batch dimensions for stateless layers, JIT accelerates by about 10%.
  • Training acceleration: MergeTB and CuPy are critical for layer-by-layer propagation because they accelerate stateless and stateful layers, respectively.
  • Backend trade-offs: CuPy is slower than JIT when T is small and tiny-kernel calling overhead is minor.
  • Training acceleration: LBL + MergeTB + CuPy maximizes training efficiency in most cases.SpikingJelly enables this configuration by setting all modules to multistep mode and using the CuPy backend.

Comparison of SNN Frameworks

SpikingJelly is compared with SNN frameworks across GPU support, neuromorphic datasets and chips, learning rules, deep learning, community activity, and module design. Its distinguishing features include broad neuromorphic support, biologically plausible learning rules, flexible step modes, and efficient state management.

  • GPU support: SpikingJelly introduces a CuPy backend that uses CUDA-level acceleration and is more efficient than naive PyTorch-based modules.
  • Framework capabilities: SpikingJelly integrates a full-stack neuromorphic-data solution and provides nine datasets.
  • Framework capabilities: SpikingJelly implements biologically plausible learning rules, whereas synaptic plasticity is often neglected in spiking deep-learning frameworks.
  • Framework capabilities: SpikingJelly supports surrogate learning and integrates ANN2SNN methods for deep learning.
  • Community activity: SpikingJelly accumulated 94+ publications within three years of its 2019 launch.
  • Module design: SpikingJelly stores states inside modules, simplifying SNN construction and training because modules manage their own states.
  • Module design: Step Mode Modules switch between single-step and multistep processing, enabling flexible step-by-step and layer-by-layer propagation patterns.
  • Module design: SeqToANNContainer merges time-step and batch dimensions for stateless layers and is much faster than comparable Norse and SpikingJelly wrappers.

Details of the Neuron Kernel

SpikingJelly uses CuPy and semiautomatic Python-based CUDA generation to reduce the effort of writing and deploying custom kernels. Its kernel hierarchy supports elementwise, sequence, and neuron BPTT/FPTT operations through reusable components and inheritance.

  • Code generation: Semiautomatic CUDA code generation modularizes CUDA writing and reduces development costs.Python conditionals generate different CUDA codes, while CuPy removes the need for C++ wrappers and related extension procedures.
  • Kernel hierarchy: CKernel organizes generated CUDA code into declaration, head, core, and tail components.CKernel1D extends it for elementwise 1-D tensors, while CKernel2D adds a time-step loop for sequence data.
  • Kernel hierarchy: CKernel1D lets users define parameters and core CUDA code in Python for elementwise operations such as the Heaviside function.The generated kernels can be executed from a pure Python environment.
  • Kernel hierarchy: The cfunction subpackage generates CUDA code for float and half2 data types by switching the dtype.Users therefore do not need to write separate kernels for both data types.
  • Neuron kernels: CKernel2D supports spiking-neuron FPTT/BPTT kernels, with child classes filling neuron-specific gradient equations through reserved placeholders.IFNodeBPTTKernel inherits common code from NeuronBPTTKernel and defines the required gradient functions.

Exchange Modules

SpikingJelly uses exchange modules to preserve module behavior across its environment and neuromorphic-chip toolchains. These modules support conversion to targets such as Loihi and Lynxi KA200, while adapting models to chip-specific constraints.

  • Deployment constraints: Neuromorphic-chip deployment is resource-constrained, including differences in typical precision between Lynxi KA200 and Loihi.The cited passage reports 16-bit precision for Lynxi KA200 and 8-bit precision for Loihi.
  • Exchange-module design: Exchange modules run in SpikingJelly and convert to the target chip toolchain.Each chip toolchain has an exchange package, including lynxi_exchange for Lynxi KA200 and lava_exchange for Loihi.
  • Exchange-module design: Exchange modules match the behavior of corresponding target-toolchain modules, supporting lossless conversion.The paper gives CubaLIFNode in lava_exchange and slayer.neuron.cuba.Neuron as a corresponding pair.
  • Loihi deployment: BlockContainer provides quantized synapses and neurons with the same behavior as Lava-DL Block modules.SNNs built with BlockContainer convert to Lava-DL with identical outputs for the same inputs.
  • KA200 deployment: KA200 deployment simplifies synapses and neurons to retain core functions, removes GPU-only CuPy functions, and compiles model files for inference.The model files include JSON network-shape definitions and BIN network weights.

Neuromorphic Datasets

SpikingJelly integrates neuromorphic datasets and supports an end-to-end workflow for building, training, evaluating, and deploying SNN applications. Its demonstrations cover gesture recognition, robotic control, ImageNet classification, brain-likeness analysis, and Loihi deployment.

  • Neuromorphic deployment: SpikingJelly models can be converted to Lava-DL and exported through HDF5 for compilation and execution on Loihi or Loihi2 systems or simulators.The deployment workflow converts the SNN, exports it from Lava-DL to HDF5, reconstructs it in Lava, and compiles it for Loihi.
  • Gesture recognition: 85.42% test accuracy was achieved on DVS Gesture by an 8-bit-quantized SNN built from SpikingJelly modules.The application used quantization-aware training, Adam optimization, and cosine annealing over 128 epochs.
  • Scope of demonstrations: The applications demonstrate SpikingJelly’s full-stack ability to process neuromorphic data, build and train SNNs, and deploy them rather than target state-of-the-art accuracy.The DVS Gesture demonstration used small networks without complex training tricks, resulting in 85.42% accuracy.
  • Robotic control: 88.45% APR was achieved by SpikingJelly’s spiking actor network across Ant-v3, HalfCheetah-v3, Hopper-v3, and Walker2d-v3.PopSAN achieved 71.94% APR and had slower training speed in the comparison.

Performance Comparison on Recurrent Structure

Recurrent SNN benchmarking shows that framework speed depends on the simulation time-step regime. Norse is faster at small T, whereas SpikingJelly is faster at large T and remains fastest during inference.

  • Acceleration constraints: Multistep-mode acceleration cannot be applied to the recurrent SEW ResNet because recurrent connections occur in all four stages and require single-step mode.Hybrid propagation patterns can still support acceleration in partially recurrent SNNs.
  • Training speed: At T = 2, 4, and 8, the recurrent SEW ResNet-18 training speed rank is Norse > SpikingJelly > SNNTorch.The benchmark compares a fully recurrent SEW ResNet-18 across the three frameworks.
  • Inference speed: 0.57× execution times were required by SpikingJelly for inference relative to Norse and SNNTorch.Norse and SNNTorch had close inference execution times in the comparison.
  • Overall comparison: SpikingJelly, Norse, and SNNTorch achieve close training performance on fully recurrent structures, with the speed advantage shifting by T.Norse is faster when T is small, while SpikingJelly is faster when T is large and during inference.

Limitations and Future Development Plans

The paper identifies unfinished support for timing-based representations, sparse acceleration, and fully automatic CUDA generation as key development areas. These limitations reflect implementation complexity, dense spike storage, and developer-effort constraints.

  • CUDA code generation: SpikingJelly’s semiautomatic CUDA code generation only partly addresses the developer effort required for forward and backward propagation.The proposed direction is fully automatic generation of CUDA code from Python definitions of neuronal dynamics.
  • Timing-based SNNs: Timing-based SNN support remains primitive, despite timing-based representations being sparser and more memory efficient than activation-based representations.The paper describes timing-based spikes through firing times tf rather than binary spike tensors S.
  • Sparse acceleration: Sparse acceleration is limited because interlayer spikes are stored in dense tensors rather than fully exploiting spike sparsity.Sparse libraries such as cuSPARSE and Minkowski Engine are identified as potential ways to reduce memory use and accelerate simulation when firing rates are sufficiently low.
  • Sparse acceleration: Exploration of sparse acceleration was paused after version 0.0.0.0.4 because of limited developer time and effort.The paper states that future work may restart sparse-acceleration development.
  • CUDA code generation: Complex neuronal dynamics can trigger numerous CUDA kernel calls, slowing SNN training especially when T is large.The paper notes that even the simple IF neuron suffers from this issue, while fused CUDA kernels can improve simulation speed.
  • Community and roadmap: Research interest in spiking deep learning increased markedly from 2019, while the framework’s development roadmap remains incomplete in several areas.The paper connects the conference-publication trend with SpikingJelly’s role in supporting the research community and lists further toolkit development as worthwhile.
Loading 2310.16620v1…