Nystromformer：一种基于 Nyström 方法的自注意力近似算法

1. 摘要

Transformer 已经成为广泛自然语言处理任务中的强大工具。推动 Transformer 展现出卓越性能的一个关键组件是 self-attention 机制，它对每个 token 编码了其他 token 的影响或依赖关系。虽然 self-attention 机制具有诸多优势，但其在输入序列长度上的二次复杂度限制了其在较长序列上的应用 —— 这是当前社区积极研究的一个主题。

为了解决这一限制，我们提出了 Nystromformer —— 一个在序列长度方面具有良好可扩展性的模型。我们的方法基于将 Nyström 方法应用于标准 self-attention 的近似，从而将其复杂度降低为 $O(n)$。Nystromformer 的可扩展性使其能够应用于包含上千个 token 的长序列。

我们在 GLUE 基准和 IMDB 评论等下游任务中，以标准序列长度进行了评估，发现 Nystromformer 的性能与标准 self-attention 可比，有时甚至略优。在长序列任务的 Long Range Arena（LRA）基准中，Nystromformer 相较于其他高效 self-attention 方法表现更优。

我们的代码已开源：https://github.com/mlpen/Nystromformer

2. 引言

基于 Transformer 的模型，如 BERT（Devlin 等，2019）和 GPT-3（Brown 等，2020），在自然语言处理（NLP）领域取得了巨大成功，在机器翻译（Vaswani 等，2017）、自然语言推理（Williams, Nangia, and Bowman，2018）、文本复述（Dolan and Brockett，2005）、文本分类（Howard and Ruder，2018）、问答系统（Rajpurkar 等，2016）以及许多其他任务中都达到了最先进的性能（Peters 等，2018；Radford 等，2018）。

Transformer 的一个核心特性是 self-attention 机制（Vaswani 等，2017），它允许每个 token 的表示来自所有其他 token 的组合。Self-attention 支持整个序列中 token 对之间的交互，并被证明非常有效。

尽管 self-attention 机制具有上述优点，它也成为了效率上的主要瓶颈，因为其在内存和时间上的复杂度为 $O(n^2)$，其中 $n$ 是输入序列的长度。这导致在训练大型 Transformer 模型时需要非常高的内存和计算资源。例如，训练一个 BERT-large 模型（Devlin 等，2019）在单张 Tesla V100 GPU 上需要 4 个月（相当于在一个 4x4 TPU pod 上训练 4 天）。此外，$O(n^2)$ 的复杂度使得训练包含长序列（如 $n=2048$）的大型 Transformer 成本高昂，难以承受。

为了解决这一挑战，最近一些工作提出了在处理长序列输入时避免二次成本的策略。例如，Dai 等（2019）提出了在内存和计算效率之间进行权衡的方案；Child 等（2019）、Kitaev, Kaiser 和 Levskaya（2019）将 self-attention 的复杂度分别降低为 $O(n\sqrt{n})$ 和 $O(n\log n)$；而 Shen 等（2018b）、Katharopoulos 等（2020）和 Wang 等（2020）通过不同的近似策略将复杂度降低为 $O(n)$，每种方法都有其优点和局限性。

在本文中，我们提出了一种在内存和时间上都为 $O(n)$ 的 self-attention 近似方法。我们提出的模型 Nystromformer 能够在输入序列长度 $n$ 上实现线性扩展。这一性能的实现得益于著名的 Nyström 方法，我们将其重新设计用于 self-attention 的近似。具体而言，我们的 Nystromformer 算法使用 landmark（或 Nyström）点 来重构 self-attention 中的 softmax 矩阵，从而避免了直接计算 $n\times n$ 的 softmax 矩阵。我们证明，这种方式能够很好地近似真实的 self-attention。

为了评估我们的方法，我们采用 Transformer 的迁移学习设定，模型首先在大规模语料库上以语言建模目标进行预训练，然后使用监督数据在目标任务上进行微调（Devlin 等，2019；Liu 等，2019；Lewis 等，2020；Wang 等，2020）。仿照 BERT 的做法（Devlin 等，2019；Liu 等，2019），我们在 English Wikipedia 和 BookCorpus（Zhu 等，2015）上使用掩码语言建模目标预训练我们的模型，并观察到与 BERT 基线相似的表现。

随后，我们在 GLUE 基准（Wang 等，2018）和 IMDB 评论数据集（Maas 等，2011）上微调预训练模型，并在准确率和效率两个方面将我们的结果与 BERT 进行比较。在所有任务中，我们的模型在保持与原始 BERT 相近性能的同时，显著提升了速度。

