Reformer: The Efficient Transformer（2020-1-13）

不负韶华ღ

已于 2022-10-14 19:44:45 修改

阅读量554

点赞数

分类专栏：深度学习（NLP）文章标签： transformer 深度学习机器学习

于 2022-10-10 12:22:40 首次发布

本文链接：https://blog.csdn.net/weixin_49346755/article/details/127239543

版权

深度学习（NLP）专栏收录该内容

44 篇文章 3 订阅

订阅专栏

模型介绍

现如今NLP领域的预训练模型实在是太大了，从最开始的显存装不下，到内存装不下，再到如今的硬盘装不下，让多少人望而却步，大模型就非得这么耗存储吗？有没有优化手段呢？针对长序列的Transformer训练问题，Reformer给出了一种存储极致压缩的方案。

Reformer主要涉及了四处创新点：

使用Axial Positional Embedding来减小位置编码矩阵
提出了基于局部敏感性哈希算法(Locality Sensitive Hashing, LSH)的self-attention机制
提出可以使用分段的方式来处理全连接层(FFN)的计算过程，每次只计算一部分，从而不需要将整个序列加载到内存
使用可逆(Reversible)残差连接来代替传统残差连接，好处是在前向过程中前层的中间值和输出都不需要存储了，只保留最后一层的输出

模型改进

Axial Positional Embedding

Axial Positional Embedding(APE)可用于输入是多维数组(Tensor)的场景，不过好在NLP的输入基本都是一维的序列，理解起来更容易些，我们假设位置向量维度是 $d$ ，输入序列的长度是 $L$ 。并且 $L = m * n$ ，我们可以将序列排列成矩阵 $P$ ，如下图所示，则矩阵元素 $P_{i,j}$ 对应原始序列中第 $(i - 1) * m + j$ 个位置(注: $i$ 和 $j$ 的下标都是从1开始)。
在这里插入图片描述
我们为矩阵 $P$ 的每一行都创建一个向量 $r_i$ ，同时也为每一列都创建一个向量 $c_i$ ，向量维度分别是 $d_1$ 、 $d_2$ ，这样 $P_{ij}$ 就关联了两个向量 $r_i,c_i)$ ，如果我们再令 $d_1+d_2=d$ ，就可以让 $concat(r_i,c_i)$ 表示 $P_{ij}$ 的位置向量，维度是 $d$ ，这就是一维场景下的APE。

我们再来分析下参数量，传统位置向量矩阵大小是 $L * d$ ，使用APE方法会得到两个矩阵，大小分别是 $m*d_1$ 、 $m*d_2$ ，我们再看下enwik8-64K（序列长度是64K）的例子， $L = 10000, d = 512$ ，如果 $m=253,n=253,d_1=256,d_2=256$ ，参数量差不多从3千万减少到了13万！

分段FFN

全连接层一直是神经网络中的参数量大户，Transformer也不例外，
在这里插入图片描述
$W_1$ 和 $W_2$ 大小均是 $d_{model}*d_{ff}$ ，注意Transformer中的FFNN全称是Position-wise Feed-Forward Networks，重点就是这个position-wise，区别于普通的全连接网络，这里FFN的输入是序列中每个位置上的元素，而不是整个序列，所以每个元素完全可以独立计算，最极端节省内存的做法是遍历序列，每次只取一个元素得到FFN的结果，但是这样做时间消耗太大，“分段”的含义就是做下折中，将序列分成N段，也就是N个子序列，每次读取一个子序列进行FFN计算，最后将N份的结果拼接：
在这里插入图片描述
分段FFN只是一种计算上的技巧，计算结果和原始FFN完全一致，所以不会影响到模型效果，好处是不需要一次性将整个序列 $batch\_size,L,d_{model})$ 读入内存，劣势当然是会增加额外的时间开销了。

可逆(Reversible)残差连接

可逆残差网络它最主要的特点是每一层的激活值都可以通过下一层的激活值计算得到，因此不需要保持大量的后向传播的参数，从而达到节约显存的目的。

可逆神经网络按照通道将神经网络输入 $x$ 分成两部分，表示为 $x_1$ 和 $x_2$ ， $F$ 和 $G$ 结果相同，就是普通的残差块，可逆残差块的输出是 $concat(y_1,y_2)$ 。

在这里插入图片描述
在反向过程中， $x_1$ 和 $x_2$ 可以通过以下方法得到：

在这里插入图片描述
下图是可逆残差块的结构：

在这里插入图片描述

Transformer中也含有残差连接，自然也可以改成可逆残差块：

在这里插入图片描述
做法和可逆残差网络完全相同： $Y_1=X_1+Attention(X_2),Y_2=X_2+FeedForward(Y_1)$ 。

Reformer如何分割得到 $X_1$ 和 $X_2$ ，在第一层之前，将输入copy了一份，也就是 $X=X_1=X_2$ 。其余层的 $X_1$ 与 $X_2$ 则对应 $Y_1$ 和 $Y_2$ 。在前向计算过程中，每一个残差块i输出 $Y_{i,1}$ 和 $Y_{i,2}$ ，再作为下一层输入计算得到下一层的 $Y_{i+1,1}$ 和 $Y_{j+1,2}$ 就可以删除 $Y_{i,1}$ 和 $Y_{i,2}$ 。这样，即使Reformer有 N层，在前向过程中我们也只需保留最后一层的 $Y_{n,1}$ 和 $Y_{n,2}$ 即可。

