STEMS: Spatial-Temporal Mapping Tool For Spiking Neural Networks

TL;DR

This paper introduces STEMS, a Spatial-Temporal Mapping Exploration Tool for Spiking Neural Networks, to minimize data movement and energy by optimizing intra-layer loop nests and inter-layer schedules across spatial and temporal dimensions. By modeling SNN neuron states and time, STEMS enables hybrid schedules that combine layer fusion and time batching, achieving up to $12\times$ off-chip data-movement reductions and up to $5\times$ energy savings on event-based vision benchmarks, with neuron-state optimization yielding up to $20\times$ memory reductions and improved accuracy. The approach is demonstrated on two benchmarks (SEW-ResNet-18 and RED-LIF) and scaled to SEW-ResNet-152, showing strong data-movement and energy benefits and revealing insights such as early blocks benefiting from neuron-state removal and hybrid schedules being most effective. Overall, STEMS provides a practical framework to deploy energy-efficient SNNs on memory-constrained hardware, enabling scalable, hybrid ANN-SNN models for event-based vision tasks.

Abstract

Spiking Neural Networks (SNNs) are promising bio-inspired third-generation neural networks. Recent research has trained deep SNN models with accuracy on par with Artificial Neural Networks (ANNs). Although the event-driven and sparse nature of SNNs show potential for more energy efficient computation than ANNs, SNN neurons have internal states which evolve over time. Keeping track of SNN states can significantly increase data movement and storage requirements, potentially losing its advantages with respect to ANNs. This paper investigates the energy effects of having neuron states, and how it is influenced by the chosen mapping to realistic hardware architectures with advanced memory hierarchies. Therefore, we develop STEMS, a mapping design space exploration tool for SNNs. STEMS models SNN's stateful behavior and explores intra-layer and inter-layer mapping optimizations to minimize data movement, considering both spatial and temporal SNN dimensions. Using STEMS, we show up to 12x reduction in off-chip data movement and 5x reduction in energy (on top of intra-layer optimizations), on two event-based vision SNN benchmarks. Finally, neuron states may not be needed for all SNN layers. By optimizing neuron states for one of our benchmarks, we show 20x reduction in neuron states and 1.4x better performance without accuracy loss.

Figures (23)

• Figure 1: Leaky Integrate-and-Fire (LIF) neuron. Colors highlight its 3 features: Leakage, integration, and fire-and-reset. The membrane potential ($V_{mem}$), which is its only state, accumulates input current and controls firing.
• Figure 2: Computation graph of one SNN layer during one discrete timestep. The neuron states are used and updated during activation.
• Figure 3: Illustration of (a) Layer-by-layer (LBL) schedule versus (b) layer-fused schedule for a 3-layer ANN block with 1D convolution connection (for simplicity). Each layer has output dimensions of size K x H x W, where K is number of channel and H and W are the height and width respectively. Layer-by-layer schedule can maximize reuse of weights, while layer-fused schedule reduces the size of intermediate features.
• Figure 4: Overview of the STEMS framework.
• Figure 5: Illustration of different inter-layer schedules: (a) single-timestep layer-by-layer (ST-LBL) schedule, (b) time-batched layer-by-layer (TB-LBL) schedule, (c) single-timestep layer-fused (ST-LF) schedule, and (d) time-batched layer-fused (TB-LF) schedule for a 3-layer 3-timestep SNN block with 1D convolution connection (for simplicity). Some arrows are removed for readability. Each layer has output dimensions of size T x K x H x W, where K is number of channel, T is the number timesteps, and H and W are the height and width respectively. Time batching increases reuse of memory over time (temporal memory reuse) at the expense of larger intermediate feature sizes, while layer-fusion decreases reuse of memory in space (spatial memory reuse) but has smaller intermediate feature sizes.