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

• URL: http://arxiv.org/abs/2603.08146v1
• Date: Mon, 09 Mar 2026 09:25:52 GMT
• Title: Training event-based neural networks with exact gradients via Differentiable ODE Solving in JAX
• Abstract summary: Eventax is a framework for training spiking neural networks.<n>It combines differentiable numerical ODE solvers with event-based spike handling.<n>It supports a wide range of neuron models, loss functions, and network architectures, which can be easily extended.
• Score: 0.869928033942254
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Physics-Modeled Neural Networks [0.0]
  We introduce emphDynamical Physics-Modeled Neural Networks (DynPMNNs)<n>DynPMNNs are a continuous-time deep learning architecture in which each hidden layer is defined as the solution of an ordinary differential equation.
  arXiv  Detail & Related papers  (2026-05-05T07:29:48Z)
• Fractional Spike Differential Equations Neural Network with Efficient Adjoint Parameters Training [63.3991315762955]
  Spiking Neural Networks (SNNs) draw inspiration from biological neurons to create realistic models for brain-like computation.<n>Most existing SNNs assume a single time constant for neuronal membrane voltage dynamics, modeled by first-order ordinary differential equations (ODEs) with Markovian characteristics.<n>We propose the Fractional SPIKE Differential Equation neural network (fspikeDE), which captures long-term dependencies in membrane voltage and spike trains through fractional-order dynamics.
  arXiv  Detail & Related papers  (2025-07-22T18:20:56Z)
• Physics-informed reduced order model with conditional neural fields [4.5355909674008865]
  This study presents the conditional neural fields for reduced-order modeling (CNF-ROM) framework to approximate solutions of parametrized partial differential equations (PDEs)<n>The approach combines a parametric neural ODE for modeling latent dynamics over time with a decoder that reconstructs PDE solutions from the corresponding latent states.
  arXiv  Detail & Related papers  (2024-12-06T18:04:33Z)
• P-SpikeSSM: Harnessing Probabilistic Spiking State Space Models for Long-Range Dependency Tasks [1.9775291915550175]
  Spiking neural networks (SNNs) are posited as a computationally efficient and biologically plausible alternative to conventional neural architectures.<n>We develop a scalable probabilistic spiking learning framework for long-range dependency tasks.<n>Our models attain state-of-the-art performance among SNN models across diverse long-range dependency tasks.
  arXiv  Detail & Related papers  (2024-06-05T04:23:11Z)
• Efficient and Flexible Neural Network Training through Layer-wise Feedback Propagation [49.44309457870649]
  Layer-wise Feedback feedback (LFP) is a novel training principle for neural network-like predictors.<n>LFP decomposes a reward to individual neurons based on their respective contributions.<n>Our method then implements a greedy reinforcing approach helpful parts of the network and weakening harmful ones.
  arXiv  Detail & Related papers  (2023-08-23T10:48:28Z)
• Continuous time recurrent neural networks: overview and application to forecasting blood glucose in the intensive care unit [56.801856519460465]
  Continuous time autoregressive recurrent neural networks (CTRNNs) are a deep learning model that account for irregular observations. We demonstrate the application of these models to probabilistic forecasting of blood glucose in a critical care setting.
  arXiv  Detail & Related papers  (2023-04-14T09:39:06Z)
• NeuralStagger: Accelerating Physics-constrained Neural PDE Solver with Spatial-temporal Decomposition [67.46012350241969]
  This paper proposes a general acceleration methodology called NeuralStagger. It decomposing the original learning tasks into several coarser-resolution subtasks. We demonstrate the successful application of NeuralStagger on 2D and 3D fluid dynamics simulations.
  arXiv  Detail & Related papers  (2023-02-20T19:36:52Z)
• A memory-efficient neural ODE framework based on high-level adjoint differentiation [4.063868707697316]
  We present a new neural ODE framework, PNODE, based on high-level discrete algorithmic differentiation. We show that PNODE achieves the highest memory efficiency when compared with other reverse-accurate methods.
  arXiv  Detail & Related papers  (2022-06-02T20:46:26Z)
• Training Feedback Spiking Neural Networks by Implicit Differentiation on the Equilibrium State [66.2457134675891]
  Spiking neural networks (SNNs) are brain-inspired models that enable energy-efficient implementation on neuromorphic hardware. Most existing methods imitate the backpropagation framework and feedforward architectures for artificial neural networks. We propose a novel training method that does not rely on the exact reverse of the forward computation.
  arXiv  Detail & Related papers  (2021-09-29T07:46:54Z)
• Multipole Graph Neural Operator for Parametric Partial Differential Equations [57.90284928158383]
  One of the main challenges in using deep learning-based methods for simulating physical systems is formulating physics-based data. We propose a novel multi-level graph neural network framework that captures interaction at all ranges with only linear complexity. Experiments confirm our multi-graph network learns discretization-invariant solution operators to PDEs and can be evaluated in linear time.
  arXiv  Detail & Related papers  (2020-06-16T21:56:22Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.