[Paper Review] AutoST: Training-free Neural Architecture Search for Spiking Transformers

Ziqing Wang, Qidong Zhao+3 more|arXiv (Cornell University)|Jul 1, 2023

Advanced Memory and Neural ComputingEngineering3 citations

TL;DR

AutoST proposes a training-free neural architecture search (NAS) method for Spiking Transformers that uses Floating-Point Operations (FLOPs) as a performance predictor to bypass the non-differentiability and high sparsity of SNNs. It achieves state-of-the-art accuracy on CIFAR-10, CIFAR-100, and CIFAR10-DVS, outperforming prior manually and automatically designed SNNs by up to 9.10% on benchmark datasets.

ABSTRACT

Spiking Transformers have gained considerable attention because they achieve both the energy efficiency of Spiking Neural Networks (SNNs) and the high capacity of Transformers. However, the existing Spiking Transformer architectures, derived from Artificial Neural Networks (ANNs), exhibit a notable architectural gap, resulting in suboptimal performance compared to their ANN counterparts. Manually discovering optimal architectures is time-consuming. To address these limitations, we introduce AutoST, a training-free NAS method for Spiking Transformers, to rapidly identify high-performance Spiking Transformer architectures. Unlike existing training-free NAS methods, which struggle with the non-differentiability and high sparsity inherent in SNNs, we propose to utilize Floating-Point Operations (FLOPs) as a performance metric, which is independent of model computations and training dynamics, leading to a stronger correlation with performance. Our extensive experiments show that AutoST models outperform state-of-the-art manually or automatically designed SNN architectures on static and neuromorphic datasets. Full code, model, and data are released for reproduction.

Motivation & Objective

• To address the performance gap between Spiking Transformers and their ANN counterparts due to architectural mismatches.
• To overcome the limitations of existing NAS methods in SNNs, which rely on backpropagation and suffer from non-differentiable spikes and high sparsity.
• To develop a training-free NAS framework that rapidly identifies high-performance Spiking Transformer architectures without requiring model training.
• To establish FLOPs as a reliable, training-free performance predictor for SNNs, independent of gradients and sparsity.
• To demonstrate superior generalization and efficiency on both static and neuromorphic datasets through extensive evaluation.

Proposed method

• AutoST employs a search space of Spiking Transformer architectures with configurable depth, width, and embedding dimensions.
• It evaluates candidate architectures using FLOPs as a performance proxy, avoiding backpropagation and activation pattern analysis.
• The method leverages the strong correlation between FLOPs and actual model performance in SNNs, circumventing issues from non-differentiable spikes and sparse activations.
• Architecture candidates are generated and ranked based on FLOPs, with no training or gradient computation required.
• The final architecture is selected based on FLOPs ranking, which correlates strongly with downstream accuracy across datasets.
• The approach is validated through ablation studies and comparisons with state-of-the-art SNN models on CIFAR-10, CIFAR-100, and CIFAR10-DVS.

Experimental results

Research questions

• RQ1 Can FLOPs serve as an effective, training-free performance predictor for Spiking Transformers, despite SNNs' non-differentiable and sparse nature?
• RQ2 Does a training-free NAS approach outperform existing training-based NAS and manual architecture design in Spiking Transformer performance?
• RQ3 How do architectural choices such as depth, width, and timestep affect performance in Spiking Transformers?
• RQ4 What architectural patterns emerge from AutoST search, and how do they differ from conventional SNN designs?
• RQ5 Can AutoST generalize effectively across static and neuromorphic datasets, including those with overfitting challenges?

Key findings

• AutoST models outperform state-of-the-art SNNs on CIFAR-10, achieving 74.54% top-1 accuracy with 34.44M parameters, surpassing Spikformer-10-512 by 0.86%.
• On CIFAR-100, AutoST -base achieves 74.54% accuracy, outperforming the best existing SNN by 4.34%.
• On the neuromorphic CIFAR10-DVS dataset, AutoST -tiny achieves 81.6% accuracy, outperforming the previous SOTA by 7.2% while using fewer parameters.
• The AutoST -small model (14.68M parameters) achieves 71.02% accuracy, exceeding Spikformer-8-384 (16.81M parameters) by 0.78%.
• The search process favors broader, shallower architectures, which mitigate quantization error accumulation and improve accuracy.
• A timestep of 4 is identified as optimal, balancing accuracy and efficiency, with performance increasing up to this point.