ST-BPTT: A Memory-efficient BPTT SNN Training Approach through Gradient-Contribution-Driven Time-Step Selection

Spiking Neural Networks (SNNs) offer the advantage of mimicking the brain’s asynchronous and event-driven computation, enabling energy efficiency and robustness in information processing. Back-Propagation Through Time (BPTT) is a training approach for SNNs, which unfolds the computational graph across both time-step and space directions, and it is widely recognized for its high-accuracy training capabilities. However, the BPTT training approach encounters a significant challenge stemming from its high hardware memory consumption, necessitated by the storage of data from all time-steps. We introduce STBPTT, a memory-efficient BPTT-based SNN training approach that selects significant time-steps for gradient back-propagation computation. A gradient-contribution-driven strategy is proposed to help select proper time-steps. We have developed a PyTorch plugin compatible with SNN training frameworks to realize ST-BPTT. Evaluated across several SNN networks and diverse datasets, ST-BPTT can reduce memory consumption by up to $6.6 \times$ compared to the baseline, and by $1.9 \times$ compared to state-of-the-art BPTT-based training methods, with minimal impact on accuracy.