Transformer可以不需要Softmax？Kakao提出了UFO-ViT，性能高，计算量还小

写在前面

视觉Transformer已成为计算机视觉任务的重要模型之一。虽然它们优于之前的卷积神经网络，但其复杂度是与输入特征长度N呈二次关系。

为了解决这个问题，在本文中，作者提出了UFO-ViT(Unit Force Operated Vision Trnasformer) ，通过消除非线性来减少自注意力操作的计算量。通过只修改Self-Attention的几行代码，UFO-ViT能够实现线性的复杂度而不降低性能。在图像分类和密集预测任务上，本文提出的UFO-ViT高于大多数基于Transformer的模型。

▊ 1. 论文和代码地址

UFO-ViT: High Performance Linear Vision Transformer without Softmax

论文地址：https://arxiv.org/abs/2109.14382
代码地址：未开源
复现地址：https://github.com/xmu-xiaoma666/External-Attention-pytorch#30-UFO-Attention-Usage

2. Motivation

最近，Transformer的一些研究成果在视觉任务和自然语言处理任务方面都有了许多突破的进展。通过消除基于CNN架构的归纳偏置，视觉Transformer在大数据集上显示出更好的可扩展性。在最近的研究中，基于Transformer的结构刷新了图像分类、目标检测和语义分割的SOTA性能。

基于Transformer的模型已经显示出与早期的SOTA模型相比有竞争力的性能。但是，尽管Transformer取得了巨大的成功，但使用自注意力的模式有众所周知的缺点。一是自注意力机制的时间和显存复杂度与输入长度N成二次关系。 计算自注意力时，与value矩阵相乘，提取两两之间的全局关系。

对于需要高分辨率的任务，例如目标检测或分割，这就是一个很关键的问题。如果宽度和高度加倍，那么自注意需要16倍的资源来计算。另一个主要问题是数据效率 。ViT必须使用大量的额外数据集进行预训练，以达到与基于CNN结构类似的性能。

在本文中，作者提出了采用简单的自注意机制来避免这两个缺陷，该方法为UFO-ViT(Unit Force Operated Vision Trnasformer) ，利用矩阵乘法的结合律，消除了softmax函数。它不是近似，而是类似于一种低秩近似的方法。早期在Performer和Efficient Attention中出现过。

在本文中里，作者使用另一种机制，即XNorm ，通过两步L2归一化来删除softmax函数。利用上述结合律，可以计算线性复杂度的自注意力。与以前的线性逼近方法不同，在多个视觉任务的基准上，与SOTA模型相比，UFO-ViT取得了更高或有竞争力的结果。

3. 方法

本文模型的结构如上图所示。它是一个卷积层，UFO模块和MLP层的混合结构。在本节中，作者将阐述如何取代softmax函数并确保线性的复杂度 。

3.1. Theoretical Interpretation

对于一个输入，传统的自注意机制表示为：

其中，代表Attention操作。

由于Softmax的非线性，不能被分解为。本文的方法利用了结合律消除softmax来先计算。由于应用identity反而会导致退化，所以作者使用了简单的约束来避免这个问题。

本文方法称为交叉归一化（cross-normalization，XNorm） ，计算如下：

其中，为可学习的参数，h为嵌入维数。XNorm在两个维度上操作：即的空间维度和的通道维度，因此被称为交叉归一化 。

使用关联法则，key和value先相乘，query再相乘。如上图所示，这两个乘法的复杂度都是，所以这个过程的计算量是线性的。(下表展示了本文方法与其他方法复杂度的比较)

3.2. XNorm

Replace softmax to XNorm

在本文的方法中，key和value直接相乘，会生成h个聚类：

XNorm直接用于输出和query，如下所示：

其中，代表输入。最后，进行投影和加权求和：

