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

Lukas König; Manuel Kuhn; David Kappel; Anand Subramoney

cs.LGMar 9, 2026

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

Authors: Lukas König, Manuel Kuhn, David Kappel, Anand Subramoney

Organizations: CITEC, Bielefeld University, Bielefeld, Germany · Department of Computer Science, Royal Holloway, University of London, Egham, United Kingdom

Abstract

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.

Figures & tables

Library	Time scale	Backend(s)	State evolution	Spike detection	Event gradients	Supported models
Eventax (ours)	continuous	JAX / Equinox / Diffrax	numerical ODE solver	root-finding	backprop through solver	model-agnostic (no closed forms required)
SNNAX [ 1 ]	discrete	JAX / Equinox	discrete (Euler)	grid threshold	surrogate	model-agnostic
Norse [ 8 ]	discrete	PyTorch	discrete (Euler)	grid threshold	surrogate / closed form (LIF)	model-agnostic; optional LIF closed form
snnTorch [ 2 ]	discrete	PyTorch	discrete (Euler)	grid threshold	surrogate	model-agnostic
Holberg & Salvi [ 7 ]	continuous	JAX / Equinox / Diffrax	numerical SDE solver	root-finding	backprop through solver	model-agnostic (single model demonstrated)
EventProp [ 4 ]	continuous	PyTorch / C++	analytical (closed-form)	root-finding	closed-form	LIF only

TABLE I: Comparison of discrete-time and event-based gradient-based SNN optimization libraries.

Fig. 2: Schematic of the multi-compartment neuron model and corresponding ODEs (Equations 2 - 5 ).

X_{1}

and

X_{2}

are the inputs.

Fig. 3: Diffrax event handling. The ODE solver advances the state

y

until the event condition

g(t,y)

changes sign. A root-finder then locates the exact event time

t_{\text{event}}

between the last two solver steps, and the state is integrated from the last step to this exact event time. Gradients are backpropagated directly through the solver steps, while the dependence of

t_{\text{event}}

on the model parameters is handled via the implicit function theorem.

Parameter	Symbol	LIF	QIF a	EIF	Izhikevich
Membrane time constant	$\tau_{\text{mem}}$	$20.0\text{\,}\mathrm{ms}$	$20.00\text{\,}\mathrm{ms}$	$20.0\text{\,}\mathrm{ms}$	–
Synaptic time constant	$\tau_{\text{syn}}$	$5.00\text{\,}\mathrm{ms}$	$5.00\text{\,}\mathrm{ms}$	$5.00\text{\,}\mathrm{ms}$	$3.0\text{\,}\mathrm{ms}$
Spike threshold	$\vartheta$ or $v_{\text{peak}}$	1.00	$1.00$	$2.98$	$30.0\text{\,}\mathrm{mV}$
Exponential threshold	$v_{T}$	–	–	1.00	–
Reset voltage	$v_{\text{reset}}$ or $c$	0.00	$0.00$	0.00	$-65.00\text{\,}\mathrm{mV}$
Leak reversal	$E_{L}$	–	–	0.00	–

TABLE II: Neuron model parameters. All models use a fixed ODE solver step size of

0.1\text{\,}\mathrm{ms}

(Euler). The AMOS wrapper restricts neurons to at most one spike per trial.

Parameter	Yin–Yang	MNIST
Network architecture	5–50–3 a	784–200–10 a
Batch size	256	1024
Epochs	100	100
Optimizer	AdamW b	AdamW b
Learning rate (TTFS)	0.005	0.005
Learning rate (state-based)	0.005 (LIF)	0.005

TABLE III: Training configuration for Yin–Yang and MNIST.

Fig. 5: Delayed-XOR task encoding (0, 1)

Neuron	Loss	Accuracy [%]	Max Firing a	Dead Neurons b	Fire Count c	TTFS [ $\text{\,}\mathrm{ms}$ ] d
LIF	TTFS (AMOS) e	$91.30\pm 0.94$ (90.90)	$1.00\pm 0.00$ (1.00)	$2.75\pm 0.96$ (2.50)	$33.42\pm 2.00$ (33.06)	$15.82\pm 0.96$ (15.71)
	Max	$94.50\pm 1.21$ (95.20)	$3.94\pm 0.88$ (3.88)	$9.00\pm 5.24$ (8.00)	$67.34\pm 13.65$ (67.93)	–
	Integral	$96.68\pm 0.41$ (96.70)	$3.45\pm 0.37$ (3.49)	$14.40\pm 5.81$ (15.00)	$56.87\pm 11.10$ (58.57)	–
	Exp-Integral	$97.08\pm 0.36$ (97.10)	$3.88\pm 0.09$ (3.88)	$9.00\pm 3.67$ (7.00)	$63.36\pm 12.30$ (60.28)	–
