哈希算法、爱因斯坦求和约定，这是2020年的注意力机制

在 Transformer 完全采用注意力机制之后，注意力机制有有了哪些改变？哈希算法、Head 之间的信息交流都需要考虑，显存占用、表征能力都不能忽视。

注意力机制是非常优美而神奇的机制，在神经网络「信息过载」的今天，让 NN 学会只关注特定的部分，无疑会大幅度提升任务的效果与效率。借助注意力机制，神经机器翻译、预训练语言模型等任务获得了前所未有的提升。

但与此同时，注意力机制也面临着重重问题，首先就是参数量太大，这有点类似于全连接层，注意力机制需要考虑两两之间的所有连接。我们可以看到，完全用注意力机制的模型，参数量轻轻松松破个亿，而卷积这类参数共享的运算，参数量一般也就几百万。

如果只是参数量大也就算了，能完整利用这些参数，正是说明模型的表征能力非常强。但问题在于，Transformer 采用的 Multi-head Attention 并没有充分利用参数的表征能力。举个例子，Transformer 中每一个注意力 Head 都是相互独立的，它们之间没有信息交流，因此谷歌最近提出的 Talking-Head 就旨在解决这个问题。

本文从原 Multi-head Attention 出发，探索 Reformer 如何用哈希算法大量降低显存需求，探索 Talking-Head 如何强化全注意力机制的表征能力 。

多头注意力：开始的地方

Transformer 因在大型预训练语言模型中的优秀性能而被世人所熟知。这一类模型已广泛应用于多种预训练语言模型中，如 BERT、GPT-2等。当然，在广泛应用之余，人们也在思考 Transformer 存在的缺陷，并进行弥补。

• 论文：Attention Is All You Need

• 论文链接：https://arxiv.org/abs/1706.03762 (https://arxiv.org/abs/1706.03762)

提到 Transformer，人们总会想到这类模型最突出的几个特征：点乘注意力、多头注意力、残差连接、Mask 机制等。在最初介绍 Transformer 的论文中，Waswani 等引入了这些优秀的特性，使 Transformer 独树一帜。

优秀的特性，问题的根源

在 Transformer 中，首先引入的是注意力点乘机制。它将输入分为查询（Query）、键（Key）和值（Value），通过以下公式计算。其相当于根据序列元素之间的相似性，确定每一个元素都应该关注哪些信息。

而另一方面，仅计算一次注意力，不足以捕捉所有的文本结构特征。因此，Transformer 的提出者想到类似于 CNN 中的 multi-channel 的方法：让点乘注意力机制「重复」多次。这样一来，Transformer 的模型复杂度显著上升，捕捉了更多的语义和语言特征，这种方法便是我们熟悉的多头注意力。

图 2：（左）点乘注意力；（右）多头注意力，有多个注意力层并行。

多头注意力的公式如下所示，每一个 Head 都是一个注意力层。为了提高差异性，每一个 Head 的输入张量都会先经过它特有的线性变换，相当于希望该 Head 注意特定的某个方面。

尽管 Transformer 因为这些特性而产生了很好的性能，但其固有的缺陷也随之而生。特别需要注意的就是多头注意力机制了。在设计之初，为了让不同的注意力头能够捕捉语言模型中的不同特征，不论是 BERT 还是GPT-2的设计者都将增加了注意力头的数量。

但是，由于 Query、Key 和 Value 在训练中会进行多次线性变换，且每个头的注意力单独计算，因此产生了大量的冗余参数，也需要极大地显存。

怎样才能降低计算复杂度呢？研究者们在这个问题上思考了很多方法，本文中可以看到采用哈希算法、参数压缩等的方法。这些方法有的是强化多头注意力机制。有的则是利用哈希算法的高效性，替换掉原本的架构，使模型性能更加优秀。

哈希改造下的注意力机制

Transformer 因为 Multi-Head Attention 而成功，也因为它而受到重重限制。在 2017 年 Transformer 提出以后，18 年到 19 年已经有很多研究在改进这种注意力机制，这里介绍的是一种「哈希大法」，提出该方法的 Reformer 已经被接收为 ICLR 2020 的 Oral 论文。

ICLR 2020 程序主席评论道：「Reformer 这篇论文提出的新型注意力机制有效减少序列长度的复杂度，同时新机制也减少了存储需求。实验证明它们非常有效，该论文可能对研究社区会产生显著影响。尽管有一些次要的观点，但所有评审者都一致同意接收该论文。」

• 论文：Reformer: The Efficient Transformer

• 论文地址：https://openreview.net/pdf?id=rkgNKkHtvB

注意力机制到底哪不对？

我们先回到注意力机制，其中非常重要的运算是 Query 与 Key 这两个张量之间的矩阵乘法，其代表着余弦相似性。具体而言，如果 Query 表示翻译领域的中文序列，Key 表示英文序列，那么矩阵乘法则表示在翻译成某个中文词时，需要注意哪些英文词。

但问题在于，Query 与 Key 的张量维度都是 [batch size, length, d_model]，乘完后的维度是 [batch size, length, length]。要是序列长度 length 不长也就罢了，但如果上下文跨度很广，需要很长的序列，也就是说 length 远远大于隐藏层大小 d_model。那可就麻烦了，研究者表示要是序列长度达到 64K，注意力机制计算一个样本也需要 16GB 的显存（Float32）。

你可能会说，我不一次性计算 Query 中的所有序列，而每次只算一步，例如一个中文词，计算它与所有英文词之间的相似性。这样不就没问题了么？这是没问题的，但并没有实质上的改变。

哈希来帮忙

