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

Feb 27, 2026

Bullet Trains: Parallelizing Training of Temporally Precise Spiking Neural Networks

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.

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

Ask PDF

Jul 16, 2026

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

Continuous-time spiking neural networks (SNNs) provide an event-driven framework for temporal computation, computational neuroscience, and neuromorphic hardware. However, training deep continuous-time SNNs is severely constrained by the memory required for exact spike-time computation, which evaluates and retains candidate firing times over intervals determined by presynaptic spike ordering. Here we introduce a memory-efficient training framework based on differentiable spike-time discretization (DSTD) for leaky integrate-and-fire neurons with general membrane and synaptic time constants. DSTD maps irregular presynaptic spikes onto differentiable weighted events at fixed time points, replacing the input-dependent candidate dimension with

M

fixed time intervals while accurately approximating continuous-time membrane-potential dynamics. This reduces candidate-related activation memory from

O(N_{\mathrm{out}}N_{\mathrm{in}})

O(N_{\mathrm{out}}M)

in the case of time-to-first-spike (TTFS) coding, where

N_{\mathrm{in}}

and

N_{\mathrm{out}}

denote the numbers of presynaptic and postsynaptic neurons, respectively. We further introduce synfire-chain-inspired temporal regularization that organizes layer-wise firing windows, mitigates dead-neuron failures, and enables pipeline-like processing. In dense LIF layers, DSTD reduced peak memory consumption by up to approximately 100-fold and training time by up to approximately 20-fold compared with exact spike-time computation. Together, these methods allowed us to train 9-layer convolutional SNNs on CIFAR-10 and 20-layer convolutional SNNs on Fashion-MNIST on a single GPU.

Yusuke Sakemi, Tomoya Takeuchi, Takeo Hosomi +1

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

Spiking Neural Networks LIF Neurons

Ask PDF

Dec 5, 2025

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

Spiking neural networks (SNNs), central to computational neuroscience and neuromorphic machine learning (ML), require efficient simulation and gradient-based training. While AI accelerators offer promising speedups, gradient-based SNNs typically implement sparse spike events using dense, memory-heavy data-structures. Existing exact gradient methods lack generality, and current simulators often omit or inefficiently handle delayed spikes. We address this by deriving gradient computation through spike event queues, including delays, and implementing this into memory-efficient, autodifferentiable spike event queues (ADSEQ). These are benchmarked across CPU, GPU, TPU, and LPU platforms. We find that queue design strongly shapes performance. CPUs, as expected, perform well with traditional tree-based or FIFO implementations, while GPUs excel with ring buffers for smaller simulations, yet under higher memory pressure prefer more sparse data-structures. TPUs seem to favor an implementation based on sorting intrinsics. Selective spike dropping provides a simple performance-accuracy trade-off, which could be enhanced by future autograd frameworks adapting diverging primal/tangent data-structures.

Lennart P. L. Landsmeer, Amirreza Movahedin, Said Hamdioui +1

Department of Computer Engineering, Delft University of Technology, Netherlands · NeuroComputingLab, Department of Neuroscience, Erasmus Medical Center, Netherlands

Automatic Differentiation Spiking Neural Networks

Ask PDF

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