TLDR: STAS is a new framework for Spiking Transformers that significantly improves energy efficiency and accuracy. It addresses the limitations of traditional SNNs by co-designing a static input module (I-SPS) to ensure temporal stability and a dynamic computation module (A-SSA) for two-dimensional token pruning. This approach leads to substantial energy savings (up to 45.9%) and enhanced accuracy on image classification tasks like CIFAR-10, CIFAR-100, and ImageNet.
Spiking Neural Networks (SNNs) are gaining attention for their potential to be much more energy-efficient than traditional Artificial Neural Networks (ANNs). However, they often come with a significant drawback: high latency and computational overhead. This is mainly because SNNs operate in a multi-timestep manner, processing information over several discrete steps, which can be slow and resource-intensive.
While various methods have been developed to make SNNs more efficient, these approaches often focus on specific aspects—like reducing spatial redundancy (pruning across layers), temporal redundancy (adjusting timesteps), or architecture-specific optimizations. The problem is that these solutions tend to be fragmented, lacking a unified strategy.
The concept of Adaptive Computation Time (ACT) offers a promising foundation for a unified approach. ACT allows neural networks to dynamically adjust their computational effort based on the input’s complexity, halting processing once a stable representation is achieved. However, applying ACT directly to SNN-based Vision Transformers (ViTs) faces two major hurdles. First, ACT’s core principle relies on temporal similarity, which is often violated in SNNs due to varying spike inputs at each timestep. Second, the static architecture of current SNN-based ViTs isn’t inherently suited for dynamic computation.
To overcome these challenges, researchers have proposed a novel framework called STAS (Spatio-Temporal Adaptive Computation Time for Spiking Transformers). STAS introduces a clever co-design approach, simultaneously optimizing both the static architecture and the dynamic computation policy. This integrated strategy is key to unlocking the full potential of SNNs.
STAS introduces a new architectural component called the Integrated Spike Patch Splitting (I-SPS) module. This module is designed to solve the problem of temporal dissimilarity by creating a unified and consistent input representation from multi-timestep spike signals. By doing so, I-SPS establishes the crucial temporal stability that was previously missing in SNN-based ViTs.
Building on this newfound stability, STAS then employs its Adaptive Spiking Self-Attention (A-SSA) module. The A-SSA module acts as a unified dynamic halting mechanism, performing two-dimensional token pruning. This means it can intelligently decide which parts of the data (tokens) are essential and can halt the processing of less informative tokens across both spatial (different parts of the image) and temporal (different timesteps) dimensions, thereby reducing unnecessary computations.
Also Read:
- Enhancing Spiking Neural Networks with Multiple Interacting Learning Mechanisms
- Enhancing Neural Network Robustness on Analog Compute-In-Memory Hardware with Extended Straight-Through Estimation
The effectiveness of STAS has been rigorously validated on popular image classification datasets, including CIFAR-10, CIFAR-100, and ImageNet, using spiking Transformer architectures like Spikformer and Spikingformer. The results are impressive: STAS significantly reduces energy consumption by up to 45.9% on CIFAR-10, 43.8% on CIFAR-100, and 30.1% on ImageNet. Crucially, it achieves these energy savings while simultaneously improving accuracy over existing state-of-the-art models. This demonstrates that STAS offers a superior trade-off between accuracy and energy efficiency for SNN-based ViTs. For more details, you can read the full research paper here.
