NLP（五）：Transformer及其attention机制
输入“/”快速插入内容
NLP（五）：Transformer及其attention机制
NLP（五）：Transformer及其attention机制
2024年8月20日修改
作者：紫气东来
Transformer 模型自 2017 年 6 月在论文《Attention Is All You Need》被提出以来，已经成为 NLP 领域中的首选模型。Transformer 抛弃了 RNN 的顺序结构，采用了 Self-Attention 机制，使得模型可以并行化训练，而且能够充分利用训练资料的全局信息，加入 Transformer 的 Seq2seq 模型在 NLP 的各个任务上都有了显著的提升。本文将试图从多角度更加清晰地讲解 Transformer 的运行原理。
Transformer 由且仅由 self-Attention 和 Feed Forward Neural Network 组成。Transformer中包括了编码器和解码器各 6 层，总共 12 层的 Encoder-Decoder。Transformer 中的核心机制就是 Self-Attention。Self-Attention 机制的本质来自于人类视觉注意力机制。
attention详解
输入是一个序列， x1,x2...，xn x^{1},x^{2}...，x^{n} ，进行embedding之后形成 a1,a2...，a4 a^{1},a^{2}...，a^{4} ；然后每个 xi x^{i} 都会经过三个矩阵处理（ WQ W^{Q} ， WK W^{K} 和 WV W^{V} ），得到三个向量 Q,V Q,V 和 K K 。
Q:query(tomatchothers)Qi=WQai Q : query (to ~match ~others) ~~~~~~~Q^{i}=W^{Q} a^{i}
K:key(tobematched)Ki=WKai K: key (to ~be ~matched)~~~~~~~K^{i}=W^{K} a^{i}
V:value(informationtobeextracted)Vi=WVai V: value(information ~to ~be ~extracted)~~~~~~V^{i}=W^{V} a^{i}
Attention最核心的公式如下，下面将逐步剖析这一公示的含义
Attention(Q,K,V)=softmax(QKTdk)V \operatorname{Attention}(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V
QKT QK^T 的含义
我们知道，对于向量 xi x_i 和 xj x_j ，其内积 <xi,xj> <x_i,x_j> 表征一个向量在另一个向量上的投影，即两个向量的相似关系。 那么对于矩阵 X X ， XXT XX^T 是一个方阵，以行向量的角度理解，里面保存了每个向量与自己和其他向量进行内积运算的结果，即与其他向量的相似程度。 使用query和每个key进行计算，得到就是相应的attention。
\sqrt{d_k} 与softmax
d_k 是 q 和 k 的维度，主要是因为随着维度增加，点乘的结果也会变大，结果太大也会对softmax的结果产生影响，所以需要除以一个系数。
下面简要予以证明：
假设向量 q 和 k 的各个分量是互相独立的随机变量，维度为 d_k ，均值是0，方差是1，即
\begin{aligned} &\mathrm{E}\left[q_{i}\right]=\mathrm{E}\left[k_{i}\right]=0 \\ &\operatorname{var}\left[q_{i}\right]=\operatorname{var}\left[k_{i}\right]=1 \end{aligned} 其中 i \in\left[0, d_{k}\right]
则有
\begin{aligned} \mathrm{E}[q \cdot k] &=\mathrm{E}\left[\sum_{i=1}^{d_{k}} q_{i} k_{i}\right] \\ &=\sum_{i=1}^{d_{k}} \mathrm{E}\left[q_{i} k_{i}\right] \\ &=\sum_{i=1}^{d_{k}} \mathrm{E}\left[q_{i}\right] \mathrm{E}\left[k_{i}\right] \\ &=0 \end{aligned}
\begin{aligned} \operatorname{var}[q \cdot k] &=\operatorname{var}\left[\sum_{i=1}^{d_{k}} q_{i} k_{i}\right] \\ &=\sum_{i=1}^{d_{k}} \operatorname{var}\left[q_{i} k_{i}\right] \\ &=\sum_{i=1}^{d_{k}} \operatorname{var}\left[q_{i}\right] \operatorname{var}\left[k_{i}\right] \\ &=\sum_{i=1}^{d_{k}} 1 \\ &=d_{k} \end{aligned}
则 var\left(\frac{q \cdot k}{\sqrt{d}_{k}}\right)=\frac{d_{k}}{\left(\sqrt{d}_{k}\right)^{2}}=1
将方差控制为1，也就有效地控制了softmax反向的梯度消失的问题 。
接着把这些分数经过一个Softmax函数，Softmax可以将分数归一化，这样使得分数都是正数并且加起来等于1， 这些分数决定了当前词向量，对其他所有位置的词向量分别有多少的注意力。
\operatorname{Softmax}\left(z_{i}\right)=\frac{e^{z_{i}}}{\sum_{c=1}^{C} e^{z_{c}}}
加权求和
得到每个词向量的分数后，将分数分别与对应的Value向量相乘。这种做法背后的直觉理解就是：对于分数高的位置，相乘后的值就越大，我们把更多的注意力放到了它们身上；对于分数低的位置，相乘后的值就越小，这些位置的词可能是相关性不大的。
多头模型
由于不同的 Attention 的权重侧重点不一样，所以将这个任务交给不同的 Attention 一起做，最后取综合结果会更好，有点像 CNN 中的 Kernel。在Transformer的论文中指出，将 Q、K、V 通过一个线性映射后，分成 h 份，对每份进行 Scaled Dot-Product Attention 效果更好， 再把这几个部分 Concat 起来，过一个线性层的效果更好， 可以综合不同位置的不同表征子空间的信息 。
即
O=\operatorname{MultiHead}(Q, K, V)=W^{O} \cdot \text { Concat }\left(\begin{array}{l} \operatorname{softmax}\left(\frac{Q^{1}\left(K^{1}\right)^{\top} }{\sqrt{d_{m}}}\right) V^{1}\\ \operatorname{softmax}\left(\frac{ Q^{2}\left(K^{2}\right)^{\top} }{\sqrt{d_{m}}}\right)V^{2} \end{array}\right)
多头attention模型的实现：
代码块
class MultiHead(nn.Module):
def __init__(self, n_head, model_dim, drop_rate):
super().__init__()
self.head_dim = model_dim // n_head
self.n_head = n_head
self.model_dim = model_dim
self.wq = nn.Linear(model_dim, n_head * self.head_dim)
self.wk = nn.Linear(model_dim, n_head * self.head_dim)
self.wv = nn.Linear(model_dim, n_head * self.head_dim)
self.o_dense = nn.Linear(model_dim, model_dim)
self.o_drop = nn.Dropout(drop_rate)
self.layer_norm = nn.LayerNorm(model_dim)
self.attention = None
def forward(self, q, k, v, mask, training):
# residual connect
residual = q
# linear projection
key = self.wk(k) # [n, step, num_heads * head_dim]
value = self.wv(v) # [n, step, num_heads * head_dim]
query = self.wq(q) # [n, step, num_heads * head_dim]
# split by head
query = self.split_heads(query) # [n, n_head, q_step, h_dim]
key = self.split_heads(key)
value = self.split_heads(value) # [n, h, step, h_dim]
context = self.scaled_dot_product_attention(query, key, value, mask) # [n, q_step, h*dv]
o = self.o_dense(context) # [n, step, dim]
o = self.o_drop(o)
o = self.layer_norm(residual + o)
return o
def split_heads(self, x):
x = torch.reshape(x, (x.shape[0], x.shape[1], self.n_head, self.head_dim))
return x.permute(0, 2, 1, 3)
def scaled_dot_product_attention(self, q, k, v, mask=None):
dk = torch.tensor(k.shape[-1]).type(torch.float)
score = torch.matmul(q, k.permute(0, 1, 3, 2)) / (torch.sqrt(dk) + 1e-8) # [n, n_head, step, step]
if mask is not None:
# change the value at masked position to negative infinity,
# so the attention score at these positions after softmax will close to 0.
score = score.masked_fill_(mask, -np.inf)
self.attention = softmax(score, dim=-1)
context = torch.matmul(self.attention, v) # [n, num_head, step, head_dim]
context = context.permute(0, 2, 1, 3) # [n, step, num_head, head_dim]
context = context.reshape((context.shape[0], context.shape[1], -1))
return context # [n, step, model_dim]