Transformer模型原理图解用Python从头实现一个简易版如果你已经对深度学习尤其是自然语言处理领域有所涉猎那么“Transformer”这个名字对你来说一定如雷贯耳。它早已不是2017年那篇论文里的一个新颖架构而是成为了驱动当今绝大多数顶尖AI模型的核心引擎。从ChatGPT到Midjourney其背后或多或少都有Transformer的影子。然而对于许多开发者来说理解Transformer常常停留在“自注意力”、“编码器-解码器”这些术语的表面。看懂了论文里的公式却依然不知道这些矩阵乘法是如何一步步构建起一个能理解语言的智能体的。这正是本文想要解决的问题。我们不满足于仅仅复述理论而是选择了一条更具挑战性也更有收获的路径用Python从零开始亲手搭建一个简易但完整的Transformer模型。我们将抛开那些庞大的深度学习框架仅依赖NumPy和一点PyTorch的基础张量操作将论文中的每一个模块都可视化为可运行的代码。这个过程就像在组装一台精密的钟表你会亲眼看到每一个齿轮自注意力头、前馈网络、层归一化是如何咬合最终让指针模型输出开始转动的。无论你是希望夯实基础以便更好地调参还是纯粹对模型的内在机理感到好奇这篇手把手的图解指南都将为你提供一次深度的、触及本质的学习体验。1. 从核心思想到代码蓝图理解Transformer的设计哲学在动手写第一行代码之前我们需要先跳出具体的网络层从更高的视角审视Transformer为何能成功。传统的RNN和LSTM在处理序列时是“串行”的必须一个接一个地处理单词这导致了训练缓慢和难以捕捉长距离依赖。Transformer的颠覆性在于其“并行”和“全局”的视野。并行性来源于它对循环的彻底抛弃。模型不再需要等待前一个时间步的计算结果整个序列的所有单词可以同时被处理。这极大地利用了现代GPU的并行计算能力使得训练超大规模模型成为可能。全局性则归功于其核心——自注意力机制。对于一个句子中的任意一个单词自注意力机制允许它直接与句子中的所有其他单词包括它自己建立联系并计算出一个“注意力分数”这个分数决定了在编码当前单词时应该“关注”其他单词的程度。这种设计让模型能够轻松捕捉“The animal didnt cross the street because it was too tired”中“it”指代“animal”这样的长距离依赖关系。提示自注意力机制的本质是一种“信息检索”过程。对于序列中的每个元素查询它去整个序列键值对存储中检索最相关的信息并加权聚合起来。基于这些思想Transformer的经典架构主要由编码器Encoder和解码器Decoder堆叠而成。为了完成我们的简易实现我们将聚焦于最核心的编码器部分并构建一个仅包含一个编码器层的微型Transformer。这足以让我们洞察所有关键组件的工作原理。我们的代码蓝图将按以下模块顺序构建词嵌入与位置编码将离散的单词转换为富含位置信息的连续向量。自注意力机制实现模型理解上下文关系的核心引擎。前馈神经网络对自注意力输出进行非线性变换增加模型表达能力。层归一化与残差连接稳定深度网络训练的关键技巧。整合构建完整的编码器层将以上模块组装起来。我们将为每个模块提供清晰的数学图解和对应的Python实现确保每一步都可验证、可运行。2. 构建基石词嵌入与位置编码Transformer模型处理的输入是数字。因此我们的第一步是将文本序列例如[hello, world]转换为一系列数字向量。这包含两个步骤词嵌入和位置编码。2.1 词嵌入从单词到向量空间的一个点词嵌入层本质上是一个查找表。假设我们的词汇表大小为vocab_size每个单词对应一个唯一的整数索引如“hello”- 102“world”- 304。嵌入层是一个形状为(vocab_size, d_model)的矩阵其中d_model是我们模型的隐藏层维度例如512。通过这个矩阵我们可以将单词索引快速映射为一个d_model维的稠密向量。这个向量并非随机。在预训练中它会被优化使得语义相似的单词如“国王”和“王后”在向量空间中的位置也接近。在我们的简易实现中我们将随机初始化这个嵌入矩阵。import torch import torch.nn as nn import numpy as np class EmbeddingLayer(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) def forward(self, x): # x 的形状: (batch_size, seq_len) - 包含单词索引的整数张量 # 输出形状: (batch_size, seq_len, d_model) return self.embedding(x) # 示例用法 vocab_size 10000 # 假设词汇表有1万个单词 d_model 512 # 模型隐藏维度 embed_layer EmbeddingLayer(vocab_size, d_model) # 假设一个批次包含2个句子每个句子5个词 batch_indices torch.LongTensor([[12, 45, 23, 9, 78], [1, 34, 56, 90, 21]]) embedded_output embed_layer(batch_indices) print(f嵌入后张量形状: {embedded_output.shape}) # 输出: torch.Size([2, 5, 512])2.2 位置编码为序列注入顺序信息自注意力机制本身是“无序”的它对输入序列的排列是不变的。为了利用序列的顺序信息Transformer引入了位置编码。它将每个位置第1个词第2个词...也编码成一个d_model维的向量然后直接加到该位置的词嵌入向量上。原始论文使用了一组固定的、基于正弦和余弦函数的编码公式[ PE_{(pos, 2i)} \sin(pos / 10000^{2i/d_{model}}) ] [ PE_{(pos, 2i1)} \cos(pos / 10000^{2i/d_{model}}) ]其中pos是位置i是维度索引。这种设计使得模型能够轻松学习到相对位置关系例如位置 posk 的编码可以表示为位置 pos 编码的线性函数。class PositionalEncoding(nn.Module): def __init__(self, d_model, max_seq_len5000): super().__init__() # 创建一个形状为 (max_seq_len, d_model) 的位置编码矩阵 pe torch.zeros(max_seq_len, d_model) position torch.arange(0, max_seq_len).unsqueeze(1) # (max_seq_len, 1) div_term torch.exp(torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) # 偶数维度用sin pe[:, 1::2] torch.cos(position * div_term) # 奇数维度用cos # 注册为缓冲区使其随模型移动但不参与梯度更新 self.register_buffer(pe, pe.unsqueeze(0)) # (1, max_seq_len, d_model) def forward(self, x): # x 的形状: (batch_size, seq_len, d_model) seq_len x.size(1) x x self.pe[:, :seq_len] return x # 将嵌入和位置编码组合 embedding embed_layer(batch_indices) pos_encoder PositionalEncoding(d_model) final_input pos_encoder(embedding) print(f加入位置编码后的张量形状: {final_input.shape}) # 仍为 torch.Size([2, 5, 512])至此我们得到了一个既包含语义信息来自词嵌入又包含顺序信息来自位置编码的输入表示可以送入Transformer的核心层进行处理。3. 核心引擎自注意力机制的实现与图解自注意力是Transformer的灵魂。它允许序列中的每个位置“查看”序列中的所有位置包括自身并根据相关性动态地聚合信息。让我们将其分解为可操作的步骤。3.1 计算查询、键和值首先对于输入序列X形状为(batch_size, seq_len, d_model)我们通过三个不同的线性变换层为其生成三组新的向量查询Query、键Key和值Value。查询Q可以理解为当前正在处理的单词提出的“问题”。键K序列中所有单词提供的“答案索引”。值V序列中所有单词提供的“实际答案内容”。这三个变换的权重矩阵W_Q,W_K,W_V的形状都是(d_model, d_k)或(d_model, d_v)。通常为了简便我们设d_k d_v d_model / h其中h是注意力头的数量。在单头注意力中我们常设d_k d_v d_model。def self_attention_single_head(X, W_Q, W_K, W_V): 单头自注意力计算 X: 输入张量形状 (batch_size, seq_len, d_model) W_Q, W_K, W_V: 权重矩阵形状 (d_model, d_k) 返回: 注意力输出形状 (batch_size, seq_len, d_v) batch_size, seq_len, d_model X.shape d_k W_Q.shape[1] # 线性变换得到 Q, K, V Q torch.matmul(X, W_Q) # (batch_size, seq_len, d_k) K torch.matmul(X, W_K) # (batch_size, seq_len, d_k) V torch.matmul(X, W_V) # (batch_size, seq_len, d_v) # 计算注意力分数: Q * K^T scores torch.matmul(Q, K.transpose(-2, -1)) # (batch_size, seq_len, seq_len) # 缩放防止点积过大导致softmax梯度消失 scores scores / (d_k ** 0.5) # 应用softmax得到注意力权重每一行和为1 attn_weights torch.softmax(scores, dim-1) # (batch_size, seq_len, seq_len) # 用注意力权重加权聚合 Value 向量 output torch.matmul(attn_weights, V) # (batch_size, seq_len, d_v) return output, attn_weights # 示例手动计算单头注意力 batch_size, seq_len, d_model 2, 5, 512 d_k d_v 64 X_demo torch.randn(batch_size, seq_len, d_model) W_Q_demo torch.randn(d_model, d_k) W_K_demo torch.randn(d_model, d_k) W_V_demo torch.randn(d_model, d_v) output_demo, weights_demo self_attention_single_head(X_demo, W_Q_demo, W_K_demo, W_V_demo) print(f单头注意力输出形状: {output_demo.shape}) print(f注意力权重矩阵形状: {weights_demo.shape})3.2 多头注意力并行化的信息聚合单头注意力只从一个“视角”去理解序列关系。为了捕捉更丰富的关系例如语法关系、指代关系、语义关系等Transformer使用了多头注意力。简单来说就是并行地运行多个独立的注意力头每个头都有自己的W_Q, W_K, W_V参数学习关注序列中不同类型的信息。最后将所有头的输出拼接起来再通过一个线性变换W_O进行融合。组件作用输入形状输出形状线性变换 (h个头)生成每个头的 Q, K, V(batch, seq_len, d_model)h 个 (batch, seq_len, d_k/d_v)缩放点积注意力计算注意力并加权h 个 (batch, seq_len, d_k/d_v)h 个 (batch, seq_len, d_v)拼接 (Concat)合并所有头的输出h 个 (batch, seq_len, d_v)(batch, seq_len, h * d_v)输出线性层 (W_O)融合多头信息(batch, seq_len, h * d_v)(batch, seq_len, d_model)class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads 0, d_model 必须能被 num_heads 整除 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads self.d_v d_model // num_heads # 将生成 Q, K, V 的线性层合并最后再分割 self.W_QKV nn.Linear(d_model, 3 * d_model) # 输出融合层 self.W_O nn.Linear(d_model, d_model) def forward(self, x): batch_size, seq_len, _ x.shape # 1. 线性变换并分割出 Q, K, V qkv self.W_QKV(x) # (batch_size, seq_len, 3*d_model) qkv qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.d_k) qkv qkv.permute(2, 0, 3, 1, 4) # (3, batch_size, num_heads, seq_len, d_k) Q, K, V qkv[0], qkv[1], qkv[2] # 每个形状: (batch_size, num_heads, seq_len, d_k) # 2. 计算缩放点积注意力 (使用矩阵广播并行计算所有头) scores torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5) # (batch_size, num_heads, seq_len, seq_len) attn_weights torch.softmax(scores, dim-1) # 3. 加权求和 context torch.matmul(attn_weights, V) # (batch_size, num_heads, seq_len, d_k) # 4. 拼接多头输出 context context.transpose(1, 2).contiguous() # (batch_size, seq_len, num_heads, d_k) context context.reshape(batch_size, seq_len, self.d_model) # (batch_size, seq_len, d_model) # 5. 输出线性变换 output self.W_O(context) return output # 测试多头注意力层 mha MultiHeadAttention(d_model512, num_heads8) test_input torch.randn(2, 10, 512) # 批次2序列长度10维度512 mha_output mha(test_input) print(f多头注意力层输出形状: {mha_output.shape}) # 应与输入一致: torch.Size([2, 10, 512])通过多头注意力模型能够同时从多个子空间学习信息其表达能力远超单头注意力。你可以将attn_weights可视化观察不同注意力头分别关注了序列的哪些部分。4. 前馈网络与稳定化技巧自注意力层的输出捕捉了上下文信息但缺乏非线性变换能力。因此每个编码器层在自注意力之后都会接一个前馈神经网络。4.1 前馈网络简单的非线性变换这个FFN非常简单由两个线性变换和一个激活函数通常是ReLU组成。它的作用是对每个位置的特征进行独立的、相同的变换。[ FFN(x) \max(0, xW_1 b_1)W_2 b_2 ]其中第一个线性层将维度从d_model扩展到d_ff例如2048第二个线性层再投影回d_model。这种“放大再缩小”的结构有助于学习更复杂的特征表示。class PositionwiseFeedForward(nn.Module): def __init__(self, d_model, d_ff): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.activation nn.ReLU() def forward(self, x): # x 形状: (batch_size, seq_len, d_model) return self.linear2(self.activation(self.linear1(x))) # 测试前馈网络 ffn PositionwiseFeedForward(d_model512, d_ff2048) ffn_output ffn(mha_output) print(f前馈网络输出形状: {ffn_output.shape}) # torch.Size([2, 10, 512])4.2 残差连接与层归一化深度网络的稳定器Transformer通常很深例如12层编码器。在深度网络中梯度消失/爆炸和训练不稳定的问题很常见。Transformer借鉴了ResNet的思想在每个子层自注意力、前馈网络周围使用了残差连接并紧接着进行层归一化。残差连接将子层的输入x直接加到其输出F(x)上即output x F(x)。这确保了梯度可以直接回传缓解了梯度消失问题让网络能够轻松学习恒等映射。层归一化对单个样本的所有特征维度进行归一化与批归一化不同。它稳定了激活值的分布加速了训练收敛。这两个操作通常组合成一个Add Norm模块。class AddNorm(nn.Module): 残差连接后接层归一化 def __init__(self, d_model, dropout0.1): super().__init__() self.norm nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, sublayer_output): x: 子层输入 sublayer_output: 子层输出 # 先残差连接再层归一化 (原始论文顺序) return self.norm(x self.dropout(sublayer_output)) # 示例在一个自注意力子层周围应用 Add Norm layer_norm nn.LayerNorm(d_model) add_norm_module AddNorm(d_model) # 假设 sublayer_input 是自注意力层的输入 sublayer_input torch.randn(2, 10, 512) # 假设 sublayer_output 是自注意力层的输出 sublayer_output mha(sublayer_input) # 应用 Add Norm output_after_add_norm add_norm_module(sublayer_input, sublayer_output) print(fAdd Norm 后输出形状: {output_after_add_norm.shape})5. 组装与测试构建完整的Transformer编码器层现在我们已经拥有了所有必要的组件。让我们将它们组装成一个完整的Transformer编码器层并创建一个简易的Transformer模型来测试其前向传播过程。一个标准的编码器层包含以下顺序多头自注意力层 Add Norm前馈网络层 Add Normclass TransformerEncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.feed_forward PositionwiseFeedForward(d_model, d_ff) self.norm1 AddNorm(d_model, dropout) self.norm2 AddNorm(d_model, dropout) def forward(self, x): # 第一个子层多头自注意力 attn_output self.self_attn(x) x self.norm1(x, attn_output) # Add Norm # 第二个子层前馈网络 ff_output self.feed_forward(x) x self.norm2(x, ff_output) # Add Norm return x class SimpleTransformer(nn.Module): 一个简易的Transformer仅包含一个编码器层 def __init__(self, vocab_size, d_model, num_heads, d_ff, max_seq_len, dropout0.1): super().__init__() self.embedding EmbeddingLayer(vocab_size, d_model) self.pos_encoding PositionalEncoding(d_model, max_seq_len) self.encoder_layer TransformerEncoderLayer(d_model, num_heads, d_ff, dropout) # 一个简单的输出层例如用于分类任务 self.output_layer nn.Linear(d_model, vocab_size) def forward(self, src_tokens): # src_tokens: (batch_size, seq_len) # 1. 嵌入和位置编码 x self.embedding(src_tokens) x self.pos_encoding(x) # 2. 通过编码器层 x self.encoder_layer(x) # 3. 取第一个位置[CLS] token或做池化用于分类这里简单取平均 pooled_output x.mean(dim1) # (batch_size, d_model) # 4. 输出层 logits self.output_layer(pooled_output) # (batch_size, vocab_size) return logits # 实例化并测试模型 vocab_size 10000 d_model 512 num_heads 8 d_ff 2048 max_seq_len 128 model SimpleTransformer(vocab_size, d_model, num_heads, d_ff, max_seq_len) # 模拟一个批次的输入数据 batch_size 4 seq_length 20 dummy_input torch.randint(0, vocab_size, (batch_size, seq_length)) # 前向传播 output model(dummy_input) print(f模型输入形状: {dummy_input.shape}) print(f模型输出形状 (分类logits): {output.shape}) print(模型构建成功前向传播完成)运行以上代码如果一切顺利你将看到模型成功接收了形状为(4, 20)的输入4个句子每句20个词并输出了形状为(4, 10000)的logits。这标志着你已经从零搭建起了一个功能完整的Transformer编码器核心。当然一个完整的Transformer模型还包括解码器层、掩码多头注意力、最终的线性层和softmax等。但通过这个从嵌入到编码器的完整流程你已经掌握了Transformer最核心、最精髓的部分。理解了这个简易版的实现再去阅读Hugging Face的Transformers库源码或是研究BERT、GPT的架构就会有一种“庖丁解牛”般的通透感。真正的掌握始于亲手将理论变为一行行运行的代码。