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

The unsustainable rise in energy cost from increasingly capable deep learning systems spurs computer architecture innovation beyond conventional deep learning accelerators such as GPUs. However, a novel computer architecture presents a problem: much of deep learning research has been optimized for conventional computer architectures, and the extent to which modern deep learning models can unlock improved efficiency on a novel computer architecture is not well understood. In this work, we demonstrate for the first time that a State Space Model (SSM) can achieve substantial efficiency improvement when mapped to Loihi 2, a state-of-the-art neuromorphic research chip, versus a Jetson Orin Nano GPU (Jetson). Specifically, we benchmark our SSM on sMNIST, psMNIST, and sCIFAR online token-by-token inference and find approximately 1000x increased energy efficiency and 75x improved latency and throughput on Loihi 2 with a decrease in accuracy of less than one to three percentage points compared to the full precision implementation on Jetson. We comprehensively tailor our implementation to Loihi-specific features and constraints, such as the co-location of memory and compute as well as fixed precision arithmetic. Our results elucidate how SSMs meaningfully bridge conventional and neuromorphic hardware via their dual nature: SSMs can operate in an offline mode using convolution or scan, which is efficient on a GPU, or in an online mode as a recurrent network, which we show is efficient on Loihi 2. This work provides a foundation for performant sequence models on neuromorphic hardware, potentially unlocking substantial improvements in latency-sensitive or energy-limited online inference applications, such as speech enhancement or vision for robotic control.