SNNAX -- Spiking Neural Networks in JAX
Jamie Lohoff Jan Finkbeiner Emre Neftci


IR[IR]Intrinsic Rewards and Motivation
PPO[PPO]Proximal Policy Optimization
RL[RL]Reinforcement Learning
AC[AC]Arrenhius & Current
AD[AD]Automatic Differentiation
AER[AER]Address Event Representation
AEX[AEX]AER EXtension board
AMDA[AMDA]“AER Motherboard with D/A converters”
ANN[ANN]Artificial Neural Network
API[API]Application Programming Interface
BP[BP]Back-Propagation
BPTT[BPTT]Back-Propagation-Through-Time
BM[BM]Boltzmann Machine
CAVIAR[CAVIAR]Convolution AER Vision Architecture for Real-Time
CCN[CCN]Cooperative and Competitive Network
CD[CD]Contrastive Divergence
CG[CG]Computational Graph
CMOS[CMOS]Complementary Metal–Oxide–Semiconductor
CNN[CNN]Convolutional Neural Network
COTS[COTS]Commercial Off-The-Shelf
CPU[CPU]Central Processing Unit
CV[CV]Coefficient of Variation
CTC[CTC]connectionist temporal classification
DAC[DAC]Digital–to–Analog
DBN[DBN]Deep Belief Network
DCLL[DECOLLE]Deep Continuous Local Learning
DFA[DFA]Deterministic Finite Automaton
DFA[DFA]Deterministic Finite Automaton
divmod3[DIVMOD3]divisibility of a number by 3
DPE[DPE]Dynamic Parameter Estimation
DPI[DPI]Differential-Pair Integrator
DSP[DSP]Digital Signal Processor
DVS[DVS]Dynamic Vision Sensor
EDVAC[EDVAC]Electronic Discrete Variable Automatic Computer
EIF[EI&F]Exponential Integrate & Fire
EIN[EIN]Excitatory–Inhibitory Network
EPSC[EPSC]Excitatory Post-Synaptic Current
EPSP[EPSP]Excitatory Post–Synaptic Potential
eRBP[eRBP]Event-Driven Random Back-Propagation
FPGA[FPGA]Field Programmable Gate Array
FSM[FSM]Finite State Machine
FZJ[FZJ]Foschungszentrum Jülich
GPU[GPU]Graphical Processing Unit
HAL[HAL]Hardware Abstraction Layer
HH[H&H]Hodgkin & Huxley
HMM[HMM]Hidden Markov Model
HNF[HNF]Helmholtz Nano-Facility
HW[HW]Hardware
hWTA[hWTA]Hard Winner–Take–All
ID[ID]Implicit Differentiation
IF2DWTA[IF2DWTA]Integrate & Fire 2–Dimensional WTA
IF[I&F]Integrate & Fire
IFSLWTA[IFSLWTA]Integrate & Fire Stop Learning WTA
INCF[INCF]International Neuroinformatics Coordinating Facility
INRC[INRC]Intel Neuromorphic Research Community
INM[INM]Institute of Neuroscience and Medicine
INI[INI]Institute of Neuroinformatics
IO[IO]Input-Output
IoT[IoT]internet of things
IPSC[IPSC]Inhibitory Post-Synaptic Current
IPU[IPU]Intelligence Processing Unit
ISI[ISI]Inter–Spike Interval
JFLAP[JFLAP]Java - Formal Languages and Automata Package
JIT[JIT]just-in-time compilation
JSC[JSC]Jülich Supercomputing Center
LIF[LIF]Linear Integrate and Fire
LSM[LSM]Liquid State Machine
LTD[LTD]Long-Term Depression
LTI[LTI]Linear Time-Invariant
LTP[LTP]Long-Term Potentiation
LTU[LTU]Linear Threshold Unit
LSTM[LSTM]long short-term memory
MAML[MAML]Model Agnostic Meta Learning
MCMCMarkov Chain Monte Carlo
MSEMean-Squared Error
NAS[NAS]Neural Architecture Search
NHML[NHML]Neuromorphic Hardware Mark-up Language
NMDA[NMDA]NMDA
NME[NE]Neuromorphic Engineering
PC[PC]Predictive Coding
PCB[PCB]Printed Circuit Board
PGI[PGI]Peter Grünberg Institute
PRC[PRC]Phase Response Curve
PSC[PSC]Post-Synaptic Current
PSP[PSP]Post–Synaptic Potential
RI[KL]Kullback-Leibler
RRAM[RRAM]Resistive Random-Access Memory
RBM[RBM]Restricted Boltzmann Machine
RTRL[RTRL]Real-Time Recurrent Learning
ROC[ROC]Receiver Operator Characteristic
RSA[RSA]Representational Similarity Analysis
RDA[RDA]Representational Dissimilarity Analysis
RDM[RDA]Representational Dissimilarity Matrix
RNN[RNN]Recurrent Neural Network
SAC[SAC]Selective Attention Chip
SCD[SCD]Spike-Based Contrastive Divergence
SCX[SCX]Silicon CorteX
SG[SG]Surrogate Gradient
SGD[SGD]Surrogate Gradient Descent
SRM[SRM]Spike Response Model
SNN[SNN]Spiking Neural Network
STDP[STDP]Spike Time Dependent Plasticity
SW[SW]Software
sWTA[SWTA]Soft Winner–Take–All
TPU[TPU]Tensor Processing Unit
VHDL[VHDL]VHSIC Hardware Description Language
VLSI[VLSI]Very  Large  Scale  Integration
WTA[WTA]Winner–Take–All
XML[XML]eXtensible Mark-up Language
SIMD[SIMD]Single Instruction Multiple Data
MIMD[MIMD]Multiple Instruction Multiple Data
UCI[UCI]University of California Irvine
Op[Op]Operation
ISA[ISA]Instruction Set Architecture
MLP[MLP]Multilayer Perceptrons
XLA[XLA]Accelerated Linear Algebra
ML[ML]machine learning
MLF[MLF]machine learning framework

