注意力机制的细节

注意力机制对近几年的深度学习产生了深远影响，从谷歌的《Attention is all you need》开始，注意力机制广泛应用在计算机视觉和自然语言处理等领域。

本文首先引入注意力机制，然后从运作过程与计算逻辑等方面讲解自注意力。自注意力是注意力的特殊呈现，本质上是一个东西，它是Transformer等模型架构的核心模块，其重要性不言而喻。

注意力机制

我们先从一个例子开始，由浅入深，由表及里地介绍注意力机制的组成部件及其运作过程。

举个栗子

例：开学班级票选班长，有我、小孙、小明、小红和小花5人参与匿名投票，每人可投100张票，结束后我共获投100票。现在我想估计这些票中分别有多少来自他们4人。一个简单粗暴的方法是：从 我 这个主体出发，将注意力放在 审视我与每个人的亲密程度 上面，然后 以得到的亲密度作为衡量标准 来 估计票数 。

现假设我得到了与每个人的亲密度如下(亲密度 $a$ 的值用数字1-5表示，越大表示关系越好，可能投的票数越多)：

$a_{我,小孙}=4、a_{我,小明}=3、a_{我,小红}=2、a_{我,小花}=1$

亲密度不好直接量化票数，因此可对其做 $softmax$ 处理以获得 与亲密度对应的票数的概率分布 ，该分布是各项非负且和为1的权重：

${\alpha }_{我,小孙}=0.64、{\alpha }_{我,小明}=0.24、{\alpha }_{我,小红}=0.09、{\alpha }_{我,小花}=0.03$

按上述权重，可以估算出100张票分别来自4人的票数，比如因为我和小孙最亲密，因此理所当然他投给我的票最多。

另外，其他人也会参与竞选获得投票，比如小明作为主体，依此方法可以估计他所获得的来自其他4人的票数。

计算过程如图1所示：

从举例引出注意力

可以令：

主体为 $q$ ：$我=q^{我}$

亲密程度为 $k$ ： $小孙=k_1、小明=k_2、小红=k_3、小花=k_4$ ，比如我与小孙的亲密度：$a_{我,小孙}=a(q^{我},k_1)$ ，权重为 ${\alpha }_{我,小孙}=\alpha (q^{我},k_1)$

持有票数为 $v$ ： $小孙=v_1、小明=v_2、小红=v_3、小花=v_4$

则所得票数 $f$ 可表示为：

$f(q^{我})=f(q^{我},(k_1,v_1),...,(k_n,v_n))=\alpha (q^{我},k_1)v_1+\alpha (q^{我},k_2)v_2+\alpha (q^{我},k_3)v_3+\alpha (q^{我},k_4)v_4={\sum }_{i=1}^n\alpha (q^{我},k_i)v_i(1)$

其中 $n=4$ ， $\alpha (q^{我},k_i)=softmax(a(q^{我},k_i))=\frac{exp(a(q^{我},k_i))}{{\sum }_{j=1}^{n}exp(a(q^{我},k_j))}(2)$

注意力机制的构成要素

从上式 $(1)、(2)$ 可抽象出 $(3)、(4)$ ：

$f(q)=f(q,(k_1,v_1),(k_2,v_2),...,(k_n,v_n))={\sum }_{i=1}^n\alpha (q,k_i)v_i(3)$

式 $(3)$ 等号最右边其实是一个加权平均，引入注意力机制后称为 注意力汇聚 。 其中 $q$ 称为 查询 ，$(k,v)$ 称为 键值对 ，一般成对出现。

$\alpha (q,k_i)=softmax(a(q,k_i))=\frac{exp(a(q,k_i))}{{\sum }_{j=1}^{n}exp(a(q,k_j))}(4)$

式 $(4)$ 中 $\alpha (q,k_i)$ 是某个查询 $q$ 与某个键 $k$ 的 注意力权重 ，一个查询 $q$ 与一组键 $k$ 的注意力权重是各项非负且和为1的概率分布，从数学上看， $\alpha (q,k_i)$ 仅仅是 $a(q,k_i)$ 通过 $softmax$ 转换而来的，故而重点在于如何得到 $a(q,k_i)$ 。

