提示: 这篇文章链接和公式特别多，🤗 宝子们可以戳 阅读原文 获得最佳阅读体验哟！
Reformer 如何在不到 8GB 的内存上训练 50 万个词元 Kitaev、Kaiser 等人于 20202 年引入的 Reformer 模型 是迄今为止长序列建模领域内存效率最高的 transformer 模型之一。
最近，人们对长序列建模的兴趣激增，仅今年一年，就涌现出了大量的工作，如 Beltagy 等人的工作 (2020) 、Roy 等人的工作 (2020) 、Tay 等人的工作 以及 Wang 等人的工作 等等。长序列建模背后的动机是，NLP 中的许多任务 (例如 摘要、问答 ) 要求模型处理更长的序列，这些序列长度超出了 BERT 等模型的处理能力。在需要模型处理长输入序列的任务中，长序列模型无需对输入序列进行裁剪以避免内存溢出，因此已被证明优于标准的 BERT 类模型 ( 见 Beltagy 等人 2020 年的工作)。
Reformer 能够一次处理多达 50 万个词元，从而突破了长序列建模的极限 (具体可参见本 笔记本)。相形之下，传统的 bert-base-uncased 模型最长仅支持 512 个词元。在 Reformer 中，标准 transformer 架构的每个部分都经过重新设计，以最小化内存需求，并避免显著降低性能。
内存的改进来自于 Reformer 作者向 transformer 世界引入的 4 大特性:
Reformer 自注意力层 - 如何在不受限于本地上下文的情况下高效地实现自注意力机制？
分块前馈层 - 如何更好地对大型前馈层的时间和内存进行权衡？
可逆残差层 - 如何聪明地设计残差架构以大幅减少训练中的内存消耗？
轴向位置编码 (Axial Positional Encodings) - 如何使位置编码可用于超长输入序列？
本文的目的是 深入 阐述 Reformer 的上述四大特性。虽然这四个特性目前是用在 Reformer 上的，但其方法是通用的。因此，读者不应被此束缚，而应该多思考在哪些情况下可以把这四个特性中的某一个或某几个应用于其他的 transformer 模型，以解决其问题。
下文四个部分之间的联系很松散，因此可以单独阅读。
Reformer 已集成入 🤗Transformers 库。对于想使用 Reformer 的用户，建议大家阅读本文，以更好地了解该模型的工作原理以及如何正确配置它。文中所有公式都附有其在 transformers 中对应的 Reformer 配置项 ( 例如 config.<param_name> )，以便读者可以快速关联到官方文档和配置文件。
注意 : 轴向位置编码 在官方 Reformer 论文中没有解释，但在官方代码库中广泛使用。本文首次深入阐释了轴向位置编码。
1. Reformer 自注意力层 Reformer 使用了两种特殊的自注意力层: 局部 自注意力层和 LSH (Locality Sensitive Hashing，局部敏感哈希， LSH ) 自注意力层。
在介绍新的自注意力层之前，我们先简要回顾一下传统的自注意力，其由 Vaswani 等人在其 2017 年的论文 中引入。
本文的符号及配色与 《图解 transformer》 一文一致，因此强烈建议读者在阅读本文之前，先阅读《图解 transformer》一文。
重要 : 虽然 Reformer 最初是为了因果自注意力而引入的，但它也可以很好地用于双向自注意力。本文在解释 Reformer 的自注意力时，将其用于 双向 自注意力。
全局自注意力回顾 Transformer 模型的核心是 自注意力 层。现在，我们回顾一下传统的自注意力层，这里称为 全局自注意力 层。首先我们假设对嵌入向量序列
config.hidden_size ， 即
简而言之，全局自注意力层将
softmax 计算最终输出
下图给出了
，
请注意，本文所有示意图都假设 batch_size 和 config.num_attention_heads 为 1。为了便于稍后更好地解释 LSH 自注意力 ，我们还在图中标记出了一些向量， 如
config.num_attention_heads > 1)。如需了解多头注意力，建议读者参阅 《图解 transformer》。
敲个重点，对于每个输出向量
这也是为什么 bert-base-cased 的 config.max_position_embedding_size 只有 512 的原因。
局部自注意力 局部自注意力 是缓解
config.local_chunk_length ， 即
继续以
，
假设
，
可以看出，我们对每个块分别执行了注意力操作
，
，
，
例如 ，我们的例子中的
一个简单的补救措施是用 config.local_num_chunks_before ( 即
config.local_num_chunks_after ( 即
其中
好吧，这个公式看起来有点复杂，我们稍微分析一下。在 Reformer 的自注意力层中，
我们注意到这里有一个循环关系，因此第一个块也可以关注最后一个块。我们再次图解一下这种增强的局部关注算法。我们先按块找到其对应的窗口，并在其上应用自注意力，然后仅保留中心输出段作为本块的输出。
最后，将相应的输出串接到
请注意，在实现局部自注意力时，为了计算效率，我们并不会像图中一样先计算全部输出并随后 丢弃 一部分。图中红叉所示的地方仅用于说明，实际并不会产生计算行为。
这里需要注意的是，扩展每个分块自注意力函数的输入向量可以使得 每个 输出向量
，
，
，
内存消耗上的降低也是显而易见的:
这种增强的局部自注意力比普通的局部自注意力架构更好，但仍然存在一个主要缺陷，因为每个输入向量只能关注预定义大小的局部上下文。对于不需要 transformer 模型学习输入向量之间的远程依赖关系的 NLP 任务 ( 例如 语音识别、命名实体识别以及短句子的因果语言建模) 而言，可能不是一个大问题。但还有许多 NLP 任务需要模型学习远程依赖关系，因此局部自注意力在这些任务下可能会导致显著的性能下降， 如 :
问答 : 模型必须学习问题词元和相关答案词元之间的关系，这些词元很可能并不相邻;
多项选择 : 模型必须将多个答案词元段相互比较，这些答案词元段通常隔得比较远;
摘要 : 模型必须学习长序列的上下文词元和较短的摘要词元序列之间的关系，而上下文和摘要之间的相关关系很可能无法通过局部自注意力来捕获。
局部自注意力本身很可能不足以让 transformer 模型学习输入向量 (词元) 彼此之间的相关关系。
因此，Reformer 额外采用了一个近似全局自注意力的高效自注意力层，称为 LSH 自注意力 。
LSH 自注意力 鉴于我们已经了解了局部自注意力的工作原理，下面我们继续尝试一下可能是 Reformer 中最具创新性的算法改进: LSH 自注意力 。
LSH 自注意力的设计目标是在效果上接近全局自注意力，而在速度与资源消耗上与局部自注意力一样高效。
LSH 自注意力因依赖于 Andoni 等人于 2015 年提出的 LSH 算法 而得名。
LSH 自注意力源于以下洞见: 如果
我们展开讲讲: 设
首先，Reformer 的作者注意到共享查询投影和键投影:
m 组位置索引，其中装的是属于同一个桶的所有查询向量:
第
簇
config.num_buckets ， 即
对每个索引
其次，作者利用 LSH 算法将查询向量聚类到预定义的
即
还用前面的例子，我们有:
接着，可以注意到，将所有查询向量聚类至
我们用例子再解释一下，假设在 config.num_buckets=4 ， config.lsh_chunk_length=4 时重排输入向量
、
、
、
对每个输入向量，仅需在簇内进行自注意力计算即可，因此每个输入向量对应的输出向量可计算如下:
我们再次图解一下该过程:
可以看出，自注意力函数的运算矩阵大小各不相同，这种情况比较麻烦，因为 GPU 和 TPU 无法高效并行处理不同尺寸的矩阵运算。
为了进一步解决高效计算的问题，可以借鉴局部注意力的方法，对重排后的输入进行分块，以使每个块的大小均为 config.lsh_chunk_length 。通过对重排后的输入进行分块，一个桶可能会被分成两个不同的块。为了解决这个问题，与局部自注意力一样，在 LSH 自注意力中，每个块除了自身之外还关注其前一个块 config.lsh_num_chunks_before=1 ( config.lsh_num_chunks_after 通常设置为 0)。这样，我们就可以大概率确保桶中的所有向量相互关注
总而言之，对于所有块
其中
这里，我们对上图中的重排向量
最后，将输出
这里还要提到的一个重要特征是，可以通过并行运行 LSH 自注意力 config.num_hashes (即
config.num_hashes > 1 ，对于每个
打完收工！至此，我们了解了 LSH 自注意力在 Reformer 中是如何工作的。
说回内存复杂度，该方法有两个可能的瓶颈点: 点积所需的内存:
config.num_buckets 进行分解，详见 此处。
我们快速总结一下:
我们希望利用 softmax 运算仅对极少数键向量赋予重要权重的先验知识来对全局注意力进行近似。
如果键向量等于查询向量，这意味着 对于每个 查询向量
这种关系是对称的，也就是说，如果
我们对输入按簇进行重排，并对重排后的输入计算局部自注意力，最后将输出重新恢复为原顺序。
假如 向量
基准测试 Transformers 最近增加了基准测试相关的代码，你可参阅 此处 以获取更详细的说明。
为了展示局部 LSH 自注意力可以节省多少内存，我们在不同的 local_attn_chunk_length 和 lsh_attn_chunk_length 上对 Reformer 模型 google/reformer-enwik8 上进行了基准测试。你可以从 此处 找到更详细的有关 google/reformer-enwik8 模型的默认配置和用法信息。
我们先进行一些必要的导入和安装。
#@title Installs and Imports # pip installs 首先，我们测试一下在 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)"Reformer" ], no_speed=True, no_env_print=True)HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1279.0, style=ProgressStyle(description…输入序列越长，输入序列和峰值内存使用之间的平方关系
对使用全局注意力的 google/reformer-enwik8 模型而言，序列长度超过 16K 内存就溢出了。
现在，我们使用模型的默认参数以使能 局部 LSH 自注意力。
config = ReformerConfig.from_pretrained("google/reformer-enwik8" )"Reformer" ], no_speed=True, no_env_print=True)1 / 1不出所料，对于较长的输入序列，使用局部 LSH 自注意力机制的内存效率更高，对于本文使用的 11GB 显存 GPU 而言，模型直到序列长度为 32K 时，内存才耗尽。
2. 分块前馈层 基于 transformer 的模型通常在自注意力层之后会有一个非常大的前馈层。该层可能会占用大量内存，有时甚至成为模型主要的内存瓶颈。Reformer 论文中首次引入了前馈分块技术，以用时间换取内存。
Reformer 中的分块前馈层 在 Reformer 中， LSH 自注意力层或局部自注意力层通常后面跟着一个残差连接，我们可将其定义为 transformer 块 的第一部分。更多相关知识，可参阅此 博文。
Transformer 块 第一部分的输出，称为 归范化自注意力 输出，可以记为
在我们的例子中，输入
Transformer 块 的第二部分通常由两个前馈层
敲重点！在数学上，前馈层在位置
从图中可以看出，所有输入向量
我们再观察一下前馈层的输出维度，看看有没有啥有意思的事情。在 Reformer 中，
config.feed_forward_size ， 即
config.hidden_size ， 即
Reformer 作者观察到
为了更好地感受维度的差异，我们将本文例子中的矩阵
很明显，张量
config.chunk_size_feed_forward 为
假设
当块大小为 1 时，必须完整存储在内存中的唯一张量是大小为
config.hidden_size 。而中间张量只需要存储大小为
最后，重要的是要记住， 分块线性层 与传统的完整线性层相比，其输出在数学上是等效的，因此可以应用于所有 transformer 线性层。因此，在某些场景下，可以考虑使用 config.chunk_size_feed_forward 在内存和速度之间进行更好的权衡。
bert-base-uncased 为例，其中间维度
config.num_attention_heads 为 1，因此假设自注意力层的输出大小为 config.hidden_size 。
读者也可以在 🤗Transformers 的 相应文档 中找到有关分块线性/前馈层的更多信息。
基准测试 我们测试一下使用分块前馈层可以节省多少内存。
#@title Installs and Imports # pip installs Building wheel for transformers (setup.py) ... [?25l[?25hdone首先，我们将没有分块前馈层的默认 google/reformer-enwik8 模型与有分块前馈层的模型进行比较。
config_no_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8" ) # no chunk "google/reformer-enwik8" , chunk_size_feed_forward=1) # feed forward chunk "Reformer-No-Chunk" , "Reformer-Chunk" ], no_speed=True, no_env_print=True)1 / 2有趣的是，分块前馈层似乎在这里根本没有帮助。原因是 config.feed_forward_size 不够大，所以效果不明显。仅当序列长度较长 (4096) 时，才能看到内存使用量略有下降。
我们再看看如果将前馈层的大小增加 4 倍，并将注意力头的数量同时减少 4 倍，从而使前馈层成为内存瓶颈，此时峰值内存情形如何。
config_no_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8" , chunk_size_feed_forward=0, num_attention_{h}eads=2, feed_forward_size=16384) # no chuck "google/reformer-enwik8" , chunk_size_feed_forward=1, num_attention_{h}eads=2, feed_forward_size=16384) # feed forward chunk "Reformer-No-Chunk" , "Reformer-Chunk" ], no_speed=True, no_env_print=True)1 / 2现在，对于较长的输入序列，可以看到峰值内存使用量明显减少。总之，应该注意的是，分块前馈层仅对于具有很少注意力头和较大前馈层的模型才有意义。
3. 可逆残差层 可逆残差层由 N. Gomez 等人 首先提出并应用在 ResNet 模型的训练上以减少内存消耗。从数学上讲，可逆残差层与 真正的 残差层略有不同，其不需要在前向传播期间保存激活，因此可以大大减少训练的内存消耗。
Reformer 中的可逆残差层 我们首先研究为什么模型训练比推理需要更多的内存。
在模型推理时，所需的内存差不多等于计算模型中 单个 最大张量所需的内存。而在训练模型时，所需的内存差不多等于所有可微张量的 总和 。
如果读者已经理解了深度学习框架中的自动微分的工作原理，对此就比较容易理解了。多伦多大学 Roger Grosse 的这些 幻灯片 对大家理解自动微分很有帮助。
简而言之，为了计算可微函数 ( 如 一层) 的梯度，自动微分需要函数输出的梯度以及函数的输入、输出张量。虽然梯度是可以动态计算并随后丢弃的，但函数的输入和输出张量 ( 又名 激活) 需要在前向传播过程中被保存下来，以供反向传播时使用。
我们具体看下 transformer 模型中的情况。Transformer 模型是由多个 transformer 层堆叠起来的。每多一个 transformer 层都会迫使模型在前向传播过程中保存更多的激活，从而增加训练所需的内存。 我们细看一下 transformer 层。Transformer 层本质上由两个残差层组成。第一个残差层是第 1) 节中解释的 自注意力 机制，第二个残差层是第 2) 节中解释的 线性层 (或前馈层)。
使用与之前相同的符号，transformer 层的输入 即
我们仍以
比如 ，要计算自注意力块
此时，可逆残差层就有用了。它的想法相对简单: 残差块的设计方式使得不必保存函数的输入和输出张量，而在反向传播期间就轻松地对二者进行重新计算，这样的话在前向传播期间就无需将这些张量保存在内存中了。
这是通过两个输入流
、
、
以
可以看出，输出
、
假设我们知道
，
，
算
出
最后一个 可逆 transformer 层的输出
，
注意 : 最近，主要的深度学习框架都支持了梯度检查点技术，以允许仅保存某些激活并在反向传播期间重计算尺寸较大的激活 (Tensoflow 代码见 此处，PyTorch 代码见 此处)。对于标准可逆层，这仍然意味着必须为每个 transformer 层保存至少一个激活，但通过定义哪些激活可以动态重新计算，能够节省大量内存。
LSH 自注意力 或 局部自注意力 层的维度为
测试基准 为了测量可逆残差层的效果，我们将增加模型层数的同时比较 BERT 和 Reformer 的内存消耗。
#@title Installs and Imports # pip installs 我们把标准 bert-base-uncased BERT 模型的层数从 4 增加到 12 ，同时测量其所需内存。
config_4_layers_bert = BertConfig.from_pretrained("bert-base-uncased" , num_hidden_layers=4)"bert-base-uncased" , num_hidden_layers=8)"bert-base-uncased" , num_hidden_layers=12)"Bert-4-Layers" , "Bert-8-Layers" , "Bert-12-Layers" ], training=True, no_inference=True, no_speed=True, no_env_print=True)HBox(children=(FloatProgress(value=0.0, description='Downloading', max=433.0, style=ProgressStyle(description_…可以看出，BERT 层数每增加 1，其所需内存就会有超 400MB 的线性增长。
config_4_layers_reformer = ReformerConfig.from_pretrained("google/reformer-enwik8" , num_hidden_layers=4, num_hashes=1)"google/reformer-enwik8" , num_hidden_layers=8, num_hashes=1)"google/reformer-enwik8" , num_hidden_layers=12, num_hashes=1)"Reformer-4-Layers" , "Reformer-8-Layers" , "Reformer-12-Layers" ], training=True, no_inference=True, no_speed=True, no_env_print=True)1 / 3另一方面，对于 Reformer 而言，每增加一层所带来的内存增量会显著减少，平均不到 100MB。因此 12 层的 reformer-enwik8 模型比 12 层的 bert-base-uncased 模型的内存需求更少。
4. 轴向位置编码 Reformer 使得处理超长输入序列成为可能。然而，对于如此长的输入序列，仅存储标准位置编码权重矩阵就需要超过 1GB 内存。为了避免如此大的位置编码矩阵，官方 Reformer 代码引入了 轴向位置编码 。
重要: 官方论文中没有解释轴向位置编码，但通过阅读代码以及与作者讨论我们很好地理解了它。
Reformer 中的轴向位置编码 Transformer 需要位置编码来对输入序列中的单词顺序进行编码，因为自注意力层 没有顺序的概念 。位置编码通常由一个简单的查找矩阵
第 i 个 输入向量上，即
即 词元) 位于位置
config.max_position_embeddings ( 即
config.hidden_size ( 即
假设
此处，我们仅展示位置编码
即 高度) 为 4。
想象一下，我们想要在长度最长为 0.5M 个词元，输入向量维度 config.hidden_size 为 1024 的序列上训练 Reformer 模型 (请参阅 此笔记本)。其对应的位置嵌入的参数量为
在将模型加载到内存中或将其保存在硬盘上时，所需要的内存是很大且很没必要的。
Reformer 作者通过将 config.hidden_size 维度一分为二，并巧妙地对
config.axis_pos_shape 设置为一个含有两个值的列表:
config.axis_pos_embds_dim 设置为一个含有两个值
大家可以将对
config.axis_pos_shape = [7, 7] 分解:
三个直立矩形棱柱分别对应于编码向量
或称 高度) 为 config.hidden_size=4 的向量切割成两个部分: 大小为
可以看到，我们已将嵌入向量切为
蓝色 ) 和
黄色 ) 两个部分。现在对 子 向量 即 图中宽度，并将其沿列 ( 又名 深度) 扩展。相反，对 子 向量
本例中，
轴向位置编码 。
下图针对我们的例子对轴向位置编码进行了更详细的说明。
现在应该很清楚如何仅根据维度为
这里的关键是，轴向位置编码能够从设计上确保向量
为了证明位置编码矩阵的尺寸得到了大幅减小，假设我们为 Reformer 模型设置了参数 config.axis_pos_shape = [1024, 512] 以及 config.axis_pos_embds_dim = [512, 512] ，且该模型支持的最长输入序列长度为 0.5M 词元。此时，生成的轴向位置编码矩阵的参数量仅为
如需更简洁、更数学化的解释，请参阅 此处 的 🤗Transformers 文档。
基准测试 最后，我们对传统位置嵌入与 轴向位置嵌入 的峰值内存消耗进行比较。
#@title Installs and Imports # pip installs 位置嵌入仅取决于两个配置参数: 输入序列允许的最大长度 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) # disable axial positional embeddings "google/reformer-crime-and-punishment" , axial_pos_embds=True, axial_pos_embds_dim=(64, 192), axial_pos_shape=(512, 1024)) # enable axial positional embeddings print ("Default Positional Encodings" )print (20 *'-' )print (f"Positional embeddings shape: {model.embeddings.position_embeddings}" )print (f"Num parameters of model: {model.num_parameters()}" )print (20 *'-' + '\n\n' )print ("Axial Positional Encodings" )print (20 *'-' )print (f"Positional embeddings shape: {model.embeddings.position_embeddings}" )print (f"Num parameters of model: {model.num_parameters()}" )print (20 *'-' + '\n\n' )HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1151.0, style=ProgressStyle(description…理解了相应的理论后，读者应该不会对轴向位置编码权重的形状感到惊讶。
从结果中可以看出，对于需要处理如此长输入序列的模型，使用标准位置编码是不切实际的。以 google/reformer-crime-and-punishment 为例，仅标准位置编码自身参数量就超过 100M。轴向位置编码可以将这个数字减少到略高于 200K。
最后，我们比较一下推理所需内存。
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)1 / 2可以看出，在 google/reformer-crime-and-punishment 模型上，使用轴向位置嵌入可减少大约一半的内存需求。
🤗 宝子们可以戳 阅读原文 查看文中所有的外部链接哟！
英文原文: https://hf.co/blog/reformer
原文作者: Patrick von Platen
译者: Matrix Yao (姚伟峰)，英特尔深度学习工程师，工作方向为 transformer-family 模型在各模态数据上的应用及大规模模型的训练推理。
