不用Softmax好像也可以？多项式激活函数替代方案出炉

本文挑战了传统观念，即 Transformer 中的softmax注意力主要因为能生成注意力分配的概率分布而有效。相反，作者从理论上证明，其成功在于在训练过程中能隐式地正则化注意力矩阵的Frobenius范数。

然后作者探索了其他能正则化注意力矩阵的Frobenius范数的激活函数，证明了某些多项式激活函数可以实现这一效果，使其适合注意力基础架构。

实验结果显示，这些激活函数在各种计算机视觉和语言任务中的性能可与softmax相当或更好，表明除了softmax之外，注意力机制还有新的可能性。

1 Introduction

变形器架构已成为各种领域（如自然语言处理（NLP）；计算机视觉；机器人学等领域的最先进模型架构。变形器架构的关键组件是softmax注意力模块，使得变形器在输出生成过程中评估单个输入元素的重要性。这一特性使得变形器在训练过程中能够有效地关注多样化的输入元素，从而有效地捕获序列数据中的空间依赖关系。与传统的循环神经网络（RNNs）和卷积神经网络（CNNs）不同，变形器能够在不显著降低性能的情况下扩展到大型数据集。这一特性使他们成为处理大规模机器学习任务的理想架构。

softmax 广泛认可其在注意力机制中的有效性，原因在于其能够生成满足三个关键条件的注意力矩阵：(i)非负性，(ii)行归一化后和为1，(iii)稀疏性。普遍认为，非负性保证了注意力权重保持正数，有助于模型为各种输入元素赋予重要性。归一化约束确保所有输入元素的注意力权重之和为1，使权重可解释为概率。

此外，稀疏性有助于模型关注少数几个关键元素，从而提高效率。有人认为，这些属性对于使注意力机制能够关注输入序列的相关部分，同时有效过滤无关细节至关重要。然而，这种关注方法已经变得有些教条，因为它主要受到经验结果的影响，而没有理论基础。尽管在几项研究中探索了其他激活方式，但 softmax 注意力仍然主导地位，主要原因在于其可解释性。

在本文中，作者通过提出软max的效力源于其在对训练时注意矩阵的Frobenius范数进行隐式正则化，从而防止注意权重变得过大或过小，来质疑这一观点。然后作者推导出一个理论框架，该框架产生的多项式激活故意违反之前提到的三个条件之一，但能够在对注意权重的Frobenius范数进行训练时进行正则化。作者的发现表明，这些激活可以实现与softmax相当或甚至优于softmax的各种视觉和自然语言处理（NLP）任务，尽管它们似乎与作者理解到的注意力相矛盾。

作者建议读者，本文并未遵循通常追求创建最先进的 Transformer 架构以在基准数据集上实现最先进结果的做法。相反，作者的重点是批判性地审视softmax注意力，以确定其有效性是否是真正的可解释性结果，还是一个更细微的正则化机制的结果。通过挑战既定的观点，作者旨在揭示 Transformer 架构的更深层次洞察，可能引领更广泛的应用和更好的理解。

然而，作者在多个基于 Transformer 的任务上验证了作者的理论，包括图像分类、分割、目标检测和自然语言处理，通常实现的结果与softmax注意力的结果相匹配或超过。作者的主要贡献是：

作者质疑软max在注意力机制中广泛接受的有效性仅源于其产生归一化稀疏注意力权重的能力这一观念。相反，作者从理论上证明软max对注意力具有规范化作用，并认为这在软max的成功中发挥了关键作用。

作者探索了故意偏离传统softmax注意力条件的激活。这些激活在训练过程中发现可以规范注意力矩阵的Frobenius范数，类似于softmax，并在各种视觉和NLP任务上表现出可比的或优越的性能。

2 Related Work

许多研究探讨了 Transformer 中注意机制的替代激活方法。沈等人（2023年）研究了ReLU激活，发现在具有长序列的任务（如文档翻译）中，它们超过了softmax。