$a(q,k_i)$ 称为 注意力评分函数 ，可以看作查询与键的某种相似度(比如上例中的亲密度)。它决定了如何对查询 $q$ 与键 $k$ 进行计算以抽取两者的相关性并得到注意力权重。

若以自然语言模型为例，输入样本通常是包含多个词元的序列，其中每个词元用一个经过词嵌入的向量表示，查询与键值对就是取自词元；当注意力评分函数 $a(q,k)$ 选的是 缩放点积注意力(scaled dot-product attention) 时，就是对两个向量 $q、k$ 执行点积运算，结果得到的是一个标量；再经 $softmax$ 转换成注意力权重；使用注意力权重对值 $v$ 进行加权平均，得到的就是考虑了所有输入的输出，一般是一个与输入同维度的向量。

以上就是注意力机制的运作过程。

自注意力机制

注意力机制中的查询、键值对可以来自相同或不同的对象，以中-英NMT为例，假设：

源句为 ‘深度学习’ ，词元化(忽略bos等特殊符号)后为： ‘深’ ‘度’ ‘学’ ‘习’ 。

目标句为 ‘deep learning’ ，词元化后为： ‘deep’ ‘learning’ 。

若以目标句作为查询 $q$ 、源句作为键值对 $k、v$ ，则查询、键值对就是来自不同的对象，可以得到 2×4=8 个注意力权重(目标句2个词元对应2个q、源句4个词元对应4个键值对k,v，每个q与每个k会得到一个注意力权重)，该注意力权重蕴含着源句每个词元对翻译结果生成的目标句每个词元的贡献或影响程度信息。这也是transformer中解码器与编码器之间注意力的计算方式。

这里顺便提一下，上面提到注意力评分函数有多种选择，其实也是要看查询 $q$ 与键值对 $k、v$ 的维度是否相同，对于查询、键值对维度相同的，使用缩放点积注意力评分函数简单又高效(因为点积运算要求维度相同)；但对于查询、键值对来自不同对象这种情况，很可能存在查询与键值对维度不同的情况，这时就不能简单使用缩放点积注意力评分函数了，而要寻求其他注意力评分函数的帮助，如 加性注意力评分函数 。

查询与键来自不同对象的计算图如下：

若源句同时作为查询 $q$ 和键值对 $k、v$ ，则查询、键值对就是来自相同的对象，这就是自注意力！这也是transformer中解码器、编码器内部各自注意力的计算方式。

自注意力不仅是指序列级的，词元也会和其自身做注意力计算，比如 $\alpha (q_1,k_1)$。

上述源句的自注意力可以得到 4×4=16 个注意力权重(源句4个词元对应4个q与4个键值对k,v)。自注意力的查询 $q$ 与键值对 $k、v$ 的维度是相同的，因此一般常用 缩放点积注意力评分函数 。

自注意力计算图如下：

自注意力的细节

本节探讨的虽是自注意力的细节，其实与一般注意力大同小异。

下图是自注意力机制的整个运作过程：从词元向量组成的序列作为注意力输入开始，通过注意力参数矩阵转换得到 $q、k、v$ ，接着使用注意力评分函数得到词元互相间的注意力权重，最后得到注意力汇聚后的输出。

这里先指定一些对象名称，方便在下面各步表示数据形状，深度学习中一些关键节点计算前后的数据形状是很重要的，需要重点关注。

输入词嵌入

首先明确一点，自注意力中的 $q、k、v$ 并不是原始的输入，事实上，原始输入序列经某种词元化方法后拆分成多个词元，再经词嵌入将词元转换成 $d$ 维向量表示：

用代码实现词嵌入。

from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"
import torch
from torch import nn
from torch.nn import functional as F
import math
from matplotlib import pyplot as plt
from matplotlib_inline import backend_inline#词嵌入
X=torch.tensor([0,1,2,3]) #输入序列(词元化后序列中的每个词元用数字索引表示)，这里 0:深 1:度 2:学 3:习 ，词元数或步数为num_steps
print('原始输入序列形状：',X.shape) #这里输入序列只有一个样本，因此加上的批量batch_size维度为1
print('原始输入序列：',X)
vocab_size=X.shape[1] #词表大小
d=6 #词嵌入维度embedding = nn.Embedding(vocab_size, d) #词嵌入接口函数
X=embedding(X) #词嵌入后的词元向量表示(上图里的输入I)print('词嵌入后序列形状：',X.shape) #形状：(batch_size,steps,d)
print('词嵌入后序列：',X)