最后，我们在 Long Range Arena（LRA）基准中的长序列任务上评估了我们的模型。Nystromformer 在多个任务中相较于近期高效 self-attention 方法（包括 Reformer（Kitaev 等，2019）、Linformer（Wang 等，2020）和 Performer（Choromanski 等，2020））表现更优，平均准确率高出约 3.4%。我们认为，这项工作是实现资源高效 Transformer 的重要一步。

3. 相关工作

我们简要回顾了与本研究相关的高效 Transformer、线性化 Softmax 核函数及 Nyström 类方法的研究。

高效 Transformer。
已有多种策略被提出以提升 Transformer 的内存效率，包括权重剪枝（Michel, Levy 和 Neubig，2019）、权重分解（Lan 等，2020）、权重量化（Zafrir 等，2019）和知识蒸馏（Sanh 等，2019）。此外，(Clark 等，2019) 中提出了新的预训练目标；(Lample 等，2019) 中使用了 product-key attention；而 (Dai 等，2019) 提出的 Transformer-XL 模型展示了如何减少总体计算需求。

在 (Child 等，2019) 中，作者通过稀疏分解 attention 矩阵，将长序列生成建模的复杂度从二次降低到 $O(n\sqrt{n})$；在 (Kitaev, Kaiser 和 Levskaya，2019) 中，Reformer 模型利用 locality-sensitive hashing（LSH）将复杂度进一步降低到 $O(n\log n)$，其前提是假设 key 和 query 需要相同，从而减少了整体的点积操作次数。

最近，(Wang 等，2020) 提出的 Linformer 模型利用 Johnson-Lindenstrauss 引理中的随机投影，将复杂度降至 $O(n)$，引入了线性投影步骤。(Beltagy, Peters 和 Cohan，2020) 提出的 Longformer 模型通过局部窗口 attention 和任务驱动的全局 attention，在处理长文档时实现了 $O(n)$ 的复杂度；BIGBIRD（Zaheer 等，2020）则采用稀疏 attention 机制。

还有其他提升优化器效率的方法，如 micro-batching（Huang 等，2019）和梯度检查点（gradient checkpointing，Chen 等，2016）。与我们的方法并行发展的是 Performer 模型（Choromanski 等，2020），其利用正交随机特征来近似 softmax attention 核函数，并将复杂度降低到 $O(n)$。

线性化 Softmax。
在 (Blanc 和 Rendle，2018) 中，提出了一种自适应采样 softmax 方法，结合核函数采样以加速训练。该方法在每次训练步骤中仅对部分类别进行采样，基于线性点积近似。在 (Rawat 等，2019) 中，Random Fourier Softmax（RF-softmax）方法利用随机傅里叶特征从近似 softmax 分布中高效采样，用于归一化嵌入表示。

(Shen 等，2018b；Katharopoulos 等，2020) 提出的 Transformer 中 softmax attention 的线性化方法，是通过启发式方式将 key 和 query 分离，基于线性点积进行近似。虽然这种方法非常有趣，但在某些情况下对 self-attention 中 softmax 矩阵的近似误差可能较大。另外，(Bello，2021) 中提出的 lambda layers 也可以被视为一种高效的相对位置 attention 机制。

Nyström 类方法

Nyström 类方法通过对矩阵的列进行采样，以实现对原始矩阵的近似。Nyström 方法（Baker，1977）最初是为了解析带有简单求积公式的积分方程而提出的离散化方法，至今仍被广泛用于通过给定的列子集来近似 kernel 矩阵（Williams 和 Seeger，2001）。许多变种已被提出以改进基础的 Nyström 近似方法，例如：

• 使用 k-means 的 Nyström 方法（Zhang, Tsang 和 Kwok，2008；Zhang 和 Kwok，2010）
• 随机化 Nyström 方法（Li, Kwok 和 Lu，2010）
• 带有谱偏移（spectral shift）的 Nyström 方法（Wang 等，2014）
• 使用伪 landmark 的 Nyström 方法、原型方法（Wang 和 Zhang，2013；Wang, Zhang 和 Zhang，2016）
• fast-Nys 方法（Si, Hsieh 和 Dhillon，2016）
• MEKA 方法（Si, Hsieh 和 Dhillon，2017）
• 集成 Nyström 方法（ensemble Nystrom，Kumar, Mohri 和 Talwalkar，2009）

在 (Nemtsov, Averbuch 和 Schclar，2016) 中，Nyström 方法被扩展以处理一般矩阵（而非对称矩阵）。(Musco 和 Musco，2017) 提出了 RLS-Nystrom 方法，该方法通过递归采样加速 landmark 点的选取。(Fanuel, Schreurs 和 Suykens，2019) 开发了 DAS（确定性自适应采样）和 RAS（随机自适应采样）算法，以增强 landmark 选择的多样性。