Multivariate Second-Order p-Poincaré Inequalities
    Tara TrauthweinDepartment of Mathematics, University of Luxembourg and Department of Statistics, University of Oxford, tara.trauthwein@stats.ox.ac.uk. The author was supported by the Luxembourg National Research Fund (PRIDE17/1224660/GPS) and by the UK Engineering and Physical Sciences Research Council (EPSRC) grant (EP/T018445/1).
    
=================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
    SNN 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 SNN such as flexible automatic differentiation and just-in-time compilation.
    We evaluate and compare SNNAX to other commonly used ML frameworks used for programming SNNs.
    We provide key performance metrics, best practices, documented examples for simulating SNN in SNNAX, and implement several benchmarks used in the literature.


BSTcontrol 



§ INTRODUCTION

SNN simulators leveraging parallel processing have become indispensable tools to rapidly evaluate and test their performance.
Furthermore, as most neuromorphic hardware are based on spiking neurons, such simulators have become essential to evaluate the performance of prototypes at large scale.
Many state-of-the-art SNN are commonly trained with gradient-based learning; typically facilitated through a modern machine learning (ML) framework like PyTorch or JAX <cit.> and parallel computing accelerators like GPUs and TPUs.
These frameworks are highly modular and extensible and thus a popular choice for a dynamic and growing field such as neuromorphic computing.
Furthermore, ML frameworks are optimized for performance, allowing for efficient training of potentially large-scale SNN models on modern hardware accelerators.
A host of recent work leveraged ML frameworks for simulating SNN <cit.>.

Here, we report SNNAX, our JAX-based library for simulating SNN that is built on Equinox <cit.>, a thin neural network and numerical computation library.
We chose JAX as the underlying ML framework because it includes several features that are essential for algorithmic exploration while offering high execution performance. 
Exploring novel learning algorithms is essential to address challenges arising from backpropagation-based methods, e.g. by incorporating results from synaptic plasticity <cit.>. 
This type of algorithmic exploration is enabled by JAX's extensive automatic differentiation library, its Just-In-Time (JIT) compiler based on XLA, and a paradigm that separates functions and parameters (functional programming).
We compare the effectiveness of SNNAX and similar SNN libraries provide notebooks implementing the standard benchmarks inspired by the Neurobench initiative <cit.>.



