SATA: Spatial Autocorrelation Token Analysis for Enhancing the Robustness of Vision Transformers

TL;DR

SATA presents a training-free robustness enhancement for Vision Transformers by exploiting spatial autocorrelation among token features. It computes a local spatial descriptor via Moran's I, splits and merges tokens in later transformer blocks using a two-set scheme and bipartite matching, and feeds the refined tokens into the FFN. Empirical results on ImageNet-1K and robustness benchmarks achieve state-of-the-art accuracy (e.g., 94.9% top-1) and strong performance on ImageNet-A, ImageNet-R, and ImageNet-C without retraining, while reducing FFN computation. The approach is modular and transferable to pre-trained ViTs and potentially to other transformer-based domains, offering both accuracy gains and improved efficiency.

Abstract

Over the past few years, vision transformers (ViTs) have consistently demonstrated remarkable performance across various visual recognition tasks. However, attempts to enhance their robustness have yielded limited success, mainly focusing on different training strategies, input patch augmentation, or network structural enhancements. These approaches often involve extensive training and fine-tuning, which are time-consuming and resource-intensive. To tackle these obstacles, we introduce a novel approach named Spatial Autocorrelation Token Analysis (SATA). By harnessing spatial relationships between token features, SATA enhances both the representational capacity and robustness of ViT models. This is achieved through the analysis and grouping of tokens according to their spatial autocorrelation scores prior to their input into the Feed-Forward Network (FFN) block of the self-attention mechanism. Importantly, SATA seamlessly integrates into existing pre-trained ViT baselines without requiring retraining or additional fine-tuning, while concurrently improving efficiency by reducing the computational load of the FFN units. Experimental results show that the baseline ViTs enhanced with SATA not only achieve a new state-of-the-art top-1 accuracy on ImageNet-1K image classification (94.9%) but also establish new state-of-the-art performance across multiple robustness benchmarks, including ImageNet-A (top-1=63.6%), ImageNet-R (top-1=79.2%), and ImageNet-C (mCE=13.6%), all without requiring additional training or fine-tuning of baseline models.

Figures (11)

• Figure 1: Visual comparison of class token attention maps and spatial autocorrelation score maps across three layers of Deit-Base/16 pre-trained on ImageNet-1K. The clean image is sourced from the 'goldfinch' class of the ImageNet-1K dataset, and its corresponding corrupted version, with maximum severity (5), is sourced from the ImageNet-C ImageNet-C dataset. $\delta$ represents the cosine similarity between either the attention maps or the spatial autocorrelation score maps of a corrupted image and its corresponding clean version at each block.
• Figure 2: (a) Comparison between conventional ViT block and the augmented ViT with SATA (b) Overall architecture of the proposed SATA module. (c) Bipartite Matching.
• Figure 3: Plotting the variations of $\mu_{\mathbf{s}}$, $|\mathbf{\hat{s}}|$, and the lower and upper bounds across different blocks of the ViT.
• Figure 4: (a) Ablation study on $\gamma$. (b) Ablation study on $\alpha$. We set $\gamma$ and $\alpha$ to 0.7 and 1.0, respectively, for all experiments throughout this paper. Ablation studies are conducted on the SATA-B model using the ImageNet-1K dataset. Dashed lines for both graphs represent the baseline Deit-Base/16 top-1 accuracy.
• Figure 5: (a) Cosine similarity between the clean and corrupted versions of the class token attention map and spatial autocorrelation scores across different blocks of SATA-B. Results are averaged across various types of image corruptions and severity levels on ImageNet-C ImageNet-C. (b) Visualisation of token splitting for a pair of clean and noisy images. Notably, the selected tokens for each set are similar for both clean and corrupted inputs.