大模型开发基础(一）自注意力机制和多头注意力机制

玉木子

1376人浏览 · 2025-08-14 09:51:44

玉木子 · 2025-08-14 09:51:44 发布

一、自注意力机制

1 简单介绍

自注意力机制(Self-Attention Mechanism)，用于捕捉序列数据内部依赖关系的关键技术,在NLP和CV中非常重要，是transformer重要组成部分。

自注意力机制的产生因为RNN、LSTM在处理长序列时有梯度小时、计算效率低等局限。

为了解决序列建模中各种问题，引入了自注意力机制。

2 基本概念

自注意力机制，就是将单个序列关联起来以计算同一序列的表示的注意机制，也就是找大自己和所在句子的所有词之间的关联关系。注意力机制与自身还有关注对象都哟i关系，并且只关注输入本身或者只关注对象本身。

自注意力的目标就是让每个位置的表示能够根据整个序列中其他位置的信息进行加权融合，从而捕获实体之间的相互关系。

3 QKV

自注意力机制就是引入查询向量（Query）、键向量（key）、值向量（Value）概念来实现序列中个元素之间的信息交互和依赖建模。

Q：也就是Query，表示当前查询者的位置；

K：Key，表示被查询者的身份，也就是所有位置给出的标签；

V：Value，表示被查询者实际信息。

序列中的每个Token都有Q、K、V。

举一个例子，我们在书店里，想找到《人生海海》这本书。我们的查询，Q就是“人生海海”。工作人员就是注意力机制，根据“人生海海”这个查询找到有这四个字的所有书籍。

而找到的书都有书名，这就相当于我们的键，也就是K。然后对找到的书和查询“人数海海”比较，这时候书的封面、作者、出版日期、故事梗概等等，这就是所谓的值，也就是V。

等到我们找到了与Q完全匹配的K，就会返回V，也就是《人生海海》这本书。

4 实现过程

自注意力机制通过计算输入数据中每个位置与其他位置的相关性，来调整每个位置的信息表示，依赖于三个核心概念：查询向量Query、键向量Key、值向量Value，对输入的X进行三次线性变换，得到三个矩阵。

些百分比便是归一化之后的注意力得分，根据得分计算出查询向量和其它向量的关联关系。

4.1 查询向量

Query(Q)是自注意机制的询问者，每个输入都会生成一个查询向量，表示当前词的需求。

用于与键向量计算相似度，也就是通过点积方式，确定当前词与其他词的相关性。

生成方式：通过一个权重矩阵将输入数据，如词向量映射到查询空间。

$Q=XWq$

Wq 是可学习权重矩阵，维度为 d×dk，dk是超参数，表示查询向量的维度。

Wq=torch.randn(512,512)#随机生成自注意力机制的Query映射的权重矩阵
Q=torch.matmul(embedding_out,Wq)

torch.size=(7,512)

4.2 键向量

Key，表示其他词的信息，供查询向量匹配。每个输入都会生成一个键向量，表示其能够提供的信息内容。

与查询向量计算点积，生成注意力权重。点积越大，表示它们的相关性越强。

生成方式：通过一个权重矩阵将输入数据，如词向量映射到空间。

$K=XWk$

Wk是可学习权重矩阵，维度为d*dk，dk是超参数，表示键向量的维度。

Wk=torch.randn(512,512)
K=torch.matmul(embedding_out,Wk)

torch.size=(7,512)

4.3 值向量

Value，也就是V，值向量包含了每个输入实际的信息内容，相关性决定了信息被聚焦的程度。

使用值向量基于注意力得分进行加权求和，生成最终的输出表示。

生成方式：通过一个权重矩阵将输入数据如词向量，映射到值空间。

$V=XWv$

Wv是可学习权重矩阵，维度为d*dv，dv是超参数，表示键向量的维度。

Wv=torch.randn(512,512)
V=torch.matmul(embedding_out,Wv)

tips:一般dv和dk通常相等。

4.4 注意力得分

使用点积来计算查询向量和键向量之间的相似度，除以缩放因子 $\sqrt{dk}$ 来便面数值过大，使得梯度更新稳定更新。

补充点积：

对于在平面上的两个向量，其点积被定义为它们模长与夹角余弦值的乘积，也可以定义为它们对应坐标乘积之和。这两种方式得到的结果是相同的，即：

u=(x1,x2),v=(y1,y2)

u·v=|u|·|v|·cos<u,v>

u·v=x1y1+x2y2

注意力得分矩阵： $Attention(Q,K)=\frac{QK_{}^{T}}{\sqrt{dk}}$

sco=torch.matmul(Q,K.transpose(0,1))/math.sqrt(512)

解读：假设Q、K的初始形状是：

Q: [seq_len, d_model]（例如 [10, 512]）

K: [seq_len, d_model]（例如 [10, 512]）

我们需要得到每个查询与所有键的点积，也就是[seq_len,seq_len]。

但 K 的原始形状是 [seq_len, d_model]，需要先转置为 [d_model, seq_len]。torch.matmul(Q, K^T): 输出形状 [seq_len, seq_len]，即注意力分数矩阵。

tips:注意力得分矩阵维度是 n×n，其中 nn 是序列的长度。每个元素 (i,j)(i, j) 表示第 ii 个元素与第 jj 个元素之间的相似度。

比如：假设“注意力机制”生成注意力得分矩阵维度5*5，（2，1）就表示意和注的相似度，（2，5）就表示意和制的相似度。

4.5 归一化

将注意力得分转换为概率分布，需按行对得分矩阵进行softmax操作，确保每行的和为1，得到的矩阵表示每个元素对其他元素的注意力权重。

$Attention Weight=softmax(\frac{QK^{_{}^{T}}}{\sqrt{dk}})$

$\hat{\alpha}_{1,i} = \frac{\exp(\alpha_{1,i})}{\sum_j \exp(\alpha_{1,j})}$ ( $\alpha_{1,i}$ :第一个词语和第i个词语之间的原始注意力得分； $\alpha_{1,i}$ ：经过归一化后的注意力得分）

    normalized_scores = F.softmax(scores, dim=1)

4.6 加权求和

通过将注意力权重矩阵与值矩阵V相乘，得到的加权值表示。

$output=Attention Weight * V=softmax(\tfrac{QK^{_{}^{T}}}{\sqrt{dk}})$

Q和K计算相似度后，经 softmax 得到注意力，再乘V，最后相加得到包含注意力的输出。

attention_result = torch.matmul(normalized_scores, V)

4.7 输出

最终就是维度n*dv的新矩阵，其中每个元素的表示都被加权了。

就是将Q和K分别计算相似度，经过softmax得到相似性概率权重即注意力，再乘以V，最后相加就是包含注意力的输出。

二、多头注意力机制

1 基本概念

将注意力机制中的Q、K、V分成多个头，每个头计算出独立的注意力结果，将所有的头的输出拼接起来，最终通过一个线性变换得到最终输出。

通常先把词向量映射为Q、K、V，再分成多个头。

2 映射权重

分头的过程通过权重矩阵映射实现，并不是直接切分。

dim=256
head_num=8
head_dim=dim//head_num
#以随机生成多头注意力机制的Q映射的权重矩阵
W_Q_list=torch.stack([torch.randn(512,d_k) for _ in range(head_num)])
#映射每个头的Q
Q_list=torch.stack([torch.matmul(embedding_out,W_Q_list[i] for i in range(hrad_num)])

3 加权求和

每个头都是独立计算的，使用自己的一套参数，得到每个头的输出。

$Oh=AhVh$ ，其中Oh是第h个头的输出。

4 输出拼接

将所有头的输出进行拼接：

$O_{\text{concat}} = [O_1, O_2, \dots, O_h] \in \mathbb{R}^{n \times h \cdot d_v}$

其中， $O_{\text{concat}}$ 是所有头拼接的结果，维度是 $n \times (h \cdot d_v)$ ，其中 hh 是头的数量，dvd_v 是每个头的值向量的维度。

5 线性变换

拼接后通过一个线性变换矩阵 $W^O$ 映射为最终输出：

$\text{Output} = O_{\text{concat}} W^O$

其中， $W^O \in \mathbb{R}^{(h \cdot d_v) \times d}$ 是可训练的权重矩阵，dd 是最终输出的维度。

import math
import torch
import torch.nn as nn
from torch.nn import functional as F
document='[un I don not like math I don not like math] '#语料库
text='I don not like math '
vocab={word : i for i,word in enumerate(text.split())}
vocab_len=len(vocab)
#构建语句的词向量
embedding=nn.Embedding(vocab_len,512)
text_embedding = embedding(torch.tensor([vocab[word] for word in text.split()]))
#进行多头操作
dim=512
head_num=8 #多头注意力的头数
head_dim=dim//head_num#每个头的维度

#第一步映射QKV
fc = nn.ModuleList([nn.Linear(dim, dim) for _ in range(3)])#3个
Q=fc[0](text_embedding)
K=fc[1](text_embedding)
V=fc[2](text_embedding)
#2、分成多个头#3*8个神经网络
#QKV的初始形状是[batch_size,seq_len,dim=512]
'''对每个头，需要将 dim=512 的输入映射到 head_dim=64 的子空间。
通过 nn.Linear(dim, head_dim)，将 512 维输入降维到 64 维（每个头的维度）。
8 个线性层会分别生成 8 个 [batch_size, seq_len, head_dim=64] 的矩阵，代表 8 个头的查询（Query）
'''
multi_head_Q_fc = nn.ModuleList([nn.Linear(dim, head_dim) for _ in range(head_num)])  # 修正：方括号
multi_head_K_fc = nn.ModuleList([nn.Linear(dim, head_dim) for _ in range(head_num)])
multi_head_V_fc = nn.ModuleList([nn.Linear(dim, head_dim) for _ in range(head_num)])
#3、shape:8 5 64

multi_head_Q = torch.stack([multi_head_Q_fc[i](Q) for i in range(head_num)])  # 修正：闭合括号
multi_head_K = torch.stack([multi_head_K_fc[i](K) for i in range(head_num)])
multi_head_V = torch.stack([multi_head_V_fc[i](V) for i in range(head_num)])

#注意力得分 8*5*5
sco_list = []
for i in range(head_num):
    score = torch.matmul(multi_head_Q[i], multi_head_K[i].transpose(-2, -1)) / math.sqrt(head_dim)
    sco_list.append(score)
#注意力拼接
sco=torch.stack(sco_list,dim=0)
sco=F.softmax(sco,dim=-1)
#计算每个头的输出
output_list = [torch.matmul(sco[i], multi_head_V[i]) for i in range(head_num)]
output_list = torch.stack(output_list, dim=0)
#每个头的输出进行拼接操作
output = torch.cat([output_list[i] for i in range(head_num)], dim=-1)
output = nn.Linear(dim, dim)(output)
print(output)

结果：
tensor([[ 0.1240, 0.1654, 0.1251, ..., -0.0815, 0.0090, -0.0352],
[ 0.1093, 0.1803, 0.1319, ..., -0.0843, 0.0128, -0.0470],
[ 0.1185, 0.1721, 0.1387, ..., -0.0827, 0.0048, -0.0449],
[ 0.1288, 0.1795, 0.1358, ..., -0.0785, 0.0145, -0.0532],
[ 0.1265, 0.1654, 0.1186, ..., -0.0739, 0.0234, -0.0415]],
grad_fn=<AddmmBackward0>)