上面代码简要实现了词嵌入过程，索引0-4分别表示 ‘深’ ‘度’ ‘学’ ‘习’ 四个词元，在词嵌入后(注意批量维度为1)每个词元用一个初始化的词向量表示(这些也是可学习的参数)，这就是自注意力的输入，本质上就是一个张量，形状为 (batch_size-小批量,num_steps-词元数,d-词嵌入维度) 。

查询与键值

直接用词嵌入后的词元向量做 $q、k、v$ 是不能表达出语义相关性的。词嵌入虽能通过学习使得词元自身具有语义，但词元之间的相关性同样也是需要学习的，因此引入注意力参数 $W^q、W^k、W^v$ 来生成 $q、k、v$ ，并通过学习来表示词元之间的相关性。

$W^q、W^k、W^v$ 是序列词元共享的一组参数，它们的形状分别为 (num_hiddens, query_size)、(num_hiddens, key_size)、(num_hiddens, value_size) ，前面说过，在transformer模型架构的自注意力中，输入经注意力操作后的输出维度不变，所以 d=num_hiddens=query_size=key_size=value_size 。

已知输入序列形状为 (batch_size,num_steps,d) ，所以与 $W^q、W^k、W^v$ 计算后：

$q$ 的形状为 (batch_size,num_steps-查询数 ,query_size) ；

$k$ 的形状为 (batch_size,num_steps-键值对数 ,key_size) ；

$v$ 的形状为 (batch_size,num_steps-键值对数 ,value_size) 。

用nn.Linear生成注意力参数 $W^q、W^k、W^v$ 。

#注意力参数
query_size, key_size, value_size, num_hiddens = d, d, d, d
W_q = nn.Linear(query_size, num_hiddens, bias=False)
W_k = nn.Linear(key_size, num_hiddens, bias=False)
W_v = nn.Linear(value_size, num_hiddens, bias=False)print('W_q形状：',W_q.weight.shape)
print('W_q：',W_q.weight.T)
print('W_k形状：',W_k.weight.shape)
print('W_k：',W_k.weight.T)
print('W_v形状：',W_v.weight.shape)
print('W_v：',W_v.weight.T)

有了输入与注意力参数，下面代码生成查询与键值对。

#生成查询与键值对
queries = W_q(X) #查询q
keys = W_k(X) #键k
values = W_v(X) #值v#自注意力中查询、键、值的形状与X相同
print('q的形状：',queries.shape) #q的形状(batch_size,num_steps-查询数,query_size)
print('k的形状：',keys.shape) #k的形状(batch_size,num_steps-键值对数,key_size)
print('v的形状：',values.shape) #v的形状(batch_size,num_steps-键值对数,values_size)queries #查看查询的值

现在来看看 $q$ 的计算细节。以词元 ‘深’ 为例，其查询 $q$ 是 $X、W^q$ 矩阵乘法生成结果的其中一行，如下图所示。

可以看出 $q$ 的最后一维取决于 $W^q$ 中的 query_size ，$k、v$ 的计算依此类推。

注意力评分

对于 $q、k、v$ ，可以使用特定的注意力评分函数计算注意力分数，Transformer的自注意力一般选择缩放点积注意力评分函数，数学公式为：

$a(q,k)=q{k}^T/\sqrt{d_k}$

其实就是 $q$ 和 $k$ 的转置这两个向量的点积，点积某种程度上可以表示两个向量的相似度，比如余弦相似度就是两个向量点积再除以它们的模，这里略去了模，影响不大。这里所有词元的 $q、k$ 将并行做点积运算，缩放点积注意力评分函数的优势就是简单高效。除以 $\sqrt{d}$ 是为了减少向量长度对点积方差的影响。

每个词元的 $q$ 都要和每个词元的 $k$ 做点积得到一个权重，上面已知：

$q$ 的形状为：(batch_size, num_steps-查询数, query_size)；

$k$ 的形状为：(batch_size, num_steps-键值对数, key_size)；

因此可以预期最后得到的注意力权重形状为：(batch_size, num_steps–4个查询, num_steps–4个键值对)。

