中文

具有$O(L)$训练内存和$O(1)$推理成本的高性能时间可逆脉冲神经网络

神经与进化计算 2024-05-28 v1

摘要

类脑脉冲神经网络(SNN)的多时间步模拟增加了训练期间的内存需求并提高了推理能耗。当前的训练方法无法同时解决训练和推理的困境。本文提出了一种新颖的SNN时间可逆架构(T-RevSNN),通过改变SNN的前向传播来共同应对训练和推理挑战。我们关闭了大多数脉冲神经元的时序动态,并在时间开启的脉冲神经元处设计多级时间可逆交互,从而实现了O(L)O(L)的训练内存。结合时间可逆特性,我们重新设计了SNN的输入编码和网络组织,以实现O(1)O(1)的推理能耗。然后,我们精细调整了基本SNN块的内部单元和残差连接,以确保稀疏时间信息交互的有效性。T-RevSNN在ImageNet上取得了优异的精度,同时内存效率、训练时间加速和推理能效分别显著提高了8.6×8.6\times2.0×2.0\times1.6×1.6\times。这项工作有望打破大规模SNN在保持高性能和低推理能耗的同时显著增加内存成本和训练时间的技术瓶颈。源代码和模型可在https://github.com/BICLab/T-RevSNN获取。

关键词

引用

@article{arxiv.2405.16466,
  title  = {High-Performance Temporal Reversible Spiking Neural Networks with $O(L)$ Training Memory and $O(1)$ Inference Cost},
  author = {JiaKui Hu and Man Yao and Xuerui Qiu and Yuhong Chou and Yuxuan Cai and Ning Qiao and Yonghong Tian and Bo XU and Guoqi Li},
  journal= {arXiv preprint arXiv:2405.16466},
  year   = {2024}
}

备注

Accepted by ICML2024