Google DeepMind｜在视觉Transformer中将Softmax替换为ReLU

​

Replacing softmax with ReLU in Vision Transformers

M Wortsman, J Lee, J Gilmer, S Kornblith
[Google DeepMind]

在视觉Transformer中将Softmax替换为ReLU

• 探索用点积激活函数(如ReLU)替换transformer注意力机制中的softmax，之前的工作发现，这种改变会降低准确率。
• 关键发现是，通过序列长度除以点积激活函数(如ReLU)，可以减轻与softmax注意力相比的准确率损失。
• 在ImageNet上的视觉transformer实验表明，除以序列长度的ReLU注意力可以达到softmax注意力的准确率。
• ReLU注意力允许在序列长度上并行化， Gather操作较少于softmax注意力。
• 文章假设序列长度缩放可以维持期望的注意力权重量级，在初始化时，期望的softmax权重量级为O(1/L)。

动机： 之前的研究发现，在将注意力softmax替换为ReLU等点对点激活时，准确率会降低。在视觉Transformer中，本文发现通过除以序列长度可以减轻这种降级。实验表明，在ImageNet-21k上训练小到大的视觉Transformer时，ReLU-attention可以在计算方面达到或接近softmax-attention的性能。 方法：探索了替代softmax操作的点对点方法，这些方法不一定输出概率分布。特别地，观察到通过序列长度除以ReLU的注意力可以在计算方面达到或接近传统softmax注意力的性能。这为并行化提供了新的机会，因为相比传统注意力，ReLU-attention在序列长度维度上可以进行更少的gather操作。 优势：提出了一种替代softmax的点对点注意力机制，通过除以序列长度可以减轻准确率降低的问题，并提供了更好的计算性能和并行化机会。

提出一种用ReLU替换softmax的点对点注意力机制，通过除以序列长度可以减轻准确率降低的问题，提供了更好的计算性能和并行化机会。

https://arxiv.org/abs/2309.08586 

​Read More