在与查询 $q$ 计算前，键 $k$ 需要转置一下。

#k的转置
keys.transpose(1,2).shape
keys.transpose(1,2) #转置最后两个维度，转置后形状为(batch_size, key_size, num_steps-键值对数)

计算注意力评分。

#计算注意力评分
attention_scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)
attention_scores.shape
attention_scores

以词元 ‘深’ 的查询 $q$ (下图queries红框里第一行)为例，其与所有词元的键 $k$ 的注意力分数的计算如下图所示，这里因为 query_size=key_size，所以可以直接点积。

注意力权重

用注意力分数经 $softmax$ 转换即可得到注意力权重，4个查询、4个键，共得到16个注意力权重值。

#注意力权重
attention_weights=nn.functional.softmax(attention_scores, dim=-1)  
attention_weights.shape
attention_weights

画出注意力权重的热力图，注意当前注意力参数是初始化未经训练的，表现不出有用的相关性。

#注意力热力图
def attention_heatmap(weights):backend_inline.set_matplotlib_formats('svg')rows, cols = weights.shape[0], weights.shape[1]fig, axes = plt.subplots(rows, cols, figsize=(3.0, 3.0), sharex=True, sharey=True, squeeze=False)for i, (row_axes, row_matrices) in enumerate(zip(axes, weights)):for j, (ax, matrix) in enumerate(zip(row_axes, row_matrices)):pcm = ax.imshow(matrix.detach().numpy(), cmap='Reds')if i == rows - 1:ax.set_xlabel('Keys')if j == 0:ax.set_ylabel('Queries')fig.colorbar(pcm, ax=axes, shrink=0.6)attention_heatmap(attention_weights.reshape((1, 1, 4, 4)))

掩蔽注意力

对于前文得到的注意力分数，有时可能不是每个分数都应当参与注意力权重的计算。

比如在语言模型中，输入序列的长度不一(不同序列的词元数量可能不同)，但为了方便小批量计算，序列会被长切短补到相同长度(相同数量的词元)，序列长度不够时会在最后面以特殊词元 pad 填补。

然而这些填补的词元没有实际意义，自然不该在注意力权重中分得一杯羹，为了达到这个目的，需将填补位置的注意力分数替换成一个很大的负值，比如-1e6，使对应位置的注意力权重为0(因为很大的负值经softmax后输出为0)，从而实现遮蔽注意力权重的效用。

下图是上文生成的注意力分数，我们现在假设原来的 ‘深’ ‘度’ ‘学’ ‘习’ 四个词元中的 ‘习’ 是用 pad 填充的，这时每个查询(查询将仍是4个，包含补充词元)对应的有实际意义的键都只有三个，因此每个查询只应使用红色虚线框里面的三个值来计算注意力权重，灰色虚线框里的值在经softmax前需替换成-1e6。

下面用代码实现掩蔽注意力。

valid_lens=torch.tensor([3,3,3,3])
attention_shape=attention_scores.shape
attention_scores=attention_scores.reshape(-1, attention_shape[-1])
mask = torch.arange((attention_shape[1]), dtype=torch.float32)[None, :] < valid_lens[:, None] #要遮蔽的位置
attention_scores[~mask] = -1e6 #遮蔽位置替换为-1e6print('遮蔽后的注意力分数：\n')
attention_scores
print('遮蔽后的注意力权重：\n')
attention_weights=nn.functional.softmax(attention_scores.reshape(attention_shape), dim=-1)
attention_weights.shape
attention_weights

下图可以看出，查询 $q$ 与填补的键 $k$ 间的注意力权重为0。

注意力汇聚

注意力汇聚就是注意力权重对值 $v$ 的加权平均，与文章开头所举例子里计算票数的方式一致。已知：

注意力权重形状：(batch_size, num_steps-查询数, num_steps-键值对数)；

值 $v$ 形状：(batch_size, num_steps-键值对数, value_size)。

print('注意力权重：\n')
attention_weights.shape
attention_weights
print('值v：\n')
values.shape
valuesprint('输出：\n')
output=torch.bmm(attention_weights, values)
output.shape
output

计算细节见下图：

汇聚后输出 $I^{\prime }$ 的形状为：(batch_size, num_steps-查询数, value_size)，最后的维度是值 $values$ 中的 value_size 。

