Search papers, labs, and topics across Lattice.
The paper introduces Eventax, a JAX-based framework for training event-based neural networks (SNNs) that combines differentiable ODE solvers with event-based spike handling to compute exact gradients. This approach overcomes the limitations of existing methods by supporting arbitrary neuron models defined by ODEs without introducing gradient bias or restricting spike-time resolution. The authors demonstrate Eventax on benchmarks like Yin-Yang and MNIST using diverse neuron models and loss functions, showcasing its flexibility and utility for prototyping and training SNNs with exact gradients, including complex multi-compartment neuron models.
Differentiable ODE solvers let you train spiking neural networks with *exact* gradients for *any* neuron model, unlocking a new level of flexibility and accuracy.
Existing frameworks for gradient-based training of spiking neural networks face a trade-off: discrete-time methods using surrogate gradients support arbitrary neuron models but introduce gradient bias and constrain spike-time resolution, while continuous-time methods that compute exact gradients require analytical expressions for spike times and state evolution, restricting them to simple neuron types such as Leaky Integrate and Fire (LIF). We introduce the Eventax framework, which resolves this trade-off by combining differentiable numerical ODE solvers with event-based spike handling. Built in JAX, our frame-work uses Diffrax ODE-solvers to compute gradients that are exact with respect to the forward simulation for any neuron model defined by ODEs . It also provides a simple API where users can specify just the neuron dynamics, spike conditions, and reset rules. Eventax prioritises modelling flexibility, supporting a wide range of neuron models, loss functions, and network architectures, which can be easily extended. We demonstrate Eventax on multiple benchmarks, including Yin-Yang and MNIST, using diverse neuron models such as Leaky Integrate-and-fire (LIF), Quadratic Integrate-and-fire (QIF), Exponential integrate-and-fire (EIF), Izhikevich and Event-based Gated Recurrent Unit (EGRU) with both time-to-first-spike and state-based loss functions, demonstrating its utility for prototyping and testing event-based architectures trained with exact gradients. We also demonstrate the application of this framework for more complex neuron types by implementing a multi-compartment neuron that uses a model of dendritic spikes in human layer 2/3 cortical Pyramidal neurons for computation. Code available at https://github.com/efficient-scalable-machine-learning/eventax.
