# Training event-based neural networks with exact gradients via Differentiable ODE Solving in JAX

arXiv: https://arxiv.org/abs/2603.08146 · HTML: https://mlsanity.com/paper/2603.08146/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 21 references that are in the map.
- First posted on arXiv 9 Mar 2026.

## Builds on (most cited first, 13 of 13)

- [On Neural Differential Equations](https://mlsanity.com/paper/2202.02435/) (2022)
- [Training Spiking Neural Networks Using Lessons From Deep Learning](https://mlsanity.com/paper/2109.12894/) (2021)
- [Long short-term memory and learning-to-learn in networks of spiking neurons](https://mlsanity.com/paper/1803.09574/) (2018)
- [Event-Based Backpropagation can compute Exact Gradients for Spiking Neural Networks](https://mlsanity.com/paper/2009.08378/) (2020)
- [Fast and energy-efficient neuromorphic deep learning with first-spike times](https://mlsanity.com/paper/1912.11443/) (2019)
- [Efficient recurrent architectures through activity sparsity and sparse back-propagation through time](https://mlsanity.com/paper/2206.06178/) (2022)
- [Loss shaping enhances exact gradient learning with Eventprop in spiking neural networks](https://mlsanity.com/paper/2212.01232/) (2022)
- [The Yin-Yang dataset](https://mlsanity.com/paper/2102.08211/) (2021)
- [Smooth Exact Gradient Descent Learning in Spiking Neural Networks](https://mlsanity.com/paper/2309.14523/) (2023)
- [Exact Gradients for Stochastic Spiking Neural Networks Driven by Rough Signals](https://mlsanity.com/paper/2405.13587/) (2024)
- [jaxsnn: Event-driven Gradient Estimation for Analog Neuromorphic Hardware](https://mlsanity.com/paper/2401.16841/) (2024)
- [SNNAX -- Spiking Neural Networks in JAX](https://mlsanity.com/paper/2409.02842/) (2024)
- [Eventprop training for efficient neuromorphic applications](https://mlsanity.com/paper/2503.04341/) (2025)

## Built on by (most cited first, 2 of 2)

- [Quadratic integrate-and-fire neurons exhibit less fragmented loss landscapes and outperform leaky integrate-and-fire neurons in spike-based gradient descent](https://mlsanity.com/paper/2606.03935/) (2026)
- [Scalable Training of Continuous-Time Spiking Neural Networks with Differentiable Spike-Time Discretization](https://mlsanity.com/paper/2607.14672/) (2026)

## Abstract (arXiv)

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.