Technical Summary: Parallel Training in Spiking Neural Networks via Dynamic Decay

1. Problem Statement

Spiking Neural Networks (SNNs) offer energy-efficient, event-driven computation suitable for neuromorphic hardware. However, training large-scale SNNs remains a significant challenge, particularly regarding parallel training efficiency on modern GPUs.

The core bottleneck is the reset mechanism inherent in standard spiking neurons (e.g., Leaky Integrate-and-Fire, LIF). When a neuron fires, its membrane potential is reset, creating complex temporal dependencies that prevent the simultaneous computation of all timesteps. Existing approaches to enable parallel training generally fall into two categories, both with limitations:

1. Retaining Reset: Methods that decouple spatial and temporal dependencies (e.g., dropping temporal dependence during backpropagation or using single-step pretraining) often compromise training quality or fail to fully utilize parallelism.
2. Removing Reset: Methods like the Parallel Spiking Neuron (PSN) remove the reset mechanism to allow parallel computation. However, these approaches often:
   • Sacrifice the serial inference capability of vanilla spiking neurons, which is crucial for efficient, stable, and long-sequence extrapolation at test time.
   • Fail to adequately preserve the functional roles of the reset mechanism (introducing nonlinearity and controlling membrane potential), leading to performance ceilings or instability.

The paper argues that a successful parallel spiking neuron must simultaneously support parallel training (for efficiency) and serial inference (for efficiency and extrapolation), while preserving or enhancing the biological functions of the reset mechanism.

2. Methodology: Dynamic Decay Spiking Neuron (DSN)

The authors propose a Dynamic Decay Spiking Neuron (DSN) based on a novel functional analysis of the reset mechanism.

2.1 Functional Analysis of Reset

The paper identifies two primary functions of the reset mechanism in vanilla spiking neurons:

1. Introducing Nonlinearity: The reset creates a nonlinear relationship between the membrane potential and the input history. Without reset, the membrane potential update is a linear combination of inputs.
2. Controlling Membrane Potential: The reset prevents "membrane potential explosion" and ceasless firing by constraining the potential within a suitable range. This is characterized as:
   • $\Delta$-short control: Limiting the influence of a large input to a short time window.
   • Long control: Ensuring the membrane potential remains bounded even under small, accumulating inputs.

The authors argue that existing parallel neurons often fail to replicate these functions effectively or lose the ability to perform serial inference.

2.2 Conditions for Parallel Training and Serial Inference

To design a neuron that supports both modes, the paper establishes three structural conditions:

1. Prefix Summarizability: The output at time $t$ depends only on the prefix of inputs $X_{1:t}$.
2. Online Updatability: The internal state can be updated recurrently using only the previous state and current input (enabling serial inference).
3. Offline Parallelizability: The sequence of states can be computed in parallel given the full input sequence (enabling parallel training).

2.3 The DSN Architecture

The DSN replaces the hard/soft reset mechanism with a dynamic decay factor, $\alpha_t$, which is derived from the input history.

• Removal of Reset: The reset term is removed from the membrane potential update equation.
• Dynamic Decay: The constant decay factor $\beta$ in standard LIF is replaced with a time-varying $\alpha_t$.
  $$H_t = \alpha_t \odot H_{t-1} + (1 - \alpha_t) \odot X_t$$
  Here, $\alpha_t$ is computed via a causal 1D convolution over the recent input history ($X_{t-k+1:t}$), passed through a Sigmoid function to ensure $0 \le \alpha_t \le 1$.
• Spike Firing: The paper utilizes integer-valued spike firing (clipping and rounding the membrane potential) to further optimize training and reduce overhead, though the method supports binary firing as well.

Key Insight: The dynamic decay $\alpha_t$ is proven to introduce nonlinearity and provide flexible, input-dependent control over the membrane potential (both short and long control), effectively replicating and enhancing the functions of the reset mechanism without breaking the conditions for parallel training.

Computation Modes:

• Training: The iterative equation is unrolled into a matrix form (using parallel scan algorithms) to compute all timesteps simultaneously.
• Inference: The model switches to the recursive form (Eq. 9), requiring only minimal state storage, thus maintaining the efficiency of serial inference.

3. Key Contributions

1. A Novel Functional View: The paper provides a systematic analysis of the reset mechanism, identifying its core functions (nonlinearity and potential control) and the limitations of existing parallel approaches that fail to preserve these functions or serial inference capabilities.
2. Design under Insightful Guidelines: The authors propose the Dynamic Decay Spiking Neuron (DSN), which satisfies the three conditions for parallel training and serial inference while implementing the functions of reset more flexibly than hard/soft reset mechanisms.
3. Generality and Performance: The method demonstrates consistent effectiveness across diverse network architectures (CNNs, Transformers, State Space Models) and five task categories (image classification, neuromorphic event processing, time-series forecasting, language modeling, and reinforcement learning).

4. Experimental Results

The authors conducted comprehensive experiments to validate DSN:

• Training Efficiency: On 16k-length sequences, DSN achieved a 25.6$\times$ training speedup over the pioneering Parallel Spiking Neuron (PSN).
• Extrapolation Capability: Models trained on 2k-length sequences successfully performed stable inference on sequences up to 30k length, whereas other parallel methods (like Masked PSN) failed beyond the training length.
• Task Performance:
  • Image Classification: Achieved state-of-the-art (SOTA) results on Sequential CIFAR-10/100 and competitive results on ImageNet and CIFAR10-DVS among parallel spiking neurons.
  • Time-Series Forecasting: Outperformed baseline SNNs and Transformers on traffic flow and solar power datasets across various prediction lengths.
  • Reinforcement Learning: Achieved SOTA performance on IDP-v4, Hopper-v4, and Walker2d-v4 benchmarks with reduced variance compared to existing SNN-based RL methods.
  • Language Modeling: Achieved the lowest perplexity (28.50) on WikiText-103 among spiking language models.
• Energy Consumption: Despite the added complexity of dynamic decay, DSN exhibited a lower spike firing rate than vanilla LIF and existing parallel neurons. Consequently, the estimated total energy consumption was slightly lower than LIF and significantly lower than PSN.

5. Significance and Claims

The paper claims that DSN offers a general and powerful alternative to the reset mechanism, resolving the trade-off between training efficiency and inference capability. By decoupling the need for reset from the mechanism itself and replacing it with a learnable, input-dependent dynamic decay, the authors enable SNNs to scale efficiently on modern GPUs while retaining the biological plausibility and inference efficiency of spiking neurons.

The work positions DSN as a foundational step toward training large-scale, high-performance spiking models that can scale alongside large foundation models, addressing the critical bottleneck of parallel training in the era of large models. The authors emphasize that their approach does not merely approximate existing neurons but enhances the functional capabilities of spiking neurons while adhering to the structural constraints required for both parallel and serial computation.