AI基本结构14,15-Attention
2026/8/4 16:01:31 网站建设 项目流程

注意力块

下三角矩阵

根据softmax函数的特性,可以通过设置矩阵的上三角部分值为负无穷(-inf),然后对矩阵进行softmax,即可设置上三角矩阵的值为0,且保证行向量相加为1;

import os os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" import torch import torch.nn as nn import torch.nn.functional as F torch.manual_seed(12046) sequence_len = 64 device = 'cuda' if torch.cuda.is_available() else 'cpu' # softmax对于-inf的计算结果 a = torch.tensor([1, 2, float('-inf')]).float() print(F.softmax(a)) # K: (B, T, H) # Q: (B, T, H) # K @ Q.transpose(-2, -1): (B, T, T) scores = torch.randn(2, 4, 4) print(scores) # 定义下三角矩阵 tril = torch.tril(torch.ones(4, 4)) s = scores.masked_fill(tril == 0, float('-inf')) print(s)

方差敏感性

若行向量数据的方差过大,在softmax后会导致极端取值,即部分维度值极大,其他维度趋近于0.

故需要在矩阵相乘等操作后,对矩阵进行方差/标准差对齐操作;即除以方差除以行向量的维度(标准差除以根号维度)

# softmax对方差的敏感性 x = torch.randn(1, 8) print(x.std()), print(F.softmax(x, dim=-1)) print(F.softmax(1000 * x, dim=-1)) # 对齐分数的方差变化 B, T, H =32, 100, 10 K = torch.randn(B, T, H) Q = torch.randn(B, T, H) scores = K @ Q.transpose(-2, -1) / H ** 0.5 print(scores.std())

注意力模块与单向自注意力

def attention(query, key, value, dropout, mask=None): # query, key, value: (B, T, H) # mask: (T, T) # output: (B, T, H) B, T, H = query.shape scores = query @ key.transpose(-2, -1) / H ** 0.5 if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) w_att = F.softmax(scores, dim=-1) # (B, T, T) out = w_att @ value # (B, T, H) return out class MaskedAttention(nn.Module): # 单向自注意力 def __init__(self, emb_size, head_size): # emb_size: C, head_size: H super().__init__() self.key = nn.Linear(emb_size, head_size, bias=False) self.query = nn.Linear(emb_size, head_size, bias=False) self.value = nn.Linear(emb_size, head_size, bias=False) # 定义下三角矩阵 self.register_buffer('tril', torch.tril(torch.ones(sequence_len, sequence_len))) self.dp = nn.Dropout(0.4) def forward(self, x): # x: (B, T, C) # out: (B, T, H) B, T, C = x.shape k = self.key(x) # (B, T, H) q = self.query(x) # (B, T, H) v = self.value(x) # (B, T, H) mask = self.tril[:T, :T] out = attention(q, k, v, self.dp, mask) return out m = MaskedAttention(3, 4) x = torch.randn(5, 10, 3) print(m(x).shape)

函数实现了标准的缩放点积注意力机制,输入包括查询(query)、键(key)、值(value)三个张量,以及可选的掩码(mask)和dropout操作。函数首先计算query和key的点积得分,并通过缩放因子√H(H是特征维度)进行归一化处理,防止梯度消失或爆炸。

若提供了掩码,将掩码中值为0的位置替换为负无穷,确保softmax后对应位置的权重接近0。随后对得分矩阵应用softmax函数,得到注意力权重w_att,最终通过权重与value的加权求和得到输出。输出张量保持与输入相同的形状(B, T, H),其中B是批次大小,T是序列长度,H是特征维度。

多头注意力

class MaskedMultiHeadAttention(nn.Module): def __init__(self, emb_size, head_size): super().__init__() # 计算单头注意力的个数 n_head = emb_size // head_size heads = [MaskedAttention(emb_size, head_size) for _ in range(n_head)] self.heads = nn.ModuleList(heads) self.proj = nn.Linear(emb_size, emb_size) self.dp = nn.Dropout(0.4) def forward(self, x): # x: (B, T, C) # out: (B, T, C) out = torch.concat([h(x) for h in self.heads], dim=-1) # (B, T, C) out = self.dp(self.proj(out)) # (B, T, C) return out

(1)MaskedMultiHeadAttention 的作用

MaskedMultiHeadAttention 实现了 Transformer 中的掩码多头自注意力机制(Masked Multi-Head Self-Attention),它是在单头注意力 MaskedAttention 基础上的扩展。相比单个注意力头只能从一个特征空间学习信息关联,多头注意力通过多个独立的注意力头同时关注不同位置之间的关系,从多个角度提取序列特征。

(2)多头结构构建

在初始化函数中,首先根据嵌入维度 emb_size 和单个注意力头维度 head_size 计算注意力头数量:n_head=head_size/emb_size

随后创建多个 MaskedAttention 实例,并通过 nn.ModuleList 保存。每个注意力头都会独立计算输入序列中的注意力关系,学习不同类型的特征表示。

(3)前向传播过程

输入 x 的形状为 (B, T, C),其中 B 表示批量大小,T 表示序列长度,C 表示特征维度。模型首先将输入分别送入多个注意力头:torch.concat([h(x) for h in self.heads], dim=-1)

每个注意力头输出一部分特征,最后沿特征维度拼接,重新组合成完整的序列表示,使模型能够同时融合多个注意力空间的信息。

(4)线性投影层作用

