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

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

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

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

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

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