搞懂 Vision Transformer 原理和代码，看这篇技术综述就够了（十八）

本文目录

23 Reformer：高效处理长序列的Transformer (ICLR 2020)
(来自UC Berkeley, Google Research)
23.1 Reformer原理分析

24 Linformer: 低秩矩阵逼近实现新的 Self-Attention (Arxiv)
(来自Facebook AI)
24.1 Linformer原理分析

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制，不采用 RNN 的顺序结构，使得模型可以并行化训练，而且能够拥有全局信息。

Transformer现在已经成了诸多NLP模型的标配，但是它的问题是，模型太大，不能处理较长的序列。在上篇文章 搞懂 Vision Transformer 原理和代码，看这篇技术综述就够了（六） 中我们提到了2种轻量化Transformer模型的DeFINE，DeLighT。它们分别代表2种轻量化Transformer的方式，即：

• DeFINE：One-hot vector → word embedding 的过程轻量化。
• DeLighT：Transformer Block内部轻量化。

它们的共同特点是都是从Transformer结构的外部入手，通过在外部添加一些模块以减少Transformer内部的embedding dimension来减少计算量。但是都没有深入探究如何从Self-attention的角度来减小模型的复杂度。本文介绍的2篇文章就是针对这个问题，即 ：如何从Self-attention的内部设计入手，来减小模型的复杂度。

原版Transformer的Self-attention的计算复杂度与序列长度 成平方关系 。而本文介绍的2种方法可以分别把Self-attention的复杂度降低到：

• Reformer： 。

• Linformer： 。

23 Reformer：高效处理长序列的Transformer (ICLR 2020)

论文名称：Reformer: The Efficient Transformer

论文地址：

https://arxiv.org/pdf/2001.04451.pdf

23.1 Reformer原理分析：

传统attention的计算方法是：

式中 的维度应为 。所以 的维度是 。假设实验的序列长度是 ，当batch size为1时， 是个 的矩阵，假设使用32 bit存储，也需要16GB的内存。

为了减小内存的使用我们可以每次计算一个 。虽然这种做法不够高效，但是它的内存使用是与序列长度线性相关的。

每一个 token 作为 query 时，要把序列中所有 token 当成 key 去计算注意力，再在所有 token 上加权得到当前 token 的一个表示，但我们知道注意力一般是非常稀疏的，权重就集中于少数几个 token 上，那不如只在这几个 token 上计算权重并加权，这样就大大减少了 self attention 里 的计算量和内存占用量。

也就是说内存耗费的真正原因是要去计算 。而其实我们真正关心的是 。而softmax的取值主要被其中较大的元素主导，因此对 的每个向量 ，只需要关注 中哪个向量最接近 。比如说如果 的长度是 ，对于每个 ，我们只需要关注其中跟 距离最近的32或64个 。

如下图1所示，对于每一个 ，我们只需要找出黑点对应的 即可。

所以现在的问题是：如何找到与每个 这距离最近的这32或64个 ？或者说怎么才知道那少数几个 token 是哪几个？

假如要完全靠注意力计算出来才能得到的话，怎么可能在计算注意力之前就知道哪几个 token 权重大？ 是不可能，但是在 self attention 里，query 和 key 计算权重，就是简单的内积，和 query 相似的 key 权重大。模型学习到注意力，是指学习到生成正确的 query 以及 key 的表示，在计算注意力时只需要比对 query 和 key 就可以了。- 所以问题转换成，对每一个 query，我先找到相近的几个 key 计算注意力就好了。怎么找？当然不是全部算一遍取 top k，那就与我们减少计算量的初衷相悖，在这里作者用到了 Local Sensitive Hashing (LSH)，局部敏感哈希，大意就是相近的向量，映射到同一哈希值的概率较大，多个相近的、映射到同一哈希值的向量相当于装进了同一个桶里 (bucket)，那么我们只需要对每个桶里的向量计算 self attention。两个向量 ，满足 LSH 的哈希函数 能做到：

