Agent Attention: On the Integration of Softmax and Linear Attention
CoRR(2023)
摘要
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 https://github.com/LeapLabTHU/Agent-Attention.
更多查看译文
AI 理解论文
溯源树
样例
生成溯源树,研究论文发展脉络
Chat Paper
正在生成论文摘要