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
摘要
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
问问这篇 Paper
智能体会读完全文。
Lune 把这篇 Paper 索引到了每一个公式,引用它的顶会 Paper 也一样。你提问,回答直接引用原文。
它引用的顶会 Paper27
- An Image is Worth 16x16 Words: Transformers for Image Recognition at ScaleAlexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn 等ICLR 2021 · 被引用 21,477 次
- CutMix: Regularization Strategy to Train Strong Classifiers With Localizable FeaturesSangdoo Yun, Dongyoon Han, Sanghyuk Chun, Seong Joon Oh 等ICCV 2019 · 被引用 5,843 次
- Deep Residual Learning in Spiking Neural NetworksWei Fang, Zhaofei Yu, Yanqi Chen, Tiejun Huang 等NeurIPS 2021 · 被引用 857 次
- Incorporating Learnable Membrane Time Constant to Enhance Learning of Spiking Neural NetworksWei Fang, Zhaofei Yu, Yanqi Chen, Timothée Masquelier 等ICCV 2021 · 被引用 731 次
- TrivialAugment: Tuning-free Yet State-of-the-Art Data AugmentationSamuel G. Müller, Frank HutterICCV 2021 · 被引用 384 次
相关 Paper
- Skipper: Enabling efficient SNN training through activation-checkpointing and time-skippingSonali Singh, Anup Sarma, Sen Lu, Abhronil Sengupta 等MICRO 2022 · 被引用 13 次
- Advancing Training Efficiency of Deep Spiking Neural Networks through Rate-based BackpropagationChengting Yu, Lei Liu, Gaoang Wang, Erping Li 等NeurIPS 2024 · 被引用 14 次
- Online Training Through Time for Spiking Neural NetworksMingqing Xiao, Qingyan Meng, Zongpeng Zhang, Di He 等NeurIPS 2022 · 被引用 121 次
- Towards Memory- and Time-Efficient Backpropagation for Training Spiking Neural NetworksQingyan Meng, Mingqing Xiao, Shen Yan, Yisen Wang 等ICCV 2023 · 被引用 84 次
- Online Stabilization of Spiking Neural NetworksYaoyu Zhu, Jianhao Ding, Tiejun Huang, Xiaodong Xie 等ICLR 2024 · 被引用 13 次
