ADSEQ: A delay-aware autograd-compatible framework for spike-event delivery in SNNs
This work addresses the need for efficient and general gradient-based training of SNNs with delays, providing a memory-efficient autodifferentiable framework for neuromorphic computing researchers.
The paper introduces ADSEQ, a framework for efficient spike-event delivery in SNNs that supports gradient computation with delays, and benchmarks it across CPU, GPU, TPU, and LPU platforms, finding that queue design significantly impacts performance with different hardware preferring different implementations.
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.