The Devil in Linear Transformer

文章标题：线性 Transformer 中的“魔鬼”
作者/机构：Zhen Qin¹, Xiaodong Han¹, Weixuan Sun²,³, Dongxu Li², Lingpeng Kong⁴,⁵, Nick Barnes², Yiran Zhong⁴
¹SenseTime Research, ²Australian National University, ³OPPO Research Institute, ⁴Shanghai AI Laboratory, ⁵The University of Hong Kong

A1 主要贡献

本文旨在解决现有基于核（kernel-based）的线性 Transformer 相较于标准 Transformer 性能下降的问题。作者通过深入分析，指出现有线性 Transformer 存在两个核心缺陷，并提出了相应的解决方案，构建了一个名为 TRANSNORMER 的新型线性 Transformer 模型。

核心问题与研究目标:
1. 性能差距：尽管线性 Transformer 在时空复杂度上具有优势（线性复杂度），但在实际应用中，其性能通常劣于标准（vanilla）Transformer，这限制了其广泛采用。
2. 根源探究：本文的目标是识别并解决导致这种性能差距的根本原因。

关键发现与创新点:
作者识别出两个导致性能下降的关键问题：
1. 梯度无界 (Unbounded Gradients)：现有的线性 Transformer 模型大多沿用了标准 Transformer 的注意力计算公式，其中包含了对注意力分数进行缩放（scaling）的操作。作者从理论上证明，这种缩放策略在线性 Transformer 中会导致梯度无界。无界的梯度会使得模型训练过程中的收敛不稳定。
2. 注意力稀释 (Attention Dilution)：标准 Transformer 的 Softmax 注意力机制倾向于关注局部（邻近）的 token。相比之下，线性 Transformer 的注意力分数往往会平庸地分布在整个序列上，即使在模型的浅层也是如此（如图2所示）。这种现象被作者称为“注意力稀"”，它导致模型无法有效保留重要的局部信息，从而损害性能。

为解决上述问题，本文提出了 TRANSNORMER 模型，其主要贡献如下：
1. 提出 NORMATTENTION：为了解决梯度无界问题，作者提出了一种新的线性注意力机制。该机制移除了注意力矩阵的缩放操作，并在注意力层之后增加了一个归一化操作（如 LayerNorm 或 RMSNorm）。这种设计在理论和实验上都被证明可以稳定训练过程中的梯度，从而实现更稳健的收敛。
2. 提出 DIAGATTENTION：为了缓解注意力稀释问题，作者设计了一种对角注意力（diagonal attention）机制。该机制通过将注意力计算限制在 token 的邻近区域内（块状对角模式），强制模型在浅层网络中关注局部上下文信息，从而模仿标准 Transformer 的行为。这种对角注意力可以被高效计算，保持了模型的整体线性时空复杂度。
3. 构建 TRANSNORMER 架构：该模型结合了以上两种注意力机制。在模型的浅层使用 DIAGATTENTION 来捕获局部信息，在深层使用 NORMATTENTION 来整合全局信息并保持训练稳定。这种混合架构充分利用了两种机制的优点。

核心成果:
- TRANSNORMER 在多个标准 NLP 任务上展示了卓越的性能，包括文本分类（GLUE）、语言建模（WikiText-103）以及具有挑战性的长序列基准测试（Long-Range Arena）。
- 在这些任务中，TRANSNORMER 不仅显著优于现有的线性 Transformer 变体，甚至在某些情况下超越了标准 Transformer。
- 同时，TRANSNORMER 保持了线性 Transformer 的时空效率优势，在处理长序列时速度更快、内存占用更小，如图1所示。

图1：在具有挑战性的 Long-Range Arena 基准测试中，与标准 Transformer 和其他竞争方法相比，TRANSNORMER 的内存占用（圆圈大小）更小，并且在速度（x轴）和综合得分（y轴）上表现出明显优势。
图1：在具有挑战性的 Long-Range Arena 基准测试中，与标准 Transformer 和其他竞争方法相比，TRANSNORMER 的内存占用（圆圈大小）更小，并且在速度（x轴）和综合得分（y轴）上表现出明显优势。

A3 背景知识/关键 Observation/设计原则

2 背景与相关工作

