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

Abstract

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.