SpikingSSMs: Learning Long Sequences with Sparse and Parallel Spiking State Space Models

TL;DR

This paper introduces SpikingSSMs, a novel spiking neural network architecture that effectively learns long sequences by combining state space models with sparse, parallel neuronal dynamics, achieving high efficiency and competitive performance.

Contribution

The work develops a hierarchical integration of neuronal dynamics with state space models and proposes a surrogate network for accelerated training, enabling long sequence learning with sparsity and efficiency.

Findings

Achieves 90% network sparsity on long-range tasks.

Surpasses existing spiking LLMs on WikiText-103 with smaller size.

Maintains competitive performance with state-of-the-art SSMs.

Abstract

Known as low energy consumption networks, spiking neural networks (SNNs) have gained a lot of attention within the past decades. While SNNs are increasing competitive with artificial neural networks (ANNs) for vision tasks, they are rarely used for long sequence tasks, despite their intrinsic temporal dynamics. In this work, we develop spiking state space models (SpikingSSMs) for long sequence learning by leveraging on the sequence learning abilities of state space models (SSMs). Inspired by dendritic neuron structure, we hierarchically integrate neuronal dynamics with the original SSM block, meanwhile realizing sparse synaptic computation. Furthermore, to solve the conflict of event-driven neuronal dynamics with parallel computing, we propose a light-weight surrogate dynamic network which accurately predicts the after-reset membrane potential and compatible to learnable thresholds, enabling orders of acceleration in training speed compared with conventional iterative methods. On the long range arena benchmark task, SpikingSSM achieves competitive performance to state-of-the-art SSMs meanwhile realizing on average 90\% of network sparsity. On language modeling, our network significantly surpasses existing spiking large language models (spikingLLMs) on the WikiText-103 dataset with only a third of the model size, demonstrating its potential as backbone architecture for low computation cost LLMs.