标准 Transformer 及其自注意力机制。标准 Transformer 的核心是自注意力机制，它对查询（Q）、键（K）和值（V）矩阵进行操作。这些矩阵是通过对输入 $X \in R^{n \times d}$ 进行线性投影得到的：


其中 $n$ 是输入长度，$d$ 是隐藏维度。输出 $O \in R^{n \times d}$ 的计算公式为：

其中 Softmax(·) 步骤导致了与输入长度相关的二次时空复杂度，这使得标准 Transformer 难以扩展到长输入序列。

高效 Transformer 的分类。为了解决上述效率问题，研究人员探索了多种高效 Transformer。这些方法大致可分为两类：基于模式（pattern based）的方法和基于核（kernel based）的方法。基于模式的方法，如【【29】Big bird: Transformers for longer sequences，2020，NeurIPS】，【【2】Longformer: The long-document transformer，2020，arXiv】，【【23】Sparse sinkhorn attention，2020，ICML】，【【12】Reformer: The efficient transformer，2020，arXiv】，【【6】Generating long sequences with sparse transformers，2019，arXiv】，通过手工设计或可学习的掩码模式来稀疏化注意力计算。

基于核的方法（线性 Transformer）。基于核的方法采用核函数来分解 Softmax 注意力，从而将理论时空复杂度降低到线性。本文为简洁起见，将这类方法称为线性 Transformer。在这些方法中（如【【7】Rethinking attention with performers，2020，arXiv】，【【11】Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML】，【【17】Random feature attention，2020，ICLR】，【【18】cosformer: Rethinking softmax in attention，2022，ICLR】，【【32】Linear complexity randomized self-attention mechanism，2022，arXiv】，【【28】Linformer: Selfattention with linear complexity，2020，arXiv】），核函数 $φ(·)$ 将查询和键映射到它们的隐藏表示。线性注意力的输出可以重写为：


这里通过先计算键和值的乘积来避免产生 $n \times n$ 的二次复杂度矩阵。

现有线性 Transformer 的局限性。现有方法主要区别在于核函数的设计。例如，Choromanski 等人【【7】Rethinking attention with performers，2020，arXiv】和 Katharopoulos 等人【【11】Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML】使用激活函数 1 + elu 来处理查询和键。Wang 等人【【28】Linformer: Selfattention with linear complexity，2020，arXiv】假设注意力矩阵是低秩的。Peng 等人【【17】Random feature attention，2020，ICLR】和 Zheng 等人【【32】Linear complexity randomized self-attention mechanism，2022，arXiv】在受限的理论边界下近似 Softmax。Qin 等人【【18】cosformer: Rethinking softmax in attention，2022，ICLR】根据 Softmax 函数的经验属性提出了一种线性替代方案。这些方法都致力于近似或改变 Softmax 算子，同时试图保留其性质。然而，与标准 Transformer 相比，这些方法通常以性能换取效率，导致任务表现变差。本文认为，导致这种性能差距有两个根本原因。

3 线性注意力中的“魔鬼”

本节通过为梯度无界问题提供理论证据，并用经验结果展示注意力稀释的负面影响，来阐述 TRANSNORMER 的设计原则。

3.1 梯度无界

对线性 Transformer 训练梯度的分析。很少有关于线性 Transformer 的工作分析其训练过程中的梯度。我们的第一个关键观察是，基于核的线性注意力存在梯度无界的问题，导致训练过程中的收敛不稳定。以下我们强调主要的理论结果，详细推导见附录D。

注意力矩阵的统一形式。考虑一个自注意力模块，无论是标准的还是线性的，其注意力矩阵 $P \in R^{n \times n}$ 都可以用以下统一形式表示：


其中，标准注意力和线性注意力在 token 间相似度 $s_{ij}$ 的计算上有所不同。在标准注意力中，$s_{ij}$ 计算为：

而对于线性注意力，$s_{ij}$ 可以使用核函数 $φ$ 分解为：

梯度推导与对比。根据上述定义，注意力矩阵 $P$ 的梯度推导如下：


对于标准注意力，偏导数 $∂p_{ij}/∂s_{ik}$ 是：

并且它是有界的：

然而，对于线性注意力，我们有：

并且：

