英文版说明： reformer-deep-dive
自从第一篇《Attention is All You Need》论文在NLP社区掀起了Transformer热潮以来，我们似乎一直在不懈地追求更大的模式。 2019年夏天，NVIDIA发表了他们的MegatronLM论文—— 83亿参数。 2020年2月，微软再次增加赌注，发表了关于拥有170亿个参数的Turing-NLG的博客文章。
了解在增加参数数量和训练数据时，这些模型能达到的程度是值得的。 我很高兴能够拥有这些资源进行大规模实验的公司已经这样做了。 但是，相比之下，在如何使Transformer体系结构更有效率方面，投资太少了。
“reformer : theefficienttransformer”来自jsdds Kitaev、ukasz Kaiser和Anselm Levskaya，与过去两年的“越大越好”趋势相反，2020年本文主要关注自我注意操作如何随着数组的长度扩展，提出了可以将来自更长上下文的信息合并到语言模型中的替代注意机制。
通过Reformer对Transformer的变更，可以通过单一的加速器对长度为64000的数组进行注意操作，与MegatronLM和TuringNLP中1024的上下文长度相反。 两个模型都采用模型并行管线来复制大量参数。
自我验证审查在深入探讨Reformer体系结构之前，让我们先简要回顾一下自我验证的形成过程，以获取整合长上下文时遇到的困难的背景知识。
为了简单起见，尽管在实践中使用了多头注意力，但我们只会讨论与单头的点积注意力。
如果想更深入地回顾自我验证的工作原理，我强烈推荐Alexander Rush的Annotated Transformer和Jay Alammar的Illustrated Transformer。
让我们把自助分成三个部分来解释
query-key-valueprojectionquery/keymatrixmultiplysoftmaxweightedsumofvaluesquery-key-value projection
在此阶段，每个token的当前隐藏状态通过线性投影分解为三个部分。
queries=NP.mat mul (query _ weights，hidden ) query_bias
kys=NP.matmul(key_weights，hidden ) key_bias
values=NP.matmul(value_weights，hidden ) value_bias
查询/密钥矩阵多项
self-attention操作的核心——一个矩阵乘法计算我们的keys和queries之间两个相似度的得分。
投影后，将queries和keys相乘以计算两个或两个相似度。 这是通过矩阵乘法实现的。
qk_agreement=NP.matmul(Queries，NP.swapaxes(keys，-1，-2) )
在你的keys和queries是形状(batch，sequence_length，hidden_size )的dtdmj的情况下，矩阵乘法运算的输出是形状) batch，sequence_length，sequence
这种乍一看无关紧要的矩阵乘法是这种自适应操作的计算复杂性问题的根源。 对于序列长度的线性增加，计算输出所需的乘法次数将按平方增加，因为必须为输入中的每对可能的token计算相似性。 该o(L )的复杂性意味着序列长度超过1024的token使用原始的transformer结构是不现实的。 实际上，由BERT及其继承人RoBERTa选择的上下文长度只有512。
softmaxweightedsumofvalueskey/value合作矩阵项除以用于消除hidden size参数对注意力分布影响的缩放因子sqrt(hidden_size )。 对于每个query，我们在所有keys上计算了softmax，使矩阵的每行和1——的新隐藏状态的大小与序列长度无关。 最后，将注意矩阵乘以values矩阵，为每个token生成新的隐藏表示。
attention _ weights=soft max (qk _ agreement/qk _ agreement.shape [-1 ] )
attention _ outputs=NP.mat mul (attention _ weights，values ) )。
计算复杂性-如上所述，解决方案
点积注意力方式非常好用，允许任意的token在我们的上下文中从任何其他的token中聚合信息，这种灵活性是有代价的，一个不幸的O (L²)计算复杂度。有几篇论文提出了帮助解决这种计算复杂性的transformer的变体。"Generating Long Sequences With Sparse Transformers”建议使用成对的注意力操作和精心选择的注意力模式来分解注意操作。"Transformer-XL: Attentive Language Models blddg a Fixed-Length Context"引入了一种循环机制，允许整合来自比自注意力操作的上下文更大的距离的信息。
The Reformer"Reformer: The Efficient Transformer"的作者采用了一种完全不同的方法来处理序列长度问题。首先，他们观察到学习不同的keys和queries的投影并不是严格必要的。他们丢弃了query投影，并将注意力权重替换为key的函数。
有点令人惊讶的是，尽管他们从注意力模块中移除了一些参数，他们的模型在enwiki8上的性能并没有下降。
现在，注意力块不再包含queries的单独投影，我们只有key和value对。然而，计算key的协同矩阵(通过将每个key与其他key进行比较)仍然是非常昂贵的。
不幸的是我们可能并没有利用好所有的这些计算。softmax的输出通常由几个关键元素控制 — 其余的往往在噪声中消失。我们在计算softmax的时候，并不一定需要那些注意力权重很小的token。
在编写传统软件时，我们总是会遇到这个问题。如果我们想找到与给定key对应的value，我们通常不会遍历所有key的列表并检查每个key是否匹配。相反，我们使用散列映射数据结构来执行O(1)的查找，而不是O(n)比较。
方便的是，向量空间的ckdch映射确实存在类似的情况，它被称为“局部敏感ckdch”(LSH)。正是基于这种方法，Reformer的论文的作者们希望产生一个transformer的替代方案，以避免使用点积注意力的平方复杂性。
局部敏感ckdch (LSH)局部敏感ckdch是一组将高维向量映射到一组离散值(桶/集群)的方法。它最常用来作为近似最近邻搜索的一种方法，用于近似的重复检测或视觉搜索等应用。
局部敏感ckdch方法尝试将高维空间中相近的向量以高概率分配到相同的ckdch。有效的ckdch函数有很多种，最简单的可能是随机投影。
lsh_proj = np.random.randn(hidden_size, hash_size)
hash_value = np.sign(np.dot(x, lsh_proj.T))
换句话说，我们选择一个随机的向量集合，观察输入向量在每个向量上的投影是正的还是负的，然后使用这个二值向量来表示给定向量的预期存储区。下图说明了LSH投影矩阵“u”中单个向量的处理过程。绿色的正号表示与向量u点积为正的点，而红色的负号表示与向量u点积为负的点。
Reformer的论文选择了局部敏感ckdch的angular变体。它们首先约束每个输入向量的L2范数(即将向量投影到一个单位球面上)，然后应用一系列的旋转，最后找到每个旋转向量所属的切片。
该图演示了一个用4个桶进行3轮ckdch的设置。下面的图中的向量映射到了同一个bucket，因为它们的输入很接近，而上一张图中的向量映射到第一个和最后一个bucket。
找到给定的向量选择之后属于哪个桶也可以看成是找到和输入最一致的向量 —— 下面是Reformer的代码：
random_rotations = np.random.randn(hidden_dim, n_buckets // 2)
rotated_vectors = np.dot(x, random_rotations) rotated_vectors =
np.hstack([rotated_vectors, -rotated_vectors]) buckets =
np.argmax(rotated_vectors, axis=-1)
在为每个token计算一个桶之后，将根据它们的桶对这些token进行排序，并将标准的点积注意力应用到桶中的token的块上。
有了足够多的桶，这就大大减少了所有的给定的token需要处理的token的数量 —— 在实验中，Reformer的论文运行的模型被配置为使用128块大小的块。因此，LSH操作将昂贵的key协同矩阵乘法的上下文大小限制为更易于管理的值。
我们现在的时间复杂度为O (L*log(L)) ，而不是时间复杂度成正比O (L²)， 这允许我们把注意力操作扩展到更长的序列的时候不会由于运行时间而受到影响。
因为这个分桶过程是随机的，所以Reformer有选择地多次运行这个过程，以减少两个在输入空间很近的向量被随机地放在不同的桶中的可能性。当所有的事情都做了之后，你就有了一个完全替代标准的多头注意力的方法，它可以与计算完整的注意力矩阵相媲美。
内存复杂度不幸的是，实现更好的时间复杂度只是问题的一半。如果我们将新的LSH注意力块替换为标准的多头注意力，并尝试输入新长度的信息，我们将很快认识到系统中的下一个瓶颈 — 内存复杂性。
即使我们已经非常小心地最小化了注意力操作的计算复杂度，我们仍然必须将所有的key和value存储在内存中，更糟糕的是，在训练期间，我们需要缓存激活以计算参数更新。
Reformer论文使用了序列长度为64k的enwiki8语言建模数据集来做实验，隐藏单元的大小为1024，层数为12层，这意味着存储key和value需要2 * 64000 * 1024 * 12 = ~ 1.5B个浮点数，大约是6GB的内存。使用这种内存使用方式，我们将无法在训练期间使用大的批处理大小，从而影响我们的运行时间。
一个选择是实现gradient checkpoint来帮助限制我们的内存使用。允许我们减少内存使用，只存储从正向传递中的关键的激活，剩余的在反向传递中重新计算。因此，我们可以选择只在key和value投影之前存储隐藏状态，而不是存储key和value，然后第二次重新投影隐藏状态来计算梯度。
不幸的是，这使我们的后向传递的成本增加了一倍，因此我们能够支持更大的批处理大小所获得的好处将通过重新计算得到部分缓解。更重要的是，即使我们选择只存储输入的一小部分，存储单个层的激活需要250MB的空间，这意味着我们很难在12GB的GPU上支持超过12个样本的批处理大小。
RevNets幸运的是，我们还有其他方法来减少内存使用。RevNet。
RevNets有个非常聪明的计算技巧，通过以一种特定的方式构造每一层，使内存使用与网络深度保持一致。每一层分为两个部分，X₁和X₂，前向计算如下：
def forward_pass(x1, x2, Wf, Wg): """ Need an extra node in the computational graph because the gradient of the loss with respect to z1 # differs from the gradient of loss with respect to y1x1: one half of layer input x2: other half of layer input Wf: weights that parameterize function f Wg: weights that parameterize function g """ z1 = x1 + f(Wf, x2) y2 = x2 + g(Wg, z1) y1 = z1可视化一下，看起来就是这样：
由于该层的特定结构，我们可以编写一个自定义函数参数更新，这意味着我们不需要缓存任何激活来计算我们的后向传播。类似于使用梯度检查点，我们仍然需要做一些冗余计算。然而，由于每一层的输入都可以很容易地从它的输出中构造出来，我们的内存使用不再随网络中层数的增加而增加。
在实践中，Reformer定义f(x)是LSH注意力块，g (x)是标准的前向块，来自transformer结构。
有了RevNet架构，我们只需要在内存中存储单层的激活，就可以在训练期间使用更大的批处理大小！现在我们不再受训练期间激活的内存占用的限制，我们可以利用LSH注意力块改进时间复杂度。
重要的是，语言模型的loss不会因为可逆层结构而降低。
这些变化实现起来并不容易 —— 很明显jsdds Kitaev, Łukasz Kaiser和Anselm Levskaya付出巨大的努力在平衡时间和内存。
总的来说，这些变化使得序列长度的扩展成为可能。虽然结果是初步的，但在enwiki8上的实验表明，在语言建模任务上，Reformer可以与它的重量级前辈竞争。
总结
位置敏感ckdch的注意力和可逆层构成了Reformer的蓝图，非常高兴可以看到基于transformer的结构选择去优化和处理长序列的问题，而不是简单的扩展之前的工作。
其他参考：
Reformer：一个高效的 Transformer
图解 Reformer
对Reformer的深入解读
花式attention Transformer