在注意力中重新思考Softmax：分解非线性，这个线性transformer变体实现多项SOTA

Transformer 在自然语言处理、计算机视觉和音频处理方面取得了巨大的成功。作为其核心组件之一，softmax 注意力有助于捕获长程依赖关系，但由于序列长度的二次空间和时间复杂度而禁止其扩展。钻研者通常采用核（Kernel）格式通过逼近 softmax 算子来降低复杂度。然而，由于近似误差，它们的功能在不同的工作 / 语料库中有所不同，与普通的 softmax 注意力相比，其功能会下降。钻研者认为 softmax 算子是主要障碍，而对 softmax 高效而准确的逼近很难实现，因此很自然地提出一个问题：我们能否用线性函数代替 softmax 算子，同时保持其关键属性？通过对 softmax 注意力的深入钻研，钻研发现了影响其经验功能的两个关键性质：(i) 注意力矩阵的非负性；(ii) 一种非线性重重加权方案，可以聚集注意力矩阵分布。这些发现揭示了当前格式的一些新见解。例如，线性 transformer 使用指数线性单元激活函数来实现属性 (i)。然而，由于缺乏重重加权（re-weighting ）方案，表现不佳。

本文中，来自商汤、上海人工智能实验室等机构的钻研者提出了一种称为 COSFORMER 的线性 transformer，它能同时满足上述两个特性。具体来说，在计算相似度分数之前，该钻研将特征传递给 ReLU 激活函数来强制执行非负属性。通过这种方式使得模型避免聚合负相关的上下文信息。此外，该钻研还鉴于余弦距离重加权体制来稳定注意力权值。这有助于模型放大局部相关性，而局部相关性通常包含更多自然语言工作的相关信息。在语言建模和文本理解工作的大量实考证明 COSFORMER 格式的有效性，并且在长序列 Long-Range Arena 基准上实现了 SOTA 功能，这一结果很好地证明了 COSFORMER 在建模长序列输入方面的强大能力。

论文地址：https://arxiv.org/pdf/2202.08791.pdf

格式钻研者表示，COSFORMER 的关键思路在于将不可分解非线性 softmax 操作交换为拥有可分解非线性重加权体制的线性操作。该模型适用于随机注意力和交叉注意力，并且输入序列长度拥有线性时间和空间复杂度，从而在建模长程依赖中显示出强大的能力。一般而言，我们可以选择任意相似度函数来计算注意力矩阵，以下公式（2）所示

下图 2 为原始自注意力（左）与线性注意力（右）的计算示意图。

大多数现有线性 transformers 想要找到针对 softmax 注意力的无偏估计。然而，钻研者实证发现，这些格式对采样率的选择非常敏感，一旦采样率过高，则会变得不稳定。并且，考虑到近期偏向，门体制可以用来更好地探索较新的上下文。另一类工作则试图直接用线性操作交换 softmax。钻研者提出了 softmax 的一种新交换，不仅可以在一系列工作中实现与 softmax 相当甚至更好的功能，而且拥有线性空间和时间复杂度。钻研者经过实证确认了 softmax 操作的两种可能在功能表现中发挥重要作用的关键特性，其一是确保注意力矩阵 A 中的所有值是非负的，其二是提供一种非线性重加权体制来聚集注意力连接的分布并稳定训练。为了考证这些假设，钻研者设计了以下表 1 中的初步钻研。

通过将的结果与 softmax 相比较，钻研者发现，拥有 softmax 重加权的模型收敛速度更快，到下游工作的泛化表现更好。 COSFORMER 鉴于以上发现，钻研者提出了模型 COSFORMER，它完全丢弃了 softmax 归一化，同时拥有非负性和重加权体制。COSFORMER 包含两个主要的组件，分别为线性投影核和 cos-Based 重加权体制。首先来看线性投影核。鉴于上文公式 2 中注意力的通用格式，钻研者将线性相似度定义以下公式（6）所示。