由于 $|s_{ik}|^{-1} = |φ(q_i)φ(k_k)^T|^{-1}$ 可以任意大，线性注意力的梯度没有上界。另一方面，我们也可以证明线性注意力的梯度没有下界。

命题 3.1。对于任意 $M > 0$，存在 $q_i, k_j \in R^d, j = 1, . . . , n$，使得：


在我们的初步研究中，梯度无界导致了优化过程更不稳定和更差的收敛结果。

3.2 注意力稀释

标准注意力的局部性。标准注意力机制倾向于强调邻近的 token 是一个已知的特性（【【25】One-vs-each approximation to softmax for scalable estimation of probabilities，2016，arXiv】，【【18】cosformer: Rethinking softmax in attention，2022，ICLR】）。然而，这个特性并没有直接继承到线性 Transformer 的变体中。

量化注意力稀释问题。为了量化注意力稀释问题，我们引入了一个名为“局部累积注意力分数”（locally accumulated attention score）的度量，它衡量在特定 token 的局部邻域内分布了多少注意力分数。对于一个长度为 $N$ 的输入序列，考虑一个以 token $x_i$ 为中心、总长度为 $r \cdot N$ 的局部邻域 $\{x_{start}, ..., x_i, ..., x_{end}\}$，其中 $r$ 是相对于总输入长度的比例，token $x_i$ 的局部累积注意力分数为 $l(i, r, N) = p_{i,start} + ... + p_{i,end}$。更高的分数表明该注意力层更关注局部邻域，而较低的分数则倾向于表明存在注意力稀释问题，即分数更均匀地分布在局部和远距离的 token 上。例如，$l(i, 0.4, N) = 0.6$ 意味着第 $i$ 个 token 周围 40% 的邻居贡献了 60% 的注意力分数。

实验观察。在图 2(a) 中，我们比较了标准 Transformer 和线性 Transformer 的局部累积注意力分数（y轴），其中邻域大小的比例（x轴）在变化。我们展示了在整个序列上每个位置的平均分数。可以看出，标准模型曲线下的面积明显大于线性模型的面积。这证明了标准注意力在局部上更集中，而线性 Transformer 存在注意力稀释问题。图 2(b) 进一步从定性上支持了这一点，其中标准模型的注意力图比线性模型的更集中。

图2：(a) 不同 Transformer 变体的局部累积注意力分数比较。x轴表示邻域大小相对于输入长度的比例；y轴表示中心 token 在该邻域内的累积注意力分数。标准 Transformer 模型的曲线增长更陡峭，表明注意力分数更集中。我们的模型极大地缓解了线性模型的注意力稀释问题。(b) 模型浅层注意力矩阵的定性比较。提出的 TRANSNORMER 产生的模式与原始标准 Transformer 更相似，有助于更好地捕捉局部-全局语言上下文，而线性模型明显受到注意力稀释问题的困扰，在浅层被远距离的 token 分散了注意力。
图2：(a) 不同 Transformer 变体的局部累积注意力分数比较。x轴表示邻域大小相对于输入长度的比例；y轴表示中心 token 在该邻域内的累积注意力分数。标准 Transformer 模型的曲线增长更陡峭，表明注意力分数更集中。我们的模型极大地缓解了线性模型的注意力稀释问题。(b) 模型浅层注意力矩阵的定性比较。提出的 TRANSNORMER 产生的模式与原始标准 Transformer 更相似，有助于更好地捕捉局部-全局语言上下文，而线性模型明显受到注意力稀释问题的困扰，在浅层被远距离的 token 分散了注意力。

A2 方法细节

基于上述观察，我们提出了一种名为 TRANSNORMER 的新型线性 Transformer 网络，旨在解决当前线性 Transformer 的两个局限性。其整体架构如图 3 所示。

4.1 整体架构

结合不同注意力机制的优势。标准注意力在注意力稀释方面问题较小，而线性注意力在处理长序列时更高效和可扩展。这启发我们设计一种方法，通过结合使用这些机制来利用两者的优点。具体来说，我们的网络包含两种类型的注意力：用于模型早期阶段的 DIAGATTENTION 和用于后期阶段的 NORMATTENTION。前者解决注意力稀释问题，后者旨在稳定训练梯度。值得注意的是，通过适当地重塑输入，对角注意力可以以线性的时空复杂度高效计算，从而保持了整体的线性复杂度。

