思、肖清、一鸣报道

哈希算法、爱因斯坦求和约定,这是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<d_k 且 h<d_v,那么交谈注意力机制的线性映射计算量,将分别小于对应多头注意力机制中的 n·m·dk ·hk 与 n·m·dv ·hv。

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

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

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

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


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

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


理论多头注意力reformer哈希算法注意力机制
5
相关数据
权重技术

线性模型中特征的系数,或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0,则相应的特征对模型来说没有任何贡献。

神经机器翻译技术

2013 年,Nal Kalchbrenner 和 Phil Blunsom 提出了一种用于机器翻译的新型端到端编码器-解码器结构 [4]。该模型可以使用卷积神经网络(CNN)将给定的一段源文本编码成一个连续的向量,然后再使用循环神经网络(RNN)作为解码器将该状态向量转换成目标语言。他们的研究成果可以说是神经机器翻译(NMT)的诞生;神经机器翻译是一种使用深度学习神经网络获取自然语言之间的映射关系的方法。NMT 的非线性映射不同于线性的 SMT 模型,而且是使用了连接编码器和解码器的状态向量来描述语义的等价关系。此外,RNN 应该还能得到无限长句子背后的信息,从而解决所谓的「长距离重新排序(long distance reordering)」问题。

自注意力技术

自注意力(Self-attention),有时也称为内部注意力,它是一种涉及单序列不同位置的注意力机制,并能计算序列的表征。自注意力在多种任务中都有非常成功的应用,例如阅读理解、摘要概括、文字蕴含和语句表征等。自注意力这种在序列内部执行 Attention 的方法可以视为搜索序列内部的隐藏关系,这种内部关系对于翻译以及序列任务的性能非常重要。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

伪代码技术

伪代码,又称为虚拟代码,是高层次描述算法的一种方法。它不是一种现实存在的编程语言;它可能综合使用多种编程语言的语法、保留字,甚至会用到自然语言。 它以编程语言的书写形式指明算法的职能。相比于程序语言它更类似自然语言。它是半形式化、不标准的语言。

注意力机制技术

我们可以粗略地把神经注意机制类比成一个可以专注于输入内容的某一子集(或特征)的神经网络. 注意力机制最早是由 DeepMind 为图像分类提出的,这让「神经网络在执行预测任务时可以更多关注输入中的相关部分,更少关注不相关的部分」。当解码器生成一个用于构成目标句子的词时,源句子中仅有少部分是相关的;因此,可以应用一个基于内容的注意力机制来根据源句子动态地生成一个(加权的)语境向量(context vector), 然后网络会根据这个语境向量而不是某个固定长度的向量来预测词。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

神经网络技术

(人工)神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型,那时候研究者构想了「感知器(perceptron)」的想法。这一领域的研究者通常被称为「联结主义者(Connectionist)」,因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型,它们都是前馈神经网络:卷积神经网络(CNN)和循环神经网络(RNN),其中 RNN 又包含长短期记忆(LSTM)、门控循环单元(GRU)等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习,但也有一些为无监督学习设计的变体,比如自动编码器和生成对抗网络(GAN)。

映射技术

映射指的是具有某种特殊结构的函数,或泛指类函数思想的范畴论中的态射。 逻辑和图论中也有一些不太常规的用法。其数学定义为:两个非空集合A与B间存在着对应关系f,而且对于A中的每一个元素x,B中总有有唯一的一个元素y与它对应,就这种对应为从A到B的映射,记作f:A→B。其中,y称为元素x在映射f下的象,记作:y=f(x)。x称为y关于映射f的原象*。*集合A中所有元素的象的集合称为映射f的值域,记作f(A)。同样的,在机器学习中,映射就是输入与输出之间的对应关系。

余弦相似性技术

余弦相似性通过测量两个向量的夹角的余弦值来度量它们之间的相似性。0度角的余弦值是1,而其他任何角度的余弦值都不大于1;并且其最小值是-1。从而两个向量之间的角度的余弦值确定两个向量是否大致指向相同的方向。两个向量有相同的指向时,余弦相似度的值为1;两个向量夹角为90°时,余弦相似度的值为0;两个向量指向完全相反的方向时,余弦相似度的值为-1。这结果是与向量的长度无关的,仅仅与向量的指向方向相关。余弦相似度通常用于正空间,因此给出的值为0到1之间。

桶排序技术

桶排序或所谓的箱排序,是一个排序算法,工作的原理是将数组分到有限数量的桶里。每个桶再个别排序(有可能再使用别的排序算法或是以递归方式继续使用桶排序进行排序)。桶排序是鸽巢排序的一种归纳结果。当要被排序的数组内的数值是均匀分配的时候,桶排序使用线性时间。

查询技术

一般来说,查询是询问的一种形式。它在不同的学科里涵义有所不同。在信息检索领域,查询指的是数据库和信息系统对信息检索的精确要求

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用,如语音识别,机器翻译,词性标注,句法分析和资讯检索。由于字词与句子都是任意组合的长度,因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题),也使得在语料库中估算字串的机率变得很困难,这也是要使用近似的平滑n元语法(N-gram)模型之原因。

GPT-2技术

GPT-2是OpenAI于2019年2月发布的基于 transformer 的大型语言模型,包含 15 亿参数、在一个 800 万网页数据集上训练而成。据介绍,该模型是对 GPT 模型的直接扩展,在超出 10 倍的数据量上进行训练,参数量也多出了 10 倍。在性能方面,该模型能够生产连贯的文本段落,在许多语言建模基准上取得了 SOTA 表现。而且该模型在没有任务特定训练的情况下,能够做到初步的阅读理解、机器翻译、问答和自动摘要。

暂无评论
暂无评论~