IceFormer | 加速 Transformer 模型即插即用，无需重新训练，更快的推理速度，CPU 上也能部署 LLM ！

AIGC 先锋科技

发布于 2024-07-08 13:15:09

1370

发布于 2024-07-08 13:15:09

文章被收录于专栏：AIGC 先锋科技

现有基于Transformer的模型的一个局限性是，由于它们的自注意力操作具有二次时间复杂度和空间复杂度，因此无法处理非常长的序列作为输入。当Transformer部署在仅配备CPU的硬件平台上时，这个问题尤为严重。为了解决这个问题，作者提出了一种在推理时加速自注意力的新方法，该方法可以与预训练的Transformer模型即开即用地工作，而无需重新训练。作者使用作者的方法来加速各种长序列Transformer，包括一个基于LLaMA 2的领先LLM，在各种基准测试上展示了

2.73\times-7.63\times

的速度提升，同时保持了原始预训练模型

98.6\%-99.6\%

的准确率。代码可在作者的项目网站https://yuzhenmao.github.io/IceFormer/上获取。

1 Introduction

Transformer推动了NLP领域的巨大进步，例如大型语言模型（LLM）如GPT-4和LLaMA 2。越来越大的LLM被应用于极长的输入序列，这使许多激动人心的应用成为可能，如长篇内容创作、扩展对话以及大型文档搜索和分析。尽管可以用昂贵的硬件加速器（例如GPU）来训练LLM，但它们需要在普通设备上部署，而这些设备可能只配备了CPU。

然而，由于计算成本高（Dice & Kogan, 2021），目前在CPU上部署LLM是一个挑战。一个重要的计算瓶颈来自于Transformer核心的自注意力机制——其时间和空间复杂度都是输入序列长度的二次方。在LLM的背景下，这个问题更加严重，因为LLM经常用于非常长的序列。

为了处理长输入序列，已经进行了大量研究以减少自注意力的二次时间复杂度——这些方法统称为“高效Transformer”。然而，许多方法无法满足LLM的需求，因此难以应用于LLM。

理想的LLM加速方法应满足四个标准：

（1）无需重新训练 - 给定训练LLM的巨大计算开销，该方法不应要求模型重新训练；

（2）通用性 - 该方法应适用于各种LLM，而不仅仅是那些内置特定约束的模型；

（3）高准确度 - 该方法不应引入大的近似误差，因为LLM具有许多注意力层，因此早期层的错误可能会累积；

（4）快速推理 - 该方法应实现比现有方法更快的测试时性能。

同时满足所有这些标准是困难的，据作者所知，没有现有方法能够做到这一点。例如，具有固定注意力模式的Transformer，如Longformer（Beltagy et al., 2020），在使用之前需要重新训练模型。Reformer（Nikita et al., 2020）要求对键进行归一化——大多数预训练模型不满足这一要求。Nystronformer（Xiong et al., 2021）和LARA（Zheng et al., 2022）不支持因果 Mask ，这在LLM中很常见。低秩方法如Performer（Ch Choromanski et al., 2020）引入了大量的近似误差，尤其是当它们没有重新训练/微调时。

在本文中，作者提出了一种加速方法，作者称之为IceFormer，因为其能够在无需重新训练的情况下直接应用于冻结模型，并同时满足上述四个标准。

具体来说，IceFormer

无需重新训练
可应用于大多数LLM
可以准确逼近普通注意力
与现有方法相比，实现了显著更快的推理速度。

作者在图1中说明了作者的方法与Transformer的对比。如图所示，Transformer计算了每个 Query

q_{i}

和键

k_{j}

可能组合的注意力权重

a_{ij}

（第1阶段）并彻底枚举每个 Query 的所有值向量

v_{j}

（第2阶段）。相比之下，作者的方法利用了注意力矩阵的稀疏性，只计算最高的注意力权重并仅枚举与它们相关的值向量。

在所有三个基准测试中，IceFormer展示了比现有方法快得多的推理速度，同时与Transformer相比几乎没有准确度损失。在LRA基准测试中，IceFormer平均实现了相对于Transformer的

7.63\times

加速，同时保持了

98.6\%

的准确度。与每个任务上准确度相当的最好高效Transformer相比，IceFormer平均快

3.04\times

。在ZeroSCROLLS基准测试中，IceFormer平均实现了相对于基于LLaMA 2的领先LLM的

2.73\times

加速，同时保持了

99.6\%

的准确度。

2 Related Work

高效的 Transformer 可以从两个维度进行分类：方法类型和重新训练要求。第一个维度上有基于稀疏性和低秩性的方法。第二个维度上分为可以不重新训练直接应用于常见预训练 Transformer 的方法和不可以的方法。