§ RELATED WORK

Due to neuromorphic computing borrowing from both Neuroscience and AI, there is a significant functional overlap between SNN "brain" simulators like NEST<cit.>, Brian2 <cit.> and GeNN <cit.> and SNN libraries for machine learning frameworks.
While these brain simulators enable distributed, highly scalable simulations on supercomputers, their main design principles are biological accuracy and reproducibility and are thus currently unsuited to explore novel AI applications.
Missing in such frameworks are traceable dynamic variables that enables the computation of gradients and thus the efficient training of SNN.


Thus, most SNN libraries are built around an automatic differentiation (AD) library like PyTorch or JAX that allows for easy and efficient gradient computations.
Popular library choices for PyTorch include snnTorch, Norse and Spiking Jelly <cit.>.
snnTorch leverages PyTorch's own compiler to accelerate the otherwise impractically slow training of large-scale SNN in Python.
Norse also builds on PyTorch but instead of discretizing incoming spikes into a fixed temporal grid, it leverages an event-driven simulation scheme.
Both capitalize on many of the already implemented ML features like optimizers, data loading and popular neural network layers such as convolutions.

Spiking Jelly leverages custom CUDA kernels for the spiking neuron dynamics implemented through the CuPy interface to further accelerate the ML workloads.
It supports both layer-by-layer and step-by-step execution of the neuron models, thereby giving a lot of flexibility to the user.
SINABS <cit.> is another example for a PyTorch-based library that is functionally similar to snnTorch, but with a focus on neuromorphic hardware.
Gradient computations can be further accelerated with the EXODUS extension <cit.> which uses implicit differentiation and custom CUDA kernels for the SRM model.



However, these features only support also only feed-forward architectures and a limited class of neurons.

A similar neuromorphic-hardware focused library is SLAYER <cit.>, which includes a functional simulation of the Intel Loihi and custom CUDA kernels for the computation of delay gradients.

JAX is a popular alternative to PyTorch which leverages a functional programming paradigm. 
Treating many of its features as function transformations, JAX puts emphasis on composability and uses the model's source code as a template to generate new source code that supports JIT compilation, vectorization or auto-differentiability.
JAX itself does not provide convenience functions and neural network building blocks, but there exist several neural network libraries with different design paradigms such as Haiku, Flax and Equinox.

While not strictly necessary to create a SNN library as neatly demonstrated by Rockpool, they can massively streamline implementation and improve user experience.

Rockpool <cit.> is one of the earliest effort to leverage JAX for efficient SNN simulation.
Currently, Rockpool is used as a training platform for the accompanying Xylo SNN inference platform, and so its published functionality is limited to linear layers and LIF neuron dynamics.

Likewise, JAXSNN falls into this line of hardware-centric SNN libraries. 
JAXSNN is designed with a time-continuous approach to SNNs in mind, exploiting event-prop <cit.> and time-to-first-spike methods <cit.> for gradient computation which tie seamlessly with the accompanying BrainScaleS-2 platform to enable hardware-in-the-loop training.
For this reason, it is not a general-purpose training framework that supports arbitrary connectivity and neuron types but rather focuses on maximally enhancing performance on the underlying platform.

Haiku and Flax <cit.> are frameworks that reimplement several elements of the JAX API such as batching, internal state management and random number generation.