图3：提出的 TRANSNORMER 的架构概览。在早期阶段，我们利用 DIAGATTENTION，其中注意力仅在块内部计算以强制邻域关注。在后期阶段，NORMATTENTION 帮助以线性复杂度获得更稳定的梯度。
图3：提出的 TRANSNORMER 的架构概览。在早期阶段，我们利用 DIAGATTENTION，其中注意力仅在块内部计算以强制邻域关注。在后期阶段，NORMATTENTION 帮助以线性复杂度获得更稳定的梯度。

4.2 NORMATTENTION

移除缩放操作的问题。正如第3节所证明的，线性 Transformer 中的缩放操作（即公式4中的分母）由于梯度无界而阻碍了优化。为了解决这个问题，我们建议移除线性 Transformer 中的缩放操作。然而，如表1所示，直接移除缩放操作会导致严重的性能下降，因为前向传播中的注意力图变得无界。因此，需要一种替代方案来在线性注意力中同时约束前向传播时的注意力图和反向传播时的梯度。
表1：带缩放操作的线性注意力的消融实验。直接移除缩放操作（即公式4中的分母）会导致显著的性能下降。我们的归一化策略取得了更好的结果。

提出 NORMATTENTION。我们提出的解决方案简单而有效。对于一个线性注意力，无缩放的注意力可以表示为：


我们凭经验发现，可以对这个注意力应用任意的归一化来约束它，这便得到了我们的 NORMATTENTION：

其中 XNorm 可以是 Layernorm【【1】Layer normalization，2016】或 RMSNorm【【30】Root Mean Square Layer Normalization，2019，NeurIPS】等。我们在实验中使用了 RMSNorm，因为它比其他选项稍快。

NORMATTENTION 的梯度有界性证明。可以证明 NORMATTENTION 的梯度是有界的：


其中 $L$ 是损失函数，$\epsilon$ 是 RMSNorm 中使用的小常数，$d$ 是嵌入维度，并且

梯度稳定性的实验验证。为了展示 NORMATTENTION 的梯度稳定性，我们将其在每次训练迭代中梯度的相对标准差与其他线性 Transformer 和标准 Transformer 进行了比较。具体来说，我们在 WikiText-103【【15】Pointer sentinel mixture models，2017，ICLR】上使用 RoBERTa 架构训练模型 5 万次迭代，并计算所有迭代梯度的相对标准差。如表2所示，现有的线性方法（【【7】Rethinking attention with performers，2020，arXiv】，【【11】Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML】）相比标准注意力有显著更高的偏差，这导致了较差的结果。NORMATTENTION 产生了更稳定的梯度，验证了我们方法的有效性。
表2：训练5万次迭代的梯度相对标准差。我们提出的 NORMATTENTION 提供了更稳定的梯度，更接近标准 Transformer。

4.3 DIAGATTENTION

在浅层使用局部注意力的重要性。为了更好地理解设计原则，我们在表3中展示了通过将线性 Transformer 的部分层替换为标准注意力，语言建模的性能得到了明显提升。结果还表明，在浅层捕获更多的局部信息比在深层更有帮助。
表3：关于注意力稀释问题的消融实验。我们在相同设置下实现了所有结构：Vanilla【【26】Attention is all you need，2017，NeurIPS】，1 + elu【【11】Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML】。

DIAGATTENTION 的设计与复杂度。为此，我们利用了基于非重叠块的策略来降低标准注意力的时空复杂度。根据图2中的观察，我们使用一种严格的对角分块模式来将注意力约束在一定范围内。由于注意力是在每个块内部计算的，我们的对角注意力的计算复杂度为 $O(nwd)$，其中 $n$ 是序列长度，$w$ 是块大小，$d$ 是特征维度。当 $d \ll n$ 时，复杂度相对于序列长度 $n$ 呈线性增长。在后续部分，我们使用 DIAGATTENTION 来指代对角注意力。

在深层使用全局注意力的必要性。我们凭经验发现，将 DIAGATTENTION 应用于模型的后期阶段会损害性能，如表9所示。这表明模型在后期层需要一个全局的视野，这也证明了我们在 TRANSNORMER 的后期层选择使用 NORMATTENTION 的合理性。

