Towards Lossless Memory-efficient Training of Spiking Neural Networks via Gradient Checkpointing and Spike Compression

Yifan Huang, Wei Fang, Zecheng Hao, Zhengyu Ma, Yonghong Tian

Keywords: Spiking Neural Network, Training Memory Optimization, Gradient Checkpointing

TL;DR: A widely applicable automatic pipeline combining spatio-temporal gradient checkpointing and spike compression that achieves up to 8× memory savings for SNN training while preserving BPTT-level accuracy and speed.

Abstract: Deep spiking neural networks (SNNs) hold immense promise for low-power event-driven computing, but their direct training via backpropagation through time (BPTT) incurs prohibitive memory cost, which limits their scalability. Existing memory-saving approaches, such as online learning, BPTT-to-BP, and reversible networks, compromise accuracy, training speed, or applicability. In this work, we propose a novel and broadly applicable pipeline for memory-efficient SNN training that preserves BPTT's accuracy. Our pipeline integrates layer-wise gradient checkpointing with lossless spike compression to eliminate internal state storage and reduce the memory cost of per-layer input spikes. We also introduce a multi-stage checkpoint adjustment strategy that adaptively refines checkpoint placement based on profiling results to further optimize memory usage and improve training speed. Wrapped in an optimization pass, the pipeline automatically restructures the computation flow before training with minimal user effort. Extensive experiments on diverse architectures and tasks demonstrate up to $8\times$ memory efficiency gains with $\le 20\%$ speed reduction and no accuracy loss. Our method provides a practical solution for efficient and scalable SNN training. Code is available at https://github.com/AllenYolk/snn-gradient-checkpointing.

Primary Area: applications to neuroscience & cognitive science

Submission Number: 7492