基于稀疏性的方法采用稀疏化注意力机制来捕捉全局信息并将其与局部注意力结果整合。一些方法旨在改善空间复杂度，而不是时间复杂度，例如相比于原始注意力机制，top-

注意力（Gupta et al., 2021）。其他方法旨在改善两者，例如Sparse Transformer（Child et al., 2019）、Longformer 和ETC。

这些模型的重大局限在于，被关注到的token是预定义的且保持静态，不适应变化的输入序列。由于原始注意力操作被允许关注任何token，这些模型必须在训练时加入它们各自预定义的关注token的约束。Reformer 通过使用局部敏感哈希（LSH）（Andoni et al., 2015）将token分组到块中，并随后只关注与每个 Query 在同一块中以及相邻块中的token，从而可以关注不同输入序列的不同token集。然而，Reformer施加了原始注意力操作中没有的两个约束：键必须归一化，且 Query 和键必须相同。

因此，Reformer必须在内部构建这些约束的情况下进行训练。结果，这些方法不能直接应用于未经修改的预训练模型；相反，模型必须在使用这些方法之前，以所需的约束重新训练。

低秩方法用低秩矩阵近似注意力权重矩阵，以减少二次时间复杂度和空间复杂度。例如Linformer（Wang et al., 2020）和Performer（Choromanski et al., 2020），它们将注意力权重矩阵分解为学习到的键和 Query 的线性特征或随机特征的瘦高矩阵和宽矩阵的乘积。然而，这些 Transformer 通常引入显著的近似误差，因为原始注意力操作产生的注意力权重矩阵，特别是在长输入序列的情况下，通常具有高秩。

因此，使用这些方法的模型必须在内部构建低秩近似，以学习对相关的近似误差健壮。结果，这些方法不能直接应用于未经修改的预训练模型；相反，模型必须在使用这些方法之前，以所需的近似重新训练。其他方法提供了更一般的方法论，可以利用标准 Transformer 预训练的权重而不需要重新训练。这些 Transformer 加速了标准注意力操作的实施，而不改变底层架构。两个例子是Nystromformer（Xiong et al., 2021）和LARA（Zheng et al., 2022），它们用单独激活的 Query 和键矩阵的乘积替换自注意力机制中的softmax结构。

Nystromformer利用Nystrom方法，而LARA结合随机注意力（RA）和随机特征注意力（RFA）（Peng et al., 2021）来重构注意力权重矩阵。在另一个例子中，H-Transformer-1D递归地将注意力权重矩阵划分为块，并截断每个非对角块的小奇异值。所有这些方法都利用了低秩近似，而不是稀疏性。

其他工作提出了针对特定硬件的优化，而不是旨在改善计算复杂度。例如FlashAttention（Dao et al., 2022），它优化了GPU内存层次之间的读写操作，以及H2O（Zhang et al., 2023），它通过KV缓存逐出策略动态保持最近的和频繁出现的token的平衡。这些策略依赖于实现，并且特定于特定的硬件平台（例如GPU）。

3 Notation and Preliminaries

在数学上，注意力操作接收三个矩阵作为输入，

\mathbf{K}\in\mathbb{R}^{m\times d},\mathbf{Q}\in\mathbb{R}^{n\times d}, \mathbf{V}\in\mathbb{R}^{m\times d^{\prime}}

，分别代表键（keys）、 Query （queries）和值（values），并输出一个矩阵

\mathbf{O}\in\mathbb{R}^{n\times d^{\prime}}

。可选地，它还可能接收一个 Mask 作为输入，

\mathbf{S}\in\mathbb{R}^{n\times m}

，其条目为0或1。

\mathbf{K}

、

\mathbf{Q}

、

\mathbf{V}

和

\mathbf{O}

的第

行，分别表示为

\mathbf{k}_{i}

、

\mathbf{q}_{i}

、

\mathbf{v}_{i}

和

\mathbf{o}_{i}

，分别代表第

个键、 Query 、值和输出。

\mathbf{S}

中第

行第

列的条目，表示为

s_{i,j}

，代表第

个 Query 是否允许关注第

个键——如果为1，则允许；如果为0，则不允许。一种常见的 Mask 方案是因果 Mask ，其中当

i\geq j

时，

s_{i,j}

为1，否则为0。键和 Query 具有相同的维度

，每个键关联一个值，因此键和值的数量相同，表示为

。

首先，注意力操作计算注意力权重矩阵

\mathbf{A}\in\mathbb{R}^{n\times m}

。其第

行第

列的条目，表示为

a_{i,j}

，用以下公式计算：