A4 实验环境

模型变体
- TRANSNORMER T1: 在 DIAGATTENTION 中使用 ReLA attention【【31】Sparse attention with linear units，2021，EMNLP】，在 NORMATTENTION 中使用 elu 作为激活函数。
- TRANSNORMER T2: 在 DIAGATTENTION 中使用 Softmax attention【【26】Attention is all you need，2017，NeurIPS】，在 NORMATTENTION 中使用 1+elu 作为激活函数。

实验任务与数据集
1. 自回归语言建模: 在 WikiText-103 数据集【【15】Pointer sentinel mixture models，2017，ICLR】上进行，评估指标为困惑度 (Perplexity, PPL)。
2. 双向语言建模: 在 WikiText-103 上预训练，然后在 GLUE 基准【【27】Glue: A multi-task benchmark and analysis platform for natural language understanding，2018，EMNLP Workshop】的多个下游任务上进行微调。
3. 长序列建模: 在 Long-Range Arena (LRA) 基准【【24】Long range arena: A benchmark for efficient transformers，2020，ICLR】上评估长程依赖建模能力和效率。

硬件配置
- 主要训练在 8块 V100 GPU 上进行。
- 速度比较实验在 一块 Nvidia A6000 GPU (48G显存) 上进行。

软件配置
- 模型在 Fairseq 框架【【16】fairseq: A fast, extensible toolkit for sequence modeling，2019，arXiv】中实现。
- LRA 基准测试在 PyTorch 中重新实现，以反映实际速度。

模型架构与训练参数
- 自回归模型: 基础模型为6个解码器层（FLASH/FLASH-quad为10层），在 WikiText-103 上训练10万步，学习率为0.005。
- 双向模型: 所有方法均基于 RoBERTa-base【【13】Roberta: A robustly optimized bert pretraining approach，2019，arXiv】，包含12个编码器层（FLASH/FLASH-quad为24层以匹配参数量）。在 WikiText-103 上预训练5万步，学习率为0.005，然后在 GLUE 上微调3个周期，学习率从 {1e-5, 3e-5, 6e-5, 1e-4} 中选择最优。
- LRA 模型: 采用与 Skyformer【【3】Skyformer: Remodel self-attention with gaussian kernel and nyström method，2021，NeurIPS】相同的配置，确保所有模型参数量相近。

对比方法
- FLASHquad【【9】Transformer quality in linear time，2022，arXiv】
- FLASH【【9】Transformer quality in linear time，2022，arXiv】
- Transformer-LS【【34】Long-short transformer: Efficient transformers for language and vision，2021，NeurIPS】
- Performer【【7】Rethinking attention with performers，2020，arXiv】
- 1+elu【【11】Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML】

A4 实验结果

5.2 实验结果

自回归语言建模 (表4)
- 实验内容: 在 WikiText-103 数据集上评估模型的困惑度(PPL)。
- 实验结果: TRANSNORMER 的两个变体 (T1, T2) 均取得了与标准注意力相当或更优的 PPL，并且显著优于所有现有的线性模型。例如，与之前在验证集和测试集上的 SOTA 线性方法相比，TRANSNORMER T2 的 PPL 分别降低了2.31和1.58。
- 分析结论: 证明了该方法在因果语言模型中的有效性。
表4：自回归语言建模的定量结果。最佳结果用粗体突出显示，次佳结果用下划线标出。PPL 指标越小越好。LS 代表 transformer-LS。

双向语言建模 (表5)
- 实验内容: 在 GLUE 基准上评估预训练模型的性能。
- 实验结果: TRANSNORMER 在平均性能上优于所有对比方法。在 SST-2, MRPC, CoLA 三个任务上全面优于所有竞争的线性方法，例如在 CoLA 任务上高出4.62分。TRANSNORMER T1 甚至显著优于标准注意力。
- 分析结论: 证明了该方法在双向语言建模中的有效性。
表5：GLUE 基准的定量结果。MNLI 报告匹配/不匹配切分的结果。MRPC 报告 F1 分数。CoLA 报告 Matthews 相关系数。所有其他任务均以准确率衡量。LS 代表 transformer-LS。最佳结果用粗体突出显示，次佳结果用下划线标出。所有指标越大越好。"-" 表示未收敛。