One of the most recently published SNN libraries is Spyx <cit.> builds on the Haiku library.
It is designed for maximum training speed on feed-forward architectures and achieves this through layer-by-layer execution and Haiku's loop unrolling feature.
This reduces the CPU-GPU communication overhead significantly and allows the batch-parallel application of stateless operations.
The downside of this method is the lack of support for arbitrary recurrent connectivity, which can be beneficial on some tasks.
Thanks to Haiku, Spyx supports a variety of neural network blocks and neuron models while still achieving similar performance to other frameworks that require the implementation of custom CUDA kernels.
It also supports the mapping of the trained parameters to neuromorphic hardware via the Neuromorphic Intermediate Representation (NIR) <cit.>.
In a similar vein, Slax <cit.> builds on Flax and is geared towards online training and algorithmic exploration.
With this goal in mind, Slax not only supports BPTT but various other approximations like OSTL, RTRL, OTTT, OTPE and FPTT.
As opposed to Spyx, Slax supports almost arbitrary connectivity thereby enabling recurrent connections.
This comes at the price of step-by-step execution of the entire network architecture as opposed to Spyx' layer-by-layer approach.

Flax and Haiku come with the downside that users interested in exploring beyond provided models are required to learn Haiku or Flax, and a much less intuitive interface compared to PyTorch.
In contrast, Equinox is a minimal library that offers many scientific computation and neural network features, while being fully compatible with all existing JAX functionalities. 
SNNAX builds on Equinox as underlying deep learning library.
Like Spyx, this allows building state-of-the-art networks with custom building blocks in a way similar to PyTorch modules while  operating seamlessly with JAX functions.
We believe that Equinox' approach is more suitable for the cutting-edge algorithmically-minded user that wishes to explore new brain-inspired algorithms and structures.
Many of the discussed libraries share a unified interface provided by the NIR initiative which enables a certain level of cross-compatibility between them thus also partly exposing the interfaces of a number of neuromorphic accelerators.

















    





































    

























































§ A MACHINE LEARNING LIBRARY FOR SNN

While developing SNNAX, we identified a set of key functionalities that a SNN library should satisfy to efficiently train modern spiking neural networks.
In the following sections, we will discuss these key functionalities and describe their implementation in SNNAX.



 §.§ Network Connectivity and Hardware Acceleration

A key element in modern deep learning is the specification of the connectivity between neurons.
This is typically realized by grouping neurons into layers whose connectivity is then described by highly regular connections, e.g. dense all-to-all connections or convolutions.
SNNAX is built on Equinox, a JAX library that already provides the core functionalities to describe the network connectivity. 
Since the corresponding mathematical operations are very costly, but also easily parallizable, SNNAX exploits JAX' XLA backend to run these on dedicated accelerators such as GPUs and TPUs.
XLA (aXelerated Linear Algebra) <cit.> is an increasingly popular ML compiler that was originally designed as Tensorflow backend, but now also serves as the backbone of JAX. 
It defines a representation of the computational graph, optimizes it, and deploys it onto the target hardware.
Furthermore, XLA is particularly amenable to extension with novel devices and computing architectures as demonstrated in <cit.>.


SNNAX and Equinox are both able to fully harness JAX' JIT compilation, making it a competitive choice for efficient SNN evaluation.
In fact, both support all of JAX' function transformation including automated and intuitive batching through vmap and device parallelism through pmap and array sharding mechanism.




 §.§ Stateful Layers and Long Sequences

SNNs are temporally evolving dynamical systems modeled through differential equations.
Discretized SNN can be viewed as a special type of RNN with multiple internal states, akin to an LSTM <cit.>. 
In its most general case of relevance to the implementation in ML frameworks, SNN dynamics can be written as:


    U^t+Δ t_i    = f(U^t_i, S^t_i, θ) + g(U^t, S^t, ϕ);   t=0,...,TΔ t 
    
        S^t_i    = Θ(U^t_i)


where U^t ∈ℝ^N× M is the internal state of the neuron which consists of M states (e.g. compartments) and S^t ∈ℝ^N is the output of other neurons activity (including inputs). 
Additionally, Θ is a discontinuous threshold function, θ and ϕ are learnable parameters of the neuron dynamics, and T is the number of time steps. N denotes the number of neurons and M denotes the number of states within each neuron.
Because the dynamics above often result from the discretization of continuous dynamics, Δ t is generally small, leading to T being large (generally T > 100) compared to RNN (generally N < 10).
Dealing with the internal states of the neurons is often tedious within JAX' functional programming paradigm.
SNNAX therefore handles the state management behind a thin layer that is still easily accessible to the user.




 §.§ Gradient-based Learning

