Your Attention Matters: to Improve Model Robustness to Noise and Spurious Correlations

TL;DR

This paper evaluates various self-attention mechanisms in Vision Transformers, demonstrating that Doubly Stochastic attention offers superior robustness to noise and spurious correlations across multiple datasets.

Contribution

It provides a comprehensive comparison of attention variants under data corruption, highlighting Doubly Stochastic attention as the most robust option.

Findings

Doubly Stochastic attention outperforms others by 0.1%-5.1% under corrupted data.

Robustness varies significantly across attention mechanisms.

Results guide better attention mechanism choices for noisy data environments.

Abstract

Self-attention mechanisms are foundational to Transformer architectures, supporting their impressive success in a wide range of tasks. While there are many self-attention variants, their robustness to noise and spurious correlations has not been well studied. This study evaluates Softmax, Sigmoid, Linear, Doubly Stochastic, and Cosine attention within Vision Transformers under different data corruption scenarios. Through testing across the CIFAR-10, CIFAR-100, and Imagenette datasets, we show that Doubly Stochastic attention is the most robust. It consistently outperformed the next best mechanism by $0.1\%-5.1\%$ when training data, or both training and testing data, were corrupted. Our findings inform self-attention selection in contexts with imperfect data. The code used is available at https://github.com/ctamayor/NeurIPS-Robustness-ViT.