SPikE-SSM: A Sparse, Precise, and Efficient Spiking State Space Model for Long Sequences Learning

TL;DR

This work proposes a sparse, precise and efficient spiking SSM framework, termed SPikE-SSM, and proposes a novel and concise neuron model incorporating reset-refractory mechanism to leverage the inherent temporal dimension for dynamic computing with biological interpretability.

Abstract

Spiking neural networks (SNNs) provide an energy-efficient solution by utilizing the spike-based and sparse nature of biological systems. Since the advent of Transformers, SNNs have struggled to compete with artificial networks on long sequential tasks, until the recent emergence of state space models (SSMs), which offer superior computational efficiency and modeling capability. However, applying the highly capable SSMs to SNNs for long sequences learning poses three major challenges: (1) The membrane potential is determined by the past spiking history of the neuron, leading to reduced efficiency for sequence modeling in parallel computing scenarios. (2) Complex dynamics of biological spiking neurons are crucial for functionality but challenging to simulate and exploit effectively in large networks. (3) It is arduous to maintain high sparsity while achieving high accuracy for spiking neurons without resorting to dense computing, as utilized in artificial neuron-based SSMs. To address them, we propose a sparse, precise and efficient spiking SSM framework, termed SPikE-SSM. For (1), we propose a boundary compression strategy (PMBC) to accelerate the inference of the spiking neuron model, enabling parallel processing for long sequence learning. For (2), we propose a novel and concise neuron model incorporating reset-refractory mechanism to leverage the inherent temporal dimension for dynamic computing with biological interpretability. For (3), we hierarchically integrate the proposed neuron model to the original SSM block, and enhance the dynamics of SPikE-SSM by incorporating trainable thresholds and refractory magnitudes to balance accuracy and sparsity. Extensive experiments verify the effectiveness and robustness of SPikE-SSM on the long range arena benchmarks and large language dataset WikiText-103, showing the potential of dynamic spiking neurons in efficient long sequence learning.

Figures (5)

• Figure 1: Main ideas of SPikE-SSM for long-sequence modeling. (Left) Overview: A parallel max-min boundary compression (PMBC) strategy is proposed to address ❶ (§ \ref{['chalg1']}); a new refractory neuron model with trainable dynamics is developed to address ❷ (§ \ref{['chalg2']}). We integrate the proposed refractory neuron with soft reset within SSMs to address ❸ (§ \ref{['chalg3']}). (Right) An example showing that the relevant information for the task at hand is often sparse in long-sequence inputs.
• Figure 2: Intuitive execution process of PMBC in Algorithm \ref{['al2']}. ESS means explicit spiking state.
• Figure 3: The SPikE-SSM block. (Left) Forward computation graph of a single SPikE-SSM layer. (Right) Comparison of SSMs. The original SSM outputs floating-point numbers, while SPikE-SSM replaces its non-linearity with the proposed refractory neuron model, which can incorporate higher-level neuronal dynamics for long sequence modeling. $D$, $N$, and $L$ represent the model dimension, SSM hidden dimension, and sequence length, respectively. SAF is the spiking activation function.
• Figure 4: Spiking rate across all layers of SPikE-SSM and SpikingSSM on sCIFAR10 and WikiText-103 datasets. The number following each legend represents the respective average spiking rate.
• Figure 5: Comparison on sCIFAR10 between SPikE-SSM-SRR with fixed $\tau$ and $\tau_r$ and that with trainable $\tau$ and $\tau_r$. $U_{th}$ and $v_{th}$ are both set to $1$ in SPikE-SSM-SRR.