具体地，为了确保全正（full positive）的注意力矩阵 A 并避免聚合负相关的信息，钻研者将 ReLU(·) 用作转换函数，从而可以有效地消除负值，以下公式（7）所示。

钻研者还重新调整了点积（dot-product）的顺序，并得到了线性复杂度中所提注意力的公式（9）。

在 softmax 注意力中引入非线性重加权体制可以聚集注意力权重的分布，因而稳定训练过程。钻研者还通过实证发现，这种做法可以惩罚远距离连接，并在某些情况下加强局部性。实际上，这类局部性偏向，也即是一大部分上下文依赖来自邻近 token，通常可以在下游 NLP 工作上发现，以下图（3）所示。

鉴于上述假设，要满足 softmax 的第二种特性需要一种可分解的重加权体制，该体制可以将近期偏向引入到注意力矩阵。钻研者提出了一种能够完美满足目标的 cos-based 重加权体制。通过上图 3（2）和（3）中比较注意力矩阵，可以看到，COSFORMER 能够比不拥有重加权体制的 COSFORMER 执行更多的局部性。拥有余弦重加权体制的模型可以定义以下公式（10）所示。

同时，COSFORMER 还可以看作一种将相对位置偏向引入到高效 transformer 中的新格式。

实验实验考证了 COSFORMER 在多种设置中的有效性。首先，该钻研使用 WikiText-103，通过自回归和双向语言模型（bidirectional）设置考证了 COSFORMER 在语言建模中的能力。通过这种方式，该钻研考证了所提出的线性注意力模块在因果和非因果情况下的有效性。自回归建模语言：该钻研采用 (Baevski & Auli, 2018) 作为基线模型，并用所提出的线性注意力模块交换自注意力模块，在 8 个 Nvidia Tesla A100 GPU 上进行训练，表 2 中报告了结果。

由结果可得，尽管基线模型是一个功能强大的标准 transformer，但它需要二次计算复杂度，但 COSFORMER 在线性计算复杂度方面明显优于基线模型。此外，该钻研在考证集上实现了与其他格式相当的困惑度（perplexity），并且在测试集上明显优于其他格式，这进一步证明了 COSFORMER 的有效性。双向语言模型：对于双向语言建模，该钻研采用 RoBERTa (Liu et al., 2019) 作为基线模型。同样的将 RoBERTa 中的自注意力模块交换为线性注意力模块，并保持其他结构不变，该钻研在 2 个 Nvidia Tesla A100 GPU 上训练双向工作。如图 4 所示，COSFORMER 在训练集和考证集上的收敛速度都比 vanilla Transformer 快，拥有可比或更小的损失值。此外，与没有重重加权的模型相比，拥有重重加权体制的 COSFORMER 变体拥有明显更好的收敛速度和最终结果，这进一步考证了鉴于余弦的距离矩阵的有效性。

其次，该钻研通过与其他现有的 Transformer 变体进行比较，钻研了 COSFORMER 在下游工作上的泛化能力。从表 3 中，我们可以看到 COSFORMER 在多个数据集上优于基线（Liu et al., 2019），此外与其他高效的 transformer 相比，COSFORMER 在所有五个下游数据集上均达到最佳或次要位置。值得注意的是，尽管 Longformer (Beltagy et al., 2020) 在 MNLI 上取得了比 COSFORMER 更好的结果，但它需要 O(Nw) 的计算复杂度。如图 1 所示，Longformer 比 COSFORMER 更慢并且需要更多的内存开销。

该钻研进一步在 long-range-arena 基准上将 COSFORMER 与其他 transformer 变体进行比较，以了解其建模远程依赖关系的能力并展示对模型效率的比较分析。如表 4 所示，总体上来说 COSFORMER 在所有工作中都取得了拥有竞争力的结果，同时在 ListOps 和文档检索方面取得了最佳功能。值得一提的是，COSFORMER 在 Long-range-arena 基准测试中取得了最好的整体分数，是仅有的两个超越 vanilla transformer 架构的模型之一。