内存高效的可逆脉冲神经网络
计算机视觉与模式识别
2023-12-14 v1 神经与进化计算
摘要
脉冲神经网络（SNN）由于在神经形态硬件上的高能效，成为人工神经网络（ANN）的潜在竞争者。然而，SNN 在训练过程中会在仿真时间步上展开。因此，SNN 比 ANN 需要更多的内存，这阻碍了更深层 SNN 模型的训练。在本文中，我们提出可逆脉冲神经网络，以减少训练期间中间激活值和膜电位的内存开销。首先，我们沿时间维度扩展可逆架构并提出可逆脉冲块，它可以通过反向过程重构计算图并在前向传播中重新计算所有中间变量。在此基础上，我们将最先进的 SNN 模型应用于可逆变体，即可逆脉冲 ResNet（RevSResNet）和可逆脉冲 Transformer（RevSFormer）。通过在静态和神经形态数据集上的实验，我们证明了我们的可逆 SNN 每张图像的内存开销不会随网络深度增加。在 CIFAR10 和 CIFAR100 数据集上，我们的 RevSResNet-37 和 RevSFormer-4-384 实现了相当的准确率，并且在模型复杂度和参数量大致相同的情况下，每张图像消耗的 GPU 内存分别比其对应模型低 3.79 倍和 3.00 倍。我们相信这项工作可以释放 SNN 训练中的内存限制，为训练超大型和超深层 SNN 铺平道路。代码可在 https://github.com/mi804/RevSNN.git 获取。
引用
@article{arxiv.2312.07922,
title = {Memory-Efficient Reversible Spiking Neural Networks},
author = {Hong Zhang and Yu Zhang},
journal= {arXiv preprint arXiv:2312.07922},
year = {2023}
}
备注
accepted by AAAI2024