尧图网站建设 
从零手写 TransformerPyTorch 逐行实现与原理拆解1. Transformer架构2. 环境准备与数据加载3. 创建词表与Tokenize4. 样本构造5. Embedding层6. 位置编码7. Transformer Block8. Multi-Head Attention 逐行拆解8.1 Q/K/V 投影8.2 拆成 4 个头8.3 注意力分数8.4 因果掩码8.5 加权求和8.6 合并多头8.7 输出投影 Wo 残差8.8 FFN 输出层本文记录一次从零手写 Transformerdecoder-onlyGPT 风格的完整过程从数据加载、tiktoken 分词、Embedding、手写位置编码到单层多头注意力与 FFN全部用 PyTorch 逐行实现并给出每一步的张量形状变化。代码与讲解一一对应适合刚接触 Transformer 想动手写一遍的同学。1. Transformer架构2. 环境准备与数据加载importtorchimporttorch.nnasnnimporttorch.nn.functionalasFimportosimportrequestsimportpandasaspd获取数据集ifnotos.path.exists(sales_textbook.txt):urlhttps://huggingface.co/datasets/goendalf666/sales-textbook_for_convincing_and_selling/resolve/main/sales_textbook.txt?downloadtruewithopen(sales_textbook.txt,wb)asf:f.write(requests.get(url).content)withopen(sales_textbook.txt,r)asf:textf.read()3. 创建词表与Tokenize创建词表这里直接使用现有的openai开源的tokenizer即可importtiktoken enctiktoken.get_encoding(o200k_base)这里的text是读取了全部的字符串然后进行tokenize之后是对每一个字母给了一个数字编号。可以看到text存储的这一段文字内容中字母、数字、标点符号转成编号之后最大的编号为199853tokenized_txtenc.encode(text)print(text[:7])print(tokenized_txt[:7])print(max(tokenized_txt))tokenized_txttorch.tensor(tokenized_txt,dtypetorch.long)train_txttokenized_txt[:int(0.9*len(tokenized_txt))]val_txttokenized_txt[int(0.9*len(tokenized_txt)):]输出Chapter[45990, 220, 16, 25, 22478, 158287, 326]1998534. 样本构造这里对全局需要有一个认识对于数据的选取训练数据和标签的选取。标签相对于训练索引在文本中是向后偏移了一位的因为需要用现有的数据对下一位进行预测。这里需要写切分的代码同时将切分后的样本按照batch进行拼接拼接之后的训练样本还不够这个时候还只是[batch,context_len]这里代表着只是把16个字母按照batch4排放在一起但是学习的时候机器并不认识这些字母所以需要将这16个字母用高维特征来表示也就是前面做的tokenize。batch_size4context_len16d_model64datatrain_txt idxstorch.randint(0,highlen(data)-context_len,size(batch_size,))x_batchtorch.stack([data[idx:idxcontext_len]foridxinidxs])y_batchtorch.stack([data[idx1:idxcontext_len1]foridxinidxs])print(x_batch)x_batch.shape输出tensor([[ 25222, 402, 2935, 1062, 6086, 3692, 4534, 71265, 2012, 5720, 11, 290, 172442, 625, 1606, 84275], [ 5892, 1989, 6028, 316, 1520, 261, 12398, 8660, 52125, 11, 472, 1023, 11747, 12486, 842, 402], [ 395, 7155, 28963, 316, 31454, 69191, 379, 290, 1432, 326, 8808, 328, 290, 17994, 7578, 13], [ 25, 37330, 289, 15193, 67384, 326, 40202, 38086, 198, 3638, 4859, 25, 54315, 13700, 326, 13888]])torch.Size([4, 16])5. Embedding层可以看到这里每一行是16个字母的tokenize后的编号现在需要用高维特征来表示这些编号让机器能够学习。同时这里不能直接使用context_len来进行初始化nn.Embedding因为这里面的高维特征是和词汇表的tokenID进行绑定的所以不能简单的对每一个token进行初始化。那就必须初始化一个词汇表最大的nn.Embedding然后用token去里面找初始化的高维特征。max_token_idtokenized_txt.max().item()embeddingsnn.Embedding(max(tokenized_txt)1,d_model)print(f词汇表大小:{embeddings.num_embeddings})print(f嵌入维度:{embeddings.embedding_dim})print(f权重矩阵形状:{embeddings.weight.shape})x_batch_embeddingembeddings(x_batch)y_batch_embeddingembeddings(y_batch)print(x_batch_embedding.shape)print(y_batch_embedding.shape)输出词汇表大小: 199854嵌入维度: 64权重矩阵形状: torch.Size([199854, 64])torch.Size([4, 16, 64])torch.Size([4, 16, 64])6. 位置编码下面给这些编码后的样本加上位置编码注意这里的位置编码是一个常数在训练过程中没有任何变化的。同时这里初始化0矩阵的时候只需要按照训练数据的后两位进行初始化即可也就是[16,64]然后根据广播机制每一个样本都可以计算到。位置编码的公式P E ( p o s , 2 i ) sin ( p o s 10000 2 i / d m o d e l ) PE_{(pos, 2i)} \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)PE(pos,2i)sin(100002i/dmodelpos)P E ( p o s , 2 i 1 ) cos ( p o s 10000 2 i / d m o d e l ) PE_{(pos, 2i1)} \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)PE(pos,2i1)cos(100002i/dmodelpos)其中这个pos是字母的位置i的话则是特征维度的位置这里比较难写。位置除以10000的2i/d_model次方这里2i的意思是i只取一半比如d_model是32i就取16个。这里生成的i在公式中就对应了是2i。position_embeddingtorch.zeros(context_len,d_model)print(position_embedding.shape)positiontorch.arange(0,context_len,dtypetorch.float).unsqueeze(1)print(position.shape)itorch.arange(0,d_model,2)print(i.shape)temptorch.pow(10000,i/d_model)position_embedding[:,0::2]torch.sin(position/temp)position_embedding[:,1::2]torch.cos(position/temp)输出torch.Size([16, 64])torch.Size([16, 1])torch.Size([32])x_inputx_batch_embeddingposition_embedding y_inputy_batch_embeddingposition_embedding Xx_input x_plotx_input[0].detach().cpu().numpy()print(Final Input Embedding of x: \n,pd.DataFrame(x_plot).iloc[0])输出Final Input Embedding of x: 0 -0.391091 1 1.948731 2 0.029039 3 2.157305 4 2.434901 ... 59 0.673963 60 1.657022 61 0.792643 62 -0.466783 63 2.021278 Name: 0, Length: 64, dtype: float327. Transformer Block到这里就获得了所有的输入接下来就是写Transformer BlockQKVX Wqnn.Linear(d_model,d_model)Wknn.Linear(d_model,d_model)Wvnn.Linear(d_model,d_model)queryWq(Q)keyWk(K)valueWv(V)num_heads4queryquery.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)keykey.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)valuevalue.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)scoresquery key.transpose(-2,-1)head_dimd_model//num_heads scoresscores/head_dim**0.5#这里是进行右上角的maskmasktorch.triu(torch.ones(context_len,context_len),diagonal1).bool()outputscores.masked_fill(mask,float(-inf))print(pd.DataFrame(output[0,0].detach().numpy()))outputF.softmax(output,-1)print(fout.shape:{output.shape})print(fvalue.shape:{value.shape})#这里是多头注意力输出[batch,num_heads,context_len,d_model/num_heads]此时的output为[batch,num_heads,context_len,context_len]value为[batch,num_heads,context_len,d_model/num_heads]Aoutput valueprint(A.shape)#输出之后需要将多头进行合并因为这个多头是特征维度的多头所以需要交换维度的位置将多头的维度交换到特征维度前面一位然后将多头的特征进行合并。AA.permute(0,2,1,3).reshape(batch_size,context_len,d_model)print(A.shape)Wonn.Linear(d_model,d_model)outputWo(A)X layer_normnn.LayerNorm(d_model)output_layernormlayer_norm(output)ffnnn.Sequential(nn.Linear(d_model,d_model*4),nn.ReLU(),nn.Linear(d_model*4,d_model))outputffn(output_layernorm)output_layernorm outputlayer_norm(output)outputF.softmax(nn.Linear(d_model,max(tokenized_txt)1)(output),-1)输出mask后的注意力分数第0头0 1 2 3 4 5 6 \ 0 0.283212 -inf -inf -inf -inf -inf -inf 1 0.342469 0.494468 -inf -inf -inf -inf -inf 2 0.414351 0.233906 0.048258 -inf -inf -inf -inf 3 0.296115 0.058448 -0.185082 -0.677303 -inf -inf -inf 4 0.086106 0.157285 0.269732 0.160495 0.076710 -inf -inf 5 0.100804 0.211951 0.174968 0.223168 0.462680 0.286964 -inf 6 1.003936 0.645493 1.333420 0.527587 0.657099 0.925585 0.722007 7 -0.439468 -0.279724 0.248188 0.020643 0.054766 -0.325878 -0.214293 8 -0.138217 0.086135 -0.399000 -0.105449 -0.023271 0.362040 -0.663289 9 -0.369189 0.067475 0.257465 -0.4069...out.shape:torch.Size([4, 4, 16, 16])value.shape:torch.Size([4, 4, 16, 16])torch.Size([4, 4, 16, 16])torch.Size([4, 16, 64])8. Multi-Head Attention 逐行拆解下面把上面的代码按模块拆开讲解。8.1 Q/K/V 投影QKVX Wqnn.Linear(d_model,d_model)Wknn.Linear(d_model,d_model)Wvnn.Linear(d_model,d_model)queryWq(Q)keyWk(K)valueWv(V)在自注意力中Query、Key、Value 都来源于同一个输入X通过三个独立的线性层分别学到不同的角色变换。这里d_model64所以每个投影矩阵的维度都是(64, 64)。8.2 拆成 4 个头num_heads4queryquery.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)keykey.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)valuevalue.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)view把(batch, 16, 64)切成(batch, 16, 4, 16)再用permute(0,2,1,3)把 head 维度提到前面变成[batch, num_heads, context_len, head_dim]最终是[4, 4, 16, 16]。多头注意力的本质把特征维度切成 h 份每份独立计算注意力。8.3 注意力分数scoresquery key.transpose(-2,-1)head_dimd_model//num_heads scoresscores/head_dim**0.5Attention ( Q , K , V ) softmax ( Q K T d k ) V \text{Attention}(Q,K,V)\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)VAttention(Q,K,V)softmax(dkQKT)V除以d k 16 4 \sqrt{d_k}\sqrt{16}4dk164是为了防止点积结果过大导致 softmax 进入饱和区。8.4 因果掩码masktorch.triu(torch.ones(context_len,context_len),diagonal1).bool()outputscores.masked_fill(mask,float(-inf))torch.triu(diagonal1)生成上三角矩阵不含对角线将其设为-inf这样 softmax 后这些位置的权重为 0。这是 decoder-only 模型的核心每个位置只能看到它自己及之前的位置不能偷看未来。8.5 加权求和outputF.softmax(output,-1)Aoutput value经过 softmax 的注意力权重形状[4, 4, 16, 16]与 Value 相乘得到加权后的特征表示。8.6 合并多头AA.permute(0,2,1,3).reshape(batch_size,context_len,d_model)permute把维度交换回[batch, context_len, num_heads, head_dim]再用reshape把后两维合并恢复成[4, 16, 64]。注意这里合并的是特征维度上的多头而不是序列维度。8.7 输出投影 Wo 残差Wonn.Linear(d_model,d_model)outputWo(A)X layer_normnn.LayerNorm(d_model)output_layernormlayer_norm(output)Wo把多头合并后的结果再投影回d_model维加上原始输入X构成残差连接然后过 LayerNorm。8.8 FFN 输出层ffnnn.Sequential(nn.Linear(d_model,d_model*4),nn.ReLU(),nn.Linear(d_model*4,d_model))outputffn(output_layernorm)output_layernorm outputlayer_norm(output)outputF.softmax(nn.Linear(d_model,max(tokenized_txt)1)(output),-1)FFN 采用经典结构d_model → 4*d_model → d_model中间用 ReLU 激活。最后接一个 Linear 映射到词表大小199854再用 softmax 得到每个位置的 next-token 概率分布。 从零手写 Transformer：PyTorch 逐行实现与原理拆解
从零手写 TransformerPyTorch 逐行实现与原理拆解1. Transformer架构2. 环境准备与数据加载3. 创建词表与Tokenize4. 样本构造5. Embedding层6. 位置编码7. Transformer Block8. Multi-Head Attention 逐行拆解8.1 Q/K/V 投影8.2 拆成 4 个头8.3 注意力分数8.4 因果掩码8.5 加权求和8.6 合并多头8.7 输出投影 Wo 残差8.8 FFN 输出层本文记录一次从零手写 Transformerdecoder-onlyGPT 风格的完整过程从数据加载、tiktoken 分词、Embedding、手写位置编码到单层多头注意力与 FFN全部用 PyTorch 逐行实现并给出每一步的张量形状变化。代码与讲解一一对应适合刚接触 Transformer 想动手写一遍的同学。1. Transformer架构2. 环境准备与数据加载importtorchimporttorch.nnasnnimporttorch.nn.functionalasFimportosimportrequestsimportpandasaspd获取数据集ifnotos.path.exists(sales_textbook.txt):urlhttps://huggingface.co/datasets/goendalf666/sales-textbook_for_convincing_and_selling/resolve/main/sales_textbook.txt?downloadtruewithopen(sales_textbook.txt,wb)asf:f.write(requests.get(url).content)withopen(sales_textbook.txt,r)asf:textf.read()3. 创建词表与Tokenize创建词表这里直接使用现有的openai开源的tokenizer即可importtiktoken enctiktoken.get_encoding(o200k_base)这里的text是读取了全部的字符串然后进行tokenize之后是对每一个字母给了一个数字编号。可以看到text存储的这一段文字内容中字母、数字、标点符号转成编号之后最大的编号为199853tokenized_txtenc.encode(text)print(text[:7])print(tokenized_txt[:7])print(max(tokenized_txt))tokenized_txttorch.tensor(tokenized_txt,dtypetorch.long)train_txttokenized_txt[:int(0.9*len(tokenized_txt))]val_txttokenized_txt[int(0.9*len(tokenized_txt)):]输出Chapter[45990, 220, 16, 25, 22478, 158287, 326]1998534. 样本构造这里对全局需要有一个认识对于数据的选取训练数据和标签的选取。标签相对于训练索引在文本中是向后偏移了一位的因为需要用现有的数据对下一位进行预测。这里需要写切分的代码同时将切分后的样本按照batch进行拼接拼接之后的训练样本还不够这个时候还只是[batch,context_len]这里代表着只是把16个字母按照batch4排放在一起但是学习的时候机器并不认识这些字母所以需要将这16个字母用高维特征来表示也就是前面做的tokenize。batch_size4context_len16d_model64datatrain_txt idxstorch.randint(0,highlen(data)-context_len,size(batch_size,))x_batchtorch.stack([data[idx:idxcontext_len]foridxinidxs])y_batchtorch.stack([data[idx1:idxcontext_len1]foridxinidxs])print(x_batch)x_batch.shape输出tensor([[ 25222, 402, 2935, 1062, 6086, 3692, 4534, 71265, 2012, 5720, 11, 290, 172442, 625, 1606, 84275], [ 5892, 1989, 6028, 316, 1520, 261, 12398, 8660, 52125, 11, 472, 1023, 11747, 12486, 842, 402], [ 395, 7155, 28963, 316, 31454, 69191, 379, 290, 1432, 326, 8808, 328, 290, 17994, 7578, 13], [ 25, 37330, 289, 15193, 67384, 326, 40202, 38086, 198, 3638, 4859, 25, 54315, 13700, 326, 13888]])torch.Size([4, 16])5. Embedding层可以看到这里每一行是16个字母的tokenize后的编号现在需要用高维特征来表示这些编号让机器能够学习。同时这里不能直接使用context_len来进行初始化nn.Embedding因为这里面的高维特征是和词汇表的tokenID进行绑定的所以不能简单的对每一个token进行初始化。那就必须初始化一个词汇表最大的nn.Embedding然后用token去里面找初始化的高维特征。max_token_idtokenized_txt.max().item()embeddingsnn.Embedding(max(tokenized_txt)1,d_model)print(f词汇表大小:{embeddings.num_embeddings})print(f嵌入维度:{embeddings.embedding_dim})print(f权重矩阵形状:{embeddings.weight.shape})x_batch_embeddingembeddings(x_batch)y_batch_embeddingembeddings(y_batch)print(x_batch_embedding.shape)print(y_batch_embedding.shape)输出词汇表大小: 199854嵌入维度: 64权重矩阵形状: torch.Size([199854, 64])torch.Size([4, 16, 64])torch.Size([4, 16, 64])6. 位置编码下面给这些编码后的样本加上位置编码注意这里的位置编码是一个常数在训练过程中没有任何变化的。同时这里初始化0矩阵的时候只需要按照训练数据的后两位进行初始化即可也就是[16,64]然后根据广播机制每一个样本都可以计算到。位置编码的公式P E ( p o s , 2 i ) sin ( p o s 10000 2 i / d m o d e l ) PE_{(pos, 2i)} \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)PE(pos,2i)sin(100002i/dmodelpos)P E ( p o s , 2 i 1 ) cos ( p o s 10000 2 i / d m o d e l ) PE_{(pos, 2i1)} \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)PE(pos,2i1)cos(100002i/dmodelpos)其中这个pos是字母的位置i的话则是特征维度的位置这里比较难写。位置除以10000的2i/d_model次方这里2i的意思是i只取一半比如d_model是32i就取16个。这里生成的i在公式中就对应了是2i。position_embeddingtorch.zeros(context_len,d_model)print(position_embedding.shape)positiontorch.arange(0,context_len,dtypetorch.float).unsqueeze(1)print(position.shape)itorch.arange(0,d_model,2)print(i.shape)temptorch.pow(10000,i/d_model)position_embedding[:,0::2]torch.sin(position/temp)position_embedding[:,1::2]torch.cos(position/temp)输出torch.Size([16, 64])torch.Size([16, 1])torch.Size([32])x_inputx_batch_embeddingposition_embedding y_inputy_batch_embeddingposition_embedding Xx_input x_plotx_input[0].detach().cpu().numpy()print(Final Input Embedding of x: \n,pd.DataFrame(x_plot).iloc[0])输出Final Input Embedding of x: 0 -0.391091 1 1.948731 2 0.029039 3 2.157305 4 2.434901 ... 59 0.673963 60 1.657022 61 0.792643 62 -0.466783 63 2.021278 Name: 0, Length: 64, dtype: float327. Transformer Block到这里就获得了所有的输入接下来就是写Transformer BlockQKVX Wqnn.Linear(d_model,d_model)Wknn.Linear(d_model,d_model)Wvnn.Linear(d_model,d_model)queryWq(Q)keyWk(K)valueWv(V)num_heads4queryquery.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)keykey.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)valuevalue.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)scoresquery key.transpose(-2,-1)head_dimd_model//num_heads scoresscores/head_dim**0.5#这里是进行右上角的maskmasktorch.triu(torch.ones(context_len,context_len),diagonal1).bool()outputscores.masked_fill(mask,float(-inf))print(pd.DataFrame(output[0,0].detach().numpy()))outputF.softmax(output,-1)print(fout.shape:{output.shape})print(fvalue.shape:{value.shape})#这里是多头注意力输出[batch,num_heads,context_len,d_model/num_heads]此时的output为[batch,num_heads,context_len,context_len]value为[batch,num_heads,context_len,d_model/num_heads]Aoutput valueprint(A.shape)#输出之后需要将多头进行合并因为这个多头是特征维度的多头所以需要交换维度的位置将多头的维度交换到特征维度前面一位然后将多头的特征进行合并。AA.permute(0,2,1,3).reshape(batch_size,context_len,d_model)print(A.shape)Wonn.Linear(d_model,d_model)outputWo(A)X layer_normnn.LayerNorm(d_model)output_layernormlayer_norm(output)ffnnn.Sequential(nn.Linear(d_model,d_model*4),nn.ReLU(),nn.Linear(d_model*4,d_model))outputffn(output_layernorm)output_layernorm outputlayer_norm(output)outputF.softmax(nn.Linear(d_model,max(tokenized_txt)1)(output),-1)输出mask后的注意力分数第0头0 1 2 3 4 5 6 \ 0 0.283212 -inf -inf -inf -inf -inf -inf 1 0.342469 0.494468 -inf -inf -inf -inf -inf 2 0.414351 0.233906 0.048258 -inf -inf -inf -inf 3 0.296115 0.058448 -0.185082 -0.677303 -inf -inf -inf 4 0.086106 0.157285 0.269732 0.160495 0.076710 -inf -inf 5 0.100804 0.211951 0.174968 0.223168 0.462680 0.286964 -inf 6 1.003936 0.645493 1.333420 0.527587 0.657099 0.925585 0.722007 7 -0.439468 -0.279724 0.248188 0.020643 0.054766 -0.325878 -0.214293 8 -0.138217 0.086135 -0.399000 -0.105449 -0.023271 0.362040 -0.663289 9 -0.369189 0.067475 0.257465 -0.4069...out.shape:torch.Size([4, 4, 16, 16])value.shape:torch.Size([4, 4, 16, 16])torch.Size([4, 4, 16, 16])torch.Size([4, 16, 64])8. Multi-Head Attention 逐行拆解下面把上面的代码按模块拆开讲解。8.1 Q/K/V 投影QKVX Wqnn.Linear(d_model,d_model)Wknn.Linear(d_model,d_model)Wvnn.Linear(d_model,d_model)queryWq(Q)keyWk(K)valueWv(V)在自注意力中Query、Key、Value 都来源于同一个输入X通过三个独立的线性层分别学到不同的角色变换。这里d_model64所以每个投影矩阵的维度都是(64, 64)。8.2 拆成 4 个头num_heads4queryquery.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)keykey.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)valuevalue.view(batch_size,context_len,num_heads,d_model//num_heads).permute(0,2,1,3)view把(batch, 16, 64)切成(batch, 16, 4, 16)再用permute(0,2,1,3)把 head 维度提到前面变成[batch, num_heads, context_len, head_dim]最终是[4, 4, 16, 16]。多头注意力的本质把特征维度切成 h 份每份独立计算注意力。8.3 注意力分数scoresquery key.transpose(-2,-1)head_dimd_model//num_heads scoresscores/head_dim**0.5Attention ( Q , K , V ) softmax ( Q K T d k ) V \text{Attention}(Q,K,V)\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)VAttention(Q,K,V)softmax(dkQKT)V除以d k 16 4 \sqrt{d_k}\sqrt{16}4dk164是为了防止点积结果过大导致 softmax 进入饱和区。8.4 因果掩码masktorch.triu(torch.ones(context_len,context_len),diagonal1).bool()outputscores.masked_fill(mask,float(-inf))torch.triu(diagonal1)生成上三角矩阵不含对角线将其设为-inf这样 softmax 后这些位置的权重为 0。这是 decoder-only 模型的核心每个位置只能看到它自己及之前的位置不能偷看未来。8.5 加权求和outputF.softmax(output,-1)Aoutput value经过 softmax 的注意力权重形状[4, 4, 16, 16]与 Value 相乘得到加权后的特征表示。8.6 合并多头AA.permute(0,2,1,3).reshape(batch_size,context_len,d_model)permute把维度交换回[batch, context_len, num_heads, head_dim]再用reshape把后两维合并恢复成[4, 16, 64]。注意这里合并的是特征维度上的多头而不是序列维度。8.7 输出投影 Wo 残差Wonn.Linear(d_model,d_model)outputWo(A)X layer_normnn.LayerNorm(d_model)output_layernormlayer_norm(output)Wo把多头合并后的结果再投影回d_model维加上原始输入X构成残差连接然后过 LayerNorm。8.8 FFN 输出层ffnnn.Sequential(nn.Linear(d_model,d_model*4),nn.ReLU(),nn.Linear(d_model*4,d_model))outputffn(output_layernorm)output_layernorm outputlayer_norm(output)outputF.softmax(nn.Linear(d_model,max(tokenized_txt)1)(output),-1)FFN 采用经典结构d_model → 4*d_model → d_model中间用 ReLU 激活。最后接一个 Linear 映射到词表大小199854再用 softmax 得到每个位置的 next-token 概率分布。