长序列竞技场 (LRA) 基准 (表6)
- 实验内容: 在 LRA 基准的五个任务上评估模型处理长序列的能力。
- 实验结果: 在所有五个任务中，TRANSNORMER 均取得第一或第二名。在总体结果上，两个变体都优于包括标准 Transformer 在内的所有其他竞争方法。
- 分析结论: 验证了模型编码长序列的能力。
表6：Long-Range Arena 基准的定量结果。最佳结果用粗体突出显示，次佳结果用下划线标出。所有指标越大越好。

5.3 速度对比 (表7)

  • 实验内容: 在不同序列长度下比较 TRANSNORMER 与其他方法的训练和推理速度（步/秒）及 GPU 显存占用。
  • 实验结果: 标准 Transformer 在序列长度超过3k时速度极慢且耗尽显存。与其他高效 Transformer 相比，TRANSNORMER 在 GPU 显存占用相当的情况下速度更快。例如，与之前的最佳线性模型 FLASH-quad 相比，TRANSNORMER 的训练速度快300%以上，推理速度快150%以上。
  • 分析结论: TRANSNORMER 在保持高性能的同时，实现了卓越的时空效率。
    表7：Long-Range Arena 基准上的速度比较。如果方法耗尽 GPU 内存，我们用破折号标记。所有指标越高越好。1K,...,5K 代表输入序列长度。

5.4 消融研究

结构设计 (表8, 表9)
- 实验内容: 探究 DIAGATTENTION 和 NORMATTENTION 在模型中的比例和顺序。
- 实验结果: 将前6层设为 DIAGATTENTION、后6层设为 NORMATTENTION 的平衡结构效果最好（表8）。在整个模型中只使用其中一种注意力机制会导致性能下降。将 DIAGATTENTION 用于早期阶段比用于后期阶段效果好得多（表9）。
- 分析结论: 验证了模型早期阶段关注邻近 token、后期阶段需要长程注意力的假设。
表8：注意力比例的消融实验。我们凭经验发现，平衡结构取得了最好的结果。我们将 DIAGATTENTION 缩写为 BlockAtt，NORMATTENTION 缩写为 NormAtt。


表9：两种所提注意力顺序的消融实验。在早期阶段使用 DIAGATTENTION 比在后期阶段使用取得了更好的结果。

FFN 模块选择 (表10)
- 实验内容: 比较传统 FFN 和 GLU 模块。
- 实验结果: GLU【【19】Glu variants improve transformer，2020，arXiv】取得了更好的结果。
- 分析结论: GLU 是更优的 FFN 选择。
表10：FFN 模块选择的消融实验。GLU 带来了更好的结果。

DIAGATTENTION 中的块大小 (表11)
- 实验内容: 探究 DIAGATTENTION 中不同块大小的影响。
- 实验结果: 随着块大小增加，性能有明显提升。
- 分析结论: 考虑到 DIAGATTENTION 的复杂度为 $O(nwd)$，选择块大小为64是在性能和计算成本之间的权衡。
表11：DIAGATTENTION 中块大小的消融实验。块大小越大，结果越好。

注意力组合方式 (表12)
- 实验内容: 探究是否应在同一层内同时使用两种注意力机制，包括顺序使用和并行使用。
- 实验结果: 在一个层内顺序使用这两种注意力机制效果不佳。并行使用它们会使计算复杂度加倍，但性能没有提升。
- 分析结论: 不应在同一层内组合使用这两种注意力。
表12：两种所提注意力组合方式的消融实验。前两行中，两种注意力层交错出现。D 代表 DIAGATTENTION，N 代表 NORMATTENTION。

A5 结论

本文识别并解决了现有线性 Transformer 模型性能不佳的两个关键问题：梯度无界和注意力稀释。

  • 针对梯度无界问题，本文提出了 NORMATTENTION，通过用归一化替代缩放操作来稳定训练梯度。
  • 针对注意力稀释问题，本文开发了 DIAGATTENTION，通过强制模型在浅层关注邻近的 token 来保留局部信息。

