STAS: Spatio-Temporal Adaptive Computation Time for Spiking Transformers

Donghwa Kang, Doohyun Kim, Sang-Ki Ko, Jinkyu Lee, Brent ByungHoon Kang, Hyeongboo Baek

Keywords: spiking neural network, vision transformer, adaptive computation time

TL;DR: STAS introduces a novel adaptive halting mechanism for Spiking Transformers that prunes redundant computations across both space (layers) and time (timesteps) to improve energy efficiency and accuracy.

Abstract: Spiking neural networks (SNNs), while energy-efficient, suffer from high latency and computational overhead, and existing dynamic computation methods to address this remain fragmented. While the principles of adaptive computation time (ACT) offer a robust foundation for a unified approach, its application to SNN-based vision Transformers (ViTs) is hindered by two core issues: the violation of its temporal similarity prerequisite and a static architecture fundamentally unsuited for its principles. To address these challenges, we propose STAS (Spatio-Temporal Adaptive computation time for Spiking transformers), a framework that co-designs the static architecture and dynamic computation policy. STAS introduces an integrated spike patch splitting (I-SPS) module to establish temporal stability by creating a unified input representation, thereby solving the architectural problem of temporal dissimilarity. This stability, in turn, allows our adaptive spiking self-attention (A-SSA) module to perform two-dimensional token pruning across both spatial and temporal axes. Implemented on spiking Transformer architectures and validated on CIFAR-10, CIFAR-100, and ImageNet, STAS reduces energy consumption by up to 45.9%, 43.8%, and 30.1%, respectively, while simultaneously improving accuracy over SOTA models.

Primary Area: applications to neuroscience & cognitive science

Submission Number: 8625