班纳吉等人（2020年）研究了softmax的泰勒级数近似，显示了在图像分类中比softmax更好的性能。王等人（2021年）提出了softmax的周期性替代方法，旨在为注意力机制提供更好的梯度，并在图像分类的简单网络中取得了比softmax更好的结果。科胡帕伊和Pirsiavash（2024年）证明了将归一化应用于线性注意力机制可以实现与softmax相当的性能。

作者的工作与这些工作不同，因为作者识别出注意力矩阵的Frobenius范数规模与输入序列长度之间存在明确的理论关系。利用这一洞察，推导出可以与softmax媲美的潜在激活方法。

3 Preliminaries and Notation

在本节中，作者通过 Transformer 块定义了 Transformer ，并设定了未来各节中使用的各种数学量的表示法。关于 Transformer 的更多信息，读者可以参阅Vaswani等人（2017年）和Dosovitskiy等人（2020年）。

Transformer 架构由 Transformer 块组成，定义如下。Transformer 块是一个映射 ，其中 和 分别为输入和输出的维度。

是一个具有残差连接的前馈 MLP(多层感知机)和 是一个注意力头。注意力头 A 定义如下：它由三个可学习的矩阵组成，Query 、键 和值 由以下公式定义： 对于输入序列 ，其中 , 。然后, 注意力头 被定义为:

在这篇论文中, 作者将关注自注意力机制, 即最常用的相似变换 （也称为点积）: 。最常用的激活函数 是softmax。这导致了注意力头的最常见形式, 由

softmax函数是矩阵软max映射，按行应用常规的softmax函数：

《1/√d》这一因素，如Vaswani等人（2017年）所述，是为了防止softmax的梯度变得过小而引入的缩放因子。在本论文的理论分析中，作者将仅使用点积相似度qk^T，并将N×N矩阵softmax(qk^T)称为_softmax自注意力矩阵。在实验部分，第5节，作者将通过实证方式在更一般的softmax注意力块上验证作者的理论框架。

对于通用Transformer架构，使用1到n的多个头。每个注意力头由方程3.3定义，然后将每个注意力头的所有输出拼接在一起，然后进入前馈层。

作者需要符号表示矩阵softmax映射的导数，该映射由方程3.4定义。

对于矩阵 , 作者可以对矩阵映射softmax在 处求导, 得到梯度线性映射 , 该映射的公式为:

给定矩阵 ，作者用 表示其 Frobenius 范数。另外，作者用 表示随机变量的期望, 具体考虑的随机变量将在上下文中明确。

4 Theoretical Analysis

Implicit regularization of Softmax

本文段提出了一种理论结果，表明softmax激活以一种指数级的方式对自注意力矩阵的Frobenius范数施加控制，而这种控制与输入序列的 Token 长度呈亚线性增长。此外，作者还证明了softmax关于自注意力矩阵的梯度也具有类似的规律性。

尽管以前的工作通过Lipschitz常数的角度分析了softmax自注意力的一致性，但作者的定理通过直接将Frobenius范数规律性与 Token 长度联系起来，提供了一个新的视角。这为作者理解自注意力激活如何随着 Token 长度的变化而调整以保持训练过程中的稳定性，尤其是基于梯度下降的算法，提供了启示。

定理4.1: 设 softmax: 为由方程3.4定义的矩阵softmax映射, 且 表示 softmax 在 处的梯度。那么，作者有以下关于Frobenius范数的界

定理4.1的关键推论是，在具有softmax自注意力机制的 Transformer 训练过程中，每个softmax自注意力矩阵的Frobenius范数始终保持在一个随着增长的价值范围内。这确保了通过自注意力矩阵权重反向传播不会导致过大的梯度。证明的关键在于，softmax内嵌的行归一化有效控制了Frobenius范数。详见附录A.1.1的详细证明。

Polynomial activations for self-attention

