智能体注意力机制：Softmax与线性注意力的融合

本文提出智能体注意力机制（Agent Attention），通过引入智能体令牌代理查询操作，巧妙融合Softmax与线性注意力，实现高效全局上下文建模。该机制将计算复杂度从二次降低至线性，适用于图像分类、目标检测、语义分割及图像生成等高分辨率任务。实验表明，智能体注意力在多种视觉Transformer模型上显著提升性能，并可直接集成至Stable Diffusion，加速生成过程并提升图像质量。

青橘MATLAB学习

964人浏览 · 2025-03-10 08:59:16

青橘MATLAB学习 · 2025-03-10 08:59:16 发布

摘要
本文提出智能体注意力机制（Agent Attention），通过引入智能体令牌代理查询操作，巧妙融合Softmax与线性注意力，实现高效全局上下文建模。该机制将计算复杂度从二次降低至线性，适用于图像分类、目标检测、语义分割及图像生成等高分辨率任务。实验表明，智能体注意力在多种视觉Transformer模型上显著提升性能，并可直接集成至Stable Diffusion，加速生成过程并提升图像质量。

关键词：智能体注意力、Softmax注意力、线性注意力、视觉Transformer、高分辨率建模、Stable Diffusion

1. 智能体注意力机制：Softmax与线性注意力的高效融合

作者：Dongchen Han*、Tianzhu Ye∗、Yizeng Han、Zhuofan Xia、Shiji Song、Gao Huang†
单位：清华大学自动化系，智能技术与系统国家重点实验室

2. 引言

Transformer模型在计算机视觉领域的广泛应用面临一个核心挑战：Softmax注意力机制的二次计算复杂度。传统全局自注意力虽具备强大的表达能力，但其高昂的计算成本限制了在高分辨率任务中的应用。现有方法（如局部窗口注意力或稀疏注意力）虽降低复杂度，但牺牲了长距离依赖建模能力。

本文提出智能体注意力机制（Agent Attention），通过引入一组**智能体令牌（Agent Tokens）**作为查询代理，将注意力计算分解为“聚合-广播”两步操作，在保留全局感受野的同时实现线性复杂度。实验证明，该机制不仅显著提升视觉任务性能，还可无缝集成至生成模型（如Stable Diffusion），加速推理并提升生成质量。

3. 智能体注意力机制设计

3.1 核心思想

智能体注意力表示为四元组 $(Q, A, K, V)$ ，其中 $A$ 为新增的智能体令牌，其数量远少于查询令牌 $Q$ 。具体流程分为两步：

智能体聚合： $A$ 作为查询，从 $K$ 和 $V$ 中聚合全局信息，生成中间特征 $V_A$ ：
$V_A = \text{Attn}^S(A, K, V) = \sigma(AK^T)V$
智能体广播： $Q$ 作为查询，从 $A$ 和 $V_A$ 中获取信息，生成最终输出：
$O^A = \text{Attn}^S(Q, A, V_A) = \sigma(QA^T)V_A$

等效性证明：上述过程可重写为广义线性注意力形式：
$O^A = \phi_q(Q) \phi_k(K)^T V$
其中 $\phi_q(Q) = \sigma(QA^T)$ ， $\phi_k(K) = (\sigma(AK^T))^T$ ，表明智能体注意力是线性注意力的高阶扩展。
在这里插入图片描述

3.2 优势分析

线性复杂度：计算复杂度从 $O(N^2)$ 降至 $O (N n)$ （ $n$ 为智能体令牌数量）。
保留全局感受野：避免局部窗口或稀疏策略的信息损失。
兼容Softmax与线性注意力：通过两步Softmax操作实现高效近似。

4. 智能体注意力模块优化

在这里插入图片描述

4.1 关键改进

智能体偏差（Agent Bias）：
在注意力计算中引入位置偏差 $B_1$ 和 $B_2$ ，增强空间感知能力：
$O^A = \sigma(QA^T + B_2) \sigma(AK^T + B_1)V$
多样性恢复模块（DWC）：
使用深度卷积（Depthwise Convolution）缓解线性注意力特征多样性不足的问题：
$O^A + \text{DWC}(V)$