当前位置: 首页 > 科技观察

只需要两个线性层就可以超越self-attention机制,清华极图团队又有新突破

时间:2023-03-15 15:47:34 科技观察

本文经AI新媒体量子位授权转载(公众号ID:QbitAI),转载请联系转载来源。5月4日,谷歌团队在arXiv上提交了一篇论文《MLP-Mixer:Anall-MLPArchitectureforVision》[1],引起了计算机视觉研究人员的热议:MLP有多大潜力??5月5日,清华大学图形学实验室Jittor团队在arXiv[2]上提交了论文《BeyondSelf-attention:ExternalAttentionusingTwoLinearLayersforVisualTasks》[2],提出了一种新的注意力机制,称为“ExternalAttention”。基于两个外部的、小型的、可学习的和共享的记忆,仅两个级联的线性层和归一化层就可以取代现有流行的学习架构中的“Self-attention”,揭示了线性层与注意力机制之间的关系。同日,清华大学软件学院丁桂光团队在arXiv上提交论文《RepMLP:Re-parameterizingConvolutionsintoFully-connectedLayersforImageRecognition》[3],展示了MLP结合重参数化技术也取得了很好的成绩。.5月6日,牛津大学的学者提交了一篇题为“DoYouEvenNeedAttention?AStackofFeed-ForwardLayersDoesSurprisinglyWellonImageNet”的论文[4],同时提出Transformer中的attention是不必要的。是的,仅使用前馈就可以在ImageNet上获得非常高的结果。FromSelf-attentiontoExternal-attention自注意力机制在自然语言处理和计算机视觉领域发挥着越来越重要的作用。对于输入的Nxd维空间的特征向量F,self-attention机制基于自身的线性变换使用Query、Key、Value特征计算自身样本中的attention,并相应更新特征:由于QKV是F的线性变换,为了简单起见,我们可以将self-attention的计算公式简写如下:这就是F对F的attention,也就是所谓的Self-attention。如果希望注意力机制能够考虑到来自其他样本的影响,那么就需要一个所有样本共享的特征。为此,我们引入了一个外部的Sxd维记忆单元M来描述所有样本最本质的特征,并用M来表示输入特征。我们称这种新的注意力机制为External-attention。我们可以发现,公式(5)(6)中的计算主要是矩阵乘法,这是一种常见的线性变换,self-attention机制因此被两个线性层和一个归一化层所取代。我们还使用了前人工作[5]中提出的Norm方法来避免因某个特征向量过大而导致注意力失效的问题。为了增强External-attention的表达能力,类似于self-attention机制,我们采用了两种不同的记忆单元。下图形象地展示了External-attention和Self-attention的区别。△图1SelfAttention和ExternalAttention的区别为什么两个线性层都能超越Self-attention?self-attention机制一个明显的缺陷就是计算量非常大,存在一定的计算冗余。通过控制记忆单元的大小,External-attention可以轻松实现线性复杂度。其次,self-attention机制只利用自身样本内的信息,忽略了不同样本之间的潜在联系,这在计算机视觉中是有意义的。例如,对于语义分割任务,不同样本中同一类别的对象应该具有相似的特征。External-attention通过引入两个外部记忆单元隐式地学习整个数据集的特征。这个想法也被应用在稀疏编码和字典学习中。图规划团队在PascalVOC测试集上可视化了注意力图和分割结果。如图2所示,可以发现使用两个线性层的外部注意力的注意力图是合理的。△图2attentionmap和分割结果的可视化从实验看ExternalAttention的效果为了证明方法的普适性,我们做了图像分类、分割、生成和点云分类分割的实验,证明了method'sEffectiveness,External-attention在大大减少计算量的同时,可以达到与当前state-of-the-art方法相当甚至更好的结果。1.图像分类2.图像语义分割(在三个数据集上)3.图像生成4.点云分类5.点云分割ExternalAttentionVSMLP-MixerGoogle的工作提出了一个小巧易用的Mixer-Layer,然后用极其丰富的实验证明只通过简单的图像块和线性层堆叠就可以达到非常好的性能,这打开了人们的想象空间。清华的ExternalAttention揭示了线性层和attention机制的内在联系,证明线性变换实际上是attention实现的一种特殊形式,如下式所示:Attention(x)=Linear(Norm(Linear(x)))Graph团队的工作和Google团队的工作都证明了线性层的有效性。值得注意的是,如果External-attention不断级联堆叠,同样是MLP的形式,可以实现纯MLP网络结构,但是External-attention使用了不同的normalization层,更符合与注意力机制。这类似于谷歌团队的工作。清华ExternalAttention的部分计算代码已经开源在Github上。未来所有的计算代码都会尽快开源。ExternalAttention部分计算代码:https://github.com/MenghaoGuo/-EANet