Transformer 系列一：Self-Attention Mechanism 自注意力机制

我在网上看到大多数文章讲transformer都会从 Attention Is All You Need 论文中那张经典的transformer架构图讲起。但是作为刚接触transformer的我，那时看到这张图其实是一头雾水，完全抓不住重点和细节。所以，我打算从transformer结构中最为重要的自注意力机制（Self-Attention Mechanism）讲起，先将transformer结构中各个重要的模块拆解，由小及大最终再将他们拼到一起形成整体的transformer架构。

什么是自注意力机制

注意力机制

先说注意力机制（Attention Mechanism）。注意力机制的灵感来源于人类视觉和认知系统。想象一下，当你在阅读一篇文章时，你的眼睛和大脑会自然地聚焦在那些最重要的部分，比如标题、关键段落或图表。这个过程就像是你的大脑在分配“注意力”，以便快速抓住文章的要点。这是因为你的潜意识认为这些部分包含的信息最为丰富和直接。注意力机制正是模仿了这种选择性关注，它通过集中关注信息的关键部分来提取出更加重要的内容。

自注意力机制

自注意力机制（Self-Attention Mechanism）是注意力机制的一种特殊形式，其独特之处在于"自"（self）这个概念。自注意力机制通常应用于处理序列数据的任务，比如音频处理（Audio）或自然语言处理（NLP）。序列数据的一个显著特点是，数据中的每个元素都可能与序列中的其他元素存在联系（即上下文关系），这与图像数据中某个像素点只与其邻近像素点有主要联系的情况不同。自注意力机制通过评估元素之间的相互关系和重要性，能够自适应地捕捉它们之间的依赖关系。

这里举一个例子来说明在自然语言处理类任务中，上下文关系对模型性能好坏影响的重要性。也可以从侧面体现出自注意力机制这种对数据内部元素之间的长程依赖关系有较强捕捉能力的方法被广泛用于自然语言处理类任务的原因。词性标注任务（Part-Of-Speech tagging，POS tagging）是NLP领域一个经典的基础任务。机器需要自动决定一句话中每一个词汇的词性，判断该词是名词还是动词还是形容词等等。现在有一个句子：

I saw a saw.

这句话的中文意思是“我看到一个锯子”，第二个 saw 是名词锯子。所以机器要知道，第一个 saw 是个动词，第二个 saw 是名词。如果我们只用一个简单的全连接网络（FC Network， Fully Connected Network）来预测每个词的词性，那由于输入都是saw，输出一定也是一样的。但实际上由于这两个单词所处的位置及其上下文所代表的关系不同，所期望的词性结果也不同。

此时，我们在FC层前引入一层自注意力机制模块，让其对输入的各个元素先进行一次处理。这里由于引入了自注意力机制，所以第一个saw和第二个saw所生成的向量就因为其所处位置不同以及所对应的上下文关系不同而产生了区别。那么这样送入FC层后所得到关于saw词性结果就可能不同了。这种机制使得模型不仅能够理解单个单词的含义，还能够理解单词在特定上下文中的含义，这对于准确进行词性标注至关重要。

自注意力机制是如何运转的

最初的token在计算词向量的时候并没有和其他token建立联系，所以自注意力层的主要作用是通过挖掘当前输入数据内部的上下文关系，来重建每个token的特征，使其包含更多的语义联系及上下文关系的信息。

输入输出

现在有一句话“我爱打网球”，将其转换成词向量的形式送入自注意力层：

• vocab_size（S）：词向量维度，即每个词（token）的特征维度，这里假设是768
• sentence_length（L）：句子长度，即句子中有几个token，这里是5
• 所以输入的一个维度为（L, S）=（sentence_length, vocab_size）= （5，768）的矩阵
• 先不用管中间过程，将输入送入自注意力层后，输出的维度是（L, M）=（5， vector_size）
• vector_size（M）：期望输出的新的特征维度，可以和原来一样也可以不一样

内部运转机制

在了解了输入输出后，我们终于可以进入到自注意力层内部一探究竟了。

这里假设输入的词向量矩阵为$I$，第一步会计算出三个新的矩阵：$Q、K、V$：
$$\begin{gathered} Q=I\cdot W^q \\ K=I\cdot W^k \\ V=I\cdot W^v \end{gathered}$$

这里用到的$W^q、W^k、W^v$都是可学习的参数矩阵，他们的维度是（S， M）=（vocab_size, vector_size）

所以最终的得到的$Q、K、V$的维度是（L，M）=（sentence_length, vector_size）

input_seq_len = 5 # 输入句子的token数(L)
input_d_model_size = 6 # 输入向量维度(S)
d_model_size = 4 # QKV向量维度(M)