将这两种机制结合而成的 TRANSNORMER 模型，融合了标准 Transformer 和线性 Transformer 的优点。实验结果表明，TRANSNORMER 在自回归和双向语言建模、文本分类任务以及具有挑战性的 Long-range Arena 基准测试中，性能均优于现有的线性 Transformer 竞争对手。

局限性与未来工作:
本文主要在自然语言处理任务中对线性 Transformer 的问题进行了全面分析和验证。然而，随着视觉 Transformer 的蓬勃发展，它们是否也存在与 NLP 线性 Transformer 相同的问题尚待探索。作者计划在未来的工作中，将该方法验证到线性的视觉 Transformer 上。

A6 附录

A 数学符号

符号定义。我们使用粗体大写字母表示矩阵(M)，粗体小写字母表示向量(m)，小写字母表示标量($m_{ij}$)。所有向量都表示为列向量，矩阵 M 的第 i 行表示为 $m_{>i}$ 或 $M_i$。我们使用 $k.k_2$ 表示 l2 范数，使用 $k.k_F$ 表示矩阵和向量的 Frobenius 范数。主要的数学符号是输入 $X \in R^{n \times d}$，Q (查询)，K (键) 和 V (值)，其形式如下：


其中 $W_Q, W_K, W_V \in R^{d \times d}$。

B 梯度上界的证明

证明目标。在本部分，我们将证明 (8) 和 (10) 中的界，我们需要证明的是：


我们采用几何平均数小于等于算术平均数的定理，即：

令 $a = p_{ik}, b = 1 - p_{ik}$ 即可完成证明。第一个界可以通过以下方式证明：

对于第二个界，我们首先使用以下事实：

C 命题 3.1 的证明

命题 3.1 的证明。对于任意 $ \epsilon > 0 $ 和核函数 $φ$，令：


则

所以

并且

根据 (10)，我们有：

令 $\epsilon \rightarrow 0^+$，则 $1/n^2, 1/n(1-1/n) \rightarrow \infty$，所以 $∂p_{ij}/∂s_{ik} \rightarrow \infty$。

D 分析各方法的梯度

在本节中，我们考虑一个单层的 Transformer，对于多层 Transformer，我们可以用归纳法证明我们的结论。我们首先引入一些数学符号。

D.1 符号

标准注意力的符号定义。在标准注意力中，我们有：


所以我们有：

线性注意力的符号定义。在线性注意力中，我们有：


虽然在线性注意力中不计算这个项，但我们从概念上讨论它。

统一的公式形式。注意以上公式可以统一为以下形式：


在标准注意力中，我们有：

在线性注意力中，我们有：

NORMATTENTION 的公式。在 NORMATTENTION 中，我们有：


其中 RMSNorm 定义如下：
定义 D.1:

梯度与映射的定义。在后续讨论中，我们定义梯度 $\nabla_M L$ 为：
定义 D.2:


其中 L 代表损失函数，M 是一个参数矩阵。
然后我们定义映射 h 为：
定义 D.3:

映射 h 具有以下性质：
命题 D.4: 对于任意 $X \in R^{n \times m}, Y \in R^{r \times m}$，我们有：

证明: 因为

所以

D.2 梯度分析
D.2.1 准备工作

梯度计算的初步定义。给定梯度 $\nabla_O L \in R^{n \times d}$，让我们计算每种情况下的 $\nabla_S L$。我们首先定义：


在开始之前，我们有以下命题，证明见附录 D.3。
命题 D.5: $c_1 < \infty$.
命题 D.6: 对于任意 $X \in R^{n \times m}$，我们有：

取 $X=V$，我们得到：

D.2.2 标准/线性注意力

统一公式下的梯度。根据 (30)，我们可以在一个公式下讨论标准注意力和线性注意力：


然后定义矩阵 $U^{(i)} \in R^{n \times n}$:

根据 (9)，在标准注意力中，我们有：

而在线性注意力中，我们有：

梯度范数的界。因为：


所以我们有：

其中在标准注意力中 $t=1$，在线性注意力中 $t=c_3$。
另一方面，根据附录 C，在线性注意力中，存在 $q_i, k_j$，使得：

则

