[Paper Review] Refiner: Refining Self-attention for Vision Transformers

Daquan Zhou, Yujun Shi+7 more|arXiv (Cornell University)|Jun 7, 2021

Advanced Neural Network Applications55 references41 citations

TL;DR

Refiner directly refines self-attention maps in Vision Transformers by attention expansion and distributed local attention, boosting data efficiency and achieving state-of-the-art results with under 100M parameters.

ABSTRACT

Vision Transformers (ViTs) have shown competitive accuracy in image classification tasks compared with CNNs. Yet, they generally require much more data for model pre-training. Most of recent works thus are dedicated to designing more complex architectures or training methods to address the data-efficiency issue of ViTs. However, few of them explore improving the self-attention mechanism, a key factor distinguishing ViTs from CNNs. Different from existing works, we introduce a conceptually simple scheme, called refiner, to directly refine the self-attention maps of ViTs. Specifically, refiner explores attention expansion that projects the multi-head attention maps to a higher-dimensional space to promote their diversity. Further, refiner applies convolutions to augment local patterns of the attention maps, which we show is equivalent to a distributed local attention features are aggregated locally with learnable kernels and then globally aggregated with self-attention. Extensive experiments demonstrate that refiner works surprisingly well. Significantly, it enables ViTs to achieve 86% top-1 classification accuracy on ImageNet with only 81M parameters.

Motivation & Objective

• Address the data-efficiency gap of Vision Transformers by refining the self-attention mechanism rather than solely changing architecture or training tricks.
• Increase diversity of attention maps and incorporate local patterns to combat over-smoothing in deep ViTs.
• Propose a simple, drop-in module (refiner) that can replace vanilla self-attention in ViT blocks.
• Demonstrate improvements on ImageNet and show generalization to NLP (GLUE) tasks.
• Provide insights into the interplay between global attention and local context in token aggregation.

Proposed method

• Introduce attention expansion: project multi-head attention maps into a higher-dimensional space to effectively increase the number of attention maps without reducing embedding dimensions.
• Use a linear projection W_A to expand A into H' attention maps, where H' > H, then aggregate with the expanded maps and later reduce back to H via a 1x1 projection.
• Apply head-wise spatial convolutions to the expanded attention maps to augment local patterns, yielding a distributed local attention (DLA) mechanism.
• Show that DLA combines global context modeling with local pattern enrichment, mitigating over-smoothing and enhancing token distinguishability.
• Replace vanilla self-attention blocks with the Refiner module to obtain Refined-ViT, a drop-in enhancement for ViT blocks.
• Demonstrate additional reduction of attention maps post-DLA to manage computational cost while preserving accuracy.

Experimental results

Research questions

• RQ1 Can refining self-attention maps through expansion and local pattern augmentation improve ViT data efficiency and accuracy?
• RQ2 Does distributed local attention provide gains over standard self-attention across various ViT architectures?
• RQ3 What is the impact of expanding and then reducing attention maps on model performance and convergence speed?
• RQ4 Do Refiner gains transfer to NLP transformers (e.g., BERT) and other vision-language or NLP benchmarks?

Key findings

• Refiner improves ViT-Base on ImageNet by 1.7% top-1 accuracy under the same training recipe with negligible memory overhead.
• Attention expansion alone improves performance as expansion ratio rises from 1 to 6, with top-1 increasing from 82.3% to 83.0% and faster convergence.
• Distributed Local Attention (DLA) consistently boosts top-1 accuracy by 1.2% to 1.7% across ViT variants, with minimal size increase.
• Refined-ViT-S achieves 83.6% top-1 on ImageNet (25M params), outperforming DeiT-S by 3.7% under the same settings.
• Refined-ViT-M reaches 85.6% top-1 (384-dim, 55M params) on 384 input, surpassing CaiT-S36 by 0.2% with less computation, and Refined-ViT-448 attains 86% with under 100M parameters, setting a new state-of-the-art among such models.
• Applying RFC (receptive field calibration) can further improve ImageNet top-1 by about 0.11% on several SOTA models without fine-tuning, and Refiner benefits extend to NLP (GLUE) tasks, improving average score by about 1% over strong baselines.