a_{i,j}=\frac{s_{i,j}\exp\left(\frac{\mathbf{q}_{i}^{\top}\mathbf{k}_{j}}{ \sqrt{d}}\right)}{\sum_{j^{\prime}=1}^{m}s_{i,j^{\prime}}\exp\left(\frac{ \mathbf{q}_{i}^{\top}\mathbf{k}_{j^{\prime}}}{\sqrt{d}}\right)} \tag{1}

然后，注意力操作以下列方式结合值和注意力权重：

\mathbf{o}_{i}=\sum_{j=1}^{m}a_{i,j}\mathbf{v}_{j} \tag{2}

注意力矩阵

\mathbf{A}

通常是很稀疏的（Nikita et al., 2020; Gupta et al., 2021），即，在

\mathbf{A}

的每一行中，只有少数几个注意力权重具有显著（大）值，而其余大部分值接近零。假设作者不知何故可以识别出对于每个 Query

\mathbf{q}_{i}

接收最高注意力权重的

个未 Mask 的键，而不计算所有键的注意力权重。那么，原始的注意力矩阵

\mathbf{A}

可以通过只计算标识键的内积来近似，这可以节省大量的时间和计算资源。

4 IceFormer: Accelerated Self-Attention for General Keys without Retraining

为了构建一个无需重新训练的通用加速方法，作者的方法不能要求修改注意力机制以改变注意力模式，或者引入新的模型参数来捕捉注意力模式中的规律性。这使得一些流行的策略变得不可行，例如具有预定义稀疏注意力模式的注意力机制。

因此，很难设计一个加速方法来利用注意力模式中已知的规律性，而又不提出重新训练的要求。因此，作者旨在设计一个不假设注意力模式中存在规律性的加速方法。

为了改进vanilla attention的

O(mn)

复杂度，作者需要自适应地识别最重要的键（即那些获得最高注意力权重的键），而不计算所有的注意力权重。这似乎是一个先有鸡还是先有蛋的问题：作者如何在不将它们与所有其他注意力权重进行比较的情况下知道哪些注意力权重是最高的？

值得注意的是，在Nikita等人（2020年）提出的归一化键的特殊情况下，这可以通过利用

General Retraining-Free Accelerated Attention

作者没有直接在原始键上应用

-NNS，而是首先将键和 Query 嵌入到一个更高维的空间中。受到Neyshabur和Srebro（2015）的启发，作者选择以下键和 Query 嵌入函数，分别表示为

T_{K}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d+1}

和

T_{Q}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d+1}

：

T_{K}(\mathbf{k}_{j}) =\left[\mathbf{k}_{j}/c\quad\sqrt{1-\|\mathbf{k}_{j}\|_{2}^{2}/c ^{2}}\right]^{\top} \tag{3}

T_{Q}(\mathbf{q}_{i}) =\left[\mathbf{q}_{i}/\|\mathbf{q}_{i}\|_{2}\quad 0\right]^{\top} \tag{4}

其中

c\geq\max_{j^{\prime}}\|\mathbf{k}_{j^{\prime}}\|_{2}

至少是所有键的最大范数。

结果表明，可以通过在键嵌入上使用 Query 嵌入执行

-NNS来识别最重要的

个键。下面作者将展示这一点：

\arg\max_{j}a_{i,j}=\arg\max_{j}\text{softmax}_{j}\left(\left\{\frac{\mathbf{ q}_{i}^{\top}\mathbf{k}_{j^{\prime}}}{\sqrt{d}}\right\}_{j^{\prime}=1}^{m}\right) \tag{5}

=\arg\max_{j}\frac{\mathbf{q}_{i}^{\top}\mathbf{k}_{j}}{\sqrt{d}} \tag{6}

=\arg\min_{j}1-2\mathbf{q}_{i}^{\top}\mathbf{k}_{j}/c\|\mathbf{q}_ {i}\|_{2}+1

(7)

=\arg\min_{j}\mathbf{q}_{i}^{\top}\mathbf{q}_{i}/\|\mathbf{q}_{i} \|_{2}^{2}-2\mathbf{q}_{i}^{\top}\mathbf{k}_{j}/c\|\mathbf{q}_{i}\|_{2}+ \mathbf{k}_{j}^{\top}\mathbf{k}_{j}/c^{2}+1-\|\mathbf{k}_{j}\|_{2}^{2}/c^{2}

(8)

=\arg\min_{j}\|\mathbf{q}_{i}/\|\mathbf{q}_{i}\|_{2}-\mathbf{k}_{ j}/c\|_{2}^{2}+1-\|\mathbf{k}_{j}\|_{2}^{2}/c^{2}

(9)

=\arg\min_{j}\|T_{Q}(\mathbf{q}_{i})-T_{K}(\mathbf{k}_{j})\|_{2}^ {2} \tag{10}

Accurate

-NNS for Accelerated Attention

