PRF: Parallel Resonate and Fire Neuron for Long Sequence Learning in Spiking Neural Networks

TL;DR

This work proposes a decoupled reset method for parallel spiking neuron training, and proposes a Parallel Resonate and Fire neuron, which leverages an oscillating membrane potential driven by a resonate mechanism from a differentiable reset function in the complex domain to enable efficient long sequence learning while maintaining parallel training.

Abstract

Recently, there is growing demand for effective and efficient long sequence modeling, with State Space Models (SSMs) proving to be effective for long sequence tasks. To further reduce energy consumption, SSMs can be adapted to Spiking Neural Networks (SNNs) using spiking functions. However, current spiking-formalized SSMs approaches still rely on float-point matrix-vector multiplication during inference, undermining SNNs' energy advantage. In this work, we address the efficiency and performance challenges of long sequence learning in SNNs simultaneously. First, we propose a decoupled reset method for parallel spiking neuron training, reducing the typical Leaky Integrate-and-Fire (LIF) model's training time from $O(L^2)$ to $O(L\log L)$, effectively speeding up the training by $6.57 \times$ to $16.50 \times$ on sequence lengths $1,024$ to $32,768$. To our best knowledge, this is the first time that parallel computation with a reset mechanism is implemented achieving equivalence to its sequential counterpart. Secondly, to capture long-range dependencies, we propose a Parallel Resonate and Fire (PRF) neuron, which leverages an oscillating membrane potential driven by a resonate mechanism from a differentiable reset function in the complex domain. The PRF enables efficient long sequence learning while maintaining parallel training. Finally, we demonstrate that the proposed spike-driven architecture using PRF achieves performance comparable to Structured SSMs (S4), with two orders of magnitude reduction in energy consumption, outperforming Transformer on Long Range Arena tasks.

Figures (12)

• Figure 1: (a) The reset mechanism prevents parallelization as the state update relies on previous spike output, causing $O(L^2)$ timing cost. (b) The proposed decoupled reset mechanism enables parallel computation. (c) Fast decay causes long-range dependencies to vanish in the membrane potential. (d) Slow decay could generate dependency over long sequence, but causes dependency ambiguity. (e) The resonate mechanism with membrane potential helps distinguish relevant inputs.
• Figure 2: Diagram of the SD-TCM.
• Figure 3: (left) Comparison of training runtime for LIF and Parallelized LIF models. (right) Comparison of sequential and parallel training runtime across different categories per batch.
• Figure 4: (left) Ablation Study on sMNIST datasets (Par. means parallel training). (right) Accuracy across different decay on psMNIST.
• Figure 5: Comparison of loss landscapes: (a) Fast decay causes gradient vanishing. (b) Moderate decay improves but keeps optimal local. (c) Further slow decay makes gradient ambiguity, hindering optimization. (d) The resonance creates a smooth gradient field, aiding efficient convergence.

Theorems & Definitions (6)

• Theorem 1
• Theorem 2
• Theorem 3
• proof
• proof
• proof