当前位置:   article > 正文

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

多头注意力机制是什么时候提出的
机器之心报道

参与:思、肖清、一鸣

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

b06446f1db2af4ff46f37801c1557638.png

注意力机制是非常优美而神奇的机制,在神经网络「信息过载」的今天,让 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),通过以下公式计算。其相当于根据序列元素之间的相似性,确定每一个元素都应该关注哪些信息。

0b85690c5c268947a16cb0d6f66271e6.png

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

d659975ae6a96e8aa621d2794af9bc92.png

图 2:(左)点乘注意力;(右)多头注意力,有多个注意力层并行。 多头注意力的公式如下所示,每一个 Head 都是一个注意力层。为了提高差异性,每一个 Head 的输入张量都会先经过它特有的线性变换,相当于希望该 Head 注意特定的某个方面。

6693ef35f6f965761456adf10e6bccca.png

尽管 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)。

46d1ba55ee5a2e99a18ad774c62d52e3.png

你可能会说,我不一次性计算 Query 中的所有序列,而每次只算一步,例如一个中文词,计算它与所有英文词之间的相似性。这样不就没问题了么?这是没问题的,但并没有实质上的改变。 哈希来帮忙 现在,假定序列长度还是 64K,对于 Query 序列中的每一个 q,正常的注意力机制需要计算它和 Key 序列 64K 个元素之间的相似度,并通过 Softmax 将相似度归一化为概率。反正都是要计算概率,且一般只有概率最高的一些元素真正对 q 有很大的贡献,那么为什么不直接找出这些元素? 之前超大的矩阵乘法会计算 Query 序列所有元素与 Key 序列所有元素之间的相似度,现在如果不通过矩阵乘法,只找每个 Query 序列「最相近」的 32 个或 64 个元素,那么显存与计算岂不是成千倍地减少? 这就是 Reformer 最核心的思想,完成查找「最相近」元素的算法即局部敏感哈希算法(Locality sensitive hashing)。 简单而言,局部敏感哈希算法(LSH)在输入数据彼此类似时,它们有很大概率映射后的哈希是一样的;而当输入数据彼此不同,它们映射后的哈希值相等概率极小。

b68371c9c582414041883963a3c290d0.png

LSH 算法根据局部敏感哈希函数族将类似的数据映射到相同的桶(Bucket)。 LSH 注意力机制 如下图 a 所示,传统注意力机制需要计算 q 与 k 之间的所有相关性,但实际上真正起作用的注意力分配是非常稀疏的。对于图 b 来说,Query 和 Key 会先通过哈希桶排序,相似的 q 与 k 会落在相同的桶内,因此直接在桶内执行注意力计算就能逼近完整的注意力机制。

08a161195e6a400c30231f14cc481a4e.png

然而图 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 表示法可改写为如下形式:

37020f6d756472452dc224355262bc79.png

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

bfa40b9c7e381eb2e168896af9901a77.png

交谈注意力机制 不同于多头注意力机制,交谈注意力机制在 softmax 运算的前后增加了两个可学习的线性映射 P_l 与 P_w,这使得多个注意力 logit 和注意力 weight 彼此之间能够相互传递信息。 于是,交谈注意力机制出现了三个不同的维度:h_k、h 和 h_v,其中 h_k 表示 Query 和 Key 的注意力头数量,h 表示 logit 和 weight 的注意力头数量,h_v 则表示值的注意力头数量。其对应伪代码表示如下,注释中标出了每个 einsum 运算所对应的计算量。

5e2da949ff26e2c714160895100d1170.png

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 具有两个三维的参数张量。 d226c481b78c6fe0ef3d53ccacea542d.png 从以下伪代码不难看出,多头注意力机制在数学上等效于,使用两个因子的乘积分别表示 GBMA 中的各参数张量。

be3ac0dc3e882e18f54fbfaf77ee5182.png

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

7f2b5f3e79d6b7741b17dab4572ec72a.png

这里的 GBMA 仅作为理论研究探讨,由于其计算量较大可能并不具备实用性。 「注意力」这个想法真的非常优雅,所以 2020 年剩下的是,我们如何才能更高效地实现「注意力」? 文为机器之心报道,转载请联系本公众号获得授权 ✄------------------------------------------------ 加入机器之心(全职记者 / 实习生):hr@jiqizhixin.com 投稿或寻求报道:content @jiqizhixin.com 广告 & 商务合作:bd@jiqizhixin.com
声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/我家自动化/article/detail/611742
推荐阅读
相关标签
  

闽ICP备14008679号