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
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× memory efficiency gains with ≤ 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 . * Corresponding authors. RELATED WORK BPTT-BASED SNN DIRECT TRAINING If simulated on discrete time steps, SNNs can be trained end-to-end as binary-activated RNNs through
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 69b8d1d1-fd06-417b-b7a8-1fade65536d9Builds on27
- An Image is Worth 16x16 Words: Transformers for Image Recognition at ScaleAlexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn et al.ICLR 2021 · 21,477 citations
- CutMix: Regularization Strategy to Train Strong Classifiers With Localizable FeaturesSangdoo Yun, Dongyoon Han, Sanghyuk Chun, Seong Joon Oh et al.ICCV 2019 · 5,843 citations
- Deep Residual Learning in Spiking Neural NetworksWei Fang, Zhaofei Yu, Yanqi Chen, Tiejun Huang et al.NeurIPS 2021 · 857 citations
- Incorporating Learnable Membrane Time Constant to Enhance Learning of Spiking Neural NetworksWei Fang, Zhaofei Yu, Yanqi Chen, Timothée Masquelier et al.ICCV 2021 · 731 citations
- TrivialAugment: Tuning-free Yet State-of-the-Art Data AugmentationSamuel G. Müller, Frank HutterICCV 2021 · 384 citations
Related papers
- Skipper: Enabling efficient SNN training through activation-checkpointing and time-skippingSonali Singh, Anup Sarma, Sen Lu, Abhronil Sengupta et al.MICRO 2022 · 13 citations
- Advancing Training Efficiency of Deep Spiking Neural Networks through Rate-based BackpropagationChengting Yu, Lei Liu, Gaoang Wang, Erping Li et al.NeurIPS 2024 · 14 citations
- Online Training Through Time for Spiking Neural NetworksMingqing Xiao, Qingyan Meng, Zongpeng Zhang, Di He et al.NeurIPS 2022 · 121 citations
- Towards Memory- and Time-Efficient Backpropagation for Training Spiking Neural NetworksQingyan Meng, Mingqing Xiao, Shen Yan, Yisen Wang et al.ICCV 2023 · 84 citations
- Online Stabilization of Spiking Neural NetworksYaoyu Zhu, Jianhao Ding, Tiejun Huang, Xiaodong Xie et al.ICLR 2024 · 13 citations
