Parallel Training in Spiking Neural Networks

TL;DR

This work reframes parallel training in Spiking Neural Networks by analyzing the reset mechanism’s functions and proposing a dynamic decay spiking neuron (DSN) that preserves the essential nonlinearity and membrane-potential control while enabling both parallel training and serial inference. DSN replaces the reset with a input-driven dynamic decay computed via a causal convolution, yielding a training-friendly, strictly sequential inference-capable architecture. Empirical results show major training speedups (up to 25.6×), robust extrapolation to long sequences, broad task and architecture generality, and favorable energy characteristics compared to prior parallel spiking neurons. The findings offer a principled path toward scalable, efficient SNNs aligned with the needs of large foundation-model-like systems.

Abstract

The bio-inspired integrate-fire-reset mechanism of spiking neurons constitutes the foundation for efficient processing in Spiking Neural Networks (SNNs). Recent progress in large models demands that spiking neurons support highly parallel computation to scale efficiently on modern GPUs. This work proposes a novel functional perspective that provides general guidance for designing parallel spiking neurons. We argue that the reset mechanism, which induces complex temporal dependencies and hinders parallel training, should be removed. However, any such modification should satisfy two principles: 1) preserving the functions of reset as a core biological mechanism; and 2) enabling parallel training without sacrificing the serial inference ability of spiking neurons, which underpins their efficiency at test time. To this end, we identify the functions of the reset and analyze how to reconcile parallel training with serial inference, upon which we propose a dynamic decay spiking neuron. We conduct comprehensive testing of our method in terms of: 1) Training efficiency and extrapolation capability. On 16k-length sequences, we achieve a 25.6x training speedup over the pioneering parallel spiking neuron, and our models trained on 2k-length can stably perform inference on sequences as long as 30k. 2) Generality. We demonstrate the consistent effectiveness of the proposed method across five task categories (image classification, neuromorphic event processing, time-series forecasting, language modeling, and reinforcement learning), three network architectures (spiking CNN/Transformer/SSMs), and two spike activation modes (spike/integer activation). 3) Energy consumption. The spiking firing of our neuron is lower than that of vanilla and existing parallel spiking neurons.

Figures (7)

• Figure 1: Illustration of a biological neuron (left) and the reset mechanism in neuronal dynamics (right).
• Figure 2: The reset mechanism serves the function of regulating the membrane potential in an input-dependent manner, which can be categorized into $\Delta$-short and long control. (a) $\Delta$-short control. Left: Hard reset enforces short control at $\Delta=1$ level, but it does not allow spatial discriminability between inputs of varying importance. Regardless of how large the membrane potential is at the current timestep, it will be forcibly reset to zero if it exceeds the threshold. Right: Soft reset extends the control duration as the input magnitude increases, which can lead to continuous spike firing and reduced temporal discriminability. That is, compared to hard reset, soft reset allows the input at the current timestep to influence several subsequent timesteps; but it also introduces the challenge of making it difficult to distinguish which inputs across different timesteps are more important. Middle: We therefore seek a balanced mechanism that adaptively determines the duration of membrane potential influence based on the input. (b) Long control. Without reset, even under a relatively small constant input sequence, e.g., {0.5}, the membrane potential of a spiking neuron would continuously accumulate, leading to infinite spike firing.
• Figure 3: Illustration of the three conditions for spiking neurons to achieve parallel training and serial inference. (a) Condition 1: Prefix summarizability. The output can only be determined by inputs from the past (the prefix), and this dependency does not change over time. (b) Condition 2: Online updatability. The internal hidden states can be recurrently updated as the sequential input arrives. (c) Condition 3: Offline parallelizability. For a fixed-length sequence, the output can be computed via parallel computation through time.
• Figure 4: Illustration of the computational process of LIF spiking neuron and reset-free spiking neuron with dynamic decay. (a) LIF neuron. The current input and the membrane potential from last timestep are integrated with a constant decay factor $\beta$. The integrated $H_t$ determines the firing of the spike $S_t$, which in turn decides whether $H_t$ is reset. (b) Dynamic decay. After replacing the reset mechanism with dynamic decay $\alpha_t$, the membrane potential can be computed both serially and in parallel.
• Figure 5: Results of training efficiency (left) and extrapolation (right). Left: Runtime of the forward and backward pass when training on sequences with lengths from 1k to 8k. Right: Perplexity (PPL, the lower the better) of extrapolation on Wikitext-103 with sequence lengths ranging from 1k to 30k, trained on sequences of length 2k.

Theorems & Definitions (16)

• Definition 3.1
• Definition 3.2
• Definition 3.3
• Proposition 4.1
• Proposition A.1
• proof
• Proposition A.2
• proof
• Lemma A.1
• proof