本文通过Locality sensitive hashing来解决这个问题。其特点是对于每个向量 ，在经过哈希函数 后，在原来的空间中挨的近的向量有更大的概率获得相同的哈希值。

这里相关领域已经有很多研究，对于不同的距离度量 ，有不同的 满足 。显然在这里我们的距离度量是 cosine 距离，对应的 哈希是球形投影，即将向量投影到一个 b 维超球面上，该球面被分成了 个象限，投影到同一象限的向量即在同一个桶中，该投影哈希具体写出来是：

式中 是尺寸为 的一个随机矩阵。 是一个 维度的行向量。注意 是个 维的向量作 之后的结果，反映了最终映射到这 个桶中的哪一个。

这个哈希过程就如图2所示：根据queries或者keys的哈希值把它们分到不同的桶里面，在图中 在1号桶， 在2号桶， 在3号桶。

接下来的一个问题是，一个桶里面，query 和 key 的数量不一定相等，比如图2就是这样。而且有可能一个桶里许多 query，没有 key。于是作者干脆 share ，即令 query 和 key 相同，都是 embedding 从同一个线性变换出来的。也就是输入经过线性变换 得到 矩阵，经过线性变换 得到 矩阵，那么 ，如图3所示。只不过 key 做了归一化操作 。

chunk 操作：接下来作者并不是让每个桶里分别做 self attention，而是做了分段，即把同一个桶里的放在一起，重新排成一个序列，然后等长切成若干个段，段内做 self attention，相邻的段也做一次 attention。这里其实有点疑问，论文的图画的非常理想，每个桶的大小差不多，可能差了一两个可以通过相邻段做 attention 来弥补，但是实际情况并不知道每个桶的大小。也许是因为 attention 本身也是学习出来的，作者这么人为设置，是不是相当于给了一个每个桶大小都趋于相同且等于段长的先验限制了 attention 的学习，如下图4所示。

Multi-round LSH attention

lsh 是有概率把距离接近的 分到一组，但也不能完全保证所有距离接近的 分到一组刚好都分到了一组中因为有概率就有误差，因此就多次 hash 来保证概率，取多轮 hash 的并集来保证相似的向量能落到同一个桶里。这里取并集而不是交集，个人理解是桶一多，hash 其实很稀疏，不相似的向量落在同一个桶的概率远小于相似的向量落在不同桶的概率。

Causal masking for shared-QK attention

正常的 transformer 在 decoder 端是要做时序掩码的，这里 lsh 把序列顺序打乱了，因此也要做对应的处理，保证时序掩码的正确性。- 值得一提的是大部分 self attention 的实现，value 包括了自身，但是在 lsh 里不能包含自身，因为 key 和 value 共享值，自身永远是最相似的。

更具体来说可分为以下5个步骤：

1. 取 ，令输入序列 。

2. 做LSH bucketing，即进行hash计算，得到每个query和key所归属的桶 (不同颜色表示不同的桶)，复杂度 。

3. 根据桶编号对query进行排序，同个桶中，按照query原本的位置进行排序，复杂度 。

4. 对于排序后的新序列，进行 chunk 拆分。

5. 对于每个query只attend自己以及自己之前的chunk，对于这些候选集中相同桶的key进行attend，复杂度 。

式中 为chunk数， 为chunk长度， 。

所以复杂度降为了 。

Reversible Transformer

这一部分的想法来自论文：The Reversible Residual Network: Backpropagation Without Storing Activations，即使用可逆残差网络代替标准残差网络。

这篇论文的出发点是：当前所有的神经网络都采用反向传播的方式来训练，反向传播算法需要存储网络的中间结果来计算梯度，而且其对内存的消耗与网络单元数成正比。这也就意味着，网络越深越广，对内存的消耗越大，这将成为很多应用的瓶颈。由于GPU的显存受限，使得网络结构难以达到最优，因为有些网络结构可能达到上千层的深度。如果采用并行GPU的话，价格既昂贵又比较复杂，同时也不适合个人研究。