-NNS问题是理论计算机科学中研究最为深入的问题之一。已经开发了多种算法，并且通常可以通过允许以一定概率出现错误来获得显著的速度提升。这类算法被称为随机算法。

在LLM的背景下，注意力层的数量通常很高，因此早期层的错误可能会累积。因此，

-NNS算法实现高准确度是至关重要的。因此，选择一个合适的

-NNS算法至关重要。

大多数

-NNS算法都是基于分桶的，它们将键值放入离散的桶中，并在包含 Query 的桶中进行搜索。另一方面，基于排序的算法比较不同键相对于 Query 的排名，并在排名较高的键中进行搜索。基于分桶的算法有效地使用了一个固定的相似度阈值，因此可能会有可变数量的键（包括零）满足该阈值；而基于排序的算法返回固定数量的键，这实际上等同于根据键的分布选择一个可变的相似度阈值，如图2所示。一个基于分桶的算法的例子是局部敏感哈希（LSH）（Indyk和Motwani，1998年），而基于排序的算法的例子是优先级DCI（Li和Malik，2017年）。如图2所示，LSH将每个键散列到一个与散列值相关联的桶中，而优先级DCI沿着随机方向对键进行排序。

为了加速注意力计算，作者认为基于排序的算法比基于分桶的算法更适合，因为注意力权重取决于不同键之间的相互比较，而不是每个键相对于固定阈值的绝对评估。因此，基于排序的算法可以更好地回忆起真正重要的键。

Fast

-NNS for Accelerated Attention

在Transformer中，注意力层中的键依赖于前一个注意力层的输出。因此，需要为每个注意力层构建一个数据库。因此，选择一个既能够快速构建和 Query 的

-近邻算法非常重要。此外，在大型语言模型（LLMs）的背景下，许多流行模型使用仅解码器架构。这种架构中的注意力层使用因果 Mask ，以防止当前生成的标记依赖于尚未生成的未来标记。这种 Mask 的注意力等同于从

-近邻算法操作的键集合中排除被 Mask 的键。因此，每次生成一个标记时，都会有一个键被解除 Mask 。与每次生成标记时都构建一个新的数据库相比，更有效的方法是逐步向

-近邻数据库中添加键。

幸运的是，优先级DCI在构建和 Query 阶段都非常高效。如果随机投影方向的数目

几乎与数据的内在维数

d^{\prime}\geq 1

一样大，且要查找的最近邻居数

很小，优先级DCI可以在大约

\tilde{O}(dk^{p/\tilde{d}}m^{1-p/\tilde{d}})

的时间内以高概率返回 Query 的确切

-最近邻居，其中

\tilde{O}(\cdot)

抑制了对数因子。其预处理非常轻量，只需要

O(dpm)

的时间。如果作者将其与普通注意力计算的复杂度

O(dmn)

进行比较，观察到不再有依赖于

的项，因此不再有与序列长度成二次依赖的关系。在后面的5.1节中，作者也实证验证了优先级DCI的效率，并发现它比其他十一个领先的

-近邻算法更快。

为了支持因果 Mask ，作者对优先级DCI的实现进行了扩展，以支持增量数据库更新。这可以高效完成，因为数据结构由排序列表组成，如果实现为二叉搜索树，插入和删除可以在

O(\log m)

的时间内完成。

5 Experiments

在本节中，作者将比较不同

-NNS算法之间的召回延迟权衡，然后分析IceFormer在LRA基准测试（Tay等人，2020年）上的性能，该基准测试是长上下文Transformers（Zhu和Soricut，2021年；Xiong等人，2021年；Zheng等人，2022年）的流行基准测试。接下来，作者将展示IceFormer应用于具有长提示作为输入的LLMs在ZeroSCROLLS基准测试（Shaham等人，2023年）和LongEval基准测试（Li等人，2023年）上的优势。为确保结果的鲁棒性，作者的实验使用了各种CPU——在LRA实验中使用了Intel(R) Core(TM) i7-6850K 6核，ZeroSCROLLS实验中使用了AMD Ryzen 9 5950X 16核，LongEval实验中使用了AMD Ryzen 9 5900X 12核。

Different

-NNS algorithms comparison

作者在图3中比较了12种

-近邻搜索（

-NNS）算法在Fashion MNIST数据集上真实最近邻的召回率以及总构建和 Query 时间，包括优先级DCI以及来自ANN基准测试（Aumuller等人，2017年）的最佳性能算法。如图所示，与其它算法相比，优先级DCI在召回率-延迟权衡方面取得了最佳效果，这证明了在需要快速构建和 Query 的作者的设置中，它是非常适用的。

Evaluation on Long Range Arena (LRA) Benchmark