在第4.1节中，作者证明了softmax隐式地正则化了自注意力矩阵的Frobenius范数。在此基础上，作者现在证明，通过缩放特定多项式激活，可以在期望值上实现对Frobenius范数的类似正则化效果，这几乎完全复制了softmax的影响。

定理4.2: 设 为独立同分布的随机变量, 分别按照 和 分布。对于 , 作者计算 矩阵 的 Frobenius 范数的期望。

通过将这种激活函数乘以，作者可以得到一个的界。

4.3推论：设4.2定理中的条件相同。

推论4.3证明了, 形式为 的激活函数在期望上提供了一种类似于softmax在自注意力矩阵上的正则化水平。推论4.2的证明可以在附录A.1.2中找到。作者接下来要证明的下一个性质与4.1中得到的梯度界类似。由于自注意力矩阵的参数由 Query 和键 （Vaswani等人, 2017）给出, 这意味着在 Transformer 训练过程中, 只有 和 矩阵是自注意力矩阵得到更新的部分。因此, 作者计算关于 和 导数的正则化结果。

定理4.4: 设 均为独立同分布随机变量, 分布分别为 , 。则当 时, 矩阵 的 参数矩阵的期望为 。

上述定理 then 表明, 如果作者以 的比例将多项式 进行缩放, 那么 导数的增长将呈现 的形式。

4.5 推论：设与定理4.4相同的条件。

类似的估计也适用于关于K矩阵的导数。4.4定理的证明可参见附录A.1.2。

4.3和4.5的推论表明, 具有形式 （其中 ）的多项式激活在应用于自注意力矩阵时, 可以达到与softmax相当的表现。

在5节中, 作者通过实证比较这些激活与softmax, 观察到它们在各种transformer任务上都优于 softmax。作者专注于 和 , 因为这些多项式显然违反了 softmax基于注意力的一些关键方面, 例如归一化行、正性和稀疏性。对于较大的 值, 性能下降是由于当 较大时, 的函数在 0 附近的梯度较小, 这导致训练困难。

5 Experiments

在本节中, 作者在各种 Transformer 任务上验证第4节中的理论。作者对第4节中提到的两个主要激活进行实证验证, 即立方多项式激活 和线性多项式 。

目标是证明, 通过适当缩放这些激活, 利用第 4 节中的理论, 作者可以实现与softmax相比具有竞争力的性能。在本节剩余部分，作者将简单地用 和 表示这些激活。

Image classification

5.1.1 ViT-Tiny on Tiny-Imagenet:

在本节中，作者测试了在Tiny-Imagenet数据集（Le和Yang，2015）上从零开始训练的ViT-Tiny架构（Steiner等人，2021）的理论（Steiner等人，2021）。

作者的第一个实验是为了测试在输入序列长度和4.3和4.5节中预测的缩放因子都使用激活函数x^3时，Top-1%的准确率如何变化。根据第4节发展的理论，当作者将X^3按1/√N进行缩放时，Frobenius范数按O(√N)进行缩放。因此，随着序列长度的减小，作者应看到对数尺度上的缩放量减少。

图1展示了该实验的结果。作者考虑了四种不同的输入序列长度，大小分别为256、64、16和8。作者运行了多种ViT-Tiny架构，这些架构的缩放形式为，其中N的范围在序列长度以下到以上。从图1可以看出，随着序列长度变小，在x轴上以对数表示的所需的缩放量减小，这验证了第4.2节中的理论。

第二项实验比较了激活函数 和 , 以及缩放版本 和 与Tiny-ViT架构在Tiny-Imagenet上的softmax。由于序列长度为 ，作者决定将 作为多项式激活函数的缩放比例。实验使用了 4 个patch大小的特征图， 3 个attention头和 12 层，如Steiner等人（2021年) 所述。表1中的结果表明, 超过了 softmax, 而未缩放版本表现较差。同样, 未缩放版本的 在竞争中表现良好，且未进行缩放时性能显著下降。