Although Neuroscience provides principled models of biological learning, the state-of-the-art for training SNN for most practical applications is still gradient descent.
There, automatic differentiation (AD) and the backpropagation algorithm are used to compute the gradient of a cost function with respect to the models parameters <cit.>.
AD allows computation of gradients up to machine precision by decomposing the models computational graph into its elemental operations whose symbolic derivatives are known and then utilizes the chain rule to assemble the functions that computes the derivatives.




  §.§.§ Temporal Credit Assignment

SNNAX maps spikes into a regular temporal grid through spike time binning, which makes it functionally similar to snnTorch and Spyx.
It then treats the discretized form of the SNNs ODE as an RNN which is evaluated using a for-loop over the time dimension with changing input and internal state but fixed computational graph per time step.
The correct implementation of the temporal processing becomes particularly relevant when considering layer-by-layer and step-by-step execution of spiking neural network architectures.
Unrolling the jax.lax.scan primitive for several steps instead of just one can be beneficial in the layer-by-layer case because it reduces the CPU-GPU communication and allows the batch-parallel application of stateless operations such as matrix-multiplications and convolutions which can incur massive runtime improvements for feed-forward architectures.
SNNAX, evaluates the time-loop over the entire network in a step-by-step manner as opposed to the layer-by-layer approach in Rockpool or Spyx.
While this is a performance bottleneck in some cases, we believe that this design choice is necessary since it enables recurrent connections across layers, a key feature of bio-plausibe SNNs that is paramount for the successful implementation of learning algorithms such as e-prop <cit.>.



Figures <ref> and <ref> show execution times of LIF-based feed-forward MLP and CNN architectures for different SNN libaries on a single RTX 4090 GPU with a batchsize of 32.


When differentiated, we arrive at a very efficient version of backpropagation through time (BPTT) <cit.>, the de-facto standard for training SNNs.




  §.§.§ Gradient Across Spikes

A spiking neuron emits a spike when its internal state reaches a certain threshold. 
Once the threshold is crossed, the state generally resets to a new value, which makes spike emission a apparent non-differentiable process. 
Several approaches exist to circumvent the non-differentiable threshold with surrogate gradient approaches currently being the most popular solution <cit.>. 
There, the ill-defined derivative of the spiking function is replaced with a surrogate function that then allows to obtain a smooth gradient for the entire network.
Since SNNAX is built on JAX' advanced AD features, it is straightforward to define these surrogates and customize them according to the users requirements and several surrogate functions commonly found in literature are already implemented.


[caption=Implementation of a feed-forward SNN in SNNAX with tools from equinox and JAX. Note the consequent use of PyTree filters throughout the entire implementation., captionpos=b, label=lst:codeexample]
# ...
import equinox as eqx
import snnax.snn as snn

model = snn.Sequential(eqx.Conv2D(2, 32, 7, 2, key=key1),
                            snn.LIF((8, 8), [.9, .8], key=key2),
                            snn.flatten(),
                            eqx.Linear(64, 11, key=key3),
                            snn.LIF(11, [.9, .8], key=key4))
# ...
# Simple batched loss function
@partial(jax.vmap, in_axes=(0, 0, 0))
def loss_fn(in_states, in_spikes, tgt_class):
    out_state, out_spikes = model(in_states, in_spikes)

    # Spikes from the last layer are summed
    pred = out_spikes[-1].sum(-1)
    loss = optax.softmax_cross_entropy(pred, tgt_class)
    return loss

# Calculating the gradient with equinox PyTree filters and
# subsequently jitting the resulting function
@eqx.filter_jit
@eqx.filter_value_and_grad
def loss_and_grad(in_states, in_spikes, tgt_class):
    return jnp.mean(loss_fn(in_states, in_spikes, tgt_class))