# 创建可学习的Wq,Wk,Wv
# 线性层不添加bias即w·x和矩阵相乘是一致的，并且线性层里的所有参数都是可学习的，所以这里直接使用线性层模拟矩阵乘法
w_q = nn.Linear(input_d_model_size, d_model_size, bias=False) # Wq
# Linear(in_features=6, out_features=4, bias=False)
w_k = nn.Linear(input_d_model_size, d_model_size, bias=False) # Wk
# Linear(in_features=6, out_features=4, bias=False)
w_v = nn.Linear(input_d_model_size, d_model_size, bias=False) # Wv
# Linear(in_features=6, out_features=4, bias=False)

# 输入 (input_seq_len, input_d_model_size)
input_x = torch.randn(input_seq_len,input_d_model_size)

print("input_x shape:", input_x.shape)
# input_x shape: torch.Size([5, 6])

Q = w_q(input_x)
K = w_k(input_x)
V = w_v(input_x)

print("Q shape:", Q.shape)
print("K shape:", K.shape)
print("V shape:", V.shape)
# Q shape: torch.Size([5, 4])
# K shape: torch.Size([5, 4])
# V shape: torch.Size([5, 4])

我们这里直接看计算公式：

这个公式总共可以被拆成四步：

第一步： $Q\cdot K$

• 一般来说，将两个向量点乘是为了计算两个向量之间的相似度。那这里的操作其实是在计算每个token和其他token的相互关系权重。在Attention中计算两个元素的相互关系或者相似度，是不会使用$I(Input)$来直接计算的，而是使用 $Q(Query)、K(Key)$来计算，这样做的好处是使用可学习的矩阵参数可以增强模型的拟合和学习能力。
• 这里的计算会得到一个$（L，L）$的矩阵，每一个slot代表了目标token对当前token的重要性或者表示两者之间的相互关系。若这个slot的值大，则在后面加权计算的时候目标token就会对当前token的信息影响权重大。

第二步：Scale $\frac{1}{\sqrt{d_k}}$

• $d_k$就是我们的vector_size，这里进行一个缩放的操作是为了防止softmax内的数值过大，从而导致其偏导数趋近于0，除以该值可以保证在训练的时候梯度稳定回传;
• 并且可以使得$Q\cdot K$的结果满足期望为0，方差为1的分布，类似于一个归一化的操作
• 拓展阅读：巴比龙：Self-attention中dot-product操作为什么要被缩放

第三步：$Softmax(\cdot )$

• 总的来说， $Softmax(\cdot )$的作用是为了保证注意力权重的非负性，同时增加非线性
• 拓展阅读：PaperWeekly：线性Attention的探索：Attention必须有个Softmax吗？

第四步： $\cdot V$

• $V(Value)$就是从token中提取出的信息，这是通过$W_v$矩阵来进行提取的。这个过程是一个加权求和的过程，通过将 $Q\cdot K$计算出来的权重叠加到$V$上，使每个token的特征不仅包含自己本身的信息还包含其余token的信息，而包含信息的多少就是两个token通过 $Q\cdot K$计算出来的权重。通过这样的操作，就将上下文的信息内容引入到了每个token的特征向量中，达到了我们期望的目的。

# Q · K^T
attention_scores = torch.matmul(Q, K.transpose(-1, -2))

# 1/sqrt(d_k)
attention_scores = attention_scores / math.sqrt(d_model_size)

print("attention_scores shape:", attention_scores.shape)
# attention_scores shape: torch.Size([5, 5])

# Softmax(·)
attention_probs = nn.Softmax(dim=-1)(attention_scores)

print("attention_probs shape:", attention_probs.shape)
# attention_probs shape: torch.Size([5, 5])

# ·V
out = torch.matmul(attention_probs, V)

print("out shape:", out.shape)
# out shape: torch.Size([5, 4])

多头注意力

为了增强拟合性能，Transformer对Attention继续扩展，提出了多头注意力（Multiple Head Attention）。对于同样的输入 $I$ ，我们定义多组不同的，比如$W^q、W^k、W^v$，每组分别计算生成不同的$Q、K、V$，最后学习到不同的参数。这有点类似于卷积神经网络的多核卷积。举个实际的例子，比如头1提取的是语义关系，头2提取的则是语法关系等。

这里举个例子来证明，不同头（head）得到的token之间的关系是不同的，也就是他们之间的权重是不同的： 最后将所有头的得到的结果进行拼接和线性映射，得到最终的输出：

---

参考文献及推荐阅读：

• Arxiv - Attention is all you need
• GitHub - datawhalechina/leedl-tutorial: 《李宏毅深度学习教程》（李宏毅老师推荐👍，苹果书🍎）
• Blog - The Annotated Transformer
• Blog - The Illustrated Transformer
• 深度学习attention机制中的Q,K,V分别是从哪来的？