QIF	TTFS	$98.62\pm 0.31$ (98.70)	$1.00\pm 0.00$ (1.00)	$1.20\pm 0.84$ (1.00)	$45.36\pm 1.35$ (45.43)	$27.27\pm 0.20$ (27.23)
	Max	$85.82\pm 18.10$ (98.60)	$1.38\pm 0.43$ (1.23)	$0.00\pm 0.00$ (0.00)	$43.26\pm 4.04$ (43.06)	–

TABLE IV: Yin-Yang classification with 50 hidden neurons (Except for multi-compartment case, where only 2 neurons were used). Test set performance at best validation epoch (mean

\pm

std, with median in parentheses) across five seeds trained for 100 epochs.

Fig. 6: Throughput of QIF models on the TTFS Yin–Yang task (NVIDIA H100). Three models were trained for various Euler step-size and hidden-size combinations, with throughput measured over 50 random batches at each batch size. Results show both forward-only and forward + backward throughput in samples per second.

Spiking Neural Networks

Spiking neural networks process information through discrete, time-dependent neural events, aiming for efficient computation and biologically inspired learning.

See all 351 papers

Bullet Trains: Parallelizing Training of Temporally Precise Spiking Neural Networks
Feb 27, 2026Todd Morrill, Christian Pehle, Anthony ZadorSpiking Neural Networks Neuromorphic Computing
Department of Computer Science, Columbia University, New York, NY, USA · Department of Neuroscience, Cold Spring Harbor Laboratory, Cold Spring Harbor, NY, USA
Scalable Training of Continuous-Time Spiking Neural Networks with Differentiable Spike-Time Discretization
Jul 16, 2026Yusuke Sakemi, Tomoya Takeuchi, Takeo Hosomi +1Spiking Neural Networks LIF Neurons
Research Center for Mathematical Engineering, Chiba Institute of Technology, Narashino, Japan · International Research Center for Neurointelligence (WPI-IRCN), The University of Tokyo, Tokyo, Japan · NEC Corporation, Kawasaki, Japan
ADSEQ: A delay-aware autograd-compatible framework for spike-event delivery in SNNs
Dec 5, 2025Lennart P. L. Landsmeer, Amirreza Movahedin, Said Hamdioui +1Automatic Differentiation Spiking Neural Networks
Department of Computer Engineering, Delft University of Technology, Netherlands · NeuroComputingLab, Department of Neuroscience, Erasmus Medical Center, Netherlands

Feb 27, 2026J/K move · Enter open · S save

Bullet Trains: Parallelizing Training of Temporally Precise Spiking Neural Networks

Todd Morrill, Christian Pehle, Anthony Zador

Department of Computer Science, Columbia University, New York, NY, USA · Department of Neuroscience, Cold Spring Harbor Laboratory, Cold Spring Harbor, NY, USA

Spiking Neural Networks Neuromorphic Computing

Open paper AskPDF arXiv

Continuous-time, event-native spiking neural networks (SNNs) operate strictly on spike events, treating spike timing and ordering as the representation rather than an artifact of time discretization. This viewpoint aligns with biological computation and with the native resolution of event sensors and neuromorphic processors, while enabling compute and memory that scale with the number of events. However, two challenges hinder practical, end-to-end trainable event-based SNN systems: 1) exact charge--fire--reset dynamics impose inherently sequential processing of input spikes, and 2) precise spike times must be solved without time bins. We address both. First, we use parallel associative scans to consume multiple input spikes at once, yielding up to 44x speedups over sequential simulation while retaining exact hard-reset dynamics. Second, we implement differentiable spike time solvers that compute spike times to machine-precision without discrete-time approximations or restrictive analytic assumptions. We demonstrate the viability of training SNNs using our solutions on four event-based datasets on GPUs.

Abstract

Spiking Neural Networks

Explore similar work

Bullet Trains: Parallelizing Training of Temporally Precise Spiking Neural Networks

Scalable Training of Continuous-Time Spiking Neural Networks with Differentiable Spike-Time Discretization

ADSEQ: A delay-aware autograd-compatible framework for spike-event delivery in SNNs