ViT神技 | 谷歌还是出手了！使用ReLU代替Softmax，精度不减，性能起飞（落地有望了）

点击下方卡片，关注「集智书童」公众号

先前的研究观察到，在将注意力Softmax替换为诸如ReLU之类的逐点激活时，准确性会下降。在视觉Transformer的背景下，作者发现通过序列长度进行除法可以缓解这种性能的下降。

作者在ImageNet-21k上训练小到大的视觉Transformer的实验表明，ReLU-Attention可以在计算的函数尺度行为方面接近或匹配传统Softmax-Attention的性能。

1 引言

Transformer架构在现代机器学习中无处不在。作为Transformer的一个核心组件，注意力包括一个Softmax，它产生了对Token的概率分布。由于指数计算和序列长度的总和使得并行化变得具有挑战性，Softmax的计算成本较高。

在本报告中，作者探讨了替代Softmax操作的逐点替代方法，这些方法不一定输出概率分布。值得一提的是，作者观察到，通过序列长度除以ReLU-Attention可以在计算的函数尺度行为方面接近或匹配传统的Softmax注意力。这个结果为并行化提供了新的机会，因为与传统的注意力相比，ReLU注意力可以在序列长度维度上进行更少的聚集操作。

2 相关工作

以前的研究已经探讨了用ReLU或 squared ReLU替代Softmax的方法。然而，这些方法没有按序列长度除法，作者在实验中发现这对于达到与Softmax相媲美的准确性是重要的。

此外，先前的研究[Robust training of neural networks using scale invariant architectures]已经替换了Softmax，同时仍然需要规范化序列长度轴上的注意力权重，以确保它们总和为1。这保留了需要聚集的缺点。

此外，还有大量文献完全删除激活函数，以便注意力是线性的，这对于长序列长度是有用的。在作者的实验中，完全删除激活函数会降低准确性。

3 方法

注意力

注意力通过两个步骤的过程来转换维Query、Key和Value 。首先，通过下式产生注意力权重α

其中ϕ通常是Softmax。接下来，使用注意力权重来计算输出α。本报告探讨了ϕ的逐点替代方法。

ReLU-Attention

作者观察到ϕ ReLU在方程1中是ϕ的一个有希望的替代方法。作者将ϕ ReLU的注意力称为ReLU-Attention。

尺度化逐点注意力

更一般地，作者的实验将探索ϕα，其中α，。

序列长度尺度

作者观察到，通过涉及序列长度 L 的尺度项对于高准确性是有益的。这种缩放在之前的工作中没有出现，这些工作删除了 Softmax。尽管序列长度缩放的主要理由是经验性的，但提供了简要的分析动机。

Transformer当前是设计为Softmax注意力的，对于其中α。这意味着α。虽然这不太可能是一个必要的条件，但ϕ ReLU确保α在初始化时是。保持这个条件可能有助于在替换Softmax时减少其他超参数的改变需求。

在初始化时，q和k的元素为，因此也为。激活函数如ReLU保留了，因此需要因子，α才能为。

4 实验

实验设置

作者的实验使用了BigVision中的ImageNet-21k和ImageNet-1k训练配置，没有修改超参数。在作者的ImageNet-21k实验中，作者训练了30个Epoch，在作者的ImageNet-1k实验中，作者训练了300个Epoch。

因此，两次训练运行使用了大约9e5左右的步数。作者使用了带有qk-layernorm的ViTs，因为在扩展模型尺寸时，先前观察到这是必要的，以防止不稳定性。然而，作者剥离了这不是作者测试的规模的重要组成部分。作者使用i21k和i1k来表示ImageNet-21k和ImageNet-1k，并通过在ImageNet-1k中取前类别中的顶级类别，而没有进行微调，来报告ImageNet-21k模型的ImageNet-1k准确性。

在评估下游任务的迁移性能时，作者使用了在三个种子上平均的10-shot线性探测。下游任务包括Caltech Birds 、Caltech-101、Stanford Cars、CIFAR-100、DTD、ColHsit、Pets和UC Merced。

主要实验

图1说明了ReLU-Attention与ImageNet-21k训练的Softmax注意力的扩展趋势相匹配。在x轴上显示了实验所需的总核心小时数。作为优势，ReLU-Attention使得在序列长度维度上进行并行化的聚合操作少于Softmax注意力。

序列长度尺度效应

图2检验了对Softmax的各种点对点替代方法的序列长度尺度效应。具体来说，作者将Softmax替换为α，其中α，。在x轴上显示α。轴显示了S/32、S/16和S/8视觉Transformer模型的准确性。

通常情况下，当α接近1时，可以获得最佳结果。由于没有明确的最佳非线性，作者在主要实验中使用了ReLU，因为它更快。

qk-layernorm的效果

作者的主要实验使用qk-layernorm，其中在计算注意力权重之前，Query和Key通过LayerNorm传递。作者默认使用qk-layernorm，因为发现在扩展模型尺寸时，它是必要的，以防止不稳定。图3显示了去除qk-layernorm的效果。结果表明，对于这些模型，qk-layernorm的影响不大，但在规模上可能会有所改变。

添加门的效果

以前的工作在删除Softmax时添加了门控单元，并且没有按序列长度缩放。具体来说，在门控注意力单元中，额外的投影产生输出，然后通过逐元素乘法组合在输出投影之前。在图4中，作者调查了是否存在门可以消除对序列长度尺度的需求。

总体上，作者观察到，无论是否存在门，最佳准确性仍然可以通过序列长度缩放来实现。

请注意，对于具有ReLU的S/8模型，门的存在使实验所需的总核心小时数增加了约9.3%。

4 参考

[1]. Replacing softmax with ReLU in Vision Transformers.

5 推荐阅读

中山大学提出MFL-YOLO | 多特征交互损失让YOLOv5再次焕发性能SOTA的光彩！

腾讯实验室提出LATR | 远超PersFormer，成就3D车道线新标杆！

微软亚研院提出DETR最强改进策略 | DETR无痛涨点Tricks集合

扫码加入👉「集智书童」交流群

（备注：方向+学校/公司+昵称）

想要了解更多：

前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉（分割、检测）、多模态、目标跟踪、NerF」

行业技术方案👉「AI安防、AI医疗、AI自动驾驶」AI模型部署落地实战👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」

欢迎扫描上方二维码，加入「集智书童-知识星球」，日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑，期待交流！

免责声明凡本公众号注明“来源：XXX（非集智书童）”的作品，均转载自其它媒体，版权归原作者所有，如有侵权请联系我们删除，谢谢。
点击下方“阅读原文”，了解更多AI学习路上的「武功秘籍」