令 $\epsilon \rightarrow 0^+$，则 $∂L/∂s_{ik} \rightarrow \infty$。这意味着线性注意力的梯度是无界的。

D.2.3 NORMATTENTION

NORMATTENTION 的梯度计算。我们首先定义 T 的第 i 行的二阶矩：


则 $∂o_{ij}/∂t_{ik}$ 如下：

注意到我们有以下上界：

我们定义矩阵 $R^{(i)} \in R^{d \times d}$ 如下：

因为

然后我们可以得到：

根据 (53)，我们有：

最终，我们得到：

梯度有界性总结。让我们总结一下前面的结果。
在标准注意力中，我们有：


在线性注意力中，存在 $q_i, k_j$，使得：

在 NORMATTENTION 中，我们有：

因此，$∂L/∂s_{ik}$ 在标准注意力和 NORMATTENTION 中是有界的，而在线性注意力中是无界的。这使得线性 Transformer 的训练不稳定。

D.3 命题的证明

命题 D.5 的证明。让我们考虑一个用于分类任务的单层 Transformer。输入为 $X \in R^{n \times d}$，标签为 $Y \in R^{n \times m}$，其中 m 是类别数， $Y_i$ 是 one-hot 向量。$f_1, f_2$ 是激活函数，这里我们以 $f_1 = f_2 = ReLU$ 为例。模型的参数是：


模型的前向传播是：

模型的反向传播是：
1. ... $\nabla_{X_4}L = Y - \text{softmax}(X_4) \in R^{n \times m}$
(a) 上界是:

2. $\nabla_{X_3}L = \nabla_{X_4} L W_4^> \in R^{n \times d}$
$X_{\text{Attention}}$ 代表标准/归一化注意力。
(a) 上界是:

3. ...
(a) 上界是:

4. $\nabla_{X_2}L = (f'_2(X_2W_2) \odot \nabla_{X_3}L) W_2^> \in R^{n \times d}$
(a) 上界是:

5. $\nabla_{X_1}L = (f'_1(X_1W_1) \odot \nabla_{X_2}L) W_1^> \in R^{n \times d_1}$
(a) 上界是:

所以传递到 $X_{\text{Attention}}$ 模块的梯度是有界的，即 $c_1 = a_3 < \infty$。□

命题 D.6 的证明。
...

E 实验配置

在本节中，我们将介绍详细的训练超参数。我们在表 F 中介绍了自回归/双向语言模型的配置。对于 LRA 基准测试，我们使用与 Skyformer 相同的配置，即使用2层 Transformer 模型，隐藏维度为64，2个注意力头，GLU 维度为85，GLU 激活函数为 Swish。对于批量大小和学习率，文本分类使用 16, 1e-4；ListOps 使用 32, 1e-4；文档检索使用 16, 2e-4；Pathfinder 使用 128, 2e-4；图像分类使用 256, 1e-4，与 Skyformer 相同。
表13：我们实验中使用的详细配置。“总批量大小”指 batch_per_gpu × update_freq × num_gpus。“注意力 dropout”仅用于标准注意力。“ALM”：自回归语言模型。“BLM”：双向语言模型。

F 可视化伪代码

在本节中，我们为图2第4列提供 Python 伪代码：

# import torch

def get_curve(w):
    n, m = w.shape
    num = 100
    P = torch.linspace(0, 1, num)
    cnts = torch.zeros(num)
    for i in range(n):
        cnt = torch.zeros(num)
        w1 = w[i].clone()
        center = i % m
        s = w1[center].item()
        L = 1
        l = center - 1
        r = center + 1
        j = 1
        l_thre = 0
        r_thre = m
        flag = 0
        while L < m and j < num:
            if (s >= P[j].item()):
                cnt[j] = L
                j += 1
                continue
            if flag == 1:
                if r != r_thre:
                    s += w1[r].item()
                    r = min(r_thre, r + 1)
                flag = 0
            else:
                if l != l_thre:
                    s += w1[l].item()
                    l = max(l_thre, l - 1)
                flag = 1
            L = min(r - l + 1, m)
        if L >= m:
            for u in range(j, num):
                cnt[u] = min(L, m)
        cnt[0] = 0
        cnts += cnt
    cnts = cnts / n / m
    plt.plot(cnts, P)
    return cnts