Spiking Neural Networks with Dynamic Time Steps for Vision Transformers

TL;DR

This paper introduces a novel training framework for Spiking Neural Networks with dynamic time steps tailored for Vision Transformers, significantly improving energy efficiency and maintaining high accuracy on image recognition benchmarks.

Contribution

It proposes a trainable, dynamic allocation of time steps per ViT module, enabling SNNs to operate efficiently with fewer time steps and high activation sparsity.

Findings

Achieved 95.97% accuracy on CIFAR10 with 4.97 time steps.

Reduced energy consumption by using only accumulate operations.

Demonstrated effectiveness on CIFAR10, CIFAR100, and ImageNet datasets.

Abstract

Spiking Neural Networks (SNNs) have emerged as a popular spatio-temporal computing paradigm for complex vision tasks. Recently proposed SNN training algorithms have significantly reduced the number of time steps (down to 1) for improved latency and energy efficiency, however, they target only convolutional neural networks (CNN). These algorithms, when applied on the recently spotlighted vision transformers (ViT), either require a large number of time steps or fail to converge. Based on analysis of the histograms of the ANN and SNN activation maps, we hypothesize that each ViT block has a different sensitivity to the number of time steps. We propose a novel training framework that dynamically allocates the number of time steps to each ViT module depending on a trainable score assigned to each timestep. In particular, we generate a scalar binary time step mask that filters spikes emitted by each neuron in a leaky-integrate-and-fire (LIF) layer. The resulting SNNs have high activation sparsity and require only accumulate operations (AC), except for the input embedding layer, in contrast to expensive multiply-and-accumulates (MAC) needed in traditional ViTs. This yields significant improvements in energy efficiency. We evaluate our training framework and resulting SNNs on image recognition tasks including CIFAR10, CIFAR100, and ImageNet with different ViT architectures. We obtain a test accuracy of 95.97% with 4.97 time steps with direct encoding on CIFAR10.