数据集和评价指标。LRA包括五个不同的任务：ListOps（Nangia和Bowman，2018），文档检索（Retrieval）（Radev等人，2013），文本分类（Text）（Maas等人，2011），CIFAR-10图像分类（Image）（Krizhevsky等人，2009）和Pathfinder（Linsley等人，2018）。具体来说，这五个任务都由最多4k个 Token 的序列组成。作者将在附录C.1中总结数据集详细信息。在这个实验中，作者遵循Tay等人（2020）的train/test划分，并报告每个任务的测试数据集分类准确度，注意力模块的平均运行时间以及推理过程中的CPU内存使用情况。 Baseline 方法。除了普通的Transformer之外，作者还与Nystromformer（Xiong等人，2021），H-Transformer-1D（Zhu和Soricut，2021），LARA（Zheng等人，2022），Reformer（Nikita等人，2020），Longformer（Beltagy等人，2020），Performer（Choromanski等人，2020）和Linformer（Wang等人，2020）进行了比较。为了与Reformer进行比较，作者根据Nikita等人（2020）训练了一个共享

\mathbf{Q}

和

\mathbf{K}

的Transformer模型。对于Longformer和Linformer，由于它们引入了额外的参数，作者在从普通Transformer加载预训练权重时随机初始化这些参数。为了公平比较，作者使用了Xiong等人（2021）在PyTorch中实现的LRA评估基准，只替换了自注意力模块，而使每个模型的其他部分与普通Transformer完全相同。

实现细节。对于每个任务，作者首先使用GPU训练一个基于普通Transformer架构的基础模型。然后，作者将普通注意力模块替换为前面提到的八种高效注意力模块之一，并直接应用预训练权重进行推理。为了确保公平比较，作者将批量大小调整为1，消除了对填充 Mask 的需求，因为作者的 Proposal IceFormer在推理过程中自动忽略填充 Mask 。需要注意的是，由于额外的共享-KQ约束，在Pathfinder任务中，作者尝试训练共享-KQ Transformer未能成功。因此，作者从后续分析中排除了相应的结果。此外，在推理过程中，作者总共使用了4个CPU线程。更详细的细节请参考附录C.2。

推理结果。理想情况下，普通Transformer（非共享-KQ）的准确性为其他七种模型（IceFormer（非共享-KQ），Nystromformer，H-Transformer-1D，LARA，Longformer，Performer和Linformer）的近似准确性提供了一个上限。共享-KQ Transformer也是如此。此外，普通Transformer的注意力模块推理时间最长，其他高效Transformer以牺牲预测准确度为代价实现更短的推理时间。表1展示了每种方法的预测准确度和注意力模块的推理时间。超参数设置列在附录C.3中。总的来说，作者提出的IceFormer在所有五个任务上始终优于所有高效Transformer，在所有任务中提供了最佳的准确度近似，同时需要最少的推理时间。这证明了作者模型的通用性和有效性。

速度与准确度权衡。对于IceFormer，增加近似的程度通常会提高模型效率，但可能导致预测性能下降。在这里，作者通过为每个任务将IceFormer返回的候选数量

从3变化到10，研究了近似程度如何影响推理速度和准确度，并在图4中呈现了结果。从图中作者可以看出，对于所有任务，当

变大时，IceFormer可以实现更高的预测准确度，但效率会降低。

内存复杂度分析。表2总结了每个方法在推理过程中的最大内存使用情况。作者使用了与表1相同的超参数，并保持批量大小为1以消除对填充 Mask 的需求。该表显示，IceFormer在所有任务上始终表现出最低的峰值内存使用量。与普通Transformer相比，IceFormer可以节省高达0.862 GB的内存。

Evaluation on Large Language Model (LLM)

作者在LLM环境中也对IceFormer进行了评估。具体来说，作者利用IceFormer来加速LLM中的提示处理过程。作者选择了Vicuna-7b-v1.5-16k，这是从LLaMA 2微调而来，并且是性能最佳的开放源码LLM之一，其上下文长度可达16K个标记，用于以下实验。关于包括IceFormer中的

-NNS中的

选择在内的更全面细节，请参阅附录E.1。

对于以下LLM实验，作者出于以下原因没有将IceFormer与Reformer、LARA和Nystromformer进行比较：Reformer要求共享键和 Query ，而预训练的LLM并非如此；Longformer只提出了一种加速Transformer编码器部分的方法，因此无法应用于仅解码器的LLM；LARA和Nystromformer将不同的标记分到不同的簇中，因此无法处理LLM中的因果 Mask ，这些LLM使用仅解码器架构。所有需要重新训练的 Baseline （Longformer、Performer和Linformer）也从比较中排除。

