去掉softmax后Transformer会更好吗？复旦&华为诺亚提出SOFT：轻松搞定线性近似

一时失言乱红尘 2022-08-28 02:44 147阅读 0赞

点击上方“计算机视觉工坊”，选择“星标”

干货第一时间送达

![a574c16c90d4949fb9f381af9c679d22.png][]

作者丨happy

编辑丨极市平台

**导读**

本文介绍了复旦大学&华为诺亚提出的一种新颖的softmax-free的Transformer—SOFT。所提SOFT显著改善了现有ViT方案的计算效率，更为关键的是：SOFT的线性复杂度可以允许更长的token序列，进而取得更佳的精度-复杂度均衡。

![23f344bdd1e37c2797bddab4b3293398.png][]

论文链接：https://arxiv.org/pdf/2110.11945.pdf

代码链接：https://github.com/fudan-zvg/SOFT

项目链接：https://fudan-zvg.github.io/SOFT/

> 本文是复旦大学&华为诺亚关于Transformer中自注意力机制复杂度的深度思考，首次提出了一种新颖的**softmax-free 的Transformer** 。本文从softmax self-attention局限性出发，分析了其存在的挑战；然后由此提出了线性复杂度的SOFT；再针对线性SOFT存在的训练问题，提出了一种具有理论保证的近似方案。所提SOFT在ImageNet分类任务上取得了比已有CNN、Transformer更佳的精度-复杂度均衡。

## Abstract ##

ViT通过图像块序列化+自注意力机制将不同CV任务性能往前推了一把。然而，自注意力机制会带来更高的计算复杂度与内存占用。在NLP领域已有不同的方案尝试采用线性复杂度对自注意力进行近似。然而，本文的深入分析表明：NLP中的近似方案在CV中缺乏理论支撑或者无效。

我们进一步分析了其局限性根因：**softmax self-attention** 。具体来说，传统自注意力通过计算token之间的点乘并归一化得到自注意力。softmax操作会对后续的线性近似带来极大挑战。基于该发现，本文首次提出了SOFT(**softmax-free transformer** )。

为移除自注意力中的softmax，我们采用高斯核函数替代点乘相似性且无需进一步的归一化。这就使得**自注意力矩阵可以通过低秩矩阵分析近似** 。近似的鲁棒性可以通过计算其MP逆(Moore-Penrose Inverse)得到。

ImageNet数据集上的实验结果表明：**所提SOFT显著改善了现有ViT方案的计算效率** 。更为关键的是：SOFT的线性复杂度可以允许更长的token序列，进而取得更佳的精度-复杂度均衡。

## Contributation ##

本文的贡献主要包含以下几点：

*  提出一种新颖的**线性空间、时间复杂度** 的**softmax-free Transformer** ；
 *  所提注意力矩阵近似可以通过**具有理论保证的矩阵分解算法** 计算得到；
 *  所提SOFT在ImageNet图像分类任务上取得了比其他ViT方案**更佳的精度-复杂度均衡** (见下图a)。

![d4f6354b6fddd16facdfe9346f697e62.png][]

## Method ##

### Softmax-free self-attention formulation ###

![d356296644735d822819478a163d5ac9.png][]

上图给出了本文所提SOFT架构示意图。我们首先来看一下该注意力模块的设计。给定包含n个token的输入序列，自**注意力旨在挖掘所有token对之间的相关性** 。

具体来说，X首先线性投影为三个维的query、key以及values：

自注意力可以表示为如下广义形式：

自注意力的关键函数包含一个非线性函数与一个相关函数。自注意力的常规配置定义如下：

虽然该softmax自注意力已成为首选且很少受到质疑，但是它并不适合进行线性化。为构建线性自注意力设计，我们引入了一种sfotmax-free自注意力函数：通过高斯核替换点乘操作。定义如下：

为保持注意力矩阵的对称性，我们设置投影矩阵相同，即。所提自注意力矩阵定义如下：

为描述的简单性，我们定义为矩阵形式：。所提自注意力矩阵S具有三个重要属性：

*  对称性
 *  所有元素均在\[0,1\]范围内；
 *  所有对角元素具有最大值1。

我们发现：**当采用无线性化的核自注意力矩阵时，transformer的训练难以收敛** 。这也就解释了：为何softmax自注意力在transformer中如此流行。

### Low-rank regularization via matrix decomposition with linear complxity ###

为解决收敛于二次复杂度问题，我们利用了矩阵分解作为带低秩正则的统一解，这就使得模型复杂度大幅下降，且无需计算全部的自注意力矩阵。

