Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

EventQueues: Autodifferentiable spike event queues for brain simulation on AI accelerators

Lennart P. L. Landsmeer, Amirreza Movahedin, Said Hamdioui, Christos Strydis

TL;DR

This work tackles efficient gradient-based training of spiking neural networks with delayed spike delivery on diverse AI accelerators. It derives generalized, autodifferentiable spike event queues and implements multiple data structures with custom gradients, enabling exact gradient propagation through delayed spikes. Through cross-platform benchmarks (CPU, GPU, TPU, Groq LPU), it demonstrates architecture-specific preferences (heaps on von Neumann, sorting on dataflow, ring buffers on GPUs) and highlights trade-offs involving memory capacity and spike dropping. The findings inform both algorithmic design and hardware development, suggesting that future autograd frameworks and specialized delay-enabled hardware can further improve performance and scalability of gradient-enabled SNNs.

Abstract

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 memory-efficient, gradient-enabled event queue structures. 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.

EventQueues: Autodifferentiable spike event queues for brain simulation on AI accelerators

TL;DR

This work tackles efficient gradient-based training of spiking neural networks with delayed spike delivery on diverse AI accelerators. It derives generalized, autodifferentiable spike event queues and implements multiple data structures with custom gradients, enabling exact gradient propagation through delayed spikes. Through cross-platform benchmarks (CPU, GPU, TPU, Groq LPU), it demonstrates architecture-specific preferences (heaps on von Neumann, sorting on dataflow, ring buffers on GPUs) and highlights trade-offs involving memory capacity and spike dropping. The findings inform both algorithmic design and hardware development, suggesting that future autograd frameworks and specialized delay-enabled hardware can further improve performance and scalability of gradient-enabled SNNs.

Abstract

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 memory-efficient, gradient-enabled event queue structures. 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.

Paper Structure

This paper contains 29 sections, 12 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Models of neurons and synapses
  4. Gradients for Spiking Neural Networks
  5. Surrogate gradients
  6. Exact gradients
  7. Existing brain simulators
  8. ML-library programming model
  9. Methods
  10. Generalized Queue Gradients
  11. EventProp
  12. Differentiable Event Queues
  13. Continuous Signal Delays
  14. Evaluation
  15. Queue Implementations
  16. ...and 14 more sections

Figures (4)

  • Figure 1: Study setup. We implement spike event queues that support gradients and benchmark how various queue structures perform on different AI accelerators.
  • Figure 2: Scaling of event queues with batch size (queue count)
  • Figure 3: A) Scaling of event queues with queue capacity. B) Architecture Effect
  • Figure 4: Spike drop rates for different lossy queue implementations given an incoming Poisson process ($\lambda$) and delay d, both in units of timesteps.