# ...
# Simple training loop
for spikes, tgt_cls in tqdm(dataloader):
    # Initializing the membrane potentials of LIF neurons
    states = model.init_states(key)
    
    # Jitting with equinox PyTree filters
    loss, grads = loss_and_grad(states, spikes, tgt_cls)

    # Update parameter PyTree with equinox and optax
    updates, opt_state = optim.update(grads, opt_state)
    model = eqx.apply_updates(model, updates)
    




  §.§.§ Advanced AD Tools

JAX' cutting-edge AD features enable a host of involved learning techniques that have proven to be helpful assets when training SNNs.
SNN training can be very memory-intensive due to the small simulation time steps necessary, often resulting in sequence lengths in the hundreds or even thousands.
When training with BPTT, this quickly saturates the memory of the underlying accelerator, and thus training large models can become prohibitively slow.
JAX, like PyTorch provides a gradient checkpointing tool that reduces memory consumption through recomputation during the backward pass.
Furthermore, JAX provides a fully-fledged forward-mode AD implementation that allows to trade memory for compute which enables the design of elaborate gradient-based learning algorithms that save on both memory and compute or might even enable online learning <cit.>.

Furthermore, JAX allows, with only very few limitations, to compute arbitrarily high orders of derivatives, thereby facilitating Meta-learning and learning-to-learn approaches which have been demonstrated to be particularly useful for few-shot learning <cit.>.
The implementation in SNNAX is straight forward in this case, by just differentiating the already differentiated code once again.



 §.§ User Interface

SNNAX, much like Equinox, has been designed with user experience in mind.
A common motive of many JAX libraries is the extensive use of PyTrees, which are simple data structures that package large assemblages of arrays and related structures.
Equinox represents parameterised functions as immutable PyTree instances in order to be compatible with JAX' functional programming paradigm.
This also allows to directly use JAX' functional transformations like jax.jit or jax.grad on calls to the model's class instance without the need for further modifications as required by Haiku or Flax.
Both frameworks require additional steps like Flax' @nn.compact decorator or Haiku's set of hk.transform primitives which convert the object-oriented model implementation into an object that is compatible with JAX' functional programming paradigm.
These additional levels of abstraction can introduce unforseeable side-effects, incur additional implementation overhead or severely limit the functionality of certain features such as the use of JAX transformations within hk.transform.
SNNAX embraces Equinox' Everything-is-a-PyTree paradigm and provides a set of additional tree primitives and management tools dedicated to the management of stateful neural network models.
In this vein, SNNAX also provides a set of convenience function such as snnax.Sequential to enable the rapid prototyping and layer-by-layer execution of feed-forward networks similar to Spyx (see Listing <ref>), while snnax.SequentialRecurrent allows the implementation of recurrent models as demonstrated in Slax.
For more elaborate models with complex recurrent connections, SNNAX provides a GraphStructure class that allows the definition of intricate feedback loops and recurrences through a graph object.
Alternatively, users can implement custom behavior by implementing their own stateful layers based on the StatefulLayer class.
The graph representation is then meticulously executed by SNNAX' execution loop while being fully differentiable.
Finally, since Equinox attempts to mimic the simplicity of the PyTorch API, SNNAX inherits many of these features as well which flattens the learning curve and reduces the codebase footprint significantly.




§ CONCLUSION

We introduced SNNAX as a new, fast and user-friendly SNN simulation library that allows to efficiently train modern SNNs.
Its design principles are in line with already existing SNN libraries and it matches their performance without the need for writing custom CUDA code while being easy to maintain and extend.









§ ACKNOWLEDGEMENTS

This work was sponsored by the Federal Ministry of Education, Germany (project NEUROTEC-II grant no. 16ME0398K and 16ME0399) and NeuroSys as part of the initiative "Clusters4Future", funded by the Federal Ministry of Education and Research BMBF (03ZU1106CB).
We also would like to thank Anil Kaya and Anurag K. Mishra for their valuable contributions to the development of the framework.


IEEEtran


