[论文解读] Enhancing the Performance of Transformer-based Spiking Neural Networks by SNN-optimized Downsampling with Precise Gradient Backpropagation

Chenlin Zhou, Han Zhang+6 more|arXiv (Cornell University)|May 10, 2023

Advanced Memory and Neural Computing被引用 5

一句话总结

本文提出 ConvBN-MaxPooling-LIF（CML），一种专为脉冲神经网络（SNN）优化的下采样模块，可在基于 Transformer 的脉冲神经网络中实现精确的梯度反向传播，克服了现有 SNN 中因梯度不精确导致的性能瓶颈。CML 实现了最先进性能，包括在 ImageNet 上达到 77.64% 的 top-1 准确率，以及在 CIFAR10-DVS 上达到 81.4%，分别优于 Spikingformer +1.79% 和 +1.16%。

ABSTRACT

Deep spiking neural networks (SNNs) have drawn much attention in recent years because of their low power consumption, biological rationality and event-driven property. However, state-of-the-art deep SNNs (including Spikformer and Spikingformer) suffer from a critical challenge related to the imprecise gradient backpropagation. This problem arises from the improper design of downsampling modules in these networks, and greatly hampering the overall model performance. In this paper, we propose ConvBN-MaxPooling-LIF (CML), an SNN-optimized downsampling with precise gradient backpropagation. We prove that CML can effectively overcome the imprecision of gradient backpropagation from a theoretical perspective. In addition, we evaluate CML on ImageNet, CIFAR10, CIFAR100, CIFAR10-DVS, DVS128-Gesture datasets, and show state-of-the-art performance on all these datasets with significantly enhanced performances compared with Spikingformer. For instance, our model achieves 77.64 $\%$ on ImageNet, 96.04 $\%$ on CIFAR10, 81.4$\%$ on CIFAR10-DVS, with + 1.79$\%$ on ImageNet, +1.16$\%$ on CIFAR100 compared with Spikingformer.

研究动机与目标

• 为解决基于 Transformer 的脉冲神经网络（SNN）中梯度反向传播不精确这一关键问题，该问题限制了模型性能。
• 设计一种专为 SNN 优化的下采样模块，以在训练过程中实现精确的梯度流动。
• 在不依赖人工神经网络（ANN）蒸馏的前提下，提升可直接训练的 SNN 性能，特别是 Spikformer 和 Spikingformer 模型。
• 在多种数据集上验证所提模块的有效性，包括静态图像和神经形态事件驱动数据。

提出的方法

• 提出新型下采样模块 ConvBN-MaxPooling-LIF（CML），结合卷积、批量归一化、最大池化和 LIF 脉冲神经元。
• 将 CML 集成到基于 Transformer 的 SNN 中，替代传统下采样层，确保与脉冲动力学的梯度兼容性。
• 理论分析证明，CML 通过在脉冲神经元膜电位中保持可微分的激活路径，实现了精确的梯度反向传播。
• 采用替代梯度方法实现 SNN 的端到端训练，CML 确保在反向传播过程中提供稳定且精确的梯度信号。
• 将 CML 应用于 Spikformer 和 Spikingformer 模型，保持与原始模型相同的训练协议和超参数。

实验结果

研究问题

• RQ1 专为 SNN 设计的下采样模块是否能提升梯度精度，从而增强模型性能？
• RQ2 所提出的 CML 模块是否能克服现有基于 Transformer 的 SNN（如 Spikformer 和 Spikingformer）中固有的梯度不精确问题？
• RQ3 CML 在静态图像和神经形态事件驱动数据集上的性能提升程度如何？
• RQ4 CML 是否能在无需依赖预训练 ANN 知识蒸馏的前提下，实现直接训练 SNN 的最先进性能？

主要发现

• 在 ImageNet 上，CML 达到 77.64% 的 top-1 准确率，相比 Spikingformer 提升 +1.79%。
• 在 CIFAR100 上，Spikingformer + CML 模型达到 81.4% 的准确率，相比原始 Spikingformer 提升 +1.16%。
• 在 CIFAR10-DVS 上，Spikingformer + CML 在 16 个时间步下达到 81.4% 的 top-1 准确率，创下直接训练 SNN 的新 SOTA 记录。
• 在 DVS128-Gesture 上，Spikingformer + CML 在 16 个时间步下达到 98.6% 的 top-1 准确率，同样创下直接训练 SNN 的新 SOTA 记录。
• 性能提升在多个数据集中保持一致，证实了 CML 在不同数据模态和网络架构下的泛化能力。