2019年11月谷歌的论文“Fast Transformer Decoding: One Write-Head is All You Need“。
Transformer神经序列模型中使用的多头注意层,是RNN的替代。虽然整个序列的并行性让这些层的训练通常快速而简单,但由于重复加载大型“Key”和“Value”张量的内存-带宽成本,增量推理(这种并行化不可能的情况下)通常运行缓慢。本文提出一种称为多查询注意的变型,其中Key和Value在所有不同的注意“头”之间共享,从而大大减少了这些张量的大小,继而减少了增量解码的内存-带宽要求。实验验证,产生的模型确实可以更快地解码,并且与基线相比只会有轻微的质量下降。
下面的代码描述了一个点积注意公式,其中权重被计算为具有不同Key的查询点积softmax:
注:代码示例使用了在TensorFlow和numpy中定义的einsum表示法,用于任意维度张量之间的广义contractions计算;在这种表示法中,一个方程命名输入和输出的张量维度;该计算在数字上等效于广播每个输入得到所有维度的并集、按分量相乘、在所有维度上求和。
Transformer的序列-序列模型[Vaswani 2017]并行使用h个不同的注意层(头部),称之为“多头注意”。h个不同层的查询向量来自输入向量对应h个不同学习的线性投影Pq。类似地,Key和Value来自m个不同输入向量对应h个学习的线性投影Pl,Pv。h层的输出本身是通过不同的学习线性投影Po,然后求和得到。为了简单起见,输入和输出向量取相同的维数d。其计算可以表示如下:
注:一个常数尺度因子在此忽略。
在实践中,将多个查询批处理在一起要高效得多。下面的代码添加了两种类型的批处理。首先,从序列中n个不同位置生成查询。这些查询都用相同Key和 Value交互。此外,同时处理一批b个不同的非交互序列。根据[Vaswani2017],在自回归模型中可以在包含值-∞的logits添加“掩码”来防止反向信息流。
在某些设置中,由于数据依赖性,无法并行处理来自多个位置的查询。一个例子是自回归语言模型中的自注意层,如Transformer[Vaswani 2017]。在每个位置产生的查询涉及在该位置之前(包括该位置)产生的KV对。在训练过程中,真值目标序列是已知的,可以高效并行实现。然而,特定位置处的自注意层输出,当训练的模型生成时,会影响下一位置处生成的token,又影响到下一个位置这个自注意层的输入。这个阻止了并行计算。下面显示了用于增量计算该自注意层的代码:
如下提出多查询注意,作为多头注意的变型。多头注意由多个并行的注意层(头)组成,对查询、Key、Value和输出进行不同的线性转换。多查询注意是相同的,只是不同的头共享一组Key和Value。(增量)多查询(自)注意的代码与上面列出的多头注意的代码相同,只是从tf.einsum方程中删除了字母“h”,其表示K、V、Pk或Pv的“头”维度。