《动手学深度学习》笔记——注意力机制、Transformer

1 注意力机制

动物需要在复杂环境下有效关注值得注意的点
心理学框架：人类根据随意线索和不随意线索选择注意点

1.1 基本概念

卷积、全连接、池化层都只考虑不随意线索。注意力机制（Attention Mechanism）则显示地考虑随意线索/意志线索（Volitional Cue），常称为查询（Query）。每个输入是一个值（Value）和不随意线索/键（Key）的对(x_i,y_i)。

通过注意力池化层（Attention Pooling）来对查询 x有偏向性地选择输入f(x)，可以表示为

f(x)=\sum\limits_i \alpha(x,x_i)y_i

其中\alpha(x,x_i)为注意力权重。

1.2 非参数注意力池化

给定询问x以及键值对数据(x_i,y_i),\ i=1,2,\cdots,n。平均池化（Average Pooling）是前文解决回归问题时常用的最简单的估计器，即

f(x)=\frac1n\sum\limits_{i=1}^n y_i

更好的方案是60年代提出的Nadaraya-Watson核回归（Nadaraya-Watson Kernel Regression），如下所示

f(x)=\sum\limits_{i=1}^n \frac{K(x-x_i)}{\sum\limits_{j=1}^n K(x-x_j)}y_i

其中K(u)为高斯核（Gaussian kernel），定义为

K(u)=\frac1{\sqrt{2\pi}}\exp(-\frac{u^2}{2})

则非参数的Nadaraya-Watson核回归可化简为

f(x)=\sum\limits_{i=1}^n \frac{\exp(-\frac12(x-x_i)^2)}{\sum\limits_{j=1}^n \exp(-\frac12(x-x_j)^2)}y_i=\sum\limits_{i=1}^n \text{softmax}(-\frac12(x-x_i)^2)y_i

1.3 带参数注意力池化

非参数的Nadaraya-Watson核回归具有一致性（Consistency）的优点： 如果有足够的数据，此模型会收敛到最优结果。在此基础上引入可以学习的w，如下所示

f(x)=\sum\limits_{i=1}^n \text{softmax}(-\frac12((x-x_i)w)^2)y_i

---

2 注意力评分函数

上述高斯核指数部分可以视为注意力评分函数（Attention Scoring Function），简称评分函数，是询问和键的相似度。注意力权重是分数的softmax结果。

2.1 高维度注意力池化

设询问\mathbf{q} \in \mathbb{R}^q，键值对(\mathbf{k}_1,\mathbf{v}_1),\cdots,(\mathbf{k}_m,\mathbf{v}_m)，其中\mathbf{k}_i\in \mathbb{R}^k,\mathbf{v}_i\in \mathbb{R}^v，则注意力池化函数f表示为

f(\mathbf{q},(\mathbf{k}_1,\mathbf{v}_1),\cdots,(\mathbf{k}_m,\mathbf{v}_m))=\sum\limits_{i=1}^m \alpha(\mathbf{q},\mathbf{k}_i)\mathbf{v}_i \in \mathbb{R}^v \\
\alpha(\mathbf{q},\mathbf{k}_i)=\text{softmax}(a(\mathbf{q},\mathbf{k}_i))=\frac{\exp(a(\mathbf{q},\mathbf{k}_i))}{\sum\limits_{j=1}^{m}\exp(a(\mathbf{q},\mathbf{k}_j))} \in \mathbb{R}

其中a(\mathbf{q},\mathbf{k}_i)为评分函数。

2.2 加性注意力

当询问和键是不同长度的矢量时，可以使用加性注意力（Additive Attention）作为评分函数。设询问\mathbf{q} \in \mathbb{R}^q和键\mathbf{k}_i\in \mathbb{R}^k，可学习的参数为\mathbf{w}_v\in \mathbb{R}^h,\mathbf{W}_k \in \mathbb{R}^{h\times k},\mathbf{W}_q \in \mathbb{R}^{h\times q}，则加性注意力的评分函数为

a(\mathbf{q},\mathbf{k}_i)=\mathbf{w}_v^\text{T}\tanh(\mathbf{W}_q\mathbf{q}+\mathbf{W}_k\mathbf{k}_i) \in \mathbb{R}

等价于将键与值合并放入一个隐藏大小为h、输出大小为1的单隐藏层MLP。

2.3 缩放点积注意力

若询问和键都是同样的长度d，则可使用缩放点积注意力（Scaled Dot-product Attention）。设\mathbf{q},\mathbf{k}_i \in \mathbb{R}^d，则评分函数为

a(\mathbf{q},\mathbf{k}_i)=\frac{\left \langle \mathbf{q},\mathbf{k}_i \right \rangle }{\sqrt{d}}

通常从小批量的角度来考虑提高效率：设有n个查询和m个键值对，其中查询和键的长度为d，值的长度为v，则查询\mathbf{Q}\in \mathbb{R}^{n\times d}、键\mathbf{K}\in \mathbb{R}^{m\times d}和值\mathbf{V}\in \mathbb{R}^{m\times v}的缩放点积注意力f为

f=\text{softmax}(a(\mathbf{Q},\mathbf{K}))\mathbf{V} \in \mathbb{R}^{n\times v} \\
a(\mathbf{Q},\mathbf{K})=\frac{\mathbf{Q}\mathbf{K}^\text{T}}{\sqrt{d}} \in \mathbb{R}^{n\times m}

可见此处直接将询问和键作内积。

---

3 自注意力

