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

2025-11-11

AI Summary

To address gradient instability in direct training of spiking neural networks (SNNs)—caused by temporal covariate shift (TCS) and learnable neuron thresholds—this paper proposes two core techniques: (1) Membrane Potential Initialization (MP-Init), which calibrates the steady-state distribution of membrane potentials across layers to mitigate TCS; and (2) Threshold-Robust Surrogate Gradient (TrSG), a surrogate gradient formulation that explicitly models threshold learnability during backpropagation to ensure gradient stability. Together, these enable robust end-to-end direct training of SNNs. Evaluated on static datasets (CIFAR-10/100) and neuromorphic benchmarks (DVS128 Gesture, NCALTECH101), our method achieves state-of-the-art classification accuracy, faster convergence, enhanced robustness to hyperparameter variations, and low-latency inference. The implementation is publicly available.

Technology Category

• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Machine Learning: Neuro-Symbolic Learning
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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

Problem

• Addresses temporal covariate shift in spiking neural networks
• Stabilizes gradient flow with learnable neuron thresholds
• Improves direct training stability for energy-efficient AI

Innovation

• Membrane Potential Initialization aligns initial state distribution
• Threshold-robust Surrogate Gradient stabilizes training gradient flow
• Two techniques address temporal covariate shift in SNNs