A Diagonal State Space Model on Loihi 2 for Efficient Streaming Sequence Processing

Svea Marie Meyer, Philipp Weidel, Plank Philipp, Leobardo Campos-Macias, Sumit Bam Shrestha, Philipp Stratmann, Jonathan Timcheck, Mathis Richter

Keywords: Neuromorphic, State-Space-Models, Hardware-Aware Machine Learning, Loihi 2, S4D

TL;DR: First ever implementation of a structured state-space model on neuromorphic hardware.

Abstract: Deep State Space Models (SSM) demonstrate state-of-the-art performance on long-range sequence modeling tasks. While the recurrent structure of SSMs can be efficiently implemented as a convolution or as a parallel scan during training, recurrent token-by-token processing cannot currently be implemented efficiently on GPUs. Here, we demonstrate efficient token-by-token inference of the SSM S4D on Intel’s Loihi 2 state-of-the-art neuromorphic processor. We compare this first-ever neuromorphic-hardware implementation of an SSM on sMNIST, psMNIST, and sCIFAR to a recurrent and a convolutional implementation of S4D on Jetson Orin Nano (Jetson). While we find Jetson to perform better in an offline sample-by-sample based batched processing mode, Loihi 2 outperforms during token-by-token based processing, where it consumes 1000 times less energy with a 75 times lower latency and a 75 times higher throughput compared to the recurrent implementation of S4D on Jetson. This opens up new avenues towards efficient real-time streaming applications of SSMs.

Submission Number: 37