作出上述选择因为在于：S为半正定矩阵，且无需后接归一化。我们将S表示为块矩阵形式：

其中，。通过上述分解，注意力矩阵可以近似表示为：

其中，表示A的MP逆。更多关于MP逆的信息建议查看原文，这里略过。

在上述公式，A和B是S通过随机采样m个token得到的子矩阵，可表示为(我们将其称之为**bottleneck token** )。然而，我们发现：随机采样对于m非常敏感。因此，我们通过利用结构先验探索两种额外的方案：

*  采用一个核尺寸为k、stride为k的卷积学习；
 *  采用一个核尺寸为k、stride为k的均值池化生成。

通过实验对比发现：**卷积层学习** **具有更好的精度** 。由于K与Q相等，因此。给定m个token，我们计算A和P：

最终，我们得到了SOFT的正则化后的自注意力矩阵：

![44f1882b95a5640d3eccb29b17d5d308.png][]

上图Algorithm1给出所提SOFT流程，它涉及到了MP逆计算。一种精确且常用的计算MP逆的方法是SVD，然而SVD对于GPU训练不友好。为解决该问题，我们采用了Newton-Raphson方法，见上图Algorithm2：一种迭代算法。与此同时，作者还给出了上述迭代可以最终收敛到MP逆的证明。对该证明感兴趣的同时强烈建议查看原文公式，哈哈。

### Instantiations ###

上面主要聚焦于**softmax-free self-attention** 模块的介绍，接下来我们将介绍如何利用SOFT模块构建Transformer模型。我们以图像分类任务为切入点，以PVT作为基础架构并引入所提SOFT模块构建最终的SOFT模型，同时还在stem部分进行了微小改动。下表给出了本文所提方案在不同容量大小下的配置信息。

![d4fa6c7daa32ca1788347475d87084df.png][]

## Experiments ##

![ad38ea2abfcf7dc64bb9dd12cb7224b0.png][]

上表对比了所提方案与现有线性Transformer模型的性能，从中可以看到：

*  相比基线Transformer，线性Transformer能够大幅降低内存占用与FLOPs，同时保持相当参数量；
 *  所提SOFT在所有线性方案中取得了最佳分类精度；
 *  所提SOFT与其他线性方案的推理速度相当，训练速度稍慢。

![b4cc902a8aecf43a25266bf354b78524.png][]

上图给出了不同方案的序列长度与内存占用之间的关系，从中可以看到：所提**SOFT确实具有线性复杂度的内存占用** 。

![37daeb42a9ddfffd6bcbd0c959ee7ea9.png][]

上表给出了所提方案与其他CNN、Transformer的性能对比，从中可以看到：

*  总体来说，ViT及其变种具有比CNN更高的分类精度；
 *  相比ViT、DeiT等Transformer方法以及RegNet等CNN方法，所提SOFT取得了最佳性能；
 *  相比PVT，所提方案具有更高的分类精度，直接验证了所提SOFT模块的有效性；
 *  相比Twins与Swin，所提SOFT具有相当的精度，甚至更优性能。

![9acf88794dd7ce8cf2560afd8c35f7de.png][]

此外，作者还在NLP任务上进行了对比，见上表，很明显：**SOFT又一次胜出** 。

本文仅做学术分享，如有侵权，请联系删文。

##  ##

**重磅！计算机视觉工坊****-学习****交流群****已成立**

扫码添加小助手微信，可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群，旨在**交流顶会、顶刊、SCI、EI等写作与投稿事宜。**

**同时**也可申请加入我们的细分方向交流群，目前主要有**ORB-SLAM系列源码学习、****3D视觉**、**CV&深度学习**、**SLAM**、**三维重建**、**点云后处理**、**自动驾驶、CV入门、三维测量、VR/AR、3D人脸识别、医疗影像、缺陷检测、行人重识别、目标跟踪、视觉产品落地、视觉竞赛、车牌识别、硬件选型、深度估计、学术交流、求职交流**等微信群，请扫描下面微信号加群，备注：”研究方向+学校/公司+昵称“，例如：”3D视觉 + 上海交大 + 静静“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进去相关微信群。**原创投稿**也请联系。

![aad755c421cb4046e4218721f0902aa7.png][]

▲长按加微信群或投稿

![92bcbc3f938bd9aae70bc4fb71d6e150.png][]

▲长按关注公众号

##  ##

