Attention 机制 • CoderXL's Blog

在先前的 RNN 文本翻译模型中，我们要求先使用 Encoder RNN 读入原文以预热 Hidden State，然后再将 Hidden State 传递给 Decoder RNN 进行自回归文本生成。
其中一个洞察就是，原文的全部信息都被压缩在最后一个时间步的 Hidden State 中了（虽然可以通过将 Output concat 到 Decoder 的输入来缓解），造成了较大的信息损失；此外，Decoder 在生成序列的时候，只有最初的 Hidden State 才能忠实反映原文的本意，后续时间步的 Hidden State 可能出现 shift.

这些问题让我们思考，能否让 Decoder 在任何时刻，都能接触到更保真、更完整的原文信息？于是便有了 Attention 机制。

古早的 Attention#

不可学习（非参）的注意力池化层，可追溯到 60 年代的工作 ↗，不是这篇笔记的主题。

这篇笔记的主题都是可以学习的注意力池化层。

朴素的 Attention#

原理#

在上述 RNN 文本翻译模型中，我的 Encoder 对每个时间步不再只计算出一个 Hidden State，而是改为计算出一个 Key 和一个 Value（也有的实践中直接使用同一个值充当 Key 和 Value，和原来的 Hidden State 差不多），分别表示这个时间步对应的输入的“特征索引”和“特征数值”。
这两个概念有些类似，都有编码当前输入特征的作用（这也是为什么有的时候会简化合并起来），但用途不同：前者的用途是给 Attention 机制提供一个索引，允许模型找到想要找的信息；后者是将具体的特征信息取出来，输入模型进行推理。

而到了 Decoder 的部分，模型每次预测之前，都会产生一个 Query（可以认为是和 Hidden State 平行的输出层），使用这个 Query 去和 Encoder 每一步产生的 Key $k^i$ （注意不再只是最后一步）进行相关度计算，得到相关度 $a^i=\operatorname{match}(q,k^i)$ ，称为注意力分数。 $\operatorname{match}$ 函数可以使用一个可训练的网络（加性 Attention），而当 $q$ 和 $k_i$ 长度相同的时候，也可以使用内积（余弦相似度）。
得到每一个 $k^i$ 关于 $q$ 的相关度 $a^i$ 之后，将所有的 $a^i$ 进行 softmax 归一化（有时候也可以不做 softmax，直接用 logits，视情况），得到注意力权重 $\hat a^i$ ，并以此为权重加权求和 $v^i$ ，得到一个 summary： $c=\sum_{i} \hat a^i v^i$ .
这个 $c$ 会作为输入，和自回归的 Input 一起进入 Decoder，供模型参考。
之后 Decoder 每次除了生成 Output，还会生成下一次的 Query 进行查询。

可以看到，Attention 机制允许模型学会主动关注一些信息，而过滤另一些。并且允许模型具有更长的上下文（每次加权之后得到的 summary 不会太大，并且信息很精准）。

Attention 机制的加入允许模型自行关注有用的信息。

除了文本翻译，Attention 还可以用在很多领域，比如语音识别、Image/Video Caption 生成。后者的大致做法是：
将 CNN 输出的 feature map 逐像素拉平成长度为 channel 的向量序列，再输入给有 Attention 机制的循环神经网络中，模型就能每次选择图片中要着重关注的区域/对象进行描述。下图为具体样例，白色蒙版对应模型生成下划线的 token 时，Attention 关注的区域。

左侧为效果不错的，右侧为效果欠佳的

公式#

Encoder 每一个时间步产生的 hidden state $h_i$ 会经过两个矩阵 $\boldsymbol{W}_k$ 和 $\boldsymbol{W}_v$ ，得到 $k_i$ 和 $v_i$ .
Decoder 当前的 hidden state 经过一个矩阵 $\boldsymbol{W}_q$ ，得到 $\boldsymbol{q}$ .
将 $\boldsymbol{q}$ 和每个 $\boldsymbol{k}_i$ 进行相似度计算 $a(\boldsymbol{q}, \boldsymbol{k}_i)$ ，其中 $a$ 为余弦相似度或者加性注意力。

采用余弦相似度#

当采用余弦相似度时， $a(\boldsymbol{q},\boldsymbol{k}_i)=\boldsymbol{k}_i\cdot \boldsymbol{q} = \boldsymbol{k}_i^T \boldsymbol{q}$ ，此时不同的 $\boldsymbol{k}_i$ （列向量）可以 concat 成一个矩阵 $\boldsymbol{K}$ ，进行并行计算： $\boldsymbol{a}(\boldsymbol{q}, \boldsymbol{K})=\boldsymbol{K}^T\boldsymbol{q}$ .

随后，每个 $a(\boldsymbol{q}, \boldsymbol{k}_i)$ 经过激活函数/归一化函数（如 softmax），得到注意力权重 $a'(\boldsymbol{q}, \boldsymbol{k}_i)$ 。然后，将 $a'$ 与 $\boldsymbol{v}_i$ 对应加权求和，得到 $\boldsymbol{b}=\sum_i a'(\boldsymbol{q}, \boldsymbol{k}_i) \boldsymbol{v}_i$ .
同样的，这一步也可以变成矩阵乘法进行并行计算，即将 $\boldsymbol{v}_i$ （列向量）concat 成一个矩阵 $\boldsymbol{V}$ ： $\boldsymbol{b}=\boldsymbol{V}\boldsymbol{a}'(\boldsymbol{q}, \boldsymbol{K})$ .

