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 <a class="link-external link-https" href="https://github.com/LeapLabTHU/InLine" rel="external noopener nofollow">this https URL</a>.

MultiMax: Sparse and Multi-Modal Attention Learning

MaxQ: Multi-Axis Query for N:M Sparsity Network

From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification

Adaptive Sparse Softmax: An Effective and Efficient Softmax Variant for Text Classification

Multi-class feature selection via Sparse Softmax with a discriminative regularization

Sparse and Structured Visual Attention

Money on the Table: Statistical information ignored by Softmax can improve classifier accuracy

Exploring the Frontiers of Softmax: Provable Optimization, Applications in Diffusion Model, and Beyond

Sparsity-Regularized Hmax For Visual Recognition

Escaping the Gradient Vanishing: Periodic Alternatives of Softmax in Attention Mechanism

Bridging the Divide: Reconsidering Softmax and Linear Attention

Softmax Dissection: Towards Understanding Intra- and Inter-class Objective for Embedding Learning

Structured Sparse Attention for End-to-end Automatic Speech Recognition.

Attention Scheme Inspired Softmax Regression

Distributed Sparse Hmax Model

A Regularized Framework for Sparse and Structured Neural Attention

ConSmax: Hardware-Friendly Alternative Softmax with Learnable Parameters

The Closeness of In-Context Learning and Weight Shifting for Softmax Regression

Max-Margin Token Selection in Attention Mechanism

Taming the Sigmoid Bottleneck: Provably Argmaxable Sparse Multi-Label Classification

Robust Softmax Regression for Multi-class Classification with Self-Paced Learning.