**3D视觉从入门到精通知识星球**：针对3D视觉领域的**视频课****程（**[**三维重建系列**][Link 1]**、**[**三维点云系列**][Link 2]**、**[**结构光系列**][Link 3]**、**[**手眼标定**][Link 4]**、**[**相机标定**][Link 5]**、激光/视觉SLAM、自动驾驶等****）、****知识点汇总、入门进阶学习路线、最新paper分享、疑问解答**五个方面进行深耕，更有各类大厂的算法工程人员进行技术指导。与此同时，星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息，打造成集技术与就业为一体的铁杆粉丝聚集区，近4000星球成员为创造更好的AI世界共同进步，知识星球入口：

学习3D视觉核心技术，扫描查看介绍，3天内无条件退款

##  ##

![28f28ffca7df34139d84fd4509976fc1.png][]

圈里有高质量教程资料、可答疑解惑、助你高效解决问题

**觉得有用，麻烦给个赞和在看~**

[a574c16c90d4949fb9f381af9c679d22.png]: /images/20220828/91fb1dda15424814a3323bc8ad057a31.png
[23f344bdd1e37c2797bddab4b3293398.png]: /images/20220828/f02cc8d313344347a1326125a42168df.png
[d4f6354b6fddd16facdfe9346f697e62.png]: /images/20220828/32cac5319f4247d0be62a0873c92506f.png
[d356296644735d822819478a163d5ac9.png]: /images/20220828/b0612a2062b4445abc13d73bf6a53fe4.png
[44f1882b95a5640d3eccb29b17d5d308.png]: /images/20220828/9d1e505d0be644b4b786a1f80fc74a48.png
[d4fa6c7daa32ca1788347475d87084df.png]: /images/20220828/42aee131eb774b84886322029d3e52f8.png
[ad38ea2abfcf7dc64bb9dd12cb7224b0.png]: /images/20220828/47fb0a69cbc2473f968fa097fc7713f5.png
[b4cc902a8aecf43a25266bf354b78524.png]: /images/20220828/6da3f29083ff44138273237ee2e87dce.png
[37daeb42a9ddfffd6bcbd0c959ee7ea9.png]: /images/20220828/be04433d0c274c1a91c3db7412d48468.png
[9acf88794dd7ce8cf2560afd8c35f7de.png]: /images/20220828/c7186c732a5a4f9f9cf3bd64e893589a.png
[aad755c421cb4046e4218721f0902aa7.png]: /images/20220828/4c492bc255a74ca18e6b1cf33182edbf.png
[92bcbc3f938bd9aae70bc4fb71d6e150.png]: /images/20220828/2b84a5f83edb4e53891e482e28a18737.png
[Link 1]: http://mp.weixin.qq.com/s?__biz=MzU1MjY4MTA1MQ==&mid=2247549738&idx=1&sn=88b4aa9e091077f107407bddc8da7acd&chksm=fbfc281ecc8ba10823e11ac0044b5a94920806c83d30d8cb90630a41ecb23ad458106b9b1f58&scene=21#wechat_redirect
[Link 2]: http://mp.weixin.qq.com/s?__biz=MzU1MjY4MTA1MQ==&mid=2247551625&idx=1&sn=4a914cfeb2d251ab25421e1e5de930c0&chksm=fbfc21bdcc8ba8abac371478a6b2b9d143a1b31fa799201c6682d9d0b6a4ed4a7da146b3c536&scene=21#wechat_redirect
[Link 3]: http://mp.weixin.qq.com/s?__biz=MzU1MjY4MTA1MQ==&mid=2247546284&idx=2&sn=d067a98047192a04f21489aca338a5d6&chksm=fbfc1e98cc8b978e5fb6b2afd1a473c89e2861374542421b9972f87a6fdfcd47eb50c443042e&scene=21#wechat_redirect
[Link 4]: http://mp.weixin.qq.com/s?__biz=MzU1MjY4MTA1MQ==&mid=2247554670&idx=2&sn=cf3323178d5d1960476208e6303a8479&chksm=fbfc3d5acc8bb44c184084dcf447a5981d297815d62bef699129338d99a0472c622f96ee71d8&scene=21#wechat_redirect
[Link 5]: http://mp.weixin.qq.com/s?__biz=MzU1MjY4MTA1MQ==&mid=2247486422&idx=1&sn=381ab793959aed0eb6f01019dea522d4&chksm=fbff20e2cc88a9f456b80e23d2482fb319b08e89da071b39597bc51942a6aa7b756a6749856c&scene=21#wechat_redirect
[28f28ffca7df34139d84fd4509976fc1.png]: /images/20220828/9f38e8bf32f74dafbb77e9a26f26b7f1.png