[论文解读] SimA: Simple Softmax-free Attention for Vision Transformers

Soroush Abbasi Koohpayegani, Hamed Pirsiavash|arXiv (Cornell University)|Jun 17, 2022

CCD and CMOS Imaging Sensors被引用 8

一句话总结

该论文提出 SimA，一种用于视觉 Transformer 的无需 Softmax 的注意力机制，通过使用查询和键矩阵的 ℓ₁-归一化替代计算成本高昂的 Softmax，实现在推理时动态计算排序，将复杂度降低至 O(ND) 或 O(D²N)，具体取决于输入分辨率。SimA 在 ImageNet、MS-COCO 和自监督基准上实现了最先进（SOTA）的准确率，同时消除了 exp(.) 操作，在边缘设备上显著提升了效率。

ABSTRACT

Recently, vision transformers have become very popular. However, deploying them in many applications is computationally expensive partly due to the Softmax layer in the attention block. We introduce a simple but effective, Softmax-free attention block, SimA, which normalizes query and key matrices with simple $\ell_1$-norm instead of using Softmax layer. Then, the attention block in SimA is a simple multiplication of three matrices, so SimA can dynamically change the ordering of the computation at the test time to achieve linear computation on the number of tokens or the number of channels. We empirically show that SimA applied to three SOTA variations of transformers, DeiT, XCiT, and CvT, results in on-par accuracy compared to the SOTA models, without any need for Softmax layer. Interestingly, changing SimA from multi-head to single-head has only a small effect on the accuracy, which simplifies the attention block further. The code is available here: https://github.com/UCDvision/sima

研究动机与目标

• 为在边缘设备上部署，降低视觉 Transformer 的计算成本，特别是 Softmax 中昂贵的 exp(.) 操作。
• 在从注意力机制中移除 Softmax 层的同时，保持或提升模型准确率。
• 基于 token 数量（N）和通道数（D），在测试时实现动态计算重排序，以优化效率。
• 探索使用单头注意力和 ReLU 激活函数替代多头和 GELU 激活函数的可行性，且准确率无显著下降。
• 提供一种数值稳定、支持半精度推理、保持排列不变性并保留 token 重要性解释能力的注意力机制。

提出的方法

• 在自注意力机制中，将 Softmax 操作替换为查询和键矩阵的 ℓ₁-归一化，再进行点积运算。
• 通过简单的矩阵乘法计算注意力：O = (Q_normalized) @ (K_normalized.T) @ V，从而无需使用 exp(.) 和 Softmax。
• 在推理时动态选择计算顺序：当 N > D 时采用 (Q @ K.T) @ V，当 D > N 时采用 Q @ (K.T @ V)，以最小化 FLOPs。
• 利用 ℓ₁-归一化确保数值稳定性，支持半精度推理且不会溢出。
• 通过将 token 视为无序集合，保持排列不变性，从而支持多源输入（如多台相机或传感器）。
• 在不修改架构的前提下，将 SimA 模块应用于 DeiT、XCiT 和 CvT 模型，同时保持标准训练协议。

实验结果

研究问题

• RQ1 无 Softmax 的注意力机制能否在视觉 Transformer 中维持最先进性能？
• RQ2 用查询和键的 ℓ₁-归一化替代 Softmax 是否能在降低计算成本的同时保持模型准确率？
• RQ3 基于 N 和 D 在推理时进行动态计算重排序，能否带来显著的效率提升？
• RQ4 在无 Softmax 的情况下，单头注意力配合 ReLU 与多头配合 GELU 相比表现如何？
• RQ5 归一化后查询和键向量的模长能否作为有意义的显著性图，用于模型解释？

主要发现

• SimA 在 ImageNet 分类、MS-COCO 目标检测与分割以及自监督学习基准上，实现了与 SOTA 模型相当的准确率。
• 在 1536×1536 图像上，由于消除了 Softmax 的 exp(.) 开销，SimA 比 DeiT 快 58%。
• 当模型维度为 8192 时，SimA 通过避免 Softmax 计算，相比 XCiT 将推理时间减少了 22%。
• 单头 SimA 达到了与多头 SOTA 模型相当的准确率，表明头数减少带来的性能损失可忽略。
• 用 ReLU 替代 GELU 仅导致准确率轻微下降，表明模型对激活函数变化具有鲁棒性。
• 归一化后查询和键向量的 ℓ₂-范数可生成有意义的显著性图，突出图像中的重要区域，而标准 DeiT 中则呈现平坦分布。