当查询、键和值来自同一组输入时称为自注意力（Self-attention）或内部注意力（Intra-attention）。

3.1 基本概念

给定序列\mathbf{x}_1,\cdots,\mathbf{x}_n,\forall \mathbf{x}_i \in \mathbb{R}^d，自注意力池化层将\mathbf{x}_i当做键、值、询问来对序列抽取特征得到\mathbf{y}_1,\cdots,\mathbf{y}_n，其中

\mathbf{y}_i=f(\mathbf{x}_i,(\mathbf{x}_1,\mathbf{x}_1),\cdots,(\mathbf{x}_n,\mathbf{x}_n)) \in \mathbb{R}^d

自注意力完全并行，最长序列为1，但对长序列计算复杂度高

3.2 位置编码

跟CNN/RNN不同，自注意力并没有记录位置信息。
位置编码（Positional Encoding）在输入表示中加入位置信息，使得自注意力能够记忆位置信息。

设输入\mathbf{X}\in\mathbb{R}^{n\times d}为长度n的序列，则使用位置编码矩阵\mathbf{P}\in\mathbb{R}^{n\times d}来输出\mathbf{X}+\mathbf{P}作为自编码输入，其中\mathbf{P}的元素如下计算

p_{i,2j}=\sin(\frac{i}{10000^{\frac{2j}{d}}}),\ p_{i,2j+1}=\cos(\frac{i}{10000^{\frac{2j}d}})

3.3 绝对、相对位置信息

观察下图中的二进制编码可知，每个数字、每两个数字和每四个数字上的比特值在第一个最低位、第二个最低位和第三个最低位上分别交替。

在二进制表示中，较高比特位的交替频率低于较低比特位， 与下面的热图所示相似，只是位置编码通过使用三角函数在编码维度上降低频率。 由于输出是浮点数，因此此类连续表示比二进制表示法更节省空间。

对于任何确定的位置偏移\delta，位置i+\delta处的位置编码可以线性投影位置i处的位置编码来表示。令\omega_i=\frac{1}{10000^{\frac{2j}{d}}}，则

\begin{bmatrix}
 \cos(\delta \omega_j) & \sin(\delta\omega_j) \\
 -\sin(\delta \omega_j)  & \cos(\delta \omega_j)
\end{bmatrix}\begin{bmatrix}
 p_{i,2j}  \\
  p_{i,2j+1}
\end{bmatrix}=\begin{bmatrix}
p_{i+\delta,2j} \\
p_{i+\delta,2j+1}
\end{bmatrix}

易知左边的投影矩阵与i无关。

---

4 Transformer

原文：Attention Is All You Need

Transformer模型完全基于注意力机制，没有任何卷积层或循环神经网络层，基于编码器-解码器架构来处理序列对。

4.1 总体架构

Transformer的编码器和解码器是基于自注意力的模块叠加而成的，源（输入）序列和目标（输出）序列的嵌入（Embedding）表示将加上位置编码，再分别输入到编码器和解码器中。

所有子层之间使用残差链接。

4.2 多头自注意力

多头自注意力（Multi-head Self-attention）对同一键、值、询问，希望抽取不同的信息，例如短距离关系和长距离关系。使用h个独立的注意力池化，合并各个头（Head）的输出得到最终输出。

设询问\mathbf{q}\in \mathbb{R}^{d_q}，键\mathbf{k}\in \mathbb{R}^{d_k}，值\mathbf{v}\in \mathbb{R}^{d_v}，头i的可学习参数\mathbf{W}_i^{(q)}\in \mathbb{R}^{p_q\times d_q},\mathbf{W}_i^{(k)}\in \mathbb{R}^{p_k\times d_k},\mathbf{W}_i^{(v)}\in \mathbb{R}^{p_v\times d_v}，则头i的输出为

\mathbf{h}_i=f(\mathbf{W}_i^{(q)}\mathbf{q},(\mathbf{W}_i^{(k)}\mathbf{k},(\mathbf{W}_i^{(v)}\mathbf{v}) \in \mathbb{R}^{p_v}

输出的可学习参数\mathbf{W}_o\in\mathbb{R}^{p_o\times hp_v}，则多头注意力的输出为

\mathbf{W}_o\begin{bmatrix}
\mathbf{h}_1 \\
\vdots \\
\mathbf{h}_h
\end{bmatrix} \in \mathbb{R}^{p_o}

有掩码的多头注意力：解码器对序列中一个元素输出时，不应该考虑该元素之后的元素，因此可通过掩码来实现。即计算\mathbf{x}_i输出时，假装当前序列长度为i。

4.3 基于位置的前馈网络

基于位置的前馈网络（Positionwise Feed-forward Network）将输入形状由(b,n,d)变换成(bn,d)，输出形状由(bn,d)变化回(b,n,d)。等价于两层1\times 1卷积层，即全连接层。

4.4 层规范化

批量规范化（详见CNN 2.5）对每个特征/通道（列）里元素进行归一化，不适合序列长度会变的NLP应用。

层规范化（Layer Normalization）对每个样本（行）里的元素进行归一化。

4.5 信息传递与输出预测

编码器到解码器的信息传递：编码器中的输出\mathbf{y}_1,\cdots,\mathbf{y}_n作为解码中第i个Transformer块中多头注意力的键和值，其询问来自目标序列。这意味着编码器和解码器中块的个数和输出维度都是一样的。

预测：预测第t+1个输出时，解码器中输入前t个预测值，在自注意力中作为键和值，其第 $t$ 个还同时作为询问。