请注意，一般使用的余弦相似度会经过 scale。 公式化地，假设 $\boldsymbol{q}$ 和 $\boldsymbol{k}$ 的维度都是 $d$ ，那么在作用了点乘之后，还会除以 $\sqrt d$ 再得到注意力分数。这样做是为了保证无论维度是多少，注意力分数的方差都恒定。此处的隐含假设是 $\boldsymbol{q}$ 和 $\boldsymbol{k}$ 的每一维都是方差相同且不变的随机变量。这是因为点乘有求和操作，而方差是线性的，因此维度数越多，点乘结果的方差也会线性地增加。此时，为了保证方差归一，要给点乘结果除以 $\sqrt d$ ，这样方差就除以了 $d$ ，保持不变。

采用加性注意力#

所谓加性注意力，就是将 $\boldsymbol{q}$ 和 $\boldsymbol{k}$ concat 起来并输入一个可学习的 MLP，最终输出一个标量作为注意力分数。加性注意力允许 $\boldsymbol{q}$ 和 $\boldsymbol{k}$ 具有不同的长度，并允许模型学习相似度的表方式。该方式也可以进行并行，但稍复杂。

多头注意力#

我们可以允许 Decoder 每个时间步产生不止一个 query，得到不止一个 $\boldsymbol{b}$ ，以获得更加全面的资讯。

具有多头注意力的 Decoder 在每个时间步会产生多个 $\boldsymbol{q}_i$ ，我们可以像上一节那样故技重施，将 $\boldsymbol{q}_i$ concat 成一个矩阵 $\boldsymbol{Q}$ 以并行计算。这样，就会得到 $\boldsymbol{A}(\boldsymbol{Q}, \boldsymbol{K}) = \boldsymbol{K}^T\boldsymbol{Q}$ ，形状为 num_keys $\times$ num_queries，进而得到 $\boldsymbol{A}'(\boldsymbol{Q},\boldsymbol{K})$ .
之后得到 $\boldsymbol{B}=\boldsymbol{V}\boldsymbol{A}'(\boldsymbol{Q}, \boldsymbol{K})$ ，形状为 v_size $\times$ num_queries.

在实践中，我们常常是将一个 qkv_size 的 query key value 先通过大矩阵 $\boldsymbol{W}_q$ $\boldsymbol{W}_k$ $\boldsymbol{W}_v$ ，再将变换后的向量平均切分成多个头（而不是先复制 qkv，再通过多个不同的矩阵 w），每个头的 dim 都会因此变小，然后每个头分别进行各自内部的 qkv 相似度加权求和，得到小的 hidden_size，再全部 concat 成一个大的 hidden_size 作为输出。

另见GPT 多头注意力 QA

Self Attention 自注意力#

自注意力旨在进一步释放并行能力。它与含有 Attention 的 RNN 模型不同，它直接抛弃了用 RNN 来顺序生成 hidden state 这一步。

从第一性原理来看，注意力机制出现之后，RNN 确实不再必要。一方面，RNN 的 hidden state 将上文极度压缩，出现失真，并且随时可能遗忘上文，而且需要使用双向 RNN 才能同时考虑上下文；另一方面，RNN 需要顺序迭代，难以并行。

自注意力的做法是，直接在原始 input 上，将每个 input $\boldsymbol{a}_i$ 通过三个可学习的矩阵 $\boldsymbol{W}_q$ $\boldsymbol{W}_k$ $\boldsymbol{W}_v$ ，得到 $\boldsymbol{q}_i$ $\boldsymbol{k}_i$ $\boldsymbol{v}_i$ 三个向量，然后将 $\boldsymbol{q}_i$ 和整个序列中其他的（有时也包括自己的） $\boldsymbol{k}_j$ 进行前述的相似度计算，再与对应的 $\boldsymbol{v}_j$ 加权求和，得到 $\boldsymbol{b}_i$ .

与之前的注意力机制不同，现在 query 也变成了“分布式的”，query 不再是由一个 Decoder 在生成时发起，而是由每一个 input 发起，得到的则是和自己有关的其他 input 的信息，这也是 “Self” 的含义。

有了 Self Attention，可以做到极佳的并行程度，从而允许扩大模型的上下文空间。

有了 Self Attention，模型没有特别重的记忆负担，因为可以随时去上下文中 query 想要的信息，并且这个 query 是全局性的，无关乎距离远近，可谓是实现了 token 之间的「天涯若比邻」（李宏毅语）。

1d-Conv、RNN、Self Attention 的上下文空间与并行能力对比

位置编码#

由于 Self Attention 是并行计算的，因此位置信息是缺失的，这点和顺序迭代的 RNN 有很大不同。因此，如果我们还想要让模型学习和位置有关的信息（譬如句子结构），就需要手动给输入加上一个位置编码。

Sinusoidal 编码#

Google 太牛逼啦！

这个初看有点复杂，但实际上是仅用一点点复杂度就带来了远超这点代价的巨大好处的极为精妙的充满数学的位置编码方式，结合了进制、高维空间旋转等巧思，一举解决了其他位置编码方式可能存在的数值不稳定、泛化性低、不能表示相对位置关系、可解释性差的问题。

解释起来有点复杂，请参考这些进行阅读：

可学习的位置编码#

Bert 等采用的方案，使用可以学习的编码方式代替 hand-crafted 的位置编码。

在图像像素位置编码这种非一维序列上前景不错。

关于在图像方面 Attention 和 CNN 的对比，可见 B 站视频杀死卷积：Transformer如何暴力接管计算机视觉 ↗.