自注意力初始输入 $I$ 的形状是(batch_size,num_steps,d)，因为 value_size=d ，因此初始输入经过自注意力计算后的输出的形状不变(词元也是一一对应)，这样设计的初衷是为了方便计算。比如在transformer中，自注意力存在于编码器与解码器中，通常transformer模型里会堆叠多个编码器/解码器块，输入输出形状保持一致可以使上一层的输出继续作为下一层的输入，使得模型的体量可以灵活控制。

多头注意力

上述可以说只是单头注意力的计算过程，因为只生成了一组 $q、k、v$，查询 $q$ 从一组键 $k$ 上仅能抽取一种相关性。这也许并不太够，我们希望能抽取更多的查询与键的相关性以使注意力包含更丰富的语义信息，有点类似卷积神经网络中的卷积核，不同卷积核能识别不同的模式。

既然一组 $q、k、v$ 能找到一种相关性，如果有多组不同的 $q、k、v$ ，自然就可以找到多种相关性，这就是多头注意力机制。

那么不同的 $q、k、v$ 如何得到？简单来说，对同一个输入 $I$ ，若想得到不同的 $q、k、v$ ，只需使用多组注意力参数 $W^q、W^k、W^v$ 即可。

下图演示了具有2个头(num_heads)的多头注意力计算过程，与上面不同点在于2个头分别对应一组注意力参数(不同颜色标示)，除了注意力参数不同外，其他的计算过程和单头注意力完全相同，只是对于每个头最后的输出 $I^{\prime }$ ，需要将它们连结起来作为最终输出。注意连结后输出的维度要和开始时的输入 $I$ 一样是 $d$ ，因此在此例中，每个头自身输出的维度都是 $\frac{d}{2}$ ，事实上，如果有 $n$ 个头，则每个头输出的维度将是 $\frac{d}{n}$ 。

来看下具体的计算过程。输入 $I$ 是不变的，注意力参数 $W^q、W^k、W^v$ 也保持不变，只是它们现在是由 $W_1^q、W_1^k、W_1^v$ 与 $W_2^q、W_2^k、W_2^v$ 这两组注意力参数拼接而成。

$queries、keys、values$ 的计算也与之前一样，只是为了方便多头注意力的并行计算，需要转换一下形状。

print('queries初始形状',queries.shape) #(batch_size, num_steps, query_size)
queriesnum_heads=2
#q形状转换
queries=queries.reshape(queries.shape[0], queries.shape[1], num_heads, -1)
queries=queries.permute(0, 2, 1, 3)
queries=queries.reshape(-1, queries.shape[2], queries.shape[3])
#k形状转换
keys=keys.reshape(keys.shape[0], keys.shape[1], num_heads, -1)
keys=keys.permute(0, 2, 1, 3)
keys=keys.reshape(-1, keys.shape[2], keys.shape[3])
#v形状转换
values=values.reshape(values.shape[0], values.shape[1], num_heads, -1)
values=values.permute(0, 2, 1, 3)
values=values.reshape(-1, values.shape[2], values.shape[3])print('queries被reshape后的形状',queries.shape) #(batch_size*num_heads, num_steps, query_size/num_heads)
queriesprint('keys被reshape后的形状',keys.shape) #(batch_size*num_heads, num_steps, query_size/num_heads)
keys

下图展示了形状转换过程：

剩余部分的计算与上面单头完全一样，只是最后每个头都有一组注意力权重，且需要连结起来。连结后的多头注意力结果与单头的结果并不相同，因为在计算注意力分数时，多个头的 $q、k$ 的维度已经被多头均分，所以得到的注意力分数已经不同了，不过本身多头注意力就是为了抽取 $q、k$ 之间的不同相关性。

#计算注意力评分
attention_scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)#注意力权重
attention_weights=nn.functional.softmax(attention_scores, dim=-1)  #输出
output=torch.bmm(attention_weights, values)#多头注意力连结
output = output.reshape(-1, num_heads, output.shape[1], output.shape[2])
output = output.permute(0, 2, 1, 3)
output = output.reshape(output.shape[0], output.shape[1], -1)
output.shape
outputattention_heatmap(attention_weights.reshape((1, 2, 4, 4)))

查看两组注意力权重的热力图，分布并不相同。