多头拼接后的结果经过 self.proj 线性层进行特征融合,将不同注意力头的信息重新映射到原始 emb_size 维度。该步骤类似于对多个注意力结果进行进一步学习,使模型能够综合不同头捕获的信息。

前馈神经网络

class FeedForward(nn.Module): def __init__(self, emb_size): super().__init__() self.ln1 = nn.Linear(emb_size, 4 * emb_size) self.ln2 = nn.Linear(4 * emb_size, emb_size) self.dp = nn.Dropout(0.4) def forward(self, x): # x: (B, T, C) out = F.gelu(self.ln1(x)) # (B, T, C) out = self.dp(self.ln2(out)) # (B, T, C) return out

上面代码用于对多头注意力输出后的特征进行进一步的非线性变换和特征提取。该模块采用两层全连接结构,首先通过 ln1 将输入维度从 emb_size 扩展到 4 * emb_size,增加模型的特征表达能力,再经过 GELU 激活函数引入非线性,使模型能够学习更加复杂的特征关系;随后通过 ln2 将维度映射回原始 emb_size,保持与输入维度一致,便于后续残差连接。最后加入 Dropout(0.4) 随机丢弃部分神经元,降低过拟合风险。与注意力层负责捕捉序列中不同位置之间的关系不同,FeedForward 主要负责对每个位置的特征进行独立的深层变换,两者结合共同提升 Transformer 的特征学习能力。

整体注意力架构

注意力块

class Block(nn.Module): def __init__(self, emb_size, head_size): super().__init__() self.l1 = nn.LayerNorm(emb_size) self.mha = MaskedMultiHeadAttention(emb_size, head_size) self.l2 = nn.LayerNorm(emb_size) self.ff = FeedForward(emb_size) def forward(self, x): # x: (B, T, C) # out: (B, T, C) # 千万不要使用+=!!! x = x + self.mha(self.l1(x)) x = x + self.ff(self.l2(x)) return x

上面代码由多头自注意力层(Masked Multi-Head Attention)和前馈神经网络层(FeedForward)两部分组成。输入 x 首先经过 LayerNorm 归一化后送入掩码多头注意力层,使模型学习序列中不同位置之间的依赖关系,并通过残差连接 x + attention 保留原始信息;随后再次经过 LayerNorm 和前馈网络进行特征变换,再通过残差连接得到最终输出。这里使用残差连接可以缓解深层网络中的梯度消失问题,提高模型训练稳定性。代码中特别强调不能使用 +=,因为 PyTorch 的自动求导机制依赖计算图,原地修改操作可能破坏计算图,导致梯度计算错误,因此应使用 x = x + ... 创建新的张量。

基于多块注意力的GPT

class CharGPT(nn.Module): def __init__(self, vs): super().__init__() self.token_emb = nn.Embedding(vs, emb_size) self.pos_emb = nn.Embedding(sequence_len, emb_size) block = [Block(emb_size, head_size) for _ in range(n_layer)] self.blocks = nn.Sequential(*block) self.l = nn.LayerNorm(emb_size) self.lm = nn.Linear(emb_size, vs) def forward(self, x): # x: (B, T) # logits: (B, T, vs) B, T = x.shape pos = torch.arange(0, T, dtype=torch.long, device=x.device) token_embeddings = self.token_emb(x) # (B, T, C) position_embeddings = self.pos_emb(pos) # ( T, C) h = token_embeddings + position_embeddings # (B, T, C) h = self.blocks(h) # (B, T, C) logits = self.lm(self.l(h)) # (B, T, vs) return logits

1)模型整体结构

CharGPT 是一个基于 Transformer 架构的字符级生成语言模型,其结构由词元嵌入(Token Embedding)+ 位置嵌入(Position Embedding)+ 多层 Transformer Block + LayerNorm + 输出层组成。模型输入为字符编号序列 (B, T),经过多层注意力机制学习字符之间的上下文关系,最终预测每个位置下一个字符的概率分布。

(2)词元嵌入与位置嵌入

token_emb 用于将离散字符编号转换为连续向量,表示每个字符本身的信息;由于 Transformer 不像 RNN 一样具有时间递归结构,无法自动感知字符顺序,因此额外加入 pos_emb 位置嵌入,为每个位置添加位置信息。两者相加后得到包含“字符内容 + 位置信息”的输入表示。

(3)Transformer Block 特征提取

处理后的特征输入多个 Block 模块,通过掩码多头自注意力学习字符之间的依赖关系,并利用前馈网络进行非线性特征变换。多个 Block 堆叠可以逐层提取更加复杂的上下文信息,同时利用残差连接和 LayerNorm 保证深层网络训练稳定。

(4)输出预测层

经过所有 Transformer Block 后,隐藏表示先通过 LayerNorm 进行归一化,再经过线性层 lm 映射到词表大小 vs,输出形状为 (B, T, vs) 的 logits。每个位置的输出表示当前位置预测下一个字符时,对所有可能字符的评分。

(5)与 RNN/LSTM 的区别

相比 RNN/LSTM 依靠隐藏状态逐步传递历史信息,CharGPT 使用自注意力机制直接建立序列中任意两个位置之间的联系,可以同时关注整个上下文,因此更擅长捕捉长距离依赖关系。同时,Transformer 可以并行处理整个序列,训练效率通常高于循环神经网络。

(6)模型前向过程总结

输入字符序列

1. token_emb 获取字符特征

2.pos_emb 加入位置信息

3.多层 Block 提取上下文关系

4.LayerNorm 稳定特征分布

5.Linear 映射到词表空间

6.输出每个位置的字符预测结果

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询