Spiking neural networks (SNNs) are potential competitors to artificial neural networks (ANNs) due to their high energy-efficiency on neuromorphic hardware. However, SNNs are unfolded over simulation time steps during the training process. Thus, SNNs require much more memory than ANNs, which impedes the training of deeper SNN models. In this paper, we propose the reversible spiking neural network to reduce the memory cost of intermediate activations and membrane potentials during training. Firstly, we extend the reversible architecture along temporal dimension and propose the reversible spiking block, which can reconstruct the computational graph and recompute all intermediate variables in forward pass with a reverse process. On this basis, we adopt the state-of-the-art SNN models to the reversible variants, namely reversible spiking ResNet (RevSResNet) and reversible spiking transformer (RevSFormer). Through experiments on static and neuromorphic datasets, we demonstrate that the memory cost per image of our reversible SNNs does not increase with the network depth. On CIFAR10 and CIFAR100 datasets, our RevSResNet37 and RevSFormer-4-384 achieve comparable accuracies and consume 3.79x and 3.00x lower GPU memory per image than their counterparts with roughly identical model complexity and parameters. We believe that this work can unleash the memory constraints in SNN training and pave the way for training extremely large and deep SNNs. The code is available at https://github.com/mi804/RevSNN.git.

翻译：脉冲神经网络（SNN）因其在神经形态硬件上的高能效特性，成为人工神经网络（ANN）的潜在竞争对手。然而，在训练过程中，SNN需要在模拟时间步长上展开，导致其比ANN消耗更多内存，这阻碍了更深层SNN模型的训练。本文提出可逆脉冲神经网络，以降低训练过程中中间激活值和膜电位的内存开销。首先，我们沿时间维度扩展可逆架构，并提出了可逆脉冲模块，该模块可重构计算图，并通过反向过程重计算前向传播中的所有中间变量。在此基础上，我们将最先进的SNN模型改造为可逆变体，即可逆脉冲残差网络（RevSResNet）和可逆脉冲变换器（RevSFormer）。通过在静态和神经形态数据集上的实验，我们证明了所提出的可逆SNN的每图像内存成本不会随网络深度增加。在CIFAR10和CIFAR100数据集上，我们的RevSResNet37和RevSFormer-4-384在模型复杂度与参数量大致相当的情况下，实现了与对照模型相当的分类精度，同时每图像的GPU内存消耗分别降低了3.79倍和3.00倍。我们相信，这项工作能够解除SNN训练中的内存限制，为训练超大规模深层SNN铺平道路。代码已在https://github.com/mi804/RevSNN.git开源。