ZeroSCROLLS结果。作者在ZeroSCROLLS基准（Shaham等人，2023年）上比较了IceFormer与vanilla Vicuna-7b-v1.5-16k模型以及应用于Vicuna-7b-v1.5-16k的H-Transformer-1D，该基准专为LLM设计，包含十个多样化的自然语言任务，这些任务需要理解长输入上下文，包括总结、问题回答、聚合情感分类和信息重新排序。每个任务的序列长度在3k到10k之间变化。作者测量了ZeroSCROLLS分数和注意力模块的延迟。表3显示，IceFormer与标准自注意力相比，速度提高了最多3.0倍，同时至少达到了vanilla未经加速模型性能的99.0%。

LongEval结果与可扩展性分析。为了更全面地分析IceFormer在LLM设置中的可扩展性，作者在LongEval基准（Li等人，2023年）上进行了额外实验，该基准旨在衡量长上下文性能，包括两个任务：提示长度从3k到16k的主题检索任务，以及提示长度从5k到16k的行检索任务。在图5中，作者展示了与不同输入提示长度对应的注意力模块平均延迟以及使用vanilla Vicuna-7b-v1.5-16k模型和IceFormer的推理准确率。从图中可以看出，IceFormer与vanilla Vicuna-7b-v1.5-16k相比，推理准确率几乎相同。值得注意的是，随着提示长度的增加，两种方法和两个任务的推理延迟都有所增加。然而，即使提示长度很长，IceFormer仍保持其可扩展性，并且一致优于vanilla Transformer。此外，随着提示长度的增加，IceFormer与vanilla Transformer之间的延迟差异变得更大，展示了IceFormer在LLM上下文中的卓越可扩展性和效率。

6 Conclusion

在本文中，作者提出了IceFormer，这是一种提高预训练Transformer在CPU上推理时间效率的新方法。值得注意的是，与其他方法不同，IceFormer不需要重新训练，也不需要对注意力机制施加特殊限制，同时实现了高准确性和快速推理。这些优点使得 IceFormer 非常适合在 CPU上部署LLM，尤其是在 LLM 需要处理非常长的序列作为输入时。在三个基准测试上的实验发现有力地证明了作者方法在减少具有双向和因果注意力机制 Transformer 的二次时间空间复杂度方面的有效性。

Proof 1

在这里，作者提供了在标准化键上进行

-最近邻搜索与确定获得最高注意力权重的键之间的数学等价的完整逐步推导。

\arg\max_{j}a_{i,j} =\arg\max_{j}\operatorname*{softmax}_{j}\left(\left\{\frac{ \mathbf{q}_{i}^{\top}\mathbf{k}_{j^{\prime}}}{\sqrt{d}}\right\}_{j^{\prime}=1 }^{m}\right) \tag{11}

=\arg\max_{j}\frac{\mathbf{q}_{i}^{\top}\mathbf{k}_{j}}{\sqrt{d}}

(12)

=\arg\min_{j}\|\mathbf{q}_{i}\|_{2}^{2}-2\mathbf{q}_{i}^{\top} \mathbf{k}_{j}+1 \tag{13}

由于对所有

j^{\prime}

都有

\|\mathbf{k}_{j^{\prime}}\|_{2}=1

，

\|\mathbf{q}_{i}\|_{2}^{2}-2\mathbf{q}_{i}^{\top}\mathbf{k}_{j}+1=\|\mathbf{q }_{i}\|_{2}^{2}-2\mathbf{q}_{i}^{\top}\mathbf{k}_{j}+\|\mathbf{k}_{j}\|_{2}^ {2}=\|\mathbf{q}_{i}-\mathbf{k}_{j}\|_{2}^{2}

，

\arg\max_{j}a_{i,j} =\arg\min_{j}\|\mathbf{q}_{i}\|_{2}^{2}-2\mathbf{q}_{i}^{\top} \mathbf{k}_{j}+1 \tag{14}

=\arg\min_{j}\|\mathbf{q}_{i}-\mathbf{k}_{j}\|_{2}^{2} \tag{15}

Proof 2

在这里，作者提供了4.1节中结果的完整逐步推导，证明了在转换后的键上进行

-最近邻搜索与识别获得最高注意力权重的键之间的数学等价性。

\arg\max_{j}a_{i,j} =\arg\max_{j}\operatorname{softmax}_{j}\left(\left\{\frac{\mathbf{q} _{i}^{\top}\mathbf{k}_{j^{\prime}}}{\sqrt{d}}\right\}_{j^{\prime}=1}^{m}\right) \tag{16}

=\arg\max_{j}\frac{\mathbf{q}_{i}^{\top}\mathbf{k}_{j}}{\sqrt{d}}

(17)

=\arg\max_{j}\mathbf{q}_{i}^{\top}\mathbf{k}_{j}