与我们的方法最相关的是 (Wang 和 Zhang，2013；Musco 和 Musco，2017)。这些方法设计用于一般矩阵近似（这与我们的设定高度一致），其策略是仅采样部分列和行。然而，直接将这些方法应用于近似 self-attention 中使用的 softmax 矩阵，并不能直接降低计算复杂度，其原因在于：

即使只访问 softmax 矩阵的一部分列或行，也必须先计算完整矩阵的所有元素，再进行 softmax 操作。而计算这些元素本身就需要二次复杂度。

尽管如此，受到“使用列的子集重构完整矩阵”这一核心思想的启发，我们提出了一种 针对 softmax 矩阵量身定制的 Nyström 近似方法，该方法具有 $O(n)$ 的复杂度，用于高效近似 self-attention。

4. 基于 Nyström 的线性 Transformer

在本节中，我们首先简要回顾 self-attention，然后讨论 Nyström 方法用于 self-attention 中 softmax 矩阵近似的基本思想，最后将该思想调整以实现我们提出的结构。

4.1 Self-Attention

什么是 self-attention？Self-attention 计算的是特征表示的加权平均，其权重与表示对之间的相似度成正比。形式上，对于一个包含 $n$ 个 token、维度为 $d$ 的输入序列：$X\in {\mathbb{R}}^{n\times d}$

通过三个矩阵 $W_Q\in {\mathbb{R}}^{d\times d_q}$、$W_K\in {\mathbb{R}}^{d\times d_k}$ 和 $W_V\in {\mathbb{R}}^{d\times d_v}$ 进行投影，从而提取特征表示 $Q$、$K$ 和 $V$，分别被称为 query、key 和 value，其中通常令 $d_k=d_q$。于是，输出表示为：

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V(1)$$

接着上文，self-attention 可以写成如下形式：
$$D(Q,K,V)=SV=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{d_q}}\right)V(2)$$
其中，softmax 表示对每一行进行归一化的 softmax 函数。因此，softmax 矩阵 $S$ 中的每一个元素都依赖于同一行中的所有其它元素。

Self-Attention 的计算代价 Self-Attention 机制需要计算 $n^2$ 个 token 对之间的相似度得分，因此在时间和内存上都具有 $O(n^2)$ 的复杂度。由于这种对输入长度 $n$ 的二次依赖，self-attention 的实际应用通常被限制在较短的序列（例如 $n<1000$）上。这一点正是我们提出一种资源高效的 self-attention 模块的主要动机。

4.2 自注意力中 Softmax 矩阵的 Nyström 近似方法

我们工作的起点是使用广泛用于矩阵近似的 Nyström 方法（Williams and Seeger 2001；Drineas and Mahoney 2005；Wang 和 Zhang 2013），以降低 Transformer 中 self-attention 的计算开销。我们基于 Wang 和 Zhang（2013）的方法，描述了使用 Nyström 方法近似 self-attention 中 softmax 矩阵的策略及其挑战：通过采样一部分列和行来近似整个矩阵。

设 self-attention 中使用的 softmax 矩阵为：$S=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_q}}\right)\in {\mathbb{R}}^{n\times n}$ ，我们可以将 $S$ 表示为：

$$S=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_q}}\right)=\left[\begin{array}{cc}{A}_S& B_S\\ F_S& C_S\end{array}\right]\text{(3)}$$

其中：

• $A_S\in {\mathbb{R}}^{m\times m}$ 是从 $S$ 中采样得到的子矩阵（包含 $m$ 行和 $m$ 列），被称为样本矩阵；
• $B_S\in {\mathbb{R}}^{m\times (n-m)}$；
• $F_S\in {\mathbb{R}}^{(n-m)\times m}$；
• $C_S\in {\mathbb{R}}^{(n-m)\times (n-m)}$。

4.3 Nyström 方法的求积技术

矩阵 $S$ 可以通过 Nyström 方法的基本求积技术进行近似。其第一步是对样本矩阵 $A_S$ 进行奇异值分解（SVD）：$A_S=U\Lambda V^T$ ，其中 $U,V\in {\mathbb{R}}^{m\times m}$ 为正交矩阵，$\Lambda \in {\mathbb{R}}^{m\times m}$ 为对角矩阵。根据 Wang 和 Zhang（2013）提出的“out-of-sample”列近似策略，可以使用从 $S$ 中抽取的 $m$ 列和 $m$ 行重建其 Nyström 近似形式：

