  Despite the success of vision transformers (ViTs), they still suffer from
significant drops in accuracy in the presence of common corruptions, such as
noise or blur. Interestingly, we observe that the attention mechanism of ViTs
tends to rely on few important tokens, a phenomenon we call token overfocusing.
More critically, these tokens are not robust to corruptions, often leading to
highly diverging attention patterns. In this paper, we intend to alleviate this
overfocusing issue and make attention more stable through two general
techniques: First, our Token-aware Average Pooling (TAP) module encourages the
local neighborhood of each token to take part in the attention mechanism.
Specifically, TAP learns average pooling schemes for each token such that the
information of potentially important tokens in the neighborhood can adaptively
be taken into account. Second, we force the output tokens to aggregate
information from a diverse set of input tokens rather than focusing on just a
few by using our Attention Diversification Loss (ADL). We achieve this by
penalizing high cosine similarity between the attention vectors of different
tokens. In experiments, we apply our methods to a wide range of transformer
architectures and improve robustness significantly. For example, we improve
corruption robustness on ImageNet-C by 2.4% while improving accuracy by 0.4%
based on state-of-the-art robust architecture FAN. Also, when fine-tuning on
semantic segmentation tasks, we improve robustness on CityScapes-C by 2.4% and
ACDC by 3.0%. Our code is available at https://github.com/guoyongcs/TAPADL.
