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

解决Transformer的固有缺陷:复旦大学等,提出线性复杂度SOFT

时间:2023-03-18 18:54:59 科技观察

VisionTransformer(ViT)借助patch-wise图像标记化和自我注意机制,在各种视觉识别任务上取得了SOTA。但是self-attention模块的使用使得类Transformer模型的空间和时间复杂度都是O(n^2)。自然语言处理领域的研究人员进行了各种尝试,使自注意力计算接近线性复杂度。近日,复旦大学、萨里大学和华为诺亚方舟实验室的研究人员进行了深入分析,发现这些尝试要么在理论上存在缺陷,要么在实验中对视觉识别无效,并进一步发现,这些方法的局限性在于:softmaxself-attention在逼近过程中仍然保持不变。具体来说,传统的自注意力是通过标准化标记特征向量之间的缩放点积来计算的。维持这个softmax操作阻碍了Transformer线性化的复杂性。基于此,本研究首次提出了一种无softmax的Transformer(softmax-freetransformer,SOFT)。为了去除self-attention中的softmax,使用高斯核函数代替点积相似度而无需进一步归一化。这使得通过低秩矩阵分解来近似完整的自注意力矩阵成为可能。近似的稳健性是通过使用Newton-Raphson方法计算其Moore-Penrose逆来实现的。在ImageNet上的大量实验表明,SOFT显着提高了现有ViT变体的计算效率。至关重要的是,对于线性复杂性,SOFT中允许使用更长的标记序列,从而在准确性和复杂性之间实现卓越的权衡。论文地址:https://arxiv.org/abs/2110.11945项目地址:https://github.com/fudan-zvg/SOFTTransformer模型存在一个瓶颈,即计算和内存使用的二次复杂度。这是自我注意机制的一个内在特征:给定一系列标记(例如,单词或图像块)作为输入,自我注意模块通过将一个标记与所有其他标记相关联来迭代学习特征表示。这导致计算(时间)和内存(空间)中长度为n的标记序列的二次复杂度O(n2),因为在推理期间需要计算和保存大小为n×n的注意力矩阵。这个问题在视觉中尤为严重:即使空间分辨率适中,标记化中的2D图像也可以产生比NLP中的更长的序列。因此,这种二次复杂性阻止了ViT模型对具有高空间分辨率的图像进行建模,而这对于视觉识别任务通常至关重要。一个自然的解决方案是通过近似来降低自注意力的计算复杂度。事实上,在NLP方面已经有很多尝试[33,5,18,38]。例如,[33]采取了一种天真的方法,通过可学习的预测来缩短键和值的长度。这种粗略的近似将不可避免地导致性能下降。相比之下,[5,17]都利用内核机制来近似softmax归一化,以线性化自注意力中的计算。[18]相反,采用哈希策略来有选择地计算最相似的对。最近,[38]使用Nystr?m矩阵分解通过多项式迭代重建完整的注意力矩阵,以近似地标矩阵的伪逆。尽管如此,softmax归一化只是在矩阵分解过程中重复进行,这在理论上是不可靠的。该研究通过实验发现,这些方法在应用于视觉时均无效(请参阅第4.2节)。本研究发现现有高效Transformer的局限性是由于使用softmaxself-attention造成的,并首次提出了不使用softmax的Transformer。更具体地说,在所有现有的Transformer(有或没有线性化)中,需要在标记特征向量之间的缩放点积之上进行softmax归一化。维持这个softmax操作会挑战任何后续的线性化工作。为了克服这个障碍,本研究提出了一种新的softmax-freeself-attention机制,命名为SOFT,在空间和时间上具有线性复杂度O(n)。具体来说,SOFT使用高斯核来定义相似度(自注意力)函数,而无需后续的softmax归一化。有了这个无softmax的注意力矩阵,研究进一步引入了一种新的低秩矩阵分解算法来对其进行逼近。通过使用Newton-Raphson方法可靠地计算矩阵的Moore-Penrose逆,理论上可以保证近似的稳健性。本研究的主要贡献包括:提出了一种具有线性空间和时间复杂度的新型softmax-freeTransformer;本研究的注意力矩阵近似是通过一种具有理论保证的新型矩阵分解算法实现的;为了评估该方法在视觉识别任务上的性能,该研究设计了一系列以SOFT作为核心自注意力组件的具有不同能力的通用骨干架构。大量实验表明,对于线性复杂度(图1b),SOFT模型可以将更长的图像标记序列作为输入。因此,给定相同的模型大小,SOFT在准确性/复杂性权衡方面优于ImageNet[9]分类上最先进的CNN和ViT变体(图1a)。下图2给出了该模型的示意图。图2:所提出的softmax-freeself-attention(SOFT)方法的示意图。P.E.:位置嵌入。虚线:线性投影。dh:每个注意力头的隐藏调光。?表示矩阵点积。作者采用了两个实验装置。在第一种设置下,对于所有方法,该研究使用相同的Tiny(表2)架构进行公平比较。也就是说,SOFT中的核心自注意力块被每个基线自己的注意力块替换,而架构的其余部分保持不变。请注意,[35]的空间缩减模块是Linformer[34]的特例。研究人员将缩减比例设置为与方法相同。使用相同的均匀采样思想,本研究将Nystr?mformer(用于NLP任务)的一维窗口平均替换为二维平均池化(用于图像)。下采样率与本研究的方法一致。另外值得一提的是,Reformer[19]没有官方代码发布,并且LocalSensitiveHashing(LSH)模块对输入令牌的长度有严格的要求,因此该方法未纳入本研究的比较.从下表1可以观察到:与Tiny架构上的Transformer相比,LinearTransformer方法在保持相似参数大小的同时大大减少了内存和FLOPs;SOFT方法在所有线性化方法中实现了最好的分类精度;该方法的推理速度与其他线性Transformer相当,训练速度略慢于Nystromformer,慢于Performer和Linformer。研究人员指出,该模型训练速度慢的主要原因是Newton-Raphson迭代,只能顺序应用以保证Moore-Penrose逆的准确性。总之,研究人员认为,增加的训练成本是值得为更高的准确性付出的代价,因为推理速度相当。该研究与最先进的替代方案进行了比较,并报告了ImageNet-1K验证集上的最高准确度。计算批量大小为1024FLOPs。从图1a和表3中得出以下观察结果:(i)总体而言,ViT及其变体比CNN产生更好的分类精度。(ii)该研究在最近基于纯视觉Transformer的方法中取得了最佳性能,包括ViT[11]和DeiT[31],以及最先进的CNNRegNet[26]。(iii)SOFT在所有变体中都优于最接近的(在架构配置中)Transformer对应物PVT[35]。由于注意力模块是主要区别,这直接验证了模型的有效性。(iv)该方法还击败了Twins,这是一种最近的ViT变体,旨在解决ViT的效率限制,并且需要更少的参数和更少的浮点计算。为了深入了解如何使用SOFT和替代方法学习注意力,图3显示了各种比较模型的注意力掩码。对于每个模型,论文中给出了前两个注意力头的输出。很明显,SOFT在捕获像素之间的局部和远程关系方面表现出稳健性和多功能性。有趣的是,虽然SOFT是在ImageNet[9]中的对象分类数据集上训练的,但它似乎能够学习同一类别实例之间共享的语义概念和实例特定特征。感兴趣的读者可以阅读论文原文了解更多研究细节。