Abstract:The attention module is the key component in Transformers. While the global attention mechanism offers high expressiveness, its excessive computational cost restricts its applicability in various scenarios. In this paper, we propose a novel attention paradigm, Agent Attention, to strike a favorable balance between computational efficiency and representation power. Specifically, the Agent Attention, denoted as a quadruple $(Q, A, K, V)$, introduces an additional set of agent tokens $A$ into the conventional attention module. The agent tokens first act as the agent for the query tokens $Q$ to aggregate information from $K$ and $V$, and then broadcast the information back to $Q$. Given the number of agent tokens can be designed to be much smaller than the number of query tokens, the agent attention is significantly more efficient than the widely adopted Softmax attention, while preserving global context modelling capability. Interestingly, we show that the proposed agent attention is equivalent to a generalized form of linear attention. Therefore, agent attention seamlessly integrates the powerful Softmax attention and the highly efficient linear attention. Extensive experiments demonstrate the effectiveness of agent attention with various vision Transformers and across diverse vision tasks, including image classification, object detection, semantic segmentation and image generation. Notably, agent attention has shown remarkable performance in high-resolution scenarios, owning to its linear attention nature. For instance, when applied to Stable Diffusion, our agent attention accelerates generation and substantially enhances image generation quality without any additional training. Code is available at <a class="link-external link-https" href="https://github.com/LeapLabTHU/Agent-Attention" rel="external noopener nofollow">this https URL</a>.

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

Agent Attention: On the Integration of Softmax and Linear Attention

Bridging the Divide: Reconsidering Softmax and Linear Attention

MultiMax: Sparse and Multi-Modal Attention Learning

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

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

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

Superiority of Softmax: Unveiling the Performance Edge Over Linear Attention

Unraveling the Gradient Descent Dynamics of Transformers

Max-Margin Token Selection in Attention Mechanism

Softmax-free Linear Transformers

Rethinking Softmax: Self-Attention with Polynomial Activations

Angle-based Softmax Loss for Face Verification

Generalized Probabilistic Attention Mechanism in Transformers

Normalized Activation Function: Toward Better Convergence

Tackling Early Sparse Gradients in Softmax Activation Using Leaky Squared Euclidean Distance

Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Transformers

Softmax Attention with Constant Cost per Token

On the Learning Dynamics of Attention Networks

A Regularized Framework for Sparse and Structured Neural Attention