AT-SNN: Adaptive Tokens for Vision Transformer on Spiking Neural Network
This work addresses energy efficiency for vision tasks using spiking neural networks, offering an incremental improvement by adapting existing adaptive computation techniques to a new architecture.
The paper tackles the problem of high power consumption in spiking neural network-based vision transformers by proposing AT-SNN, which dynamically adjusts the number of tokens during inference to reduce energy use while maintaining accuracy. For example, it achieves up to 42.4% fewer tokens than the best existing method on CIFAR-100 with higher accuracy.
In the training and inference of spiking neural networks (SNNs), direct training and lightweight computation methods have been orthogonally developed, aimed at reducing power consumption. However, only a limited number of approaches have applied these two mechanisms simultaneously and failed to fully leverage the advantages of SNN-based vision transformers (ViTs) since they were originally designed for convolutional neural networks (CNNs). In this paper, we propose AT-SNN designed to dynamically adjust the number of tokens processed during inference in SNN-based ViTs with direct training, wherein power consumption is proportional to the number of tokens. We first demonstrate the applicability of adaptive computation time (ACT), previously limited to RNNs and ViTs, to SNN-based ViTs, enhancing it to discard less informative spatial tokens selectively. Also, we propose a new token-merge mechanism that relies on the similarity of tokens, which further reduces the number of tokens while enhancing accuracy. We implement AT-SNN to Spikformer and show the effectiveness of AT-SNN in achieving high energy efficiency and accuracy compared to state-of-the-art approaches on the image classification tasks, CIFAR10, CIFAR-100, and TinyImageNet. For example, our approach uses up to 42.4% fewer tokens than the existing best-performing method on CIFAR-100, while conserving higher accuracy.