上面是torchsummary截图，forward和backward pass size就是需要保存的中间变量大小，可以看出这部分占据了大部分显存。如果不存储中间层结果，那么就可以大幅减少GPU的显存占用，有助于训练更深更广的网络。多伦多大学的Aidan N.Gomez和Mengye Ren提出了可逆残差神经网络，当前层的激活结果可由下一层的结果计算得出，也就是如果我们知道网络层最后的结果，就可以反推前面每一层的中间结果。

这样我们只需要存储网络的参数和最后一层的结果即可，激活结果的存储与网络的深度无关了，将大幅减少显存占用。

实验结果显示，可逆残差网络的表现并没有显著下降，与之前的标准残差网络实验结果基本旗鼓相当。

可逆块结构

可逆神经网络将每一层分割成两部分，通过将 channel 一分为二，做成两路，分别为 和 ，每一个可逆块的输入是 ，输出是 。其结构如下：

正向计算：

逆向计算：

其中 和 都是相似的残差函数，参考残差网络。可逆块的 stride 只能为1，也就是说可逆块必须一个接一个连接，中间不能采用其它网络形式衔接，否则的话就会丢失信息，并且无法可逆计算了，这点与残差块不一样。如果一定要采取跟残差块相似的结构，也就是中间一部分采用普通网络形式衔接，那中间这部分的激活结果就必须显式的存起来。

不用存储激活结果的反向传播

为了更好地计算反向传播的步骤，我们修改一下上述正向计算和逆向计算的公式：

尽管 和 的值是相同的，但是两个变量在图中却代表不同的节点，所以在反向传播中它们的总体导数是不一样的。 的导数包含通过 产生的间接影响，而 的导数却不受 的任何影响。

在反向传播计算流程中，先给出最后一层的激活值 和误差传播的总体导数 ，然后要计算出其输入值 和对应的导数 ，以及残差函数 和 中权重参数的总体导数，求解步骤如下：

计算开销

一个N个连接的神经网络，正向计算的理论加乘开销为N，反向传播求导的理论加乘开销为2N(反向求导包含复合函数求导连乘)，而可逆网络多一步需要反向计算输入值的操作，所以理论计算开销为4N (前向N，反向N(反向计算输入值)+2N(反向传播))，比普通网络开销约多出33%左右。但是在实际操作中，正向和反向的计算开销在GPU上差不多，可以都理解为N。那么这样的话，普通网络的整体计算开销为2N，可逆网络的整体开销为3N，也就是多出了约50%。

再回到Reformer中：

Reformer 中把两个函数 和 分别改成了Self-attention和FFN，这样就对应了 transformer 的可逆结构。

可以看到计算 时只用了上一层的激活值 ，计算 时用了上一步计算出来的 ，因此不需要存储这两个激活值。虽然节省了空间，但是激活函数需要重新算一遍，相当于用时间换空间。 原始论文用在 resnet 里，节约显存可以换得更大的 batch_size，在 transformer 中就可以用来训练更长的 sequence。

注意到Transformer的FFN中的hidden dimension 的值也相当大 (约为4K)，其依然受序列长度影响，为了减少这一部分显存占用，作者有一次采用了 chunking，因为 FFN 这里是不存在序列依赖的，完全可以拆成几段计算。为了减少这一部分显存占用，作者有一次采用了 chunking，因为 FFN 这里是不存在序列依赖的，完全可以拆成几段计算，又一次用时间换空间。

此外，对于词典较大的应用场景，作者在计算损失 log-probabilities 时也是分段的。

