[Paper Review] ReViT: Enhancing Vision Transformers Feature Diversity with Attention Residual Connections

Anxhelo Diko, Danilo Avola+2 more|arXiv (Cornell University)|Feb 17, 2024

CCD and CMOS Imaging Sensors4 citations

TL;DR

ReViT proposes a residual attention mechanism that enhances Vision Transformers by preserving low-level visual features through attention residual connections, mitigating feature collapse in deeper layers. The method improves performance across multiple image classification and object detection benchmarks by balancing local and global attention, achieving state-of-the-art results on ImageNet1k, CIFAR100, and COCO2017.

ABSTRACT

Vision Transformer (ViT) self-attention mechanism is characterized by feature collapse in deeper layers, resulting in the vanishing of low-level visual features. However, such features can be helpful to accurately represent and identify elements within an image and increase the accuracy and robustness of vision-based recognition systems. Following this rationale, we propose a novel residual attention learning method for improving ViT-based architectures, increasing their visual feature diversity and model robustness. In this way, the proposed network can capture and preserve significant low-level features, providing more details about the elements within the scene being analyzed. The effectiveness and robustness of the presented method are evaluated on five image classification benchmarks, including ImageNet1k, CIFAR10, CIFAR100, Oxford Flowers-102, and Oxford-IIIT Pet, achieving improved performances. Additionally, experiments on the COCO2017 dataset show that the devised approach discovers and incorporates semantic and spatial relationships for object detection and instance segmentation when implemented into spatial-aware transformer models.

Motivation & Objective

• To address feature collapse in deep Vision Transformers, where low-level visual features are lost due to over-globalization of self-attention.
• To improve model robustness and feature diversity by integrating residual attention connections across multi-head self-attention layers.
• To enhance performance in visual recognition tasks such as image classification and object detection by preserving local and semantic details.
• To evaluate the effectiveness of the residual attention mechanism across diverse datasets with varying complexity and class distribution.

Proposed method

• Introduces a residual attention learning mechanism that combines current and past self-attention outputs via a learnable gating variable α.
• Applies residual connections to the query and key matrices in Multi-Head Self-Attention (MHSA) layers, enabling feature propagation from shallow to deep layers.
• Uses a weighted fusion of past and current attention maps: α × (past attention) + (1−α) × (current attention), with α learned during training.
• Employs GradCAM visualization to demonstrate that ReViT preserves more localized and diverse features compared to standard ViT.
• Applies the residual module to both standard and spatial-aware transformer architectures for object detection and instance segmentation on COCO2017.
• Conducts ablation studies with varying α values to analyze the trade-off between local and global feature learning.

Figure 1: Illustration of feature maps learned from ReViT and ViT obtained using GradCAM [ 31 ] algorithm.

Experimental results

Research questions

• RQ1 How does residual attention learning affect feature diversity and representation quality in Vision Transformers?
• RQ2 What is the optimal balance between past and current attention information (via α) for maximizing performance across different datasets?
• RQ3 To what extent does ReViT improve performance on image classification benchmarks such as ImageNet1k, CIFAR10, and CIFAR100?
• RQ4 Can ReViT enhance object detection and instance segmentation by capturing semantic and spatial relationships in spatial-aware transformer models?
• RQ5 How does ReViT compare to standard ViT in terms of robustness to object position variations and inductive bias limitations?

Key findings

• ReViT achieves state-of-the-art performance on ImageNet1k, outperforming standard ViT and other SOTA models in accuracy and robustness.
• On CIFAR100 and Oxford Flowers-102, ReViT achieves peak performance at α = 0.75, with a sharp decline in accuracy as α approaches 1, indicating the importance of past attention.
• For the Oxford-IIIT Pet dataset, ReViT shows significant performance gains due to better handling of intra-class and inter-class variations when α is optimally tuned.
• On COCO2017, ReViT enhances object detection and instance segmentation by discovering and incorporating semantic and spatial relationships in spatial-aware transformers.
• The ablation study confirms that α = 0 (relying only on first-layer attention) leads to performance collapse, especially on complex datasets, highlighting the necessity of residual attention.
• GradCAM visualizations confirm that ReViT preserves more localized and diverse features than standard ViT, particularly in deeper layers, reducing feature collapse.

Figure 2: Illustration of the standard self-attention (a) and the proposed mechanism with residual attention module (b).