Agent Attention

此仓库包含 Agent Attention 的官方 PyTorch 代码及预训练模型。

• Agent Attention：关于 Softmax 注意力与线性注意力的融合

简介

注意力模块是 Transformer 中的核心组件。尽管全局注意力机制具有强大的表达能力，但其过高的计算成本限制了它在多种场景中的应用。本文提出了一种全新的注意力范式——Agent Attention，旨在在计算效率与表征能力之间取得良好平衡。具体而言，Agent Attention 以四元组 $(Q, A, K, V)$ 表示，在传统注意力模块中引入了一组额外的代理 token $A$。这些代理 token 首先作为 query token $Q$ 的“代理人”，从 $K$ 和 $V$ 中聚合信息，然后再将这些信息广播回 $Q$。由于代理 token 的数量可以远小于 query token 的数量，因此 Agent Attention 在保持全局上下文建模能力的同时，相比广泛使用的 Softmax 注意力显著提升了效率。有趣的是，我们证明了所提出的 Agent Attention 等价于一种广义的线性注意力形式。因此，Agent Attention 无缝地整合了功能强大的 Softmax 注意力和高效稳定的线性注意力。

动机

(a) 在 Softmax 注意力中，每个 query 都会从所有特征中聚合信息，导致复杂度呈二次方增长。(b) 利用注意力权重之间的冗余性，Agent Attention 使用少量的代理 token 作为 query 的“代理人”，从所有特征中捕捉多样化的语义信息，再将其传递给每个 query。

方法

我们提出的 Agent Attention 及其模块示意图。 (a) Agent Attention 使用代理 token 聚合全局信息，并将其分发到各个图像 token，从而实现了 Softmax 注意力与线性注意力的实际融合。$\rm{\sigma}(\cdot)$ 表示 Softmax 函数。在 (b) 中，我们展示了 Agent Attention 模块的信息流。作为示例，我们通过池化操作获取代理 token。随后，代理 token 用于从 $V$ 中聚合信息，而 $Q$ 则从代理特征中查询信息。此外，还采用了代理偏置和 DWC 来加入位置信息并保持特征多样性。

结果

分类

请前往 agent_transformer 文件夹查看具体文档。

• 不同模型在 ImageNet-1K 数据集上的对比结果。

• ImageNet 上的准确率-运行时间曲线。

• 分辨率提升至 ${256^2, 288^2, 320^2, 352^2, 384^2}$。

下游任务

请前往 detection 和 segmentation 文件夹查看具体文档。

AgentSD

当应用于 Stable Diffusion 时，我们的 Agent Attention 在无需任何额外训练的情况下，能够加速生成过程并显著提升图像质量。请前往 agentsd 文件夹查看具体文档。

• Stable Diffusion、ToMeSD 和我们提出的 AgentSD 的定量结果。

• Stable Diffusion、ToMeSD ($r=0.4$) 和 AgentSD ($r=0.4$) 生成的样本。

待办事项

• 分类
• 分割
• 检测
• 将 Agent Attention 应用于 Stable Diffusion

致谢

我们的代码基于 PVT、Swin Transformer、CSwin Transformer 和 ToMeSD 开发而成。

引用

如果您觉得本仓库对您有所帮助，请考虑引用我们的工作。

@inproceedings{han2024agent,
  title={Agent attention: On the integration of softmax and linear attention},
  author={Han, Dongchen and Ye, Tianzhu and Han, Yizeng and Xia, Zhuofan and Pan, Siyuan and Wan, Pengfei and Song, Shiji and Huang, Gao},
  booktitle={欧洲计算机视觉会议},
  year={2024},
}

联系方式

如有任何问题，请随时联系作者。

Dongchen Han：hdc23@mails.tsinghua.edu.cn

Tianzhu Ye：ytz20@mails.tsinghua.edu.cn

Agent Attention 快速上手指南

Agent Attention 是一种新型注意力机制，旨在平衡 Softmax 注意力强大的表达能力与线性注意力的高效计算成本。通过引入少量的“代理令牌（Agent Tokens）”来聚合全局信息并分发给查询令牌，该模型在保持全局上下文建模能力的同时，显著降低了计算复杂度。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 macOS
• Python: 3.8 或更高版本
• GPU: 支持 CUDA 的 NVIDIA GPU (推荐显存 8GB+)
• 核心依赖:
  • PyTorch >= 1.8.0
  • torchvision
  • timm
  • mmcv / mmdetection / mmsegmentation (如需运行下游任务)

国内加速建议： 推荐使用清华源或阿里源安装 Python 依赖，以加快下载速度。

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torchvision torchaudio

安装步骤

1. 克隆仓库 从 GitHub 克隆官方代码库：

   git clone https://github.com/hanc00l/Agent-Attention.git
   cd Agent-Attention
   

2. 安装依赖 安装项目所需的 Python 包。如果网络较慢，可结合国内镜像源使用：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   注：若 requirements.txt 未包含所有下游任务依赖，请根据具体任务文件夹（如 downstream/detection）中的说明额外安装 mmdet 或 mmseg。

3. 验证安装 确保能够导入核心模块（以分类任务为例）：

   python -c "from agent_transformer import create_model; print('Installation successful')"
   

基本使用

以下是最简单的使用示例，展示如何加载预训练的 Agent Attention 模型并进行推理。

1. 图像分类推理

假设您已下载 ImageNet 预训练权重（位于 agent_transformer 目录下的检查点文件），可以使用以下代码加载模型并处理一张图片：

import torch
from PIL import Image
from torchvision import transforms
from agent_transformer import create_model

# 配置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 创建模型 (以 agent_tiny 为例，可根据需要更换为 agent_small, agent_base 等)
# num_classes 需与预训练权重匹配，ImageNet 通常为 1000
model = create_model('agent_tiny', pretrained=True, num_classes=1000)
model.to(device)
model.eval()

# 预处理图像
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 加载图像
image = Image.open("your_image.jpg").convert("RGB")
input_tensor = transform(image).unsqueeze(0).to(device)

# 推理
with torch.no_grad():
    output = model(input_tensor)
    probabilities = torch.nn.functional.softmax(output[0], dim=0)

print(f"Prediction shape: {probabilities.shape}")

2. 应用于 Stable Diffusion (AgentSD)

Agent Attention 可直接加速 Stable Diffusion 的生成过程并提升质量，无需额外训练。具体用法请参考 agentsd 文件夹中的脚本。基本调用逻辑如下：

# 伪代码示例，详细参数请参阅 agentsd 目录下的文档
from agentsd import AgentSDPipeline

pipe = AgentSDPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", 
    use_agent_attn=True, 
    ratio=0.4  # 代理令牌压缩率
)
pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
image.save("generated_image.png")

提示：更多关于检测（Detection）、分割（Segmentation）及高分辨率训练的详细配置，请分别查阅仓库中 downstream/detection、downstream/segmentation 及 agent_transformer 目录下的具体文档。