Refining Attention Maps for Vision Transformer

Self-attention layers extract long-range dependencies from long sequences of input tokens via query-key interactions. Although self-attention acquires high flexibility in modeling visual data, the inductive biases of the input images are weakly captured by the query-key interactions. This requires a lot of data to compensate for the lack of inductive biases. This paper addresses the above issue by refining query-key interactions of the attention maps. The better interaction patterns are learned and structured through convolution built on the top of query-key matrix multiplication. Each query in the proposed attention attends to all the spatial locations to generate attention maps. Convolution refines the structure of attention maps, learning the local contexts of the input images. Both global and local contexts of the visual data are aggregated by the interactions between structured attention maps and value tokens. Based on the obtained refining attention map designs, this paper introduces a new Vision Transformer, dubbed RAM. To verify the efficacy of the proposed RAM, extensive experiments are conducted and evaluated on ImageNet-1K image classification task. RAM achieves superior performances than previous methods across visual tasks with faster speeds on CPU and GPU devices. Typically, RAM outperforms efficient MobileOne-S0 by 1.9% Top-1 accuracy while running 3.3⋉, 1.8⋉ faster on GPU and CPU devices.