PMSN: A Parallel Multi-compartment Spiking Neuron for Multi-scale Temporal Processing

TL;DR

This work introduces the Parallel Multi-compartment Spiking Neuron (PMSN), a neuron model that splits a single neuron into an adjustable number of interacting compartments to enable robust multi-scale temporal processing. To overcome training bottlenecks, the authors develop two temporal parallelization strategies that linearize the hidden-compartment dynamics and decouple the nonlinear output-reset, allowing efficient GPU-based training. The approach yields strong empirical gains in long-range temporal tasks, competitive static image recognition on ImageNet, and substantial simulation acceleration with manageably increased computation and memory. The combination of dynamic multi-compartment interactions and temporal parallelization offers a scalable pathway for neuromorphic SNNs to handle complex temporal patterns with improved efficiency. The work also provides gradient analysis, ablation studies, and visualization of compartment dynamics to support the interpretability and practicality of PMSN in real-world tasks.

Abstract

Spiking Neural Networks (SNNs) hold great potential to realize brain-inspired, energy-efficient computational systems. However, current SNNs still fall short in terms of multi-scale temporal processing compared to their biological counterparts. This limitation has resulted in poor performance in many pattern recognition tasks with information that varies across different timescales. To address this issue, we put forward a novel spiking neuron model called Parallel Multi-compartment Spiking Neuron (PMSN). The PMSN emulates biological neurons by incorporating multiple interacting substructures and allows for flexible adjustment of the substructure counts to effectively represent temporal information across diverse timescales. Additionally, to address the computational burden associated with the increased complexity of the proposed model, we introduce two parallelization techniques that decouple the temporal dependencies of neuronal updates, enabling parallelized training across different time steps. Our experimental results on a wide range of pattern recognition tasks demonstrate the superiority of PMSN. It outperforms other state-of-the-art spiking neuron models in terms of its temporal processing capacity, training speed, and computation cost. Specifically, compared with the commonly used Leaky Integrate-and-Fire neuron, PMSN offers a simulation acceleration of over 10 $\times$ and a 30 % improvement in accuracy on Sequential CIFAR10 dataset, while maintaining comparable computational cost.

Figures (7)

• Figure 1: Comparison of neuronal structure and dynamics between the popular single-compartment model, biological neurons, and the proposed generalized multi-compartment spiking neuron model. (a, b) The widely used Leaky Integrate-and-Fire model simplifies the biological neurons into a single unit and ignores the interaction among neuronal substructures, resulting in deficiencies in multi-scale temporal processing and slow training speed. (c) In contrast, the detailed morphologies and electrophysical properties of biological neurons improve their computational power in interactive ways, which is crucial for temporal information processing. Drawing inspiration from these, we propose a generalized multi-compartment spiking neuron model that divides a single neuron into a variable number of interconnected subunits with heterogeneous temporal properties and endows them with interactive dynamics. (d) Our proposed Parallel Multi-compartment Spiking Neuron model further extends the compartmental structure in (c). It not only supports rich neuronal dynamics essential for temporal processing, but also facilitates efficient parallel training across time.
• Figure 2: Illustration of the proposed PMSN model and its parallel implementation. (a) The PMSN can be divided into two parts: $n-1$ hidden compartments with membrane potential matrix $V_h$ (blue box), and one output compartment with membrane potential $v_s$ (green box). $I_h$ denotes the total input current to the output compartment. To accelerate the training speed, two temporal parallel strategy are introduced to unfold the recurrent computation within $V_h$ and $v_s$, respectively. (b) The proposed parallel implementation of PMSN. Each bolded symbol represents a set of states over time.
• Figure 3: The learning curves of PMSN models (with and w/o reset), PSN model families, and LIF models (with and w/o reset) on (a) Sequential CIFAR10 and (b) Sequential CIFAR100 tasks. The mean (solid lines) and standard deviations (shaded regions) are derived from three independent runs with different random seeds.
• Figure 4: The visualization of PMSN model dynamics ($n=5$). Left: The impulse response of different compartments within one PMSN neuron, indicating the multi-timescale properties of a single PMSN neuron. Each hidden compartment is characterized by its own dynamic coefficient $\lambda_i=e^{\alpha+\beta\boldsymbol{i}}$, exhibiting damped oscillation patterns after receiving inputs, while the compartment$5$ is responsible for spike generation and reset. Middle: The distribution of oscillation frequencies $\beta/2\pi$, and Right: damping coefficients $\alpha$ for different neurons in one layer, suggesting the population of PMSNs possess neuron-wise specificity to effectively integrate and preserve information across different timescales.
• Figure 5: Comparison of simulation speed ratios $t_i/{t_{PMSN}}$ among serial MSN with identical structure, LIF, PSN, SPSN, and PMSN, where $t_i$ represents the recorded runtime per propagation for model $i$.