基础学习一:Transformer架构+多头注意力机制代码及解析

📅 发布时间:2026/9/28 8:09:36
基础学习一:Transformer架构+多头注意力机制代码及解析
一. Transformer架构回顾Transformer架构是当前AI的基础架构网络结构如下我们对Transformer架构先来一次清晰的回顾我们以一个简单的例子来解释解释整个流程。首先Transformer顾名思义就是一种用于语言翻译的网络架构。假设我们想将 “I love apples” 翻译成 “我爱苹果”。我们可以分为训练阶段和推理阶段。我们观察网络结构他是一个 Encoder(左) 和 Decoder(右) 的编解码两部分结构其中在训练过程中我们将原句 “I love apples” 输入到Encoder中将翻译句 “我爱苹果” 输入到Decoder中。这个过程是原句是翻译过程的条件输入翻译句是模型在原句条件下生成的目标序列。现在我们来详细解释每一个部分。Encoder1. 输入是“I love apples” 经过一次Input Embedding得到X。Input Embedding实际上就是将 “I love apples” 进行token化分成“I” “love” “apples” 三个词。然后利用编码表将这三个词对应好三个数字 (Token ID) 即 “I”: 37 “love”:102 “apples”: 65。因此“I love apples” 这句话就转换成向量[37,102,65]这还没有完我们还需要其转换成网络可用的输入即一个3xd_model的矩阵3代表3个token。[[xxxxxxx...] d_model维 “I”[xxxxxxx...] d_model维 “love”[xxxxxxx...] d_model维 “apples”]2. 然后将X经过一次位置编码Positional Encoding即XPE 得到X_in。Positional Encoding公式如下即对X矩阵每行计算一下PE2i和2i1是来区分奇偶项。3.然后我们将X_in经过一次Multi-Head-Attention。Encoder这部分的Multi-Head-Attention是self attention即找到“I” “love” “apples”这三个词之间的相关性得到Attn具体如下首先我们将X_in经历三个线性变换W_q W_k W_v 得到Q(query) K(key) V(value)。在注意力机制中词 “I” 的查询向量 (Q) 与包括 “love” 在内的所有词的键向量 (K) 计算相关性并根据所得注意力权重对相应的值向量 (V) 进行加权聚合。具体计算公式为将除以sqrt(d_k)是为了抑制点积的方差随特征维度d_k 增大而上升避免 Softmax 进入饱和区域、使注意力分布过度尖锐并导致梯度变小。多头注意力将 d_model划分为heads个维度为 d_model / heads的子空间使各注意力头能够并行学习不同位置关系和语义模式例如词法依赖、句法结构及长距离关联随后将各头的结果拼接并投影从而获得比单头注意力更丰富的特征表示如下图所示。整个流程X_in - QKV划分 - Head划分 - 子空间Attn计算 - concat - 投影 - Attn4.AddNorm 首先将多头注意力输出 Attn 经 Dropout 后与输入 X_in 进行残差相加再通过 LayerNorm 进行归一化从而改善梯度传播并提高训练稳定性。5.将第四步得到的中间表征送入前馈网络Feed‑Forward Network, FNN执行前向传播再对 FNN 输出执行残差相加与层归一化AddNorm运算。以上是Encoder部分的将“I love apples” 进行了编码、多头注意力计算、前向传播的操作。Decoder基于Encoder输出的先验上下文训练阶段Decoder将目标译文 “我爱苹果”右移一位并拼接起始符BOS与结束符EOS得到输入序列“BOS我爱苹果EOS”以开展学习而推理阶段解码器采用自回归方式逐 token 生成输出即1.input:BOS output:我2.input:BOS我 output:爱3.input:BOS我爱 output:苹果4.input:BOS我爱苹果 output:EOS训练阶段解码器以“BOS我爱苹果”作为输入序列、以“我爱苹果EOS”作为对应标签序列通过掩码机制实现对目标序列的并行同步预测。输入和标签解决好了再通过 Output Embedding 和 Positional Encoding 后接入到Masked Multi-Head-Attention这一部分与多头注意力稍有区别的部分在于Masked就是在如下表示中X_out - QKV划分 - Head划分 - 带Masked的子空间Attn计算 - concat - 投影 - Attn带掩码Masked的自注意力主要用在Decoder / 自回归生成场景目的是防止 未来信息泄露在预测第 i 个位置的词时模型只能看到它自己以及它之前的词绝对不能看到后面还没生成的词 —— 否则训练时等于提前看到答案推理时就会失效。以序列BOS 我 爱 中国 EOS为例位置编号 0…4当前位置能看到的位置即0BOS0只看自己1我0, 1不能看到后面的 爱 中国2爱0, 1, 2不能看到 中国3中国0, 1, 2, 3…4EOS0, 1, 2, 3, 4…二、多头注意力机制代码import math import torch import torch.nn as nn class MultiheadAttention(nn.Module): # n_heads 头数 # d_model 向量维度 def __init__(self, n_heads, d_model, dropout0.1): super(MultiheadAttention, self).__init__() if d_model % n_heads ! 0: raise ValueError( fd_model ({d_model}) must be divisible by fn_heads ({n_heads}). ) self.d_model d_model self.n_heads n_heads # d_model 必须是 n_heads的整数倍 assert self.d_model % self.n_heads 0 self.d_k self.d_model // self.n_heads self.d_v self.d_k # 定义QKV矩阵 # 假如是单个头(T,d_model) - Wq_0 (d_model, d_k) - (T,d_k) - 多个头合并成完整的d_model - (T,d_model) self.W_q nn.Linear(self.d_model, self.d_model) self.W_k nn.Linear(self.d_model, self.d_model) # 假如是单个头(T,d_model) - Wv_0 (d_model, d_v) - (T,d_v) - 多个头合并成完整的d_model - (T,d_model) self.W_v nn.Linear(self.d_model, self.d_model) # W_o self.W_o nn.Linear(self.d_model, self.d_model) # dropout self.dp nn.Dropout(dropout) # query key value 来自于输入 # 如果做self attention 则 query key value X # 如果做cross attention 则 query X1 key value X2 def forward(self, query, key, value, maskNone): # 现获取一下 batch_size大小为后续做准备 batch_size query.shape[0] # [B, T, d_model] - B # 输入X分别经过W_q W_k W_v得到QKV # [B, T, d_model] - W - [B, T, d_model] Q self.W_q(query) K self.W_k(key) V self.W_v(value) # 将我们得到的QKV分解成多头的形式 # [B, T, d_model] - [B, T, n_heads, d_k] - permute - [B, n_heads, T, d_k] # -1是让pytorch自行判断填充完整 Q Q.view(batch_size, -1, self.n_heads, self.d_k).permute(0, 2, 1, 3) K K.view(batch_size, -1, self.n_heads, self.d_k).permute(0, 2, 1, 3) V V.view(batch_size, -1, self.n_heads, self.d_v).permute(0, 2, 1, 3) # 计算attention # Q:[B, n_heads, T_q, d_k] K: [B, n_heads, T_k, d_k] # Q K^T - [B, n_heads, T_q, d_k] [B, n_heads, d_k, T_k] - [B, n_heads, T_q, T_k] # 缩放一下 scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # 如果有maskmask是一个01矩阵则1的位置元素保持不变0的位置元素改为-1e10(无穷小) if mask is not None: scores scores.masked_fill(mask 0, -1e10) # attention: [B, n_heads, T_q, T_k] # softmax是按照T_k的维度计算的即按每一行的所有列计算一次softmax。 # 再经历一个dropout目的是让注意力权重始终不要依赖于其中一个key而是学习如何利用其他的key attention_weights self.dp(torch.softmax(scores, dim-1)) # 计算多头注意力的结果 # attention: [B, n_heads, T_q, T_k] V: [B, n_heads, T_k, d_v] - [B, n_heads, T_q, d_v] x torch.matmul(attention_weights, V) # 将x转成可以合并多头的形式为后续合并做铺垫 [B, n_heads, T_q, d_v] - [B, T_q, n_heads, d_v] # permute只是改变了解读方式而不是改变了存储方式contiguous是改变了存储方式 x x.permute(0, 2, 1, 3).contiguous() # [B, T_q, n_heads, d_v] - [B, T_q, d_model] x x.view(batch_size, -1, self.n_heads*self.d_k) # 再经历一个W_o output self.W_o(x) return output