图2显示了在ViT-Tiny的层2和12中, 对五个激活进行自注意力矩阵的Frobenius范数训练时的结果, 所有 Head 的平均值。 和 的范数高于softmax, 但通过将它们乘以 进行缩放, 可以将它们降低到更稳定的水平, 从而提高训练稳定性。图3显示了Jacobian的Frobenius范数, 缩放同样使范数更接近softmax, 以确保更稳定的梯度。附录A.2.2中包含其他层的进一步图表。

表格1：在Tiny-Imagenet上，软max和多项式激活方法在Top-1%精度上的对比。当应用正确的缩放比例时，立方激活方法优于软max。同样，线性激活方法只有在最优缩放时才具有竞争力。

5.1.2 Larger vision transformers on ImageNet-1k

在本次实验中，作者使用来自ImageNet-1k数据集的各种不同视觉 Transformer 进行了图像分类任务。作者发现，比例为1/14对于x^3和x都取得了最佳效果。

作者从头训练所有模型在ImageNet-1k数据集上，并在验证集上报告Top-1精度。作者使用PyTorch Paszke等人 和Timm Wightman（2019年）库来训练作者的模型，设置与He等人 和Liu等人（2021年）相似。作者检查了作者的方法与以下三种Transformer架构一起，以展示其泛化能力：

ViT：Dosovitskiy等人（2020年）的开创性工作将图像解释为一系列 Patch ，并通过NLP中使用的标准Transformer编码器对其进行处理。这种简单但可扩展的策略，在与大型数据集的预训练相结合时，表现出惊人的效果。作者使用ViT-Small，具有以下设置：Patch 大小=16，嵌入维度=384，头数=6，层数=12。此外，作者还使用ViT-Base，具有以下设置：Patch 大小=16，嵌入维度=768，头数=12，层数=12。

DeiT: Touvron等人（2021年）提出了一种基于ViT的 Transformer 模型。除了作者没有使用DeiT中强制的蒸馏 Token 外，它与ViT非常相似。作者使用的是DeiT-Small，它具有以下设置：patch size = 16，嵌入维度 = 384，head数量 = 6，层数 = 12。此外，作者还使用了DeiT-Base，它具有以下设置：patch size = 16，嵌入维度 = 768，head数量 = 12，层数 = 12。

Swin Transformer： 刘等（2021年）提出了一个层次化的特征表示，并提出了基于位移窗口的自注意力机制，该机制在视觉问题上被证明是有效和高效的。作者在第一阶段的隐藏层中使用Swin-Small（96通道）和Swin-Base（128通道）。默认为设置窗口大小M=7，每个 Head 的 Query 维度d=32，所有实验的层数为2,2,18,2。

XciT：Xiong等人（2021年）提出的视觉 Transformer 架构与标准ViT不同，它包含两个不同的组件。首先，在每个块中具有局部patch交互，包括一个深度可分离的3×3卷积后接 BatchNorm 、GELU，以及另一个深度可分离的3×3卷积。其次，它使用交叉协方差注意力，其中注意力图是从计算在 Token 和 Query 投影的token特征上的关键和 Query 协方差矩阵得到的。作者使用了XCiT-S12，其patch大小为16，以及XCiT-M24，其patch大小为24。

表格 2 中的结果如下。激活函数 在 ViT 和 Swin Transformer 上表现最好, 而 softmax 在 DeIT 架构上表现最好。在附录 A.2.1 中，可以找到不同尺度和激活函数的进一步消融实验结果。

图4展示了ViT-Small架构中, 训练时层2和层12中每个头中自注意力矩阵的Frobenius范数。通过将激活值 和 乘以 ，作者可以控制自注意力矩阵的Frobenius范数的规模，并得到与 softmax的规模相当的规模。同样, 图5展示了层2和层12中自注意力矩阵的Jacobian范数在训练时的Frobenius范数, 平均值在所有头中。通过将激活值 和 乘以 , 作者可以控制 Jacobian范数的规模, 并在训练时得到与softmax的规模相当的规模。该架构其他层在训练时的 Frobenius范数图在附录A.2.2中给出。

