Spiking Vision Transformer with Saccadic Attention

TL;DR

An innovative Saccadic Spike Self-Attention (SSSA) method is introduced that employs a novel spike distribution-based method to effectively assess the relevance between Query and Key pairs in SNN-based ViTs and develops a SNN-ViT, which achieves state-of-the-art performance with linear computational complexity.

Abstract

The combination of Spiking Neural Networks (SNNs) and Vision Transformers (ViTs) holds potential for achieving both energy efficiency and high performance, particularly suitable for edge vision applications. However, a significant performance gap still exists between SNN-based ViTs and their ANN counterparts. Here, we first analyze why SNN-based ViTs suffer from limited performance and identify a mismatch between the vanilla self-attention mechanism and spatio-temporal spike trains. This mismatch results in degraded spatial relevance and limited temporal interactions. To address these issues, we draw inspiration from biological saccadic attention mechanisms and introduce an innovative Saccadic Spike Self-Attention (SSSA) method. Specifically, in the spatial domain, SSSA employs a novel spike distribution-based method to effectively assess the relevance between Query and Key pairs in SNN-based ViTs. Temporally, SSSA employs a saccadic interaction module that dynamically focuses on selected visual areas at each timestep and significantly enhances whole scene understanding through temporal interactions. Building on the SSSA mechanism, we develop a SNN-based Vision Transformer (SNN-ViT). Extensive experiments across various visual tasks demonstrate that SNN-ViT achieves state-of-the-art performance with linear computational complexity. The effectiveness and efficiency of the SNN-ViT highlight its potential for power-critical edge vision applications.

Figures (7)

• Figure 1: Distribution of magnitudes for Q and K in ViTs within ANNs and SNNs on CIRAF100. In ANNs, Q and K exhibit similar magnitude distributions, whereas in SNNs, the magnitude differences between Q and K are pronounced.
• Figure 2: Comparison of three self-attention computation paradigms. (a) VSA employs floating-point matrix multiplication to assess the spatial correlation between Q and K, resulting in a computational complexity of $\mathcal{O}(N^2D)$. (b) SSA lacks a dedicated temporal interaction module, maintaining the same complexity as VSA. (c) In contrast, STSA introduces global spatial-temporal interactions, increasing the complexity to $\mathcal{O}(T^2N^2D)$.
• Figure 3: Overview of SSSA method. (a) SSSA consisting of two key components: cross-entropy relevance computation and saccadic spiking neurons. The latter outputs spike-driven decisions that mask $V$ in N-dimensional space. (b) training and inference process for saccadic spiking neurons. (c) the structure of the spatial relevance computation based on spike distribution. (d) the structure of SSSA-V2 on spatial relevance computation, significantly reducing computational complexity.
• Figure 4: The overall structure of SNN-ViT, mainly consisting of GL-SPS blocks and SSSA-based transformer blocks. GL-SPS block combines dilated convolution and standard convolution at different scales to facilitate multi-scale feature extraction from images. The SSSA-based block, composed of SSSA methods and Linear layers, achieves lower computational complexity.
• Figure 5: The detection results of SNN-ViT-YOLO on the NWPU-10 dataset are displayed in the first rows. SSSA attention heatmaps are showcased in the second rows.