AttnZero: Efficient Attention Discovery for Vision Transformers¶

会议: ECCV 2024
arXiv: 无
代码: https://github.com/lliai/AttnZero
领域: 模型效率 / 神经架构搜索
关键词: 线性注意力, 注意力机制搜索, Vision Transformer, 进化算法, 注意力基准

一句话总结¶

本文提出 AttnZero，首个自动发现高效注意力模块的框架，通过构建包含六类计算图和丰富算子的搜索空间、利用进化算法进行多目标搜索，自动发现了适用于多种 ViT 的线性注意力公式，在 DeiT/PVT/Swin/CSwin 上分别达到 74.9%/78.1%/82.1%/82.9% 的 ImageNet top-1 准确率，并构建了包含 2000 种注意力变体的 Attn-Bench-101 基准。

研究背景与动机¶

领域现状：Vision Transformer（ViT）已成为计算机视觉领域的主流架构，但其核心的自注意力（self-attention）机制具有 \(O(n^2)\) 的计算复杂度，其中 \(n\) 为 token 数量。这严重限制了 ViT 在高分辨率输入和实时推理场景中的应用。线性注意力（Linear Attention）通过近似 softmax 注意力来实现 \(O(n)\) 复杂度，是一种有前景的替代方案。

现有痛点：现有的线性注意力方法（如 Efficient Attention、FLatten Transformer、ReLU Attention 等）都是手工设计的，存在两个问题：(1) 手工设计依赖研究者的经验和直觉，搜索空间有限，容易遗漏更优的注意力公式；(2) 现有线性注意力相比标准 softmax 注意力通常存在明显的性能下降，尤其在 small-scale 模型上。

核心矛盾：线性注意力在理论上具有 \(O(n)\) 的量级优势，但实践中手工设计的线性注意力公式难以在效率和性能之间取得最佳平衡。巨大的设计空间意味着可能存在更优的公式组合，但人工枚举不可行。

本文目标 (1) 如何系统化地搜索最优的线性注意力公式；(2) 如何保证发现的注意力模块在不同 ViT 架构上都有良好的泛化性；(3) 如何高效地过滤大量不可行的候选方案。

切入角度：作者将注意力公式建模为计算图（computation graph），每个节点是一个操作（激活函数、归一化、二元运算），通过进化算法在计算图空间中搜索最优的线性注意力公式。关键创新是以多架构联合性能作为搜索目标，并使用程序验证加速过滤。

核心 idea：用进化算法在结构化搜索空间中自动发现比手工设计更优的线性注意力模块。

方法详解¶

整体框架¶

AttnZero 的流程分为三个阶段：(1) 搜索空间构建——定义六类计算图模板和丰富的操作符集合，组合整个线性注意力搜索空间；(2) 进化搜索——使用多目标进化算法，以候选注意力在多种 ViT（DeiT、PVT、Swin 等）上的联合精度作为适应度函数进行搜索，同时通过程序检查和拒绝协议快速过滤不可行方案；(3) 验证与迁移——将搜索到的最优注意力模块替换到各种 ViT 架构中进行完整训练和下游任务评估。

关键设计¶

结构化搜索空间（Structured Search Space）:
- 功能：定义线性注意力的所有可能公式组合
- 核心思路：搜索空间包含六种计算图类型，覆盖不同的 \(Q\)、\(K\)、\(V\) 交互模式。每种计算图中，操作节点可选择的算子包括：激活函数类（ReLU、GELU、Sigmoid、ELU+1、Softmax 等）、归一化类（L1-Norm、L2-Norm、LayerNorm、BatchNorm 等）、二元运算类（加法、乘法、逐元素最大值等）。整个搜索空间涵盖了数万种可能的注意力公式。关键约束是所有候选公式必须满足线性复杂度——即可以先计算 \(\phi(K)^T V\)（复杂度 \(O(d^2 n)\)）再与 \(\phi(Q)\) 相乘，而非先计算 \(\phi(Q)\phi(K)^T\)（复杂度 \(O(n^2 d)\)）
- 设计动机：手工设计只能探索极小的子空间，系统化的搜索空间构建确保不会遗漏有潜力的公式组合。六种计算图覆盖了 \(Q\)-\(K\) 先交互、\(K\)-\(V\) 先交互、添加偏置项等多种范式
多目标进化搜索（Multi-Objective Evolutionary Search）:
- 功能：在搜索空间中高效找到最优的注意力公式
- 核心思路：进化算法维护一个候选注意力公式的种群。每一代中：(a) 评估每个候选公式在多种 ViT 架构（DeiT-Tiny、PVT-Tiny 等）上的 proxy 性能（如短训练 epoch 的 ImageNet 精度）作为多目标适应度；(b) 通过交叉操作组合两个高适应度个体的子图产生新个体；(c) 通过变异操作随机替换节点中的操作符。多目标优化确保搜索到的注意力在不同架构上都有效，而非只对某一个架构过拟合
- 设计动机：单目标搜索容易导致搜索到的注意力模块过于特化于某一架构，多目标搜索强制发现具有通用性的注意力设计原则
程序检查与拒绝协议（Program Checking and Rejection Protocols）:
- 功能：快速过滤不可行或明显低质量的候选方案，加速搜索过程
- 核心思路：在对候选公式进行昂贵的 proxy 训练评估之前，先执行一系列廉价的检查：(a) 语法检查——验证计算图是否合法（如类型匹配、维度一致）；(b) 数值稳定性检查——在随机输入上前向传播，检测是否存在 NaN 或数值溢出；(c) 梯度检查——验证反向传播是否能正常计算梯度；(d) 复杂度检查——确认计算复杂度为线性。只有通过所有检查的候选才进入 proxy 评估阶段
- 设计动机：搜索空间中大量组合是不可行的（如数值不稳定、梯度消失等），如果对每个都进行 proxy 训练将极其浪费。预检查可以用几乎零成本过滤掉 90%+ 的无效候选

