跳到主要内容
基于 PyTorch 从零构建 Transformer 模型详解 | 极客日志
Python AI 算法
基于 PyTorch 从零构建 Transformer 模型详解 综述由AI生成 Transformer 架构通过注意力机制有效解决了传统 RNN 无法并行处理和长距离依赖的问题。基于 PyTorch 从零构建 Transformer 模型的完整流程，涵盖缩放点积注意力 (SDPA)、多头注意力机制、位置编码、编码器与解码器的具体实现，以及如何组装成完整的机器翻译模型。内容包含关键代码解析与数学原理说明，适合希望深入理解 Transformer 内部机制的开发者参考。
涅槃凤凰 发布于 2026/2/21 更新于 2026/4/25 前言
相较于传统的循环神经网络 (RNN) 和卷积神经网络 (CNN)，Transformer 的优势在于能够有效地理解输入和输出序列中元素之间的关系，尤其是在长距离依赖的情况下。与 RNN 不同，Transformer 能够并行训练，显著减少训练时间，并且能够处理大规模数据集。这种创新性的架构在大语言模型 (LLM) 如 ChatGPT、BERT 和 DeepSeek 的发展中起到了关键作用。
在 Transformer 模型之前，自然语言处理 (NLP) 主要依赖 RNN，其中包括长短期记忆 (LSTM) 网络。然而，RNN 按顺序处理信息，无法并行训练，限制了速度，并且在保持序列早期部分信息方面存在困难，难以捕捉长期依赖关系。
Transformer 架构的创新在于其注意力机制。注意力机制通过分配权重来评估序列中单词之间的关系，基于训练数据决定单词之间的语义相关性。这使得像 ChatGPT 这样的模型能够理解单词之间的关系，从而更有效地理解语言数据。为了从零开始构建 Transformer，我们首先探讨自注意力机制的工作原理，包括查询 (query)、键 (key) 和值 (value) 向量的作用，以及缩放点积注意力 (SDPA) 的计算。
注意力机制 和 Transformer
要理解机器学习中的 Transformer，首先必须理解注意力机制。该机制使 Transformer 能够识别序列元素之间的长程依赖性。通过注意力机制，Transformer 能够同时关注序列中的每个元素，理解每个单词的上下文。
以单词 "bank" 为例，说明注意力机制如何根据上下文解释词义。在句子 "I went fishing by the river yesterday, remaining near the bank the whole afternoon" 中，单词 "bank" 与 "fishing" 相关联，因为它指的是河岸的区域。相比之下，在句子 "Kate went to the bank after work yesterday and deposited a check there" 中，"bank" 与 "check" 相关联，使得 Transformer 将 "bank" 识别为金融机构。
注意力机制
注意力机制是一种用于确定序列中元素之间相互关系的方法，通过计算得分来表示一个元素与序列中其他元素的关系。在 NLP 中，这一机制有助于有意义地连接句子中的单词。
我们将构建一个由编码器和解码器组成的 Transformer 机器翻译模型，编码器将英语句子转化为捕捉其含义的向量表示，解码器则使用这些向量表示生成法语翻译。
为了将短语 "How are you?" 转换为向量表示，模型首先将其拆分为词元序列 [how, are, you, ?]。每个词元由一个 256 维的向量表示，称为词嵌入 (word embeddings)。编码器还使用位置编码 (positional encoding) 来确定词元在序列中的位置。将位置编码添加到词嵌入中，形成输入嵌入 (input embeddings)，用于计算自注意力 (self-attention)。
计算注意力有多种方法，本节将介绍最常见的方法——缩放点积注意力 (Scaled Dot Product Attention, SDPA)。该机制也称为自注意力 (self-attention)，因为算法计算一个单词如何关注序列中的所有单词，包括它自身。
在计算注意力时，查询 (query)、键 (key) 和值 (value) 的灵感来自于检索系统。假如，我们访问图书馆寻找一本书，在图书馆的搜索引擎中搜索'金融中的机器学习'，这个短语就是查询 (query)，图书馆中的书名和描述则充当键 (key)。根据查询与这些键之间的相似度，图书馆的检索系统会推荐一份书单 (值，value)。
为了计算 SDPA，输入嵌入 X 会通过三层不同的神经网络进行处理。这些层的相应权重为 W_Q、W_K 和 W_V，每个权重的维度为 256 × 256，这些权重在训练阶段从数据中学习。因此，可以通过以下公式计算查询 Q、键 K 和值 V：
Q = X × W_Q
K = X × W_K
V = X × W_V
其中，Q、K 和 V 的维度与输入嵌入 X 相匹配，均为 4 × 256。
类似于检索系统示例，在注意力机制中，使用 SDPA 方法来评估查询向量与键向量之间的相似度。SDPA 计算查询 Q 和键 K 向量的点积。较高的点积表示这两个向量之间的相似度较强，反之亦然。缩放后的注意力分数计算如下：
Attention Score (Q,K) = (Q · K^T) / √d_k
其中，d_k 表示键向量 K 的维度，在本节中为 256。通过除以 d_k 的平方根来缩放 Q 和 K 的点积，从而稳定训练过程，防止点积的值过大。
接下来，对这些注意力分数应用 softmax 函数，将它们转化为注意力权重，确保一个单词对句子中所有单词的总注意力之和为 100%。
上图展示注意力的计算过程。对于句子 "How are you?"，注意力权重形成一个 4×4 的矩阵，表示句子中每个词元与其他所有词元的关系。
最终的注意力输出通过注意力权重与值向量 V 的点积计算得出：
Attention (Q,K,V) = softmax((Q · K^T) / √d_k) · V
总结来说，整个过程从句子的输入嵌入 X 开始。注意力机制的最终输出是 Attention(Q, K, V)，该输出可以看作是对四个原始词元的上下文组合。原始词元的权重根据每个词元在上下文中的相关性变化，赋予在句子中更重要的词更高的权重。通过这一过程，注意力机制将表示孤立词元的向量转换为蕴含上下文意义的向量。
此外，Transformer 模型并不只使用一组查询、键和值向量，而是采用了多头注意力 (multihead attention) 机制。例如，256 维的查询、键和值向量可以被分割成 8 个注意力头，每个注意力头有一组维度为 32 的查询、键和值向量。每个注意力头关注输入的不同部分或方面，使得模型能够捕捉更广泛的信息。
Transformer 架构 注意力机制的概念于 2014 年提出，在《Attention Is All You Need》中得到了广泛应用。该论文专注于创建一种用于机器语言翻译的模型，称为 Transformer，它具有一个依赖于注意力机制的编码器 - 解码器结构。
以英法翻译为例。Transformer 的编码器将一个英语句子转换成表示其含义的向量表示。然后，Transformer 的解码器处理这些向量，并生成法语翻译。编码器的作用是捕捉原始英语句子的核心含义。
Transformer 中的编码器首先对英语和法语句子进行分词，采用子词分词 (subword tokenization) 技术。深度学习模型无法直接处理文本，因此在输入模型之前，词元会被索引为整数。这些词元通常会首先使用独热编码 (one-hot encoding) 表示，然后，通过词嵌入层将它们压缩成具有连续值的向量。
与循环神经网络按顺序处理数据不同，Transformer 并行处理输入数据。这种并行性提高了效率，但并不会自动识别输入的顺序。为了解决这个问题，Transformer 在输入嵌入中添加了位置编码 (positional encoding)，位置编码是分配给输入序列中每个位置的独特向量。
PositionalEncoding(pos, 2i) = sin(pos / n^(2i/d))
PositionalEncoding(pos, 2i+1) = cos(pos / n^(2i/d))
需要注意的是，每个词元的位置都由一个 256 维的向量唯一标识，这些向量值在训练过程中保持不变。在输入到注意力层之前，这些位置编码会被添加到序列的词嵌入中。
Transformer 模型的编码器如下图所示，由六个相同的层 (N = 6) 组成。每一层包含两个不同的子层，第一个子层是多头自注意力层，第二个子层是一个逐位置的全连接前馈网络。在模型架构中，每个子层都包含层归一化和残差连接。
Transformer 模型的解码器如下图所示，由六个相同的解码器层 (N = 6) 组成。每个解码器层包含三个子层：一个多头自注意力子层，一个在第一子层输出和编码器输出之间执行多头交叉注意力的子层，最后是一个前馈子层。
解码器的自注意力子层的一个关键特点是掩码机制。这种掩码机制防止模型访问序列中的未来位置，确保某个位置的预测只能依赖于之前已知的元素。这种顺序依赖性对于语言翻译或文本生成等任务至关重要。
解码过程从解码器接收一个法语输入短语开始。解码器以自回归方式运行，逐个生成输出序列的词元。在第一个时间步，它从表示句子开始的 "BOS" (句子开始词元) 开始。持续进行上述过程，解码器将每个新预测的词元添加到其输入序列中，用于后续的预测。翻译过程在解码器预测出 "EOS" (句子结束词元) 时结束。
不同类型的 Transformer Transformer 有三种类型：仅编码器 Transformer、仅解码器 Transformer 和编码器 - 解码器 Transformer。
仅编码器 Transformer 由 N 个相同的编码器层组成，能够将序列转换为抽象的连续向量表示。例如，BERT 是一个仅编码器 Transformer，可以用于文本分类。
仅解码器 Transformer 由 N 个相同的解码器层组成。例如，ChatGPT 是一个仅解码器 Transformer，可以根据提示生成文本。
编码器 - 解码器 Transformer 用于处理复杂任务，例如文本到图像生成或语音识别。编码器与解码器的结合使得模型能够有效理解复杂的输入，并生成复杂的输出。
构建编码器 本节将讨论如何构建 Transformer 中的编码器。具体而言，我们将深入探讨如何构建每个编码器层内的各个子层，并实现多头自注意力机制。
注意力机制 虽然存在多种注意力机制，本节中我们使用缩放点积注意力 (SDPA)。Transformer 模型并不是使用单一的查询、键和值向量，而是采用了多头注意力机制。将 256 维的查询、键和值向量分成 8 个头，每个头有一组 32 维的查询、键和值向量。每个头关注输入的不同部分或方面，使得模型能够捕捉更广泛的信息。
(1) 在模块 util.py 中定义 attention() 函数：
import torch
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
import numpy as np
from torch import nn
from copy import deepcopy
import math
def attention (query, key, value, mask=None , dropout=None ):
d_k = query.size(-1 )
scores = torch.matmul(query, key.transpose(-2 ,-1 )) / math.sqrt(d_k)
if mask is not None :
scores = scores.masked_fill(mask == 0 , -1e9 )
p_attn = nn.functional.softmax(scores, dim=-1 )
if dropout is not None :
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn
下图通过一个实例展示了多头注意力的工作原理。将位置编码添加到词嵌入之后，句子的嵌入是一个大小为 (1, 6, 256) 的张量。这些嵌入通过三个线性层获得查询 (Q)、键 (K) 和值 (V) 向量。然后将它们分成八个头，得到八组不同的 Q、K 和 V。将 attention() 函数应用于每一组，从而产生八个注意力输出。然后我们将这八个注意力输出拼接成一个单一的注意力输出。最后，拼接后的注意力通过另一个大小为 256 × 256 的线性层，生成最终输出。
(2) 使用 PyTorch 在 util.py 中实现多头注意力：
class MultiHeadedAttention (nn.Module):
def __init__ (self, h, d_model, dropout=0.1 ):
super ().__init__()
assert d_model % h == 0
self .d_k = d_model // h
self .h = h
self .linears = nn.ModuleList([deepcopy(nn.Linear(d_model, d_model)) for i in range (4 )])
self .attn = None
self .dropout = nn.Dropout(p=dropout)
def forward (self, query, key, value, mask=None ):
if mask is not None :
mask = mask.unsqueeze(1 )
nbatches = query.size(0 )
query, key, value = [l(x).view(nbatches, -1 , self .h, self .d_k).transpose(1 ,2 ) for l, x in zip (self .linears, (query, key, value))]
x, self .attn = attention(query, key, value, mask=mask, dropout=self .dropout)
x = x.transpose(1 ,2 ).contiguous().view(nbatches, -1 , self .h * self .d_k)
output = self .linears[-1 ](x)
return output
(3) 每个编码器层和解码器层还包含一个前馈子层，这是一个两层的全连接神经网络，旨在增强模型捕捉和学习训练数据集中复杂特征的能力。此外，神经网络独立处理每个嵌入，而不是将嵌入序列视为单一的向量。因此，我们通常称其为逐位置前馈网络。在 util.py 模块中定义 PositionwiseFeedForward() 类：
class PositionwiseFeedForward (nn.Module):
def __init__ (self, d_model, d_ff, dropout=0.1 ):
super ().__init__()
self .w_1 = nn.Linear(d_model, d_ff)
self .w_2 = nn.Linear(d_ff, d_model)
self .dropout = nn.Dropout(dropout)
def forward (self, x ):
h1 = self .w_1(x)
h2 = self .dropout(h1)
return self .w_2(h2)
PositionwiseFeedForward() 类定义了两个关键参数：d_ff (前馈层的维度) 和 d_model (模型维度的大小)。通常，d_ff 的值选择为 d_model 的四倍。在本节中，d_model 为 256，因此将 d_ff 设置为 256 x 4 = 1024。
创建编码器 (1) 为了创建编码器层，首先定义 EncoderLayer() 类和 SublayerConnection() 类：
class SublayerConnection (nn.Module):
def __init__ (self, size, dropout ):
super ().__init__()
self .norm = LayerNorm(size)
self .dropout = nn.Dropout(dropout)
def forward (self, x, sublayer ):
output = x + self .dropout(sublayer(self .norm(x)))
return output
class EncoderLayer (nn.Module):
def __init__ (self, size, self_attn, feed_forward, dropout ):
super ().__init__()
self .self_attn = self_attn
self .feed_forward = feed_forward
self .sublayer = nn.ModuleList([deepcopy(SublayerConnection(size, dropout)) for i in range (2 )])
self .size = size
def forward (self, x, mask ):
x = self .sublayer[0 ](x, lambda x: self .self_attn(x, x, x, mask))
output = self .sublayer[1 ](x, self .feed_forward)
return output
每个编码器层由两个不同的子层组成：一个是多头自注意力层；另一个是简单的逐位置全连接前馈网络。此外，这两个子层都包含了层归一化和残差连接，残差连接方法用于解决消失梯度问题。
(2) 层归一化在一定程度上类似于批归一化，将层中的观测标准化，使其均值为零，标准差为 1。在 util.py 中定义 LayerNorm() 类：
class LayerNorm (nn.Module):
def __init__ (self, features, eps=1e-6 ):
super ().__init__()
self .a_2 = nn.Parameter(torch.ones(features))
self .b_2 = nn.Parameter(torch.zeros(features))
self .eps = eps
def forward (self, x ):
mean = x.mean(-1 , keepdim=True )
std = x.std(-1 , keepdim=True )
x_zscore = (x - mean) / torch.sqrt(std ** 2 + self .eps)
output = self .a_2 * x_zscore + self .b_2
return output
(3) 通过堆叠六个编码器层创建编码器，在 util.py 模块中定义 Encoder() 类：
class Encoder (nn.Module):
def __init__ (self, layer, N ):
super ().__init__()
self .layers = nn.ModuleList([deepcopy(layer) for i in range (N)])
self .norm = LayerNorm(layer.size)
def forward (self, x, mask ):
for layer in self .layers:
x = layer(x, mask)
output = self .norm(x)
return output
其中，Encoder() 类定义了两个参数：layer，即在 EncoderLayer() 类所指定的编码器层；N，即编码器中的编码器层数量。
构建编码器 - 解码器 Transformer 接下来我们继续实现解码器。首先创建解码器层，然后，堆叠 N = 6 个相同的解码器层来形成一个解码器。
创建解码器层
多头自注意力层
来自第一个子层的输出与编码器输出之间的交叉注意力层
前馈网络
每个子层都包含层归一化和残差连接，此外，解码器部分的多头自注意力子层被掩码化，以防止当前位置关注后续的位置。
(1) 在 util.py 模块中定义 DecoderLayer() 类：
class DecoderLayer (nn.Module):
def __init__ (self, size, self_attn, src_attn, feed_forward, dropout ):
super ().__init__()
self .size = size
self .self_attn = self_attn
self .src_attn = src_attn
self .feed_forward = feed_forward
self .sublayer = nn.ModuleList([deepcopy(SublayerConnection(size, dropout)) for i in range (3 )])
def forward (self, x, memory, src_mask, tgt_mask ):
x = self .sublayer[0 ](x, lambda x: self .self_attn(x, x, x, tgt_mask))
x = self .sublayer[1 ](x, lambda x: self .src_attn(x, memory, memory, src_mask))
output = self .sublayer[2 ](x, self .feed_forward)
return output
为了说明解码器层的操作，我们继续使用上一小节的示例。解码器接受词元 ['BOS', 'comment', 'et', 'es-vous', '?']，以及来自编码器的输出 (memory) 来预测序列 ['comment', 'et', 'es-vous', '?', 'EOS']。
可以看到，掩码的下半部分被设置为 True，而掩码的上半部分被设置为 False。当这个掩码应用于注意力分数时，第一个词元在第一个时间步仅关注自身。随着过程的继续，解码器使用词元来预测下一个词元，注意力分数仅在这几个词元之间计算，从而有效地隐藏了之后的词元。
遵循这一过程，第一个子层生成的输出是一个大小为 (1, 5, 256) 的张量。这个输出随后输入到第二个子层。在第二个子层中，计算 x 与编码器部分输出 (称为 memory) 之间的交叉注意力。
第二个子层中的最终交叉注意力是通过注意力权重和值向量 V 的点积计算得到的。因此，第二个子层的输入和输出具有相同的维度 (1, 5, 256)。
创建编码器 - 解码器 Transformer (1) 解码器由 N = 6 个相同的解码器层组成。Decoder() 类在 util.py 模块中定义：
class Decoder (nn.Module):
def __init__ (self, layer, N ):
super ().__init__()
self .layers = nn.ModuleList([deepcopy(layer) for i in range (N)])
self .norm = LayerNorm(layer.size)
def forward (self, x, memory, src_mask, tgt_mask ):
for layer in self .layers:
x = layer(x, memory, src_mask, tgt_mask)
output = self .norm(x)
return output
(2) 为了创建编码器 - 解码器 Transformer，首先在 util.py 模块中定义 Transformer() 类：
class Transformer (nn.Module):
def __init__ (self, encoder, decoder, src_embed, tgt_embed, generator ):
super ().__init__()
self .encoder = encoder
self .decoder = decoder
self .src_embed = src_embed
self .tgt_embed = tgt_embed
self .generator = generator
def encode (self, src, src_mask ):
return self .encoder(self .src_embed(src), src_mask)
def decode (self, memory, src_mask, tgt, tgt_mask ):
return self .decoder(self .tgt_embed(tgt), memory, src_mask, tgt_mask)
def forward (self, src, tgt, src_mask, tgt_mask ):
memory = self .encode(src, src_mask)
output = self .decode(memory, src_mask, tgt, tgt_mask)
return output
Transformer() 类由五个关键组件构成：编码器、解码器、源语言嵌入、目标语言嵌入和生成器。
基于 Transformer 构建机器翻译模型 在本节中，我们将整合所有组件，创建一个可以在任意两种语言之间进行翻译的 Transformer 模型。
定义生成器 首先，我们在 util.py 模块中定义 Generator() 类，用于生成下一个词元的概率分布。其基本思路是为解码器附加一个输出头 (head)，用于下游任务。在本节示例中，下游任务是预测法语翻译中的下一个词元。
(1) 定义 Generator() 类，Generator() 类会为每个索引生成预测概率，这些索引对应于目标语言中的词元，这使得模型能够利用先前生成的词元和编码器的输出，以自回归的方式顺序地预测词元：
class Generator (nn.Module):
def __init__ (self, d_model, vocab ):
super ().__init__()
self .proj = nn.Linear(d_model, vocab)
def forward (self, x ):
out = self .proj(x)
probs = nn.functional.log_softmax(out, dim=-1 )
return probs
创建翻译模型 (1) 创建 Transformer 模型，用于在任意两种语言之间进行翻译。在 util.py 中定义 create_model() 函数实现此功能：
def create_model (src_vocab, tgt_vocab, N, d_model, d_ff, h, dropout=0.1 ):
attn = MultiHeadedAttention(h, d_model).to(DEVICE)
ff = PositionwiseFeedForward(d_model, d_ff, dropout).to(DEVICE)
pos = PositionalEncoding(d_model, dropout).to(DEVICE)
model = Transformer(
Encoder(EncoderLayer(d_model, deepcopy(attn), deepcopy(ff), dropout).to(DEVICE), N).to(DEVICE),
Decoder(DecoderLayer(d_model, deepcopy(attn), deepcopy(attn), deepcopy(ff), dropout).to(DEVICE), N).to(DEVICE),
nn.Sequential(Embeddings(d_model, src_vocab).to(DEVICE), deepcopy(pos)),
nn.Sequential(Embeddings(d_model, tgt_vocab).to(DEVICE), deepcopy(pos)),
Generator(d_model, tgt_vocab)).to(DEVICE)
for p in model.parameters():
if p.dim() > 1 :
nn.init.xavier_uniform_(p)
return model.to(DEVICE)
在 create_model() 函数中，使用 Encoder()、Decoder() 和 Generator() 类顺序地构建 Transformer() 类的五个关键组件。
小结
Transformer 是一种先进的深度学习模型，擅长处理序列到序列的预测任务。其优势在于能够有效理解输入和输出序列中元素之间的长距离关系。
Transformer 架构的创新在于其注意力机制。注意力机制通过分配权重来评估序列中单词之间的关系，基于训练数据确定单词之间的关联程度。
为了计算缩放点积注意力 (SDPA)，输入嵌入 X 通过三个不同的神经网络层进行处理：查询 (Q)、键 (K) 和值 (V)。我们可以按以下方式计算 Q、K 和 V：
Q = X · W_Q
K = X · W_K
V = X · W_V
SDPA 计算如下：
Attention (Q, K, V) = softmax((Q · K^T) / √d_k) · V
其中 d_k 表示键向量 K 的维度。对注意力分数应用 softmax 函数，将其转换为注意力权重，这确保了一个单词对句子中所有单词的总注意力之和为 100%。
Transformer 模型采用了多头注意力 (multihead attention) 机制，查询、键和值向量被拆分成多个头，每个头关注输入的不同部分或方面，使得模型能够捕捉更广泛的信息，并形成对输入数据更详细和更具上下文的理解。
相关免费在线工具 加密/解密文本 使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
RSA密钥对生成器 生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
Mermaid 预览与可视化编辑 基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
随机西班牙地址生成器 随机生成西班牙地址（支持马德里、加泰罗尼亚、安达卢西亚、瓦伦西亚筛选），支持数量快捷选择、显示全部与下载。 在线工具，随机西班牙地址生成器在线工具，online
Gemini 图片去水印 基于开源反向 Alpha 混合算法去除 Gemini/Nano Banana 图片水印，支持批量处理与下载。 在线工具，Gemini 图片去水印在线工具，online
curl 转代码 解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
