[论文解读] Masked Spiking Transformer

Ziqing Wang, Yuetong Fang+4 more|arXiv (Cornell University)|Oct 3, 2022

Advanced Memory and Neural Computing被引用 5

一句话总结

本文提出了一种新型SNN架构——掩码脉冲变压器（Masked Spiking Transformer, MST），通过利用人工神经网络到脉冲神经网络（ANN-to-SNN）转换方法，在图像分类任务中实现了最先进性能。通过引入随机脉冲掩码（Random Spike Masking, RSM），该方法在推理过程中剪枝冗余脉冲，使在75%掩码率下功耗降低26.8%的同时保持性能无损，展现出极高的能效比与模型鲁棒性。

ABSTRACT

The combination of Spiking Neural Networks (SNNs) and Transformers has attracted significant attention due to their potential for high energy efficiency and high-performance nature. However, existing works on this topic typically rely on direct training, which can lead to suboptimal performance. To address this issue, we propose to leverage the benefits of the ANN-to-SNN conversion method to combine SNNs and Transformers, resulting in significantly improved performance over existing state-of-the-art SNN models. Furthermore, inspired by the quantal synaptic failures observed in the nervous system, which reduces the number of spikes transmitted across synapses, we introduce a novel Masked Spiking Transformer (MST) framework that incorporates a Random Spike Masking (RSM) method to prune redundant spikes and reduce energy consumption without sacrificing performance. Our experimental results demonstrate that the proposed MST model achieves a significant reduction of 26.8% in power consumption when the masking ratio is 75% while maintaining the same level of performance as the unmasked model.

研究动机与目标

• 弥合SNN与ANN在复杂视觉任务（如ImageNet分类）中的性能差距。
• 在不牺牲准确率的前提下降低SNN的能耗，尤其针对基于Transformer的架构。
• 探索将ANN-to-SNN转换方法应用于SNN中完整自注意力机制的可行性。
• 开发一种通用的脉冲剪枝方法，在保持模型性能的同时减少冗余脉冲传输。

提出的方法

• MST框架采用ANN-to-SNN转换方法，将预训练的人工Transformer转换为脉冲版本，从而保持高准确率。
• 在推理过程中应用随机脉冲掩码（RSM），随机丢弃部分输入脉冲，以减少脉冲数量和能耗。
• RSM方法被应用于Transformer模块中的自注意力与MLP模块，以针对冗余脉冲活动。
• 在Transformer模块中，批量归一化（Batch normalization, BN）层替代层归一化（Layer normalization, LN）层，以提升训练稳定性和与SNN的兼容性。
• 通过调整掩码率（0%、50%、75%）评估能耗节省与性能下降之间的权衡。
• 该方法在多种架构（包括Swin Transformer、ResNet18和VGG16）上进行评估，证明其在Transformer之外也具备良好的泛化能力。

实验结果

研究问题

• RQ1 ANN-to-SNN转换方法能否有效应用于Transformer中的完整自注意力机制，以在SNN中实现高性能？
• RQ2 在SNN中，通过随机掩码剪枝冗余脉冲的程度如何，且不会导致模型准确率下降？
• RQ3 随机脉冲掩码（RSM）方法对不同SNN架构的功耗和推理性能有何影响？
• RQ4 RSM方法是否可泛化至Transformer以外的其他深度脉冲网络（如ResNet和VGG）？
• RQ5 在脉冲Transformer模型中，何种掩码率能最佳平衡能效与模型准确率？

主要发现

• 在CIFAR-10、CIFAR-100和ImageNet数据集上，掩码脉冲变压器（MST）相较于之前SOTA的SNN模型，分别实现了1.21%、7.3%和3.7%的top-1准确率提升。
• 在75%掩码率下，MST模型功耗降低26.8%，同时保持与未掩码模型相同的准确率。
• RSM方法使MST模型的脉冲传输减少高达74%，证明了在无性能损失下可实现显著的能效节省。
• 在ResNet18上，75%掩码率下功耗降低30.1%，但准确率下降3.6%，表明在高掩码率下存在性能权衡。
• MST模型的注意力图在不同掩码率下均保持对物体相关区域的一致关注，证实关键信息得以保留。
• RSM方法可有效泛化至非Transformer类SNN，在ResNet18和VGG16上实现中等至显著的功耗降低。