5.1.3 Visualizing self-attention with ViT-Base

作者绘制了在收敛后使用 和softmax 激活的Self-Attention矩阵的热力图, 覆盖了两个层和 8 个头。在固定大小为 128 的训练批次的平均下, 图 6 显示了层 2 , 头 8 , 突出了注意力模式的不同, 其中 包含正负值。同样, 图7的层12, 头6展示了每个激活的显著模式。总体而言, ViT-Base架构在 下表现出与softmax 显著不同的Self-Attention模式。

作者观察了自注意力矩阵如何针对图像的不同区域进行定位。作者使用了ImageNet-1k验证集的一张图像，提取了类 Token ，将其 Reshape 为（14,14）的网格，代表196个 Patch ，然后使用最近邻插值将其映射回原始图像大小。图8显示了输入图像，而图9说明了ViT-Base架构在收敛后，第12层， Head 6的自注意力矩阵对于不同激活值的关注区域。

Object Detection and Instance Segmentation

在本节中, 为了检验作者模型的迁移学习能力, 作者通过在目标检测和分割任务上微调作者的 ImageNet预训练XCiT模型来展示作者的方法。作者的实验在COCO 2017数据集（Lin等人, 2014年）上进行，该数据集有 118 K训练图像和 验证图像，包含 80 个类别。作者将XCiT架构作为Mask R-CNN（He等人，2017年）检测器的 Backbone，并与特征金字塔网络（FPN）相结合。由于XCiT固有的列状设计, 作者通过从各种层提取特征来适应FPN的兼容性, 这些特征具有一致的步长为 16 。然后, 作者将特征分辨率调整为 的步长。这种下采样通过最大池化实现, 而通过单层反向卷积实现上采样。模型使用AdamW优化器训练 36 个周期, 学习率为 , 权重衰减率为 0.05 , 批大小为 16 。在表 3 中, 作者在具有激活 和 softmax的XCiT-S12模型上进行了实验。作者发现在这个任务上, 作者无法很好地训练激活 和 , 因此只报告其他激活。

Natural language processing(NLP)

为了评估作者的方法在NLP任务上的有效性, 作者在Long Range Arena（LRA）套件中的五个基准测试（Tay等人, 2020年）上训练了模型：ListOps, 文本分类, 检索, 图像分类和 Pathfinder。作者将激活 和 与softmax进行比较, 发现 和 在自己的训练上效果不佳, 因此只呈现这些缩放激活和softmax的结果。作者的实现遵循了Xiong等人（2021年）的指南。结果汇总于表4。

6 Limitations

尽管作者的工作引入了具有挑战性的新激活，这些激活对传统的softmax方法提出了质疑，但仍需解决一些限制。作者的理论框架主要是为点积自注意力设计的，可能无法立即扩展到其他注意力机制，尽管作者的实证结果表明，在不同的架构中，作者的激活相对于softmax具有竞争力的性能。此外，作者观察到，尽管作者的激活在视觉任务上表现良好，但在NLP任务上的表现不够稳定，这表明可能需要一种更精细的理论方法来处理这些应用。

7 Conclusion

这项工作挑战了传统观点，即自注意力机制中的 Transformer 激活必须产生稀疏概率分布。

作者引入了一个理论框架，分析自注意力矩阵的Frobenius范数，该框架揭示了注意力机制中激活的关键缩放性质。

作者证明了特定多项式激活，其行为与softmax大不相同，满足这些性质。

通过在视觉和NLP任务上的广泛实验，作者证明了这些替代激活不仅与softmax竞争，而且有时可以超过softmax，为 Transformer 中的注意力机制提供了全新的视角。

参考文献

[0]. Rethinking Softmax: Self-Attention with Polynomial Activations.

公众号后台回复“数据集”获取100+深度学习各方向资源整理