(18)

=\arg\min_{j}-2\mathbf{q}_{i}^{\top}\mathbf{k}_{j}

(19)

=\arg\min_{j}2-2\mathbf{q}_{i}^{\top}\mathbf{k}_{j}/c\|\mathbf{q} _{i}\|_{2}

(20)

=\arg\min_{j}1-2\mathbf{q}_{i}^{\top}\mathbf{k}_{j}/c\|\mathbf{q} _{i}\|_{2}+1

(21)

=\arg\min_{j}\|\mathbf{q}_{i}\|_{2}^{2}/\|\mathbf{q}_{i}\|_{2}^{2 }-2\mathbf{q}_{i}^{\top}\mathbf{k}_{j}/c\|\mathbf{q}_{i}\|_{2}+\|\mathbf{k}_{j }\|_{2}^{2}/c^{2}+1-\|\mathbf{k}_{j}\|_{2}^{2}/c^{2}

(22)

=\arg\min_{j}\mathbf{q}_{i}^{\top}\mathbf{q}_{i}/\|\mathbf{q}_{i} \|_{2}^{2}-2\mathbf{q}_{i}^{\top}\mathbf{k}_{j}/c\|\mathbf{q}_{i}\|_{2}+\mathbf{ k}_{j}^{\top}\mathbf{k}_{j}/c^{2}+1-\|\mathbf{k}_{j}\|_{2}^{2}/c^{2}

(23)

=\arg\min_{j}(\mathbf{q}_{i}/\|\mathbf{q}_{i}\|_{2}-\mathbf{k}_{j }/c)^{\top}(\mathbf{q}_{i}/\|\mathbf{q}_{i}\|_{2}-\mathbf{k}_{j}/c)+1-\| \mathbf{k}_{j}\|_{2}^{2}/c^{2}

(24)

=\arg\min_{j}\|\mathbf{q}_{i}/\|\mathbf{q}_{i}\|_{2}-\mathbf{k}_{ j}/c\|_{2}^{2}+1-\|\mathbf{k}_{j}\|_{2}^{2}/c^{2}

(25)

=\arg\min_{j}\|\mathbf{q}_{i}/\|\mathbf{q}_{i}\|_{2}-\mathbf{k}_{ j}/c\|_{2}^{2}+\left(0-\sqrt{1-\|\mathbf{k}_{j}\|_{2}^{2}/c^{2}}\right)^{2}

(26)

=\arg\min_{j}\|T_{Q}(\mathbf{q}_{i})-T_{K}(\mathbf{k}_{j})\|_{2}^ {2}

(27)

=\arg\min_{j}\|T_{Q}(\mathbf{q}_{i})-T_{K}(\mathbf{k}_{j})\|_{2} \tag{28}

（

Dataset Details

在作者的LRA实验中，对于检索、文本和路径finder（

64\times 64

版本），作者直接使用了来自LRA代码库1的数据集。因为ListOps和Image原始数据集只包含短序列，所以作者使用与LRA代码库相同的代码生成了更长的ListOps样本，最大长度为4000；对于Image任务，作者使用了CIFAR-10数据集的超分辨率版本，分辨率为

64\times 64

2，而不是原始的低分辨率

32\times 32

CIFAR-10数据集。作者遵循了原始LRA论文（Tay等人，2020年）完全相同的训练/测试划分。LRA数据集的详细信息列在表4中。### 基模型配置

作者遵循先前工作（Zhu和Soricut，2021年）的实验设置来训练基模型。然而，由于作者无法使用原始设置在Image和Pathfinder数据集上成功训练基Transformer模型以达到满意的准确度，因此作者减少了这两个任务的头数和层数。每个任务的基模型详细信息在表5中概述。### Baseline 方法及提出方法的超参数

对于LARA和Nystromformer，作者通过在范围{64, 128, 256, 512, 1024}内优化调整了参数_num_landmarks_。对于H-Transformer-1D，作者通过在范围{64, 128, 256, 512, 1024}内优化调整了参数_block_size_。对于Reformer，作者调整了参数_num_hash_和_bucket_size_：作者考虑了_num_hash_的值范围为{1, 2, 4}，以及_bucket_size_的值范围为{64, 128, 256, 512, 1024}。对于需要重新训练的Longformer、Performer和Linformer，由于它们的性能较差，作者选择了导致最少近似程度的超参数值。对于IceFormer，作者在范围{3, 5, 8, 10, 15, 20}内调整了参数_top_k_。一般来说，_bucket_size_、_num_landmarks_、_block_size_或_top_k_的值越大，近似程度越不积极，意味着模型的性能越接近于原始Transformer。作者为每个模型选择了导致最佳准确度-时间权衡的超参数值，并在表6中列出。

