Transformer计算全流程:从词嵌入到概率输出的维度推演

📅 2026/7/14 13:07:49
Transformer计算全流程:从词嵌入到概率输出的维度推演
1. Transformer模型的基本框架Transformer模型由编码器Encoder和解码器Decoder两部分组成每一部分都由多个相同的层堆叠而成。编码器负责将输入序列转换为一系列连续的表示解码器则根据编码器的输出和之前的输出生成目标序列。整个过程可以看作是一个从输入到输出的数据流每一步都伴随着张量维度的精确变化。为了更好地理解这个过程我们以一个具体的例子来说明。假设输入序列是中文句子早上好对应的目标序列是英文句子good morning。我们将追踪这个输入样本在Transformer模型中的流动过程重点关注每一步矩阵运算引发的张量维度变化。2. 输入表示与词嵌入2.1 词嵌入层输入序列首先经过词嵌入层Embedding Layer将每个词转换为一个固定维度的向量表示。假设我们使用的嵌入维度是512那么对于输入序列早上好假设分词为3个token经过嵌入层后的输出维度将是(3, 512)。具体计算过程如下每个token会被映射为一个512维的向量对于长度为3的序列最终得到一个3×512的矩阵这个矩阵的行表示序列中的位置列表示特征维度# 伪代码示例 embedding_layer nn.Embedding(vocab_size, 512) input_ids [1, 2, 3] # 早上好对应的token id embedded embedding_layer(input_ids) # 输出形状(3, 512)2.2 位置编码由于Transformer没有内置的顺序信息我们需要额外添加位置编码Positional Encoding来注入序列的位置信息。位置编码与词嵌入具有相同的维度可以直接相加。位置编码的计算公式如下PE(pos,2i) sin(pos/10000^(2i/d_model)) PE(pos,2i1) cos(pos/10000^(2i/d_model))其中pos是位置i是维度索引。将位置编码加到词嵌入上后输出维度保持不变仍然是(3, 512)。# 位置编码实现示例 position torch.arange(0, max_len).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) x x pe[:x.size(0), :] # 形状保持不变3. 编码器中的自注意力机制3.1 QKV矩阵计算编码器的核心是多头自注意力机制。首先我们需要计算查询Query、键Key和值Value矩阵。这是通过将输入矩阵与三个不同的权重矩阵相乘得到的。具体计算过程输入矩阵X的形状为(3, 512)三个权重矩阵WQ、WK、WV的形状都是(512, 64)假设每个头的维度是64计算Q、K、VQ X × WQ → (3, 64)K X × WK → (3, 64)V X × WV → (3, 64)# 伪代码示例 WQ nn.Linear(512, 64) WK nn.Linear(512, 64) WV nn.Linear(512, 64) Q WQ(x) # (3, 64) K WK(x) # (3, 64) V WV(x) # (3, 64)3.2 注意力分数计算接下来计算注意力分数这是通过Q和K的点积得到的然后除以√d_kd_k是key的维度这里为64进行缩放再经过softmax归一化。计算步骤Q × K^T → (3, 3) 的注意力分数矩阵除以√64即8进行缩放应用softmax得到注意力权重scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(64) attn_weights F.softmax(scores, dim-1) # (3, 3)3.3 上下文向量计算最后用注意力权重对V进行加权求和得到上下文向量context torch.matmul(attn_weights, V) # (3, 64)3.4 多头注意力实际中我们使用多头注意力假设8个头每个头独立计算注意力然后将结果拼接起来每个头的输出是(3, 64)8个头拼接后是(3, 512)通过一个线性层进行投影# 假设有8个头 multi_head [] for _ in range(8): head self_attention(Q, K, V) # (3, 64) multi_head.append(head) multi_head torch.cat(multi_head, dim-1) # (3, 512) output linear_proj(multi_head) # (3, 512)4. 前馈神经网络自注意力层之后是一个前馈神经网络Feed Forward Network它由两个线性变换和一个ReLU激活函数组成第一个线性层将维度从512扩展到2048ReLU激活函数第二个线性层将维度从2048投影回512维度变化(3, 512) → (3, 2048) → (3, 512)self.ffn nn.Sequential( nn.Linear(512, 2048), nn.ReLU(), nn.Linear(2048, 512) ) output self.ffn(input) # 输出形状(3, 512)5. 解码器的工作原理5.1 解码器输入处理解码器的输入是目标序列如good morning同样需要经过嵌入层和位置编码。此外解码器还需要编码器的输出作为额外的输入。目标序列嵌入(2, 512)假设good morning分词为2个token添加位置编码计算解码器的Q、K、V矩阵5.2 掩码自注意力解码器的第一个注意力层是掩码自注意力防止当前位置关注到未来的信息。这通过添加一个上三角矩阵upper triangular matrix实现。# 生成掩码 mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool() scores scores.masked_fill(mask, float(-inf)) # 将未来位置设为负无穷5.3 编码器-解码器注意力解码器的第二个注意力层使用编码器的输出作为K和V解码器的输出作为Q。这使得解码器可以关注输入序列的相关部分。# enc_output是编码器输出形状为(3, 512) # dec_output是解码器当前输出形状为(2, 512) Q dec_output WQ # (2, 64) K enc_output WK # (3, 64) V enc_output WV # (3, 64) # 计算交叉注意力 scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(64) attn torch.softmax(scores, dim-1) context torch.matmul(attn, V) # (2, 64)6. 输出概率计算解码器的最终输出经过一个线性层和softmax函数转换为词汇表上的概率分布线性层将512维映射到词汇表大小如50000softmax转换为概率logits linear_layer(decoder_output) # (2, vocab_size) probs F.softmax(logits, dim-1) # 每个位置的概率分布7. 训练与推理过程7.1 训练阶段在训练时我们使用教师强制Teacher Forcing即将真实的目标序列作为解码器输入即使模型预测错误。计算预测结果与真实标签的交叉熵损失。criterion nn.CrossEntropyLoss(ignore_index0) # 忽略padding loss criterion(logits.view(-1, vocab_size), labels.view(-1))7.2 推理阶段推理时采用自回归方式每一步预测一个token并将其作为下一步的输入output [] current_input start_token # 起始token for _ in range(max_len): logits model(enc_input, current_input) next_token logits.argmax(-1)[-1] # 选择概率最高的token output.append(next_token) if next_token end_token: # 结束token break current_input torch.cat([current_input, next_token.unsqueeze(0)])8. 维度变化总结让我们总结一下主要步骤中的维度变化输入嵌入(序列长度, 512)位置编码(序列长度, 512) → 相加后维度不变自注意力Q/K/V(序列长度, 64) 每个头注意力分数(序列长度, 序列长度)上下文向量(序列长度, 64) 每个头多头拼接(序列长度, 512)前馈网络第一个线性层(序列长度, 512) → (序列长度, 2048)第二个线性层(序列长度, 2048) → (序列长度, 512)解码器输出线性投影(序列长度, 512) → (序列长度, 词汇表大小)softmax后得到每个位置的概率分布在实际项目中理解这些维度变化对于调试模型和优化性能至关重要。特别是在实现自定义的Transformer变体时确保各层之间的维度匹配是避免错误的关键。