现在，假定序列长度还是 64K，对于 Query 序列中的每一个 q，正常的注意力机制需要计算它和 Key 序列 64K 个元素之间的相似度，并通过 Softmax 将相似度归一化为概率。反正都是要计算概率，且一般只有概率最高的一些元素真正对 q 有很大的贡献，那么为什么不直接找出这些元素？

之前超大的矩阵乘法会计算 Query 序列所有元素与 Key 序列所有元素之间的相似度，现在如果不通过矩阵乘法，只找每个 Query 序列「最相近」的 32 个或 64 个元素，那么显存与计算岂不是成千倍地减少？

这就是 Reformer 最核心的思想，完成查找「最相近」元素的算法即局部敏感哈希算法（Locality sensitive hashing）。

简单而言，局部敏感哈希算法（LSH）在输入数据彼此类似时，它们有很大概率映射后的哈希是一样的；而当输入数据彼此不同，它们映射后的哈希值相等概率极小。

LSH 算法根据局部敏感哈希函数族将类似的数据映射到相同的桶（Bucket）。

LSH注意力机制

如下图 a 所示，传统注意力机制需要计算 q 与 k 之间的所有相关性，但实际上真正起作用的注意力分配是非常稀疏的。对于图 b 来说，Query 和 Key 会先通过哈希桶排序，相似的 q 与 k 会落在相同的桶内，因此直接在桶内执行注意力计算就能逼近完整的注意力机制。

然而图 b 只是理想情况，真正用于注意力机制还有很多困难要克服，比如说，一个桶内包含很多 q 而没有一个 k。为此，研究者对 b 执行了一系列转换，并使排序后的注意力权重，来自相同桶的 q 和 v 都集中在对角线上，就如同图 c 一样。注意 c 图中 Q=K，表示 Transformer 中的自注意力机制。

排好序后，d 图就可以分成不同的批量，并完成并行计算。这样不仅计算量与参数量大大减少，还能加速训练。

最后，整个 LSH注意力机制就如上左图所示。对于自注意力机制，先根据 LSH 分成不同的桶，然后对桶排序并分割为不同的批量，这些批量能够并行处理。最后只要在相同的桶内执行注意力机制就行了，整个过程会节省大量存储与计算资源。

Talking-Heads Attention

近日，来自 Google 的研究团队提出一种「交谈注意力机制」（Talking-Heads Attention），在 softmax 操作前后引入对多头注意力之间的线性映射，以此增加多个注意力机制间的信息交流。这样的操作虽然增加了模型的计算复杂度，却能够在多项语言处理问题上取得更好的效果。

• 论文：Talking-Heads Attention

• 论文地址：https://arxiv.org/abs/2003.02436

einsum 表示法

作者在论文的伪代码中使用大写字母表示张量，而小写字母表示其对应维度大小。同时作者在张量的计算中使用了 einsum 表示法，也就是爱因斯坦求和约定。它在 numpy、tensorflow、pytorch 等 Python 扩展库中均有实现。使用 einsum 表示法能够仅使用一个函数，就可以优雅地实现如点积、外积、转置、矩阵-向量乘法、矩阵-矩阵乘法等运算。

einsum 解释：https://rockt.github.io/2018/04/30/einsum

限于本文篇幅，有关 einsum 表示法的概念，感兴趣的小伙伴可参考上面这篇博客。这里举个栗子，两个矩阵的乘法运算使用 einsum 表示法可写为：

Y[a, c] = einsum(X[a, b], W[b, c])

于是，前面介绍的多头注意力机制使用 einsum 表示法可改写为如下形式：

同时，einsum 表示法还支持大于两个矩阵作为输入的运算。于是，以上伪代码可进一步精简为如下极简模式：

交谈注意力机制

不同于多头注意力机制，交谈注意力机制在 softmax 运算的前后增加了两个可学习的线性映射P_l 与 P_w，这使得多个注意力 logit 和注意力 weight 彼此之间能够相互传递信息。

于是，交谈注意力机制出现了三个不同的维度：h_k、h 和 h_v，其中 h_k 表示 Query 和 Key 的注意力头数量，h 表示 logit 和 weight 的注意力头数量，h_v 则表示值的注意力头数量。其对应伪代码表示如下，注释中标出了每个 einsum 运算所对应的计算量。

Talking-Head 是 Multi-Head 的延伸

当假设注意力机制的输入与输出维度相同时，可得到计算多头注意力机制所需进行的标量乘法运算数目为：

h·(dk +dv)·(n·dX +m·dM +n·m)

而交谈注意力机制所需的标量乘法运算数目为：

(dk ·hk +dv ·hv)·(n·dX +m·dM +n·m)+n·m·h·(hk +hv)

可以看到上式中第一项与多头注意力机制的计算量类似，而第二项是由交谈注意力中的线性映射导致的。假如 h

作者指出，由于交谈注意力机制引入了额外的小尺寸张量运算，在实际运行过程中很可能导致速度比 Transformer 更慢。

最后，多头注意力机制与交谈注意力机制都可看作一种「通用双线性多头注意力机制」（GBMA, i.e. general bilinear multihead attention）的特殊形式。我们可以从以下伪代码看出，GBMA 具有两个三维的参数张量。

从以下伪代码不难看出，多头注意力机制在数学上等效于，使用两个因子的乘积分别表示 GBMA 中的各参数张量。

而交谈注意力机制在数学上等效于，使用三个因子的乘积分别表示 GBMA 中的各参数张量。

这里的 GBMA 仅作为理论研究探讨，由于其计算量较大可能并不具备实用性。

「注意力」这个想法真的非常优雅，所以 2020 年剩下的是，我们如何才能更高效地实现「注意力」？