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

2024-08-27

arXiv.org

AI Summary

To address the challenge of simultaneously achieving strong temporal modeling capability, biological plausibility, and training efficiency in Spiking Neural Networks (SNNs) for long-sequence tasks, this paper proposes SpikingSSM—a novel spiking state space model. It innovatively integrates hierarchical dendritic neuron dynamics into the state space framework and introduces a lightweight surrogate dynamic network to reconcile event-driven sparsity with parallel gradient-based training, supporting learnable thresholds and predictive membrane potential reset. By unifying SNNs’ sparse spike mechanisms with state space models’ long-range dependency modeling, SpikingSSM enables end-to-end training via surrogate gradients. Experiments demonstrate that it matches mainstream SSMs on the Long Range Arena benchmark while achieving ~90% average sparsity; on WikiText-103, it significantly outperforms existing spiking large models using only one-third the parameters. This work establishes a new paradigm for energy-efficient, long-sequence AI.

Technology Category

• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Spiking Neural Networks
• Long-range Dependencies
• Energy-efficient Computing

Innovation

• SpikingSSMs
• Energy-Efficient Computing
• Long-Chain Information Processing