Stabilizing Direct Training of Spiking Neural Networks: Membrane Potential Initialization and Threshold-robust Surrogate Gradient

TL;DR

This paper introduces two techniques, MP-Init and TrSG, to stabilize and improve the direct training of Spiking Neural Networks, achieving state-of-the-art accuracy on various datasets.

Contribution

The paper proposes novel initialization and gradient stabilization methods specifically designed for training SNNs effectively.

Findings

Achieved state-of-the-art accuracy on static image datasets.

Effectively mitigated temporal covariate shift in SNN training.

Stabilized gradient flow with respect to neuron thresholds.

Abstract

Recent advancements in the direct training of Spiking Neural Networks (SNNs) have demonstrated high-quality outputs even at early timesteps, paving the way for novel energy-efficient AI paradigms. However, the inherent non-linearity and temporal dependencies in SNNs introduce persistent challenges, such as temporal covariate shift (TCS) and unstable gradient flow with learnable neuron thresholds. In this paper, we present two key innovations: MP-Init (Membrane Potential Initialization) and TrSG (Threshold-robust Surrogate Gradient). MP-Init addresses TCS by aligning the initial membrane potential with its stationary distribution, while TrSG stabilizes gradient flow with respect to threshold voltage during training. Extensive experiments validate our approach, achieving state-of-the-art accuracy on both static and dynamic image datasets. The code is available at: https://github.com/kookhh0827/SNN-MP-Init-TRSG