注：可逆残差块的结构和原始Transformer中的结构已经不相同了，原始Transformer是 $Y_1=X_1+Attention(X_1),Y_2=Y_1+FeedForword(Y_1)$ ，所以在反向过程中无法还原出 $X_1$ 和 $X_2$ ，也就是不可逆。

局部敏感性哈希Self-Attention

由于softmax先指数缩放再归一化的本质，使其极易受到极大值影响，特别是长序列，得到的分布几乎总是稀疏的，这种稀疏性分布是有现实含义的：序列中的某个元素一般只会和少数几个元素具有较高的相似性/关联性。

一种通用的思路就是如果我们能为每个query $q_i$ 找到最相似的 $K$ 个key，只和它们计算点积，再取softmax，只要寻找相似key集合速度够快，并且 K << L，这种方式不论是时间还是空间都是有优势的。那么问题来了，如何快速找到和 $q_i$ 最相似的key呢？

这个问题可以转化为向量检索问题，向量检索又可以分为精确(exact)检索和近似(approximate)检索两类，后者在学术界被称为**approximated nearest neighbor search(ANN)**问题，在工业界有非常广泛的应用，毕竟能提供检索服务的场景数据量一般都非常大，如何快速召回就显得格外重要，ANN问题的解决方案包括局部敏感性哈希(Locality Sensitive Hashing, LSH)、树方法和Product Quantization。

Reformer选择的就是LSH，啥是LSH呢？提起哈希函数，大家应该都不陌生，局部敏感性哈希函数是一类特殊的哈希函数，特殊在“局部敏感”上面，如果一个哈希函数具有如下的特点，则它就属于局部敏感性哈希函数：

如果有三个点 $x_1,x_2,x_3$ ，如果 $x_1$ 和 $x_2$ 相邻， $x_1$ 和 $x_3$ 相隔较远，同时 $hash(x_1)$ 和 $hash(x_2)$ 碰撞的概率要比 $hash(x_1)$ 和 $hash(x_3)$ 碰撞的概率大得多，我们就说 $ha s h ()$ 函数属于LSH，也就是它保留了原始数据之间的距离属性。

LSH有很多种，Reformer里用到的是基于随机投影(random projections)的哈希方法，实现非常简单，假设我们希望有 b 个哈希结果，也就是哈希函数对应 b 个桶，首先创建大小是 $d_{model} * \frac b 2$ 的矩阵 $R$ ，矩阵元素服从标准正态分布，向量 $x$ 哈希值等于 $a r g ma x (co n c a t (x R, - x R))$ 。

随机投影是一种降维方法，最大的特点估计就是简单，构造一个随机矩阵 $x * W$ ，完事了，这玩意真的有用？随机投影背后的数学原理是Johnson–Lindenstrauss lemma：

在这里插入图片描述
继续说回LSH self-attention，我们知道，在Transformer中， $Q, K, V$ 是通过 $XW_Q,XW_K,XW_V$ 得到的，由于经过了不同的线性映射(乘以矩阵)，即使是同一位置的 $q_i$ 和 $k_i$ 都很难保证哈希值相同，这还怎么找相似？Reformer的解决办法是让 $Q = K$ ，实验证明，这样做并不会影响模型效果，还有一个问题，当 $Q = K$ 时， $q_i$ 和 $k_i$ 相同， $k_i$ 必然属于最相似的 $K$ 个key，并且 $q_i$ 和 $k_i(q_i)$ 的点积结果一定是最大的，同样会造成softmax稀疏，所以LSH self-attention会mask掉自己。

在这里插入图片描述

在得到每个 $q_i$ 的哈希值后，如何快速得到每个 $q_i$ 的最相似的key集合？也就是和 $q_i$ 落入同一个桶内的 $Set_{q_i}=q_j,j \not =i$ ？Reformer选择了排序，也就是下图中的"Sort by LSH bucket"，将落入同一个桶的 $q$ 排列在一起，下面就可以进行点积计算了，这时候又有一个问题，有的桶元素个数多，有的桶元素个数少，如果以桶数据组成batch，明显各个batch的size不同，不便于批处理，那就分段(chunk)吧，假设序列长度 $L$ ，一共有 $b$ 个桶，则平均下来每个桶内有 $\frac L b$ 个元素，考虑到有的桶元素个数会大于 $\frac L b$ ，为了让同一个桶内的 $q_j$ 都属于 $Set_{q_i}$ ，肯定要让每一段序列长度大于 $\frac L b$ ，Reformer设置的子序列长度是 $2*\frac L b$ ，平均一个子序列包含两个桶的数据，具体在计算点积时， $q_i$ 会在它所在的子序列和前一个子序列中找同一个桶内的 $q_j$ 进行点积，然后计算softmax，得到 $v_i$ 。