Appendix D Additional Experiments on LRA

近似质量。为了评估各种高效Transformer对普通修改注意模块输出的逼近程度，作者通过计算它们注意模块输出与标准普通注意模块输出（方程2中的

\mathbf{o}_{i}

）之间的L2范数差来衡量逼近误差。使用表6中的相同超参数设置，不同高效Transformer的平均逼近误差在表7中总结。如表所示，IceFormer在所有LRA任务中一致地实现最低的逼近误差，进一步证明了其逼近效果的有效性。

表1和表2结果的视觉表示。表1和表2的结果分别以图6和图7的形式进行视觉展示。

LLM Experiment Setting

作者在第5.3节中使用vicuna-7b-v1.5-16k作为测试的LLM。它包含32个注意力层，每个注意力层包含32个注意力头，其维度等于128。其最大输入序列长度为16,384。作者观察到LLM的不同层之间存在不同程度的稀疏性，并且这种稀疏性在不同提示下保持一致。因此，在 section 5.3中的所有LLM实验中，作者将IceFormer应用于vicuna-7b-v1.5-16k中从第16层到第31层的相对稀疏层的近似。这个选择涵盖了总共16层，相当于模型总层数量的一半。

在IceFormer的每个ZeroSCROLLS基准测试和LongEval基准测试任务的

-NNS中，

定义为：

k=\max(\min(\left\lfloor n*\alpha\right\rfloor,50),30) \tag{29}

其中，

是输入标记的数量，

\left\lfloor x\right\rfloor

是下取整函数，

\alpha

是用户设置的超级参数。在ZeroSCROLLS基准测试中，作者将

\alpha

设置为4e-3用于SSFD和QMsm任务；5e-3用于GvRp、SQLAL、Qspr、Ntrv、MuSQ和BkSS任务；6e-3用于QALT和SpDg任务。在LongEval基准测试中，对于两个任务的所有设置，作者将

\alpha

设置为5e-3。

Causal Masks and Other Inference-Time Efficient Transformers

在主要论文中，作者没有将IceFormer与LARA和Nystromformer在LLMs上进行了比较。在本节中，作者将详细阐述这两种方法在因果 Mask 方面的问题。

图6：IceFormer和 Baseline 方法（vanilla Transformer, Reformer, LARA, Nystromer, H-Transformer-1D）在LRA基准测试上的推理延迟（越小越好）。

图7：在LRA基准测试上的峰值内存使用（GB）。峰值内存使用是指整个程序的总内存使用，包括优先级DCI数据库/索引的内存（越小越好）。

大多数基于随机特征的模型，如LARA和Nystromformer，将不同的标记分组成不同的簇，称为地标。为了使这些模型能够进行因果 Mask ，不仅需要在地标 Level 应用 Mask 以防止未来标记的信息泄露，还需要另外一组 Mask ，以对不同 Query 在同一地标内 Mask 不同数量的标记。后者不是原生支持的，而且特别难以实现。因此，很难将LARA和Nystromformer应用于具有因果 Mask 的模型。IceFormer + LLM的文本输出

在本节中，作者提供了在LLM（vicuna-7b-v1.5-16k）上应用IceFormer时的文本输出，如图8&9所示。作者还将在作者的补充文件夹中包含输入提示的完整信息。提示（8k标记）

一个好奇的用户与人工智能助手之间的对话。助手对用户的问题给出有帮助、详细和礼貌的回答。用户：请帮我总结以下内容：

迅速增长的人口加上数字化和全球化导致互联网用户数量不断增加。（ITU 2019）与这一发展密切相关的是电子商务业务的扩展。特别是在由当前作者所经历的由冠状病毒危机引起的隔离时期，作者在互联网上的时间比以前多了很多，这似乎是线上购物行业增长的大好机会。全世界各国的封锁对在线购买数量产生了重大影响。经济合作与发展组织汇编的统计数据显示，这一假设是正确的，2020年第一季度和第二季度（英国于2020年3月26日首次封锁，美国于2020年3月19日封锁）美国和英国的电子商务零售额分别从11.8%增长到16.1%，从20.3%增长到31.3%。在线购买数量的增加意味着更高的收入，同时也增加了对该行业的重要性。

如果广告商希望他们在OBA实践中获得成功，他们还应该试图了解更多关于用户上网的原因。根据交互式广告模型（Rogers和Thorson 2000），针对这些动机的广告可以取得更好的效果，因为消费者倾向于更关注这些广告，并且认为它们更容易被记住。如果消费者觉得OBA实践可能支持他们追求目标，他们更愿意接受OBA实践。（Kirmani和Campbell 2004）

Keyformer输出示例（8=40）