作者还额外提到，当使用了chunking和reversible layers以后用于存储激活值的memory use就与layer数无关了。这样节省的是反向传播计算梯度时用到的中间临时变量，并不会节省参数量，节省参数量在 GPU 的消耗可以通过将其转到 CPU 内存来解决，通常这样的操作得不偿失，因为在 CPU 和 GPU 之间传输数据非常耗时。

24 Linformer: 低秩矩阵逼近实现新的 Self-Attention

论文名称： Linformer: Self-Attention with Linear Complexity

论文地址：

https://arxiv.org/pdf/2006.04768.pdf

• 24.1 Linformer 原理分析：

Linformer看重的是Attention内部的轻量化。原版Transformer的Self-attention的计算复杂度与序列长度 成平方关系 。作者的目的是设计一个新的Self-attention机制，使得它的计算复杂度与序列长度 成线性关系 。注意不仅是计算量的下降，参数量也会相应减小。由此产生的线性Transformer，即Linformer，性能与标准变压器模型相当，同时具有更大的内存和时间效率。

Linformer的实现基于这样一个假设，即：self-attention是一个低秩矩阵，更准确地说，self-attention的这个随机矩阵可以用低阶矩阵近似。根据这个假设，就可以把self-attention在空间和时间复杂度降低为 操作。做法是把原来的dot-product attention通过linear projection变成小一点的attention，这些attention形成了原版self-attention的低秩分解。

为了下文叙述的方便我们先统一下attention操作的符号：

式中， 代表input embedding matrices， 为序列长度， 是embedding dimension， 是head的数量，每个head的计算方法是：

其中， 是attention操作的权重：

，一般来说 。

self-attention就是为了计算上式中的context mapping matrix ，它的计算要把2个 的矩阵作矩阵乘法，其计算复杂度是 。

Self-Attention is Low Rank

回到一开始的假设：为什么self-attention是一个低秩矩阵？

分析2个已有的预训练模型：RoBERTa-base (12层堆叠的Transformer)和RoBERTa-large (24层堆叠的Transformer)。任务是masked-language-modeling task和分类任务，数据集分别是Wiki103和IMDB。作者对模型的不同层、不同注意力头对应的矩阵 ，都进行了谱分解。

谱分解( Spectral decomposition) 又叫做 特征分解 (Eigendecomposition) ，又称谱分解（Spectral decomposition）是将矩阵分解为由其特征值和特征向量表示的矩阵之积的方法。需要注意只有对可对角化矩阵才可以施以特征分解。

可对角化矩阵是线性代数和矩阵论中重要的一类矩阵。如果一个方块矩阵 相似于对角矩阵，也就是说，如果存在一个可逆矩阵 使得 是对角矩阵，则它就被称为可对角化的。

下面是对特征分解的一个简单的介绍，熟悉的读者可以直接跳到图1。

当且仅当下式成立时， 维非零向量 是 矩阵 的特征向量。

其中 为一标量，称为 对应的特征值。也称 为特征值 对应的特征向量。在一定条件下（如其矩阵形式为实对称矩阵的线性变换），一个变换可以由其特征值和特征向量完全表述，也就是说：所有的特征向量组成了这向量空间的一组基底。为了求出特征值，令

显然，这个方程是一个齐次线性方程组。

有非零解的充分必要条件是 ，即 。

所以求特征值和特征向量的过程转化为了求 的解的过程，而这个式子也叫做矩阵 的特征方程， 是关于 的 次多项式，称为矩阵 的特征多项式。由代数基本定理，特征方程有 个解。这些解的解集也就是特征值的集合，有时也称为 “谱”(Spectrum) 。对 进行因式分解，可以得到：

其中， 。对每一个特征值 ，我们有下式成立：

对每一个特征方程，都会有 个线性无关的解。这 个向量与一个特征值 相对应。这里，整数 称为特征值 的几何重数，而整数 称为特征值 的d代数重数。这里需要注意的是几何重数与代数重数可以相等，但也可以不相等。一种最简单的情况是