在这个公式中，关系特征由patch和聚类之间的余弦相似度定义。XNorm将query和聚类限制为单位向量，原因是它通过归一化以缩小范围，可以防止它们的value抑制关系属性。然而，这种解释不足以解释为什么XNorm必须采用L2归一化形式。因此，作者又给出了另一种简单的物理解释。

Details of XNorm

带残差连接任意模块的输出公式为：

其中，n和x表示当前深度和输入图像。假设n和x是某物体的位移和时间，则上式可以重新定义为:

大多数神经网络是离散的，所以是常数。为简单起见，设，残差项表示速度。

在物理学中，胡克定律被定义为弹性矢量与位移矢量的点积。弹性力产生谐波势（上图产生谐波势的质心的示意图），这是关于的函数。从物理上讲，势能干扰了粒子通过它的运动路径。

一般可以用上述公式来近似在处的分子势能。对于多分子，可以利用弹性的线性：

假设少量的k值非常大，它的影响就会很大。如果某个粒子绕着它们运动，就不容易逃脱。在这种情况下，就会出现特征崩溃 或者关系特征被忽略 的现象。因此，本文中的XNorm强制所有向量为单位向量，以防止这种情况发生。因此，在某种程度上讲，XNorm其实就是一种约束 。

Feed-forward layers(FFN)

在注意力模块中，FFN的作用是不可忽视的。然而，FFN是静态的，因此不依赖于空间维度。在物理上，这类FFN函数可以表示谐振子方程的驱动力：

这就起到了放大器或减速器的作用，从而增强或消除与空间关系无关的特征。

3.3. UFO-ViT

基于上面提出的XNorm，作者又基于之前的视觉Transformer，来构建本文提出的Transformer模型——UFO-ViT，其结构的示意图如上图所示。

Patch embedding with convolutions

最近的研究表明视觉Transformer用卷积的Patch Embedding能够获得更好的结果，因此，在本文中，作者也采用他们的策略，在模型的早期阶段使用卷积层。

Positional encoding

在本文中，作者使用位置编码作为可学习参数。虽然位置编码对消融研究没有影响，但作者还是决定添加位置编码减小结构的差异性。

Multi-headed attention

之前的工作也证明了，multi-head的重要性，为了更好的正则化，作者将参数应用于所有头，以衡量每个head的重要性。

Local patch interaction

为自注意模块设计额外的模块来提取局部特征已经非常常见了。在本文中，作者选择了其中最简单的方法——即，3×3的深度分离卷积。

Feed-forward network

与传统的视觉Transformer一样，本文的模型同样使用了MLP，其拓展因子为4。

▊ 4.实验

4.1. Image Classification

上表展示了本文方法的训练时候的一些设置细节。

上表展示了本文方法消融实验的结果。

上表展示了本文方法的不同变体和SOTA方法对比结果。

上图展示了本文方法和SOTA方法的FLOPs/Acc的实验结果对比。

从上面的结果可以看出，本文的方法计算效率非常高，能够在更小的FLOPs下，实现比较高的精度。

4.2. Object Detection with Mask R-CNN

上表展示了本文方法用Mask R-CNN进行目标检测的结果，可以看出，本文方法相比于其他CNN和Transformer方法，性能上还是相对来说比较不错的。

▊ 5. 总结

在本文中，作者提出了一种简单的方法，使自注意力具有线性复杂度而不损失性能。通过替换softmax函数，作者利用矩阵乘法的关联律消除了二次复杂度的运算。在早期的研究中，这种类型的因子分解通常会导致性能下降。

但是在图像分类方面，UFO-ViT模型优于大多数现有的基于Transformer和CNN模型的最新结果，因此UFO-ViT具有很强的实用性。此外，UFO-ViT模型在密集的预测任务中显示出比先前SOTA方法有竞争力或更高的性能，充分展示了本文方法的优越性。

▊ 作者简介

研究领域：FightingCV公众号运营者，研究方向为多模态内容理解，专注于解决视觉模态和语言模态相结合的任务，促进Vision-Language模型的实地应用