$$\hat{S}=\left[\genfrac{}{}{0ex}{}{{\mathcal{A}}_S{\mathcal{B}}_S}{{\mathcal{F}}_S{\mathcal{A}}_S^{+}{\mathcal{B}}_S}\right]=\left[\genfrac{}{}{0ex}{}{{\mathcal{A}}_S}{{\mathcal{F}}_S}\right]{\mathcal{A}}_S^{+}\left[{\mathcal{A}}_S{\mathcal{B}}_S\right](4)$$

其中 $A_S^{+}$ 表示 $A_S$ 的 Moore-Penrose 伪逆。$C_S$ 被近似为 $F_S{A}_S^{+}{B}_S$。这说明 $n\times n$ 的矩阵 $S$ 可以通过从 $S$ 中采样 $m$ 行（即 $A_S$ 和 $B_S$）和 $m$ 列（即 $A_S$ 和 $F_S$）来构造 Nyström 近似 $\hat{S}$。

用于 Softmax 矩阵的 Nyström 近似

接下来，我们简要讨论如何用标准 Nyström 方法构造 self-attention 中 softmax 矩阵的 out-of-sample 近似。

对于给定的 query $q_i$ 和 key $k_j$，我们定义：

• $K_K(q_i)=\text{softmax}\left(\frac{q_i{K}^T}{\sqrt{d_q}}\right)$
• $K_Q(k_j)=\mathrm{softmax}\left(\frac{Q{k}_j^T}{\sqrt{d_q}}\right)$
  

其中 $K_K(q_i)\in {\mathbb{R}}^{1\times n}$，$K_Q(k_j)\in {\mathbb{R}}^{n\times 1}$。我们可以构造

$${\varphi }_K(q_i)={\Lambda }^{-\frac{1}{2}}{V}^T[K_K^T(q_i){]}_{m\times 1}$$

$${\varphi }_Q(k_j)={\Lambda }^{-\frac{1}{2}}{U}^T[K_Q(k_j){]}_{m\times 1}$$

其中 $[\cdot {]}_{m\times 1}$ 表示先计算完整的 $n\times 1$ 向量再取前 $m$ 行。
当 ${\varphi }_K(q_i)$ 和 ${\varphi }_Q(k_j)$ 可得时，标准 Nystrom 近似的 $\hat{S}$ 的元素为
$${\hat{S}}_{ij}={\varphi }_K(q_i{)}^T{\varphi }_Q(k_j),\forall i=1,\dots ,n;j=1,\dots ,n\text{\hspace{1em}(5)}$$

矩阵形式为

$$\hat{S}={\left[\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_q}}\right)\right]}_{n\times m}{A}_S^{+}{\left[\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_q}}\right)\right]}_{m\times n}\text{\hspace{1em}(6)}$$

其中 $[\cdot {]}_{n\times m}$ 表示从 $n\times n$ 矩阵中选取 $m$ 列，$[\cdot {]}_{m\times n}$ 表示从中选取 $m$ 行。该表示对应于式（4）在 self-attention 中 softmax 矩阵近似的应用。式（4）中的 ${[}_{FS}^{A_S}]$ 对应于（6）中的第一个 $n\times m$ 矩阵，$[A_S{B}_S]$ 对应于（6）中的第二个 $n\times m$ 矩阵。矩阵形式的更多细节可见补充材料。
Nystrom 近似的关键挑战。在式（4）和（6）中，即使近似只需要访问 softmax 矩阵的一部分列（即 ${[}_{FS}^{A_S}]$），仍需计算全部 $Q{K}^T$ 的元素，这是因为 softmax 是逐行归一化的。计算 softmax 矩阵中的任意元素都需要该行中所有元素指数的总和作为分母。因此，计算 ${[}_{FS}^{A_S}]$必须访问完整的 $Q{K}^T$，如图 1 所示，这使得直接应用 Nystrom 近似效率不高。

通过 Nystrom 方法实现线性化自注意力
我们现在将 Nystrom 方法应用于对完整的 softmax 矩阵 S 进行近似计算。基本思路是利用来自键 K 和查询 Q 的地标点 $\tilde{K}$ 和 $\tilde{Q}$，从而在不访问完整 $Q{K}^T$ 的情况下，导出一种高效的 Nystrom 近似。当地标点数量 $m$ 远小于序列长度 $n$ 时，我们的 Nystrom 近似在内存和时间两个方面都能实现关于输入序列长度的线性扩展。

温馨提示：
阅读全文请访问"AI深语解构" Nystromformer：一种基于 Nyström 方法的自注意力近似算法