P-SpikeSSM: Harnessing Probabilistic Spiking State Space Models for Long-Range Dependency Tasks

2024-06-05

AI Summary

Spiking neural networks (SNNs) face fundamental challenges in sequence learning: difficulty capturing long-range dependencies, limited representational capacity of conventional leaky integrate-and-fire (LIF) neurons, and computational bottlenecks arising from serial spike generation. To address these, this paper proposes the Probabilistic Spiking State Space Model (PSSM), a scalable architecture grounded in state space modeling. Its key contributions are: (1) SpikeSampler—the first stochastic spiking layer built upon SSMs; (2) a surrogate gradient function explicitly designed for stochastic spiking dynamics; (3) SpikeMixer, a module enhancing population-level communication via cross-neuron interaction; and (4) ClampFuse, a residual-connected spiking fusion mechanism with bounded activation. Evaluated on Long Range Arena, psMNIST, and Speech Commands benchmarks, PSSM achieves state-of-the-art accuracy among SNNs while exhibiting higher spike sparsity, improved computational efficiency, and superior long-range dependency modeling.

Technology Category

• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Machine Learning: Deep Neural Architectures and Foundation Models
• Natural Language Processing: Learning & Optimization for NLP

Application Category

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

Abstract

Spiking neural networks (SNNs) are posited as a computationally efficient and biologically plausible alternative to conventional neural architectures, with their core computational framework primarily using the leaky integrate-and-fire (LIF) neuron model. However, the limited hidden state representation of LIF neurons, characterized by a scalar membrane potential, and sequential spike generation process, poses challenges for effectively developing scalable spiking models to address long-range dependencies in sequence learning tasks. In this study, we develop a scalable probabilistic spiking learning framework for long-range dependency tasks leveraging the fundamentals of state space models. Unlike LIF neurons that rely on the deterministic Heaviside function for a sequential process of spike generation, we introduce a SpikeSampler layer that samples spikes stochastically based on an SSM-based neuronal model while allowing parallel computations. To address non-differentiability of the spiking operation and enable effective training, we also propose a surrogate function tailored for the stochastic nature of the SpikeSampler layer. To enhance inter-neuron communication, we introduce the SpikeMixer block, which integrates spikes from neuron populations in each layer. This is followed by a ClampFuse layer, incorporating a residual connection to capture complex dependencies, enabling scalability of the model. Our models attain state-of-the-art performance among SNN models across diverse long-range dependency tasks, encompassing the Long Range Arena benchmark, permuted sequential MNIST, and the Speech Command dataset and demonstrate sparse spiking pattern highlighting its computational efficiency.

Problem

• Develop scalable spiking models for long-range dependencies
• Introduce stochastic spike sampling for parallel computations
• Propose surrogate functions for effective SNN training

Innovation

• Stochastic SpikeSampler for parallel computations
• Surrogate function for spiking operation training
• SpikeMixer and ClampFuse for inter-neuron communication