Bridging the Divide: Reconsidering Softmax and Linear Attention

TL;DR

This paper analyzes the core differences between Softmax and linear attention in Vision Transformers, providing theoretical insights and practical methods to enhance linear attention's performance and scalability.

Contribution

It introduces theoretical analyses of linear attention's limitations and proposes ways to improve its effectiveness, bridging the gap with Softmax attention.

Findings

Linear attention is not injective, leading to semantic confusion.

Local modeling is crucial for Softmax attention's success.

Enhanced linear attention can outperform Softmax in various tasks.

Abstract

Widely adopted in modern Vision Transformer designs, Softmax attention can effectively capture long-range visual information; however, it incurs excessive computational cost when dealing with high-resolution inputs. In contrast, linear attention naturally enjoys linear complexity and has great potential to scale up to higher-resolution images. Nonetheless, the unsatisfactory performance of linear attention greatly limits its practical application in various scenarios. In this paper, we take a step forward to close the gap between the linear and Softmax attention with novel theoretical analyses, which demystify the core factors behind the performance deviations. Specifically, we present two key perspectives to understand and alleviate the limitations of linear attention: the injective property and the local modeling ability. Firstly, we prove that linear attention is not injective, which is prone to assign identical attention weights to different query vectors, thus adding to severe semantic confusion since different queries correspond to the same outputs. Secondly, we confirm that effective local modeling is essential for the success of Softmax attention, in which linear attention falls short. The aforementioned two fundamental differences significantly contribute to the disparities between these two attention paradigms, which is demonstrated by our substantial empirical validation in the paper. In addition, more experiment results indicate that linear attention, as long as endowed with these two properties, can outperform Softmax attention across various tasks while maintaining lower computation complexity. Code is available at https://github.com/LeapLabTHU/InLine.