改革者 – 推动语言建模的极限
改革者 - 语言建模推进的极限
Reformer如何使用少于8GB的RAM来训练长度为50万个标记的序列
Reformer模型是Kitaev、Kaiser等人于2020年提出的，是迄今为止用于长序列建模的最节省内存的变压器模型之一。
最近，长序列建模经历了一个兴趣激增的时期，从今年的许多论文投稿中可以看出，比如Beltagy等人（2020），Roy等人（2020），Tay等人，Wang等人等。长序列建模的动机在于许多NLP任务，例如摘要、问答，需要模型处理比BERT等模型能够处理的更长的输入序列。在需要模型处理大量输入序列的任务中，长序列模型不必切割输入序列以避免内存溢出，因此已被证明在性能上优于标准的“BERT”模型，参见Beltagy等人（2020）。
Reformer通过能够一次处理多达50万个标记的能力将长序列建模的极限推至极致，如本演示所示。作为比较，传统的bert-base-uncased模型将输入长度限制为仅512个标记。在Reformer中，标准变压器架构的每个部分都经过重新设计，以优化最小内存需求，而性能几乎没有明显下降。
内存方面的改进可以归因于Reformer作者引入的四个特性：
- Reformer自注意层 – 如何高效地实现自注意力而不受局部上下文的限制？
- 分块前馈层 – 如何在大型前馈层中取得更好的时间-内存权衡？
- 可逆残差层 – 如何通过智能残差架构大幅减少训练中的内存消耗？
- 轴向位置编码 – 如何使得位置编码适用于极大输入序列？
本博文的目标是让读者深入了解上述四个Reformer特性的每个细节。虽然解释主要集中在Reformer上，但读者也能更好地理解在其他变压器模型中每个特性何时有效。这四个部分只是松散相关的，因此可以单独阅读。
Reformer是🤗Transformers库的一部分。对于所有使用Reformer的用户，建议阅读这篇非常详细的博文，以更好地了解模型的工作原理和正确设置其配置。所有方程都附有它们在Reformer配置中的等效名称，例如config.<param_name>，以便读者能够快速参考官方文档和配置文件。
注意：轴向位置编码在官方的Reformer论文中没有解释，但在官方代码库中被广泛使用。本博文首次对轴向位置编码进行了深入解释。
1. Reformer自注意层
Reformer使用了两种特殊的自注意层：局部自注意层和局部敏感哈希（LSH）自注意层。
为了更好地介绍这些新的自注意层，我们将简要回顾Vaswani等人在2017年介绍的传统自注意层。
本博文使用与流行博文“可视化变压器”相同的符号和着色，因此强烈建议读者先阅读该博文。
重要提示：虽然Reformer最初是用于因果自注意力的，但它同样适用于双向自注意力。在本博文中，Reformer的自注意力被介绍为双向自注意力。
回顾全局自注意力
每个变压器模型的核心是自注意力层。为了回顾传统的自注意力层，我们将其称为全局自注意力层，假设我们将变压器层应用于嵌入向量序列X = x 1 , … , x n \mathbf{X} = \mathbf{x}_1, \ldots, \mathbf{x}_n X = x 1 , … , x n ，其中每个向量x i \mathbf{x}_{i} x i 的大小为config.hidden_size，即d h d_h d h 。
简而言之，全局自注意层将X投影到查询、键和值矩阵Q，K，V，并使用softmax操作计算输出Z如下：Z = SelfAttn(X) = softmax(QKT)V（为简单起见，省略了关键规范化因子和自注意权重WO）。有关完整Transformer操作的更多细节，请参阅illustrated transformer。
从可视化角度来看，对于n = 16，dh = 3，我们可以将此操作表示如下：
请注意，对于所有可视化，假定batch_size和config.num_attention_heads为1。一些向量，例如x3和其对应的输出向量z3，被标记为后续更好地解释LSH自注意力。所展示的逻辑可以轻松扩展为多头自注意力（config.num_attention_heads > 1）。建议读者参考illustrated transformer以了解多头自注意力。
重要的是要记住，对于每个输出向量zi，整个输入序列X都会被处理。内部点积张量QKT的内存复杂度为O(n^2)，通常在Transformer模型中表示内存瓶颈。
这也是为什么bert-base-cased的config.max_position_embedding_size只有512的原因。
本地自注意力
本地自注意力是减少O(n^2)内存瓶颈的明显解决方案，可以以较低的计算成本对更长的序列进行建模。在本地自注意力中，输入X = X1:n = x1, …, xn被切割成nc个块：X = [X1:lc, …, X((nc-1)*lc:nc*lc)]，每个块的长度为config.local_chunk_length，即lc，并且随后对每个块分别应用全局自注意力。
让我们以n = 16，dh = 3的输入序列再次进行可视化：
假设 lc = 4，nc = 4，分块注意力可以如下所示：
可以看到，注意力操作分别应用于每个块 X1:4，X5:8，X9:12，X13:16。这种架构的第一个缺点变得明显：一些输入向量无法访问它们的即时上下文，例如在我们的示例中，x9无法访问x8，反之亦然。这是有问题的，因为这些令牌无法学习考虑其即时上下文的词表示。
一个简单的解决方法是将每个块增加config.local_num_chunks_before（即np）个块和config.local_num_chunks_after（即na）个块，以便每个输入向量至少能够访问np个前面的输入向量和na个后面的输入向量。这也可以理解为具有重叠的分块，其中np和na定义了每个块与所有前面块和后面块的重叠量。我们将这个扩展的本地自注意力表示为：
Zloc = [Z1:lcloc, …, Z(nc – 1) * lc : nc * lcloc], 其中 Zlc * (i – 1) + 1 : lc * i loc = SelfAttn(Xlc * (i – 1 – np) + 1 : lc * (i + na))[np * lc: -na * lc], ∀ i ∈ { 1, …, nc }
好的，这个公式看起来相当复杂。让我们简化一下。在Reformer的自注意力层中，通常将na设置为0，np设置为1，因此让我们再次为i = 1写出公式：
Z1:lcloc = SelfAttn(X-lc + 1 : lc)[lc:]
我们注意到存在一个循环关系，因此第一个片段也可以参与到最后一个片段中。让我们再次详细说明这种略微改进的局部注意机制。首先，我们在每个窗口化片段内应用自注意力，并仅保留中心输出片段。
最后，相关的输出被连接到 Z loc \mathbf{Z}^{\text{loc}} Z loc ，其如下所示。
请注意，局部自注意力的实现方式非常高效，因此不会计算任何输出，并最终将其“抛弃”，正如红叉所示。
这里需要注意的是，对于每个分块自注意力函数来说，扩展输入向量允许该自注意力函数的每个单独输出向量 z i \mathbf{z}_{i} z i 学习更好的向量表示。例如，输出向量 z 5 loc ， z 6 loc ， z 7 loc ， z 8 loc \mathbf{z}_{5}^{\text{loc}}，\mathbf{z}_{6}^{\text{loc}}，\mathbf{z}_{7}^{\text{loc}}，\mathbf{z}_{8}^{\text{loc}} z 5 loc ，z 6 loc ，z 7 loc ，z 8 loc 可以考虑输入向量 X 1 : 8 \mathbf{X}_{1:8} X 1 : 8 来学习更好的表示。
内存消耗的增益非常明显：O ( n 2 ) \mathcal{O}(n^2) O ( n 2 ) 内存复杂度在每个片段上分解，因此总体渐近内存消耗降低为 O ( n c ∗ l c 2 ) = O ( n ∗ l c ) \mathcal{O}(n_{c} * l_{c}^2) = \mathcal{O}(n * l_{c}) O ( n c ∗ l c 2 ) = O ( n ∗ l c ) 。
这种改进的局部自注意力优于普通的局部自注意力架构，但仍存在一个主要缺点，即每个输入向量只能参与到预定义大小的局部上下文中。对于不需要变换器模型学习输入向量之间的长程依赖关系的自然语言处理任务（例如语音识别、命名实体识别和短句的因果语言建模），这可能不是一个大问题。但是，许多自然语言处理任务需要模型学习长程依赖关系，因此局部自注意力可能导致性能显著下降，例如：
- 问答：模型必须学习问题标记与相关答案标记之间的关系，这两者通常不在同一个局部范围内
- 多选题：模型必须比较多个答案标记片段，这些片段通常相隔较远
- 摘要：模型必须学习上下文标记的长序列与摘要标记的短序列之间的关系，而局部自注意力很可能无法捕捉到上下文和摘要之间的相关关系
- 等等…
局部自注意力本身很可能不足以使变换器模型学习到输入向量（标记）之间的相关关系。
因此，Reformer还使用了一种称为局部敏感哈希（LSH）自注意力的高效自注意力层，用于近似全局自注意力。
局部敏感哈希自注意力
好了，现在我们已经理解了局部自注意力的工作原理，我们可以来看一下Reformer中可能最具创新性的部分：局部敏感哈希（LSH）自注意力。
LSH自注意力的前提是要尽量像局部自注意力一样高效，同时近似全局自注意力。
LSH自注意力依赖于LSH算法，正如Andoni等人在2015年的研究中所介绍的那样，因此得名。
LSH自注意力的思想基于这样一个洞察：如果 n n n 很大，对于每个查询向量，应用于 Q K T \mathbf{Q}\mathbf{K}^T Q K T 注意力点积权重的softmax函数只会给很少的值向量赋予显著大于0的权重。
让我们更详细地解释一下。假设 k i ∈ K = [ k 1 , … , k n ] T \mathbf{k}_{i} \in \mathbf{K} = \left[\mathbf{k}_1, \ldots, \mathbf{k}_n \right]^T k i ∈ K = [ k 1 , … , k n ] T 和 q i ∈ Q = [ q 1 , … , q n ] T \mathbf{q}_{i} \in \mathbf{Q} = \left[\mathbf{q}_1, \ldots, \mathbf{q}_n\right]^T q i ∈ Q = [ q 1 , … , q n ] T 是键向量和查询向量。对于每个 q i \mathbf{q}_{i} q i ，计算 softmax ( q i T K T ) \text{softmax}(\mathbf{q}_{i}^T \mathbf{K}^T) softmax ( q i T K T ) 可以通过仅使用与 q i \mathbf{q}_{i} q i 具有高余弦相似度的键向量 k j \mathbf{k}_{j} k j 近似。这是因为 softmax 函数在较大的输入值上放置指数权重。目前为止，一切顺利，下一个问题是高效地找到与所有 i i i 的 q i \mathbf{q}_{i} q i 具有高余弦相似度的向量。
首先，Reformer 的作者注意到共享查询和键的投影：Q = K \mathbf{Q} = \mathbf{K} Q = K 不会影响 Transformer 模型的性能^1。现在，不再需要为每个查询向量 q i q_i q i 找到高余弦相似度的键向量，而是需要找到查询向量之间的余弦相似度。这很重要，因为查询-查询向量点积的近似具有传递性：如果 q i \mathbf{q}_{i} q i 与查询向量 q j \mathbf{q}_{j} q j 和 q k \mathbf{q}_{k} q k 具有高余弦相似度，那么 q j \mathbf{q}_{j} q j 也与 q k \mathbf{q}_{k} q k 具有高余弦相似度。因此，可以将查询向量聚类到桶中，使得属于同一桶的所有查询向量彼此具有高余弦相似度。让我们将 C m C_{m} C m 定义为第 m m m 组位置索引的集合，使得它们对应的查询向量在同一个桶中：C m = { i ∣ s.t. q i ∈ m t h c l u s t e r } C_{m} = \{ i | \text{ s.t. } \mathbf{q}_{i} \in \text{mth cluster}\} C m = { i ∣ s.t. q i ∈ m t h c l u s t e r } ，将 config.num_buckets ，即 n b n_{b} n b ，定义为桶的数量。
对于每组索引集合 C m C_{m} C m ，对应桶的查询向量的 softmax 函数 softmax ( Q i ∈ C m Q i ∈ C m T ) \text{softmax}(\mathbf{Q}_{i \in C_{m}} \mathbf{Q}^T_{i \in C_{m}}) softmax ( Q i ∈ C m Q i ∈ C m T ) 近似了具有共享查询和键投影的全局自注意力的 softmax 函数 softmax ( q i T Q T ) \text{softmax}(\mathbf{q}_{i}^T \mathbf{Q}^T) softmax ( q i T Q T ) ，其中 i i i 是 C m C_{m} C m 中的所有位置索引。
其次，作者使用 LSH 算法将查询向量聚类到预定义的桶数 n b n_{b} n b 中。LSH 算法是一个理想的选择，因为它非常高效，并且是余弦相似度的最近邻算法的近似。解释 LSH 方案超出了这个笔记本的范围，所以让我们记住对于每个向量 q i \mathbf{q}_{i} q i ，LSH 算法将其位置索引 i i i 归属于 n b n_{b} n b 个预定义的桶之一，即 LSH ( q i ) = m \text{LSH}(\mathbf{q}_{i}) = m LSH ( q i ) = m ，其中 i ∈ { 1 , … , n } i \in \{1, \ldots, n\} i ∈ { 1 , … , n } ，m ∈ { 1 , … , n b } m \in \{1, \ldots, n_{b}\} m ∈ { 1 , … , n b } 。
从视觉上来看，我们可以用以下方式来说明我们的原始示例：
第三，需要注意的是，将所有查询向量聚类到 n b n_{b} n b 个桶中后，对应的索引集合 C m C_{m} C m 可以用来按照 2 {}^2 2 进行输入向量 x 1 , … , x n \mathbf{x}_1, \ldots, \mathbf{x}_n x 1 , … , x n 的排列，从而可以分块应用共享查询-键自注意力，类似于局部注意力。
让我们用我们的示例输入向量 X = x 1 , . . . , x 16 \mathbf{X} = \mathbf{x}_1, …, \mathbf{x}_{16} X = x 1 , . . . , x 1 6 来进行澄清，并假设 config.num_buckets=4 和 config.lsh_chunk_length = 4 。从上面的图中我们可以看到，我们已将每个查询向量 q 1 , … , q 16 \mathbf{q}_1, \ldots, \mathbf{q}_{16} q 1 , … , q 1 6 分配给了其中一个聚类 C 1 , C 2 , C 3 , C 4 \mathcal{C}_{1}, \mathcal{C}_{2}, \mathcal{C}_{3}, \mathcal{C}_{4} C 1 , C 2 , C 3 , C 4 。如果我们现在按照相应的顺序对应排列的输入向量 x 1 , … , x 16 \mathbf{x}_1, \ldots, \mathbf{x}_{16} x 1 , … , x 1 6 ，我们将得到以下排列后的输入 X ′ \mathbf{X’} X ′ ：
自注意机制应该分别应用于每个聚类，以便对于每个聚类 C m \mathcal{C}_m C m ，计算相应的输出如下：Z i ∈ C m LSH = SelfAttn Q = K ( X i ∈ C m ) \mathbf{Z}^{\text{LSH}}_{i \in \mathcal{C}_m} = \text{SelfAttn}_{\mathbf{Q}=\mathbf{K}}(\mathbf{X}_{i \in \mathcal{C}_m}) Z i ∈ C m LSH = SelfAttn Q = K ( X i ∈ C m ) 。
让我们再次用我们的示例来说明这一点。
可以看到，自注意函数在不同大小的矩阵上操作，这对于在 GPU 和 TPU 上进行高效的批处理是不理想的。
为了解决这个问题，可以对排列后的输入进行与局部注意力相同的分块，使得每个块的大小为 config.lsh_chunk_length 。通过对排列后的输入进行分块，一个桶可能会被分成两个不同的块。为了解决这个问题，在 LSH 自注意力中，每个块都与它前面的块 config.lsh_num_chunks_before=1 进行自注意力，与局部自注意力相同（ config.lsh_num_chunks_after 通常设置为 0）。通过这种方式，我们可以确保一个桶中的所有向量都有很高的概率相互关注 3 {}^3 3 。
总之，对于所有的块 k ∈ { 1 , … , n c } k \in \{1, \ldots, n_{c}\} k ∈ { 1 , … , n c } ，LSH 自注意力可以表示如下：
Z ′ l c ∗ k + 1 : l c ∗ ( k + 1 ) LSH = SelfAttn Q = K ( X ′ l c ∗ k + 1 ) : l c ∗ ( k + 1 ) ) [ l c : ] \mathbf{Z’}_{l_{c} * k + 1:l_{c} * (k + 1)}^{\text{LSH}} = \text{SelfAttn}_{\mathbf{Q} = \mathbf{K}}(\mathbf{X’}_{l_{c} * k + 1): l_{c} * (k + 1)})\left[l_{c}:\right] Z ′ l c ∗ k + 1 : l c ∗ ( k + 1 ) LSH = SelfAttn Q = K ( X ′ l c ∗ k + 1 ) : l c ∗ ( k + 1 ) ) [ l c : ]
其中 X ′ \mathbf{X’} X ′ 和 Z ′ \mathbf{Z’} Z ′ 表示根据 LSH 算法重新排列的输入和输出向量。够复杂的公式了，让我们来说明一下 LSH 自注意力。
如上所示，重新排列的向量 X ′ \mathbf{X’} X ′ 被分块，对每个块应用共享的查询键自注意力。
最后，将输出 Z ′ LSH \mathbf{Z’}^{\text{LSH}} Z ′ LSH 重新按照原始的排列方式排序。
还有一个重要的特性需要提到的是，通过同时运行 LSH 自注意力 config.num_hashes 次，例如并行运行 n h n_{h} n h 次，每次使用不同的随机 LSH 哈希函数，可以提高 LSH 自注意力的准确性。通过设置 config.num_hashes > 1 ，对于每个输出位置 i i i ，计算多个输出向量 z i LSH , 1 , … , z i LSH , n h \mathbf{z}^{\text{LSH}, 1}_{i}, \ldots, \mathbf{z}^{\text{LSH}, n_{h}}_{i} z i LSH , 1 , … , z i LSH , n h ，然后将它们合并：z i LSH = ∑ k n h Z i LSH , k ∗ weight i k \mathbf{z}^{\text{LSH}}_{i} = \sum_k^{n_{h}} \mathbf{Z}^{\text{LSH}, k}_{i} * \text{weight}^k_i z i LSH = ∑ k n h Z i LSH , k ∗ weight i k 。其中 weight i k \text{weight}^k_i weight i k 表示哈希轮 k k k 的输出向量 z i LSH , k \mathbf{z}^{\text{LSH}, k}_{i} z i LSH , k 在与其他哈希轮的归一化项 softmax 计算中的重要性比较中的权重，与其 softmax 归一化项的指数成比例。其背后的直觉是，如果相应的查询向量 q i k \mathbf{q}_{i}^{k} q i k 与其所在块中的所有其他查询向量具有较高的余弦相似度，则该块的 softmax 归一化项倾向于较高，因此与具有较低 softmax 归一化项的哈希轮的输出向量相比，相应的输出向量 q i k \mathbf{q}_{i}^{k} q i k 应该是对全局注意力的更好近似，并且应该获得更多的权重。更多细节请参阅论文的附录 A 。对于我们的示例，多轮 LSH 自注意力可以如下所示。
太好了。现在我们知道了 Reformer 中 LSH 自注意力的工作原理。
关于内存复杂度，现在我们有两个竞争成为内存瓶颈的因素：点积运算：O ( n h ∗ n c ∗ l c 2 ) = O ( n ∗ n h ∗ l c ) \mathcal{O}(n_{h} * n_{c} * l_{c}^2) = \mathcal{O}(n * n_{h} * l_{c}) O ( n h ∗ n c ∗ l c 2 ) = O ( n ∗ n h ∗ l c ) 和 LSH 分桶所需的内存：O ( n ∗ n h ∗ n b 2 ) \mathcal{O}(n * n_{h} * \frac{n_{b}}{2}) O ( n ∗ n h ∗ 2 n b ) ，其中 l c l_{c} l c 是块长度。因为对于大的 n n n ，桶的数量 n b 2 \frac{n_{b}}{2} 2 n b 增长得比块长度 l c l_{c} l c 快得多，用户可以根据这里的解释再次分解桶的数量 config.num_buckets 。
让我们快速回顾一下我们上面所讲的内容：
- 我们想要用 softmax 操作只对少数关键向量赋予显著权重的知识来近似全局注意力。
- 如果关键向量等于查询向量，这意味着对于每个查询向量 q i \mathbf{q}_{i} q i ，softmax 只会对其他在余弦相似度方面相似的查询向量赋予显著权重。
- 这种关系是双向的，也就是说如果 q j \mathbf{q}_{j} q j 和 q i \mathbf{q}_{i} q i 相似，那么 q j \mathbf{q}_{j} q j 也和 q i \mathbf{q}_{i} q i 相似，因此我们可以在对一个排列输入应用自注意力之前进行全局聚类。
- 我们在排列输入上应用局部自注意力，并将输出重新排序为原始排列。
1 {}^{1} 1 作者们进行了一些初步实验证实，共享查询键自注意力的性能与标准自注意力的性能差不多。
2 {}^{2} 2 更准确地说，桶内的查询向量按照它们的原始顺序进行排序。这意味着，例如，向量 q 1 , q 3 , q 7 \mathbf{q}_1, \mathbf{q}_3, \mathbf{q}_7 q 1 , q 3 , q 7 都被散列到桶 2 中，那么桶 2 中向量的顺序仍然是 q 1 \mathbf{q}_1 q 1 ，然后是 q 3 \mathbf{q}_3 q 3 和 q 7 \mathbf{q}_7 q 7 。
3 {}^3 3 顺便提一下，作者们在查询向量 q i \mathbf{q}_{i} q i 上放置了一个掩码，防止向量自己参与注意力。因为向量与自己的余弦相似度总是高于或等于与其他向量的余弦相似度，所以共享查询键自注意力中的查询向量被强烈禁止与自己进行注意力。
基准测试
最近在 Transformers 中添加了基准测试工具-详细说明请参见此处。
为了展示使用”本地” + “LSH”自注意力可以节省多少内存，对 Reformer 模型 google/reformer-enwik8 进行了不同的 local_attn_chunk_length 和 lsh_attn_chunk_length 的基准测试。可以在此处详细检查 google/reformer-enwik8 模型的默认配置和用法。
首先，让我们进行一些必要的导入和安装。
#@title 安装和导入
# pip 安装
!pip -qq install git+https://github.com/huggingface/transformers.git
!pip install -qq py3nvml
from transformers import ReformerConfig, PyTorchBenchmark, PyTorchBenchmarkArguments首先，让我们基准测试使用全局自注意力的 Reformer 模型的内存使用情况。这可以通过设置 lsh_attn_chunk_length = local_attn_chunk_length = 8192 来实现，这样对于所有小于等于 8192 的输入序列，模型会自动切换到全局自注意力。
config = ReformerConfig.from_pretrained("google/reformer-enwik8", lsh_attn_chunk_length=16386, local_attn_chunk_length=16386, lsh_num_chunks_before=0, local_num_chunks_before=0)
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[2048, 4096, 8192, 16386], batch_sizes=[1], models=["Reformer"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config], args=benchmark_args)
result = benchmark.run()
HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1279.0, style=ProgressStyle(description…
1 / 1
Doesn't fit on GPU. CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 11.17 GiB total capacity; 8.87 GiB already allocated; 1.92 GiB free; 8.88 GiB reserved in total by PyTorch)
==================== 推理-内存-结果 ====================
--------------------------------------------------------------------------------
模型名称 批次大小 序列长度 内存（MB）
--------------------------------------------------------------------------------
Reformer 1 2048 1465
Reformer 1 4096 2757
Reformer 1 8192 7893
Reformer 1 16386 N/A
--------------------------------------------------------------------------------输入序列越长，输入序列和峰值内存使用之间的二次关系O(n^2)越明显。正如可以看到的，实际上需要一个更长的输入序列才能清楚地观察到将输入序列加倍会使峰值内存使用量增加四倍。
对于使用全局注意力的google/reformer-enwik8模型，序列长度超过16K会导致内存溢出。
现在，让我们使用模型的默认参数来激活局部和LSH自注意力。
config = ReformerConfig.from_pretrained("google/reformer-enwik8")
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[2048, 4096, 8192, 16384, 32768, 65436], batch_sizes=[1], models=["Reformer"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config], args=benchmark_args)
result = benchmark.run()
1 / 1
无法适应GPU。CUDA内存不足。尝试分配2.00 GiB（GPU 0；总容量11.17 GiB；已分配7.85 GiB；剩余1.74 GiB；PyTorch总共预留9.06 GiB）
无法适应GPU。CUDA内存不足。尝试分配4.00 GiB（GPU 0；总容量11.17 GiB；已分配6.56 GiB；剩余3.99 GiB；PyTorch总共预留6.81 GiB）
==================== 推断 - 内存 - 结果 ====================
--------------------------------------------------------------------------------
模型名称 批大小 序列长度 内存（MB）
--------------------------------------------------------------------------------
Reformer 1 2048 1785
Reformer 1 4096 2621
Reformer 1 8192 4281
Reformer 1 16384 7607
Reformer 1 32768 N/A
Reformer 1 65436 N/A
--------------------------------------------------------------------------------如预期所料，对于较长的输入序列，使用局部和LSH自注意力在内存效率上更高，因此在这个笔记本上，对于11GB RAM的GPU，模型在16K个标记处才会耗尽内存。
2. 分块前馈层
基于Transformer的模型通常在自注意力层之后并行使用非常大的前馈层。因此，这一层可以占用相当大的内存，并且有时甚至代表模型的内存瓶颈。在Reformer论文中首次引入的前馈分块是一种技术，可以通过增加时间消耗来有效地减少内存消耗。
Reformer中的分块前馈层
在Reformer中，LSH或局部自注意力层通常后跟一个残差连接，然后定义了变压器块的第一部分。有关此问题的更多详细信息，请参阅此博客。
第一部分变压器块的输出，称为规范化的自注意力输出，可以表示为Z̅ = Z + X，其中Z是Z LSH或Z loc中的Z。
对于我们的示例输入x1，…，x16，我们将规范化的自注意力输出表示如下。
现在，变压器块的第二部分通常由两个前馈层1^1和Linear int（…）组成，该层将Z̅处理为中间输出Y int，以及Linear out（…）将中间输出处理为输出Y out。这两个前馈层可以定义为
输出 Y out = 线性 out (输入 Y int) = 线性 out (线性 int ( Z ‾ ))。
现在需要记住的是，从数学上讲，前馈层在位置 y out , i 的输出 \mathbf{y}_{\text{out}, i} 仅取决于该位置的输入 \mathbf{\overline{y}}_{i}。与自注意层不同，每个输出 \mathbf{y}_{\text{out}, i} 完全独立于不同位置的所有输入 \mathbf{\overline{y}}_{j \ne i}。
让我们以 z ‾ 1 ，…，z ‾ 16 为例来说明前馈层。
从图中可以看出，所有输入向量 z ‾ i 都由相同的前馈层并行处理。
当我们观察前馈层的输出维度时，就会变得有趣。在 Reformer 中，线性 int 的输出维度被定义为 config.feed_forward_size，即 d f，而线性 out 的输出维度被定义为 config.hidden_size，即 d h。
Reformer 的作者观察到，在 transformer 模型中，中间维度 d f 通常比输出维度 2 d h 大得多。这意味着尺寸为 d f × n 的张量 Y int 占用了大量的内存，并且甚至可能成为内存瓶颈。
为了更好地了解维度之间的差异，让我们来看一下我们示例中的矩阵 Y int 和 Y out。
很明显，张量 Y int 比 Y out 占用更多的内存（确切地说，是 d f d h 的 d f 倍）。但是，是否有必要计算完整的中间矩阵 Y int？实际上并不需要，因为只有输出矩阵 Y out 是相关的。为了在内存中减少内存开销，可以将线性层的计算分成多个块进行处理。将 config.chunk_size_feed_forward 定义为 c f，分块线性层定义为 Y out = [Y out, 1:c f, …, Y out, (n – c f):n]，其中 Y out, (c f ∗ i):(i ∗ c f + i) = 线性 out (线性 int (Z ‾ (c f ∗ i):(i ∗ c f + i)))。实际上，这意味着输出是逐步计算并连接起来，避免了存储整个中间张量 Y int。
假设 c f = 1 c_{f}=1 c f = 1 是我们示例的一个假设，我们可以用以下方式说明在位置 i = 9 i=9 i = 9 的输出的增量计算：
通过以大小为 1 的块处理输入，同时需要存储在内存中的张量只有最大大小为 16 × d h 16 \times d_{h} 1 6 × d h 的输出向量 Y out \mathbf{Y}_\text{out} Y out ，大小为 d f d_{f} d f 的中间向量 y int , i \mathbf{y}_{\text{int}, i} y int , i ，以及大小为 16 × d h 16 \times d_{h} 1 6 × d h 的输入向量 Z ‾ \mathbf{\overline{Z}} Z ，其中 d h d_{h} d h 是 config.hidden_size 3 ^{3} 3 。
最后，重要的是要记住，分块线性层和传统线性层产生的输出在数学上是等价的，因此可以应用于所有变压器的线性层。在某些用例中，使用 config.chunk_size_feed_forward 可以在内存和速度之间取得更好的权衡。
1 {}^1 1 为了更简单的解释，现在省略了通常应用于 Z ‾ \mathbf{\overline{Z}} Z 的层归一化层在被前馈层处理之前。
2 {}^2 2 在 bert-base-uncased 中，例如中间维度 d f d_{f} d f 是输出维度 d h d_{h} d h 的四倍大。
3 {}^3 3 作为提醒，在这个笔记本中，为了清晰和说明的目的，假设输出 config.num_attention_heads 为 1 ，因此可以假设自注意力层的输出大小为 config.hidden_size 。
有关分块线性/前馈层的更多信息也可以在 🤗Transformers 文档中找到。
基准测试
让我们测试使用分块前馈层可以节省多少内存。
#@title 安装和导入
# 安装依赖
!pip -qq install git+https://github.com/huggingface/transformers.git
!pip install -qq py3nvml
from transformers import ReformerConfig, PyTorchBenchmark, PyTorchBenchmarkArguments
Building wheel for transformers (setup.py) ... [?25l[?25hdone首先，我们将默认的 google/reformer-enwik8 模型与具有分块前馈层的模型进行比较。
config_no_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8") # 无分块
config_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8", chunk_size_feed_forward=1) # 分块前馈
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[1024, 2048, 4096], batch_sizes=[8], models=["Reformer-No-Chunk", "Reformer-Chunk"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_no_chunk, config_chunk], args=benchmark_args)
result = benchmark.run()
1 / 2
无法适应 GPU。CUDA 内存不足。尝试分配 2.00 GiB（GPU 0；总容量 11.17 GiB；已分配 7.85 GiB；剩余 1.74 GiB；PyTorch 总共保留了 9.06 GiB）
2 / 2
无法适应 GPU。CUDA 内存不足。尝试分配 2.00 GiB（GPU 0；总容量 11.17 GiB；已分配 7.85 GiB；剩余 1.24 GiB；PyTorch 总共保留了 9.56 GiB）
==================== 推断 - 内存 - 结果 ====================
--------------------------------------------------------------------------------
模型名称 批量大小 序列长度 内存使用（MB）
--------------------------------------------------------------------------------
Reformer-No-Chunk 8 1024 4281
Reformer-No-Chunk 8 2048 7607
Reformer-No-Chunk 8 4096 N/A
Reformer-Chunk 8 1024 4309
Reformer-Chunk 8 2048 7669
Reformer-Chunk 8 4096 N/A
--------------------------------------------------------------------------------有趣的是，分块的前馈层在这里似乎没有起到任何帮助。原因是config.feed_forward_size的大小不足以产生真正的差异。只有在序列长度为4096时，才能稍微减少内存使用。
让我们看看如果将前馈层的大小增大4倍，并将注意力头的数量减少4倍，使得前馈层成为内存瓶颈，内存峰值使用情况会发生什么。
config_no_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8", chunk_size_feed_forward=0, num_attention_heads=2, feed_forward_size=16384) # 无分块
config_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8", chunk_size_feed_forward=1, num_attention_heads=2, feed_forward_size=16384) # 前馈分块
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[1024, 2048, 4096], batch_sizes=[8], models=["Reformer-No-Chunk", "Reformer-Chunk"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_no_chunk, config_chunk], args=benchmark_args)
result = benchmark.run()
1 / 2
2 / 2
==================== 推理 - 内存 - 结果 ====================
--------------------------------------------------------------------------------
模型名称 批量大小 序列长度 内存使用（MB）
--------------------------------------------------------------------------------
Reformer-No-Chunk 8 1024 3743
Reformer-No-Chunk 8 2048 5539
Reformer-No-Chunk 8 4096 9087
Reformer-Chunk 8 1024 2973
Reformer-Chunk 8 2048 3999
Reformer-Chunk 8 4096 6011
--------------------------------------------------------------------------------现在可以看到在较长的输入序列中，峰值内存使用量明显减少。总结一下，只有在具有少量注意力头和大型前馈层的模型中，分块的前馈层才有意义。
3. 可逆残差层
可逆残差层最早由N. Gomez等人引入，并用于减少训练流行的ResNet模型时的内存消耗。从数学上讲，可逆残差层与“真实”残差层略有不同，但不需要在前向传递期间保存激活值，这可以大大减少训练时的内存消耗。
Reformer中的可逆残差层
让我们首先研究为什么训练模型需要比推理模型更多的内存。
在推理模型中运行模型时，所需内存大致等于计算模型中“最大”的张量所需的内存。另一方面，在训练模型时，所需内存大致等于所有可微分张量的“总和”。
考虑到深度学习框架中自动微分的工作原理，这并不令人意外。多亏了多伦多大学的Roger Grosse的这些讲座幻灯片，我们可以更好地理解自动微分。
简而言之，为了计算可微函数（例如一个层）的梯度，自动微分需要函数输出的梯度以及函数的输入和输出张量。尽管梯度是动态计算并随后丢弃的，但函数的输入和输出张量（也称为激活值）在前向传递期间被存储。
好的，让我们将这个应用到Transformer模型上。Transformer模型包括多个所谓的Transformer层的堆叠。每个额外的Transformer层都会强制模型在前向传递期间存储更多的激活值，从而增加训练所需的内存。让我们仔细看看。一个Transformer层实质上包含两个残差层。第一个残差层表示自注意机制，如第1节所述，第二个残差层表示线性或前馈层，如第2节所述。
使用前面的符号表示法，Transformer层的输入即X \mathbf{X} X首先进行了归一化1 ^{1} 1，然后通过自注意层进行处理，得到输出Z = SelfAttn ( LayerNorm ( X ) ) \mathbf{Z} = \text{SelfAttn}(\text{LayerNorm}(\mathbf{X})) Z = SelfAttn ( LayerNorm ( X ) )。我们将这两个层简写为G G G，使得Z = G ( X ) \mathbf{Z} = G(\mathbf{X}) Z = G ( X )。接下来，将残差Z \mathbf{Z} Z添加到输入Z ‾ = Z + X \mathbf{\overline{Z}} = \mathbf{Z} + \mathbf{X} Z = Z + X，并将总和输入到第二个残差层 – 两个线性层。将Z ‾ \mathbf{\overline{Z}} Z通过第二个归一化层进行处理，然后通过两个线性层进行处理，得到Y = Linear ( LayerNorm ( Z + X ) ) \mathbf{Y} = \text{Linear}(\text{LayerNorm}(\mathbf{Z} + \mathbf{X})) Y = Linear ( LayerNorm ( Z + X ) )。我们将第二个归一化层和两个线性层简写为F F F，得到Y = F ( Z ‾ ) \mathbf{Y} = F(\mathbf{\overline{Z}}) Y = F ( Z )。最后，将残差Y \mathbf{Y} Y添加到Z ‾ \mathbf{\overline{Z}} Z，得到Transformer层的输出Y ‾ = Y + Z ‾ \mathbf{\overline{Y}} = \mathbf{Y} + \mathbf{\overline{Z}} Y = Y + Z。
让我们以 x 1 , … , x 16 \mathbf{x}_1, \ldots, \mathbf{x}_{16} x 1 , … , x 1 6 为例，来说明一个完整的Transformer层。
为了计算例如自注意力块 G G G 的梯度，需要提前知道三个张量：梯度 ∂ Z \partial \mathbf{Z} ∂ Z ，输出 Z \mathbf{Z} Z 和输入 X \mathbf{X} X 。虽然梯度 ∂ Z \partial \mathbf{Z} ∂ Z 可以在运行时计算并且之后丢弃，但是输出 Z \mathbf{Z} Z 和输入 X \mathbf{X} X 的值必须在前向传播期间计算和存储，因为在反向传播期间无法轻松地重新计算它们。因此，在前向传播期间，需要将大型张量输出（如查询-键点积矩阵 Q K T \mathbf{Q}\mathbf{K}^T Q K T 或线性层的中间输出 Y int \mathbf{Y}^{\text{int}} Y int ）存储在内存中 2 ^{2} 2 。
在这里，可逆残差层能够帮助我们。这个想法相对直接。残差块被设计成以一种方式，以便在反向传播期间可以重新计算输入和输出张量，从而在前向传播期间不需要将任何张量存储在内存中。这是通过使用两个输入流 X ( 1 ) , X ( 2 ) \mathbf{X}^{(1)}, \mathbf{X}^{(2)} X ( 1 ) , X ( 2 ) 和两个输出流 Y ‾ ( 1 ) , Y ‾ ( 2 ) \mathbf{\overline{Y}}^{(1)}, \mathbf{\overline{Y}}^{(2)} Y ( 1 ) , Y ( 2 ) 实现的。第一个残差 Z \mathbf{Z} Z 是由第一个输出流计算的 Z = G ( X ( 1 ) ) \mathbf{Z} = G(\mathbf{X}^{(1)}) Z = G ( X ( 1 ) ) ，然后添加到第二个输入流的输入中，以便 Z ‾ = Z + X ( 2 ) \mathbf{\overline{Z}} = \mathbf{Z} + \mathbf{X}^{(2)} Z = Z + X ( 2 ) 。类似地，残差 Y = F ( Z ‾ ) \mathbf{Y} = F(\mathbf{\overline{Z}}) Y = F ( Z ) 被再次添加到第一个输入流中，以便两个输出流由 Y ( 1 ) = Y + X ( 1 ) \mathbf{Y}^{(1)} = \mathbf{Y} + \mathbf{X}^{(1)} Y ( 1 ) = Y + X ( 1 ) 和 Y ( 2 ) = X ( 2 ) + Z = Z ‾ \mathbf{Y}^{(2)} = \mathbf{X}^{(2)} + \mathbf{Z} = \mathbf{\overline{Z}} Y ( 2 ) = X ( 2 ) + Z = Z 定义。
对于 x 1 , … , x 16 \mathbf{x}_1, \ldots, \mathbf{x}_{16} x 1 , … , x 1 6 ，可逆Transformer层可以如下可视化。
可以看出，输出 Y ‾ ( 1 ) , Y ‾ ( 2 ) \mathbf{\overline{Y}}^{(1)}, \mathbf{\overline{Y}}^{(2)} Y ( 1 ) , Y ( 2 ) 的计算方式与非可逆层的 Y ‾ \mathbf{\overline{Y}} Y 非常相似，但它们在数学上是不同的。Reformer的作者在一些初步实验中观察到，可逆Transformer模型的性能与标准Transformer模型的性能相匹配。与标准Transformer层的第一个可见差异是存在两个输入流和输出流 3 ^{3} 3 ，这在前向传播期间稍微增加了所需的内存。然而，两个流体结构对于在前向传播期间不需要保存任何激活是至关重要的。让我们解释一下。对于反向传播，可逆Transformer层必须计算梯度 ∂ G \partial G ∂ G 和 ∂ F \partial F ∂ F 。除了可以在运行时计算的梯度 ∂ Y \partial \mathbf{Y} ∂ Y 和 ∂ Z \partial \mathbf{Z} ∂ Z 外，还必须知道张量值 Y \mathbf{Y} Y ，Z ‾ \mathbf{\overline{Z}} Z 用于 ∂ F \partial F ∂ F ，张量值 Z \mathbf{Z} Z 和 X ( 1 ) \mathbf{X}^{(1)} X ( 1 ) 用于 ∂ G \partial G ∂ G ，以使自动微分工作。
如果我们假设已知 Y ‾ ( 1 ) , Y ‾ ( 2 ) \mathbf{\overline{Y}}^{(1)}, \mathbf{\overline{Y}}^{(2)} Y ( 1 ) , Y ( 2 ) ，从图中可以很容易地描绘出如下计算 X ( 1 ) , X ( 2 ) \mathbf{X}^{(1)}, \mathbf{X}^{(2)} X ( 1 ) , X ( 2 ) 的过程。 X ( 1 ) = F ( Y ‾ ( 1 ) ) − Y ‾ ( 1 ) \mathbf{X}^{(1)} = F(\mathbf{\overline{Y}}^{(1)}) – \mathbf{\overline{Y}}^{(1)} X ( 1 ) = F ( Y ( 1 ) ) − Y ( 1 ) 。好的，现在已知 X ( 1 ) \mathbf{X}^{(1)} X ( 1 ) ，可以通过 X ( 2 ) = Y ‾ ( 1 ) − G ( X ( 1 ) ) \mathbf{X}^{(2)} = \mathbf{\overline{Y}}^{(1)} – G(\mathbf{X}^{(1)}) X ( 2 ) = Y ( 1 ) − G ( X ( 1 ) ) 计算出 X ( 2 ) \mathbf{X}^{(2)} X ( 2 ) 。好的，现在，Z \mathbf{Z} Z 和 Y \mathbf{Y} Y 可以通过 Y = Y ‾ ( 1 ) − X ( 1 ) \mathbf{Y} = \mathbf{\overline{Y}}^{(1)} – \mathbf{X}^{(1)} Y = Y ( 1 ) − X ( 1 ) 和 Z = Y ‾ ( 2 ) − X ( 2 ) \mathbf{Z} = \mathbf{\overline{Y}}^{(2)} – \mathbf{X}^{(2)} Z = Y ( 2 ) − X ( 2 ) 完成计算。因此，总结一下，如果仅在前向传播过程中存储了最后一个可逆变换层的输出 Y ‾ ( 1 ) , Y ‾ ( 2 ) \mathbf{\overline{Y}}^{(1)}, \mathbf{\overline{Y}}^{(2)} Y ( 1 ) , Y ( 2 ) ，所有其他相关的激活值都可以通过在反向传播过程中利用 G G G 和 F F F 并传递 X ( 1 ) \mathbf{X}^{(1)} X ( 1 ) 和 X ( 2 ) \mathbf{X}^{(2)} X ( 2 ) 来推导得到。在反向传播过程中，对于每个可逆变换层，通过两次 G G G 和 F F F 的前向传播来换取不需要在前向传播过程中存储任何激活值的开销。这是个不错的交易！
注意：最近，主要的深度学习框架已经发布了代码，允许只存储某些激活值，并在反向传播过程中重新计算较大的激活值（Tensoflow 这里和 PyTorch 这里）。对于标准的可逆层来说，仍然意味着每个变换层至少需要存储一个激活值，但通过定义可以动态重新计算哪些激活值，可以节省大量内存。
1 ^{1} 1 在前两个部分中，我们省略了自注意层和线性层之前的层归一化层。读者应该知道，在输入自注意层和线性层之前，X \mathbf{X} X 和 Z ‾ \mathbf{\overline{Z}} Z 都经过了层归一化的处理。2 ^{2} 2 虽然在设计中，Q K \mathbf{Q}\mathbf{K} Q K 的维度被写作 n × n n \times n n × n ，但在 LSH 自注意或局部自注意层中，维度分别为 n × l c × n h n \times l_{c} \times n_{h} n × l c × n h 或 n × l c n \times l_{c} n × l c ，其中 l c l_{c} l c 是块长度，n h n_{h} n h 是哈希数。3 ^{3} 3 在第一个可逆变换层中，X ( 2 ) \mathbf{X}^{(2)} X ( 2 ) 被设置为等于 X ( 1 ) \mathbf{X}^{(1)} X ( 1 ) 。
基准测试
为了衡量可逆残差层的效果，我们将比较BERT和Reformer在训练中使用不同层数时的内存消耗。
#@title 安装和导入
# pip安装
!pip -qq install git+https://github.com/huggingface/transformers.git
!pip install -qq py3nvml
from transformers import ReformerConfig, BertConfig, PyTorchBenchmark, PyTorchBenchmarkArguments通过增加层数，我们来测量标准的bert-base-uncased BERT模型所需的内存。
config_4_layers_bert = BertConfig.from_pretrained("bert-base-uncased", num_hidden_layers=4)
config_8_layers_bert = BertConfig.from_pretrained("bert-base-uncased", num_hidden_layers=8)
config_12_layers_bert = BertConfig.from_pretrained("bert-base-uncased", num_hidden_layers=12)
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[512], batch_sizes=[8], models=["Bert-4-Layers", "Bert-8-Layers", "Bert-12-Layers"], training=True, no_inference=True, no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_4_layers_bert, config_8_layers_bert, config_12_layers_bert], args=benchmark_args)
result = benchmark.run()
HBox(children=(FloatProgress(value=0.0, description='Downloading', max=433.0, style=ProgressStyle(description_…
1 / 3
2 / 3
3 / 3
==================== TRAIN - MEMORY - RESULTS ====================
--------------------------------------------------------------------------------
模型名称 批次大小 序列长度 内存使用（MB）
--------------------------------------------------------------------------------
Bert-4-Layers 8 512 4103
Bert-8-Layers 8 512 5759
Bert-12-Layers 8 512 7415
--------------------------------------------------------------------------------可以看到，增加BERT的一个层会线性增加所需的内存超过400MB。
config_4_layers_reformer = ReformerConfig.from_pretrained("google/reformer-enwik8", num_hidden_layers=4, num_hashes=1)
config_8_layers_reformer = ReformerConfig.from_pretrained("google/reformer-enwik8", num_hidden_layers=8, num_hashes=1)
config_12_layers_reformer = ReformerConfig.from_pretrained("google/reformer-enwik8", num_hidden_layers=12, num_hashes=1)
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[512], batch_sizes=[8], models=["Reformer-4-Layers", "Reformer-8-Layers", "Reformer-12-Layers"], training=True, no_inference=True, no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_4_layers_reformer, config_8_layers_reformer, config_12_layers_reformer], args=benchmark_args)
result = benchmark.run()
1 / 3
2 / 3
3 / 3
==================== TRAIN - MEMORY - RESULTS ====================
--------------------------------------------------------------------------------
模型名称 批次大小 序列长度 内存使用（MB）
--------------------------------------------------------------------------------
Reformer-4-Layers 8 512 4607
Reformer-8-Layers 8 512 4987
Reformer-12-Layers 8 512 5367
--------------------------------------------------------------------------------而对于Reformer，增加一层实际上所需的内存要少得多。增加一层仅平均增加了不到100MB的内存，因此一个更大的12层reformer-enwik8模型所需的内存比一个12层的bert-base-uncased模型要少。
4. 轴向位置编码
Reformer可以处理巨大的输入序列。然而，对于如此长的输入序列，仅使用标准的位置编码权重矩阵就会使用超过1GB的存储空间来存储权重。为了防止这样大的位置编码矩阵，官方的Reformer代码引入了轴向位置编码。
重要：轴向位置编码在官方论文中没有解释，但可以通过查看代码和与作者交流来理解。
Reformer中的轴向位置编码
Transformers需要位置编码来考虑输入中单词的顺序，因为自注意力层没有顺序概念。位置编码通常由一个简单的查找矩阵E = [ e 1 , … , e n max ] \mathbf{E} = \left[\mathbf{e}_1, \ldots, \mathbf{e}_{n_\text{max}}\right] E = [ e 1 , … , e n max ]定义。然后，将位置编码向量e i \mathbf{e}_{i} e i 简单地添加到第i个输入向量x i + e i \mathbf{x}_{i} + \mathbf{e}_{i} x i + e i ，以使模型能够区分输入向量（也称为标记）是在位置i还是位置j。对于每个输入位置，模型需要能够查找相应的位置编码向量，以便E \mathbf{E} E的维度由模型可以处理的输入向量的最大长度config.max_position_embeddings，即n max n_\text{max} n max ，和输入向量的config.hidden_size，即d h d_{h} d h 定义。
假设d h = 4 d_{h}=4 d h = 4和n max = 49 n_\text{max}=49 n max = 4 9，则这样一个位置编码矩阵可以可视化如下：
在这里，我们展示了仅维度为4的位置编码e 1 \mathbf{e}_{1} e 1 ，e 2 \mathbf{e}_{2} e 2 和e 49 \mathbf{e}_{49} e 4 9 。
假设我们想要在长度为0.5M标记的序列上训练一个Reformer模型，并且输入向量config.hidden_size为1024（请参见此处的notebook）。相应的位置嵌入的大小为0.5 M × 1024 ∼ 512 M 0.5M \times 1024 \sim 512M 0 . 5 M × 1 0 2 4 ∼ 5 1 2 M个参数，相当于2GB的大小。
这样的位置编码会在将模型加载到内存中和将模型保存到硬盘时使用过多的内存。
Reformer的作者通过将config.hidden_size维度减半，并巧妙地因式分解n max n_\text{max} n max 维度，成功地大大缩小了位置编码的大小。在Transformer中，用户可以通过将config.axial_pos_shape设置为适当的两个值n max 1 n_\text{max}^1 n max 1 和n max 2 n_\text{max}^2 n max 2 ，以将n max n_\text{max} n max 分解成哪种形状。通过将config.axial_pos_embds_dim设置为适当的两个值d h 1 d_{h}^{1} d h 1 和d h 2 d_{h}^2 d h 2 ，使得d h 1 + d h 2 = d h d_{h}^1 + d_{h}^2 = d_{h} d h 1 + d h 2 = d h ，用户可以决定如何切割隐藏大小维度。现在，让我们更直观地可视化和解释。
可以将n max n_{\text{max}} n max 的因式分解想象为将维度折叠成第三个轴，如下所示，用于因式分解config.axial_pos_shape = [7, 7]：
每个立方体对应一个编码向量e 1 ，e 2 ，e 49 ，但我们可以看到49个编码向量被分成了7行，每行有7个向量。现在的想法是只使用一行7个编码向量，并将这些向量扩展到其他6行，从而重复使用它们的值。由于不鼓励在不同的编码向量中具有相同的值，每个维度（即高度）为config.hidden_size=4的向量被分成了尺寸为1的下部编码向量e down和尺寸为3的上部编码向量e up，这样下部可以沿着行维度扩展，上部可以沿着列维度扩展。让我们通过可视化更清楚地看到。
我们可以看到我们将嵌入向量切成了蓝色的e down（下部）和黄色的e up（上部）。现在对于”子”向量E down = [ e down,1 ，…，e down,49 ]，只保留了图中的第一行，即宽度为7，并沿着列维度进行扩展，即图的深度。相反，对于”子”向量E up = [ e up,1 ，…，e up,49 ]，只保留了图中的第一列，即7，并沿着行维度进行扩展。最终得到的嵌入向量e ′ i 对应于
e ′ i = [ [ e down, i % n max 1 ] T , [ e up, ⌊ i n max 2 ⌋ ] T ] T
其中n max 1 = 7，n max 2 = 7是我们的示例中的值。这些新的编码E’ = [ e’ 1 ，…，e’ n max ]被称为轴向位置编码。
以下是我们示例中更详细地说明了这些轴向位置编码。
现在应该更容易理解了，最终的位置编码向量 E′ 是如何仅从维度为 d h 1 × n max 1 的 E down 和维度为 d h 2 × n max 2 的 E up 计算而来的。
这里需要注意的关键方面是，轴向位置编码确保设计上不会有任何向量 [ e ′ 1 , … , e ′ n max ] 相等，并且编码矩阵的整体大小从 n max × d h 减少到 n max 1 × d h 1 + n max 2 × d h 2 。通过设计使每个轴向位置编码向量都不同，使得模型在学习轴向位置表示时具有更大的灵活性。
为了演示大小的显著缩减，假设我们设置 config.axial_pos_shape = [1024, 512] 和 config.axial_pos_embds_dim = [512, 512]，用于可以处理最长长度为 0.5M 个标记的 Reformer 模型。得到的轴向位置编码矩阵的大小将仅为 1024 × 512 + 512 × 512 ∼ 800 K，这大约相当于 3MB。相比之下，标准位置编码矩阵在这种情况下需要 2GB 的空间。
有关更简洁和数学密集的解释，请参阅 🤗Transformers 文档。
基准测试
最后，我们还要将传统的位置嵌入和轴向位置嵌入的峰值内存消耗进行比较。
#@title 安装和导入
# pip 安装
!pip -qq install git+https://github.com/huggingface/transformers.git
!pip install -qq py3nvml
from transformers import ReformerConfig, PyTorchBenchmark, PyTorchBenchmarkArguments, ReformerModel位置嵌入仅依赖于两个配置参数：输入序列的最大允许长度 config.max_position_embeddings 和 config.hidden_size。让我们使用一个将输入序列的最大允许长度推到 50 万个标记的模型 google/reformer-crime-and-punishment，来看使用轴向位置嵌入的效果。
首先，我们将比较轴向位置编码和标准位置编码的形状以及模型中的参数数量。
config_no_pos_axial_embeds = ReformerConfig.from_pretrained("google/reformer-crime-and-punishment", axial_pos_embds=False) # 禁用轴向位置嵌入
config_pos_axial_embeds = ReformerConfig.from_pretrained("google/reformer-crime-and-punishment", axial_pos_embds=True, axial_pos_embds_dim=(64, 192), axial_pos_shape=(512, 1024)) # 启用轴向位置嵌入
print("默认位置编码")
print(20 * '-')
model = ReformerModel(config_no_pos_axial_embeds)
print(f"位置嵌入的形状：{model.embeddings.position_embeddings}")
print(f"模型的参数数量：{model.num_parameters()}")
print(20 * '-' + '\n\n')
print("轴向位置编码")
print(20 * '-')
model = ReformerModel(config_pos_axial_embeds)
print(f"位置嵌入的形状：{model.embeddings.position_embeddings}")
print(f"模型的参数数量：{model.num_parameters()}")
print(20 * '-' + '\n\n')
HBox(children=(FloatProgress(value=0.0, description='下载', max=1151.0, style=ProgressStyle(description…
默认位置编码
--------------------
位置嵌入的形状：PositionEmbeddings(
(embedding): Embedding(524288, 256)
)
模型的参数数量：136572416
--------------------
轴向位置编码
--------------------
位置嵌入的形状：AxialPositionEmbeddings(
(weights): ParameterList(
(0): Parameter containing: [torch.FloatTensor of size 512x1x64]
(1): Parameter containing: [torch.FloatTensor of size 1x1024x192]
)
)
模型的参数数量：2584064
--------------------阅读了理论之后，轴向位置编码权重的形状对读者来说应该不会感到意外。
关于结果，可以看到对于能够处理如此长的输入序列的模型来说，使用默认的位置编码是不切实际的。在 google/reformer-crime-and-punishment 的情况下，仅标准位置编码就包含了超过一亿个参数。而轴向位置编码将这个数量降低到了20万多个。
最后，让我们也来比较一下推理时所需的内存。
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[512], batch_sizes=[8], models=["Reformer-No-Axial-Pos-Embeddings", "Reformer-Axial-Pos-Embeddings"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_no_pos_axial_embeds, config_pos_axial_embeds], args=benchmark_args)
result = benchmark.run()
1 / 2
2 / 2
==================== 推理 - 内存 - 结果 ====================
--------------------------------------------------------------------------------
模型名称 批次大小 序列长度 内存（MB）
--------------------------------------------------------------------------------
Reformer-No-Axial-Pos-Embeddin 8 512 959
Reformer-Axial-Pos-Embeddings 8 512 447
--------------------------------------------------------------------------------可以看到，在 google/reformer-crime-and-punishment 的情况下，使用轴向位置嵌入将内存需求减少了大约一半。
