Improving Vision Transformers by Overlapping Heads in Multi-Head Self-Attention

TL;DR

This paper introduces Multi-Overlapped-Head Self-Attention (MOHSA), a novel approach that overlaps attention heads in Vision Transformers to improve their performance across multiple datasets.

Contribution

The paper proposes MOHSA, a new head-overlapping technique for Vision Transformers, demonstrating significant performance improvements over standard multi-head self-attention.

Findings

MOHSA outperforms standard MHSA on multiple benchmarks.

Overlapping heads with adjacent ones enhances feature learning.

Optimal overlapping ratios vary across models and datasets.

Abstract

Vision Transformers have made remarkable progress in recent years, achieving state-of-the-art performance in most vision tasks. A key component of this success is due to the introduction of the Multi-Head Self-Attention (MHSA) module, which enables each head to learn different representations by applying the attention mechanism independently. In this paper, we empirically demonstrate that Vision Transformers can be further enhanced by overlapping the heads in MHSA. We introduce Multi-Overlapped-Head Self-Attention (MOHSA), where heads are overlapped with their two adjacent heads for queries, keys, and values, while zero-padding is employed for the first and last heads, which have only one neighboring head. Various paradigms for overlapping ratios are proposed to fully investigate the optimal performance of our approach. The proposed approach is evaluated using five Transformer models on four benchmark datasets and yields a significant performance boost. The source code will be made publicly available upon publication.