损失函数 / 训练策略¶

搜索阶段使用 proxy 训练（短 epoch 的 ImageNet 分类训练）评估候选注意力的质量。最终选定的注意力模块（如 Trial-105）被替换到目标 ViT 中进行完整训练。训练策略沿用各 ViT 原始的训练配置（如 DeiT 的 DeiT-III 训练策略、CSwin 的 EMA 训练策略），以保证公平比较。此外，作者还构建了 Attn-Bench-101——包含 2000 种注意力变体的预计算性能库，可直接查表获取某种注意力在各种 ViT 上的性能，无需重新训练。

实验关键数据¶

主实验（ImageNet-1K 分类）¶

模型	注意力类型	Top-1 Acc	相比原始提升
DeiT-Tiny	Softmax（原始）	72.2%	-
DeiT-Tiny	AttnZero	74.9%	+2.7%
PVT-Tiny	Softmax（原始）	75.1%	-
PVT-Tiny	AttnZero	78.1%	+3.0%
Swin-Tiny	Softmax（原始）	81.3%	-
Swin-Tiny	AttnZero	82.1%	+0.8%
CSwin-Tiny	Softmax（原始）	82.7%	-
CSwin-Tiny	AttnZero	82.9%	+0.2%

消融实验¶

配置	关键指标	说明
单目标搜索（仅 DeiT）	DeiT 最优但泛化差	过拟合到单一架构
多目标搜索（DeiT+PVT+Swin）	各架构均衡提升	注意力泛化性更好
无程序检查	搜索时间大幅增加	大量无效候选浪费资源
有程序检查	搜索效率提升 10x+	快速过滤不可行方案

关键发现¶

发现的 AttnZero 注意力模块不仅具有线性复杂度，在小模型上还能超越标准 softmax 注意力，这是非常反直觉的结果
不同 ViT 架构偏好的注意力公式有共性——说明存在普适性更强的注意力设计原则
Attn-Bench-101 的分析揭示了一些有趣的注意力设计洞见，如适当的归一化比激活函数选择更重要

亮点与洞察¶

这是首个将注意力机制设计建模为搜索问题的工作，开辟了自动注意力设计的新方向
多目标搜索策略确保了泛化性，搜索到的注意力模块能"即插即用"到多种 ViT 中
Attn-Bench-101 的构建为后续注意力机制研究提供了宝贵的基准资源
程序检查协议的设计实用且高效，值得其他 NAS 工作借鉴

局限与展望¶

搜索空间虽然丰富但仍是离散的计算图，未覆盖连续化的注意力参数化形式
Proxy 评估使用短 epoch 训练，与完整训练的排序可能存在偏差
目前主要验证了图像分类任务，在检测、分割等密集预测任务上的效果需要更多验证
搜索过程仍需大量 GPU 资源，对于资源受限的研究者可能不太友好
未考虑硬件感知的效率指标（如实际推理延迟），仅使用理论 FLOPs

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首个自动注意力发现框架，搜索空间和搜索策略设计出色
实验充分度: ⭐⭐⭐⭐ 多种 ViT 架构验证，Attn-Bench-101 加分
写作质量: ⭐⭐⭐⭐ 搜索空间描述清晰，实验组织合理
价值: ⭐⭐⭐⭐ 开辟了注意力自动化设计新方向，Attn-Bench 有长期价值