汤不热吧自2017年Google在论文《Attention Is All You Need》中提出Transformer架构以来，它已经彻底改变了自然语言处理乃至整个深度学习的格局。从BERT到GPT，从ViT到CLIP，几乎所有现代大型模型的核心都建立在Transformer之上。然而，许多开发者在使用Hugging Face等框架调用预训练模型时，对Transformer的内部机制仍然一知半解。本文将从零开始，用PyTorch逐层实现Transformer的核心组件，帮助你真正理解注意力机制、位置编码、多头注意力等关键概念。
一、Transformer整体架构概览
Transformer是一个基于自注意力机制的编码器-解码器架构。编码器负责将输入序列编码为连续表示，解码器则根据编码器的输出和已生成的token来生成目标序列。原始论文中的Transformer包含6个编码器层和6个解码器层的堆叠。
每个编码器层包含两个子层：多头自注意力层（Multi-Head Self-Attention）和前馈神经网络层（Feed-Forward Network）。每个子层都使用了残差连接和层归一化。解码器层则多了一个编码器-解码器交叉注意力层，并且在自注意力层中使用了掩码机制来防止看到未来的信息。
| 组件 | 编码器 | 解码器 |
|---|---|---|
| 自注意力 | 是（双向） | 是（带掩码，因果） |
| 交叉注意力 | 否 | 是 |
| 前馈网络 | 是 | 是 |
| 残差连接 | 是 | 是 |
| 层归一化 | 是 | 是 |
二、Self-Attention机制：核心思想与数学推导
自注意力机制是Transformer最核心的创新。其基本思想是：对于序列中的每个位置，计算它与其他所有位置的相关性（注意力权重），然后用这些权重对序列中所有位置的表示进行加权求和，得到该位置的新表示。这使得模型能够捕捉序列内部的长距离依赖关系，而不像RNN那样需要逐步传递信息。
2.1 Query、Key、Value的概念
自注意力机制借鉴了信息检索中Query-Key-Value的概念。对于输入序列中的每个token，我们将其投影为三个向量：Query（查询向量）、Key（键向量）和Value（值向量）。Query表示当前token想要查询的信息，Key表示该token可以被匹配的特征，Value则是该token实际承载的信息内容。
注意力分数的计算公式为：
其中d_k是Key向量的维度，除以sqrt(d_k)是为了缩放点积，防止在维度较大时点积值过大导致softmax梯度消失。
2.2 PyTorch实现Scaled Dot-Product Attention
下面我们从零实现缩放点积注意力函数：
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import torch.nn as nn
import torch.nn.functional as F
import math
class ScaledDotProductAttention(nn.Module):
def __init__(self):
super().__init__()
def forward(self, Q, K, V, mask=None):
"""
Args:
Q: (batch_size, num_heads, seq_len, d_k)
K: (batch_size, num_heads, seq_len, d_k)
V: (batch_size, num_heads, seq_len, d_v)
mask: (batch_size, 1, 1, seq_len) 或 None
Returns:
context: 加权后的上下文向量
attention_weights: 注意力权重
"""
d_k = K.size(-1)
# 计算注意力分数: Q @ K^T / sqrt(d_k)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
# 应用掩码（将需要屏蔽的位置设为负无穷）
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
# 计算注意力权重
attention_weights = F.softmax(scores, dim=-1)
# 加权求和
context = torch.matmul(attention_weights, V)
return context, attention_weights
这段代码虽然简短，但包含了几个关键细节。首先，我们使用torch.matmul进行高效的批量矩阵乘法。其次，mask操作通过masked_fill将需要屏蔽的位置替换为-1e9，这样经过softmax后这些位置的权重会趋近于零。最后，我们同时返回上下文向量和注意力权重，权重可以用于可视化或调试。
三、Multi-Head Attention：并行多视角注意力
单头注意力只能学习一种注意力模式，而多头注意力通过将Q、K、V分别投影到多个不同的子空间，让模型能够同时关注序列中不同位置和不同层面的信息。例如，一个头可能关注语法关系，另一个头关注语义相似性，还有的头关注位置邻近性。
具体做法是将维度为d_model的输入投影到h个头，每个头的维度为d_model/h，分别计算注意力后再拼接回来。计算公式为：
2
where head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)
3.1 完整的Multi-Head Attention实现
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def __init__(self, d_model, num_heads, dropout=0.1):
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
# Q, K, V 的线性投影层
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
# 输出投影层
self.W_o = nn.Linear(d_model, d_model)
# 缩放点积注意力
self.attention = ScaledDotProductAttention()
# Dropout层
self.dropout = nn.Dropout(dropout)
# 层归一化
self.layer_norm = nn.LayerNorm(d_model)
def forward(self, query, key, value, mask=None):
"""
Args:
query: (batch_size, seq_len, d_model)
key: (batch_size, seq_len, d_model)
value: (batch_size, seq_len, d_model)
mask: (batch_size, seq_len, seq_len) 或 None
"""
batch_size = query.size(0)
residual = query
# 1. 线性投影
Q = self.W_q(query) # (batch, seq_len, d_model)
K = self.W_k(key)
V = self.W_v(value)
# 2. 分割多头: (batch, seq_len, d_model) -> (batch, num_heads, seq_len, d_k)
Q = Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
K = K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
V = V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
# 3. 扩展mask维度以匹配多头
if mask is not None:
mask = mask.unsqueeze(1) # (batch, 1, seq_len, seq_len)
# 4. 计算注意力
context, attn_weights = self.attention(Q, K, V, mask)
# 5. 合并多头: (batch, num_heads, seq_len, d_k) -> (batch, seq_len, d_model)
context = context.transpose(1, 2).contiguous()
context = context.view(batch_size, -1, self.d_model)
# 6. 输出投影
output = self.W_o(context)
output = self.dropout(output)
# 7. 残差连接 + 层归一化
output = self.layer_norm(output + residual)
return output, attn_weights
这里有几个值得注意的实现细节。首先是view和transpose的组合操作：我们先将投影后的向量reshape为(batch, seq_len, num_heads, d_k)，然后通过transpose将num_heads维度移到前面，这样每个头就能独立进行矩阵乘法。其次是contiguous()的调用：transpose后张量在内存中不再是连续的，view操作需要连续张量，所以必须先调用contiguous()。最后，残差连接和层归一化被内置在模块中，这就是Post-LN结构。
四、位置编码：让模型理解序列顺序
由于自注意力机制本身是排列不变的——交换输入序列中任意两个token的位置，输出只是相应地交换，值不会改变——Transformer需要额外的位置信息来感知序列的顺序。原始论文使用正弦和余弦函数生成位置编码：
2
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
使用正弦余弦函数有几个优势：首先，对于任意偏移k，PE(pos+k)可以表示为PE(pos)的线性函数，这使得模型容易学习到相对位置关系。其次，不同频率的正弦余弦函数在不同维度上编码了不同尺度的位置信息，使模型能够同时关注局部和全局的位置模式。
4.1 位置编码的实现
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, d_model, max_len=5000, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
# 创建位置编码矩阵
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
# 计算频率项: 1 / 10000^(2i/d_model)
div_term = torch.exp(
torch.arange(0, d_model, 2, dtype=torch.float) *
(-math.log(10000.0) / d_model)
)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
# 增加batch维度: (max_len, d_model) -> (1, max_len, d_model)
pe = pe.unsqueeze(0)
self.register_buffer('pe', pe)
def forward(self, x):
"""
Args:
x: (batch_size, seq_len, d_model)
"""
# 将位置编码加到输入上
x = x + self.pe[:, :x.size(1)]
return self.dropout(x)
注意div_term的计算技巧：我们使用torch.exp和-log(10000)来替代直接的幂运算，这在数值上更稳定。另外，register_buffer将位置编码注册为模型的一部分，但它不是可学习参数，不会在反向传播中更新，并且会随模型一起移动到GPU。
除了正弦位置编码，现代Transformer模型还使用了多种其他位置编码方案：
- 可学习位置编码：将位置编码作为可学习参数，如GPT系列所用，简单但不可外推到训练时的最大长度。
- 相对位置编码：如T5使用的相对位置偏置，直接在注意力分数中加入位置偏移信息。
- 旋转位置编码（RoPE）：如LLaMA所使用，通过旋转变换将位置信息编码到Q和K中，具有良好的外推性。
- ALiBi：在注意力分数中直接加入与距离成正比的偏置项，无需位置嵌入。
五、前馈网络与编码器层组装
每个注意力子层之后是一个逐位置的前馈神经网络（FFN），它由两个线性变换和一个ReLU激活函数组成：
虽然FFN对每个位置独立处理，但它通过将维度先扩展到d_ff（通常是d_model的4倍）再压缩回来，增加了模型的非线性表达能力。下面是实现：
2
3
4
5
6
7
8
9
10
11
12
13
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)
self.layer_norm = nn.LayerNorm(d_model)
def forward(self, x):
residual = x
x = self.w_2(self.dropout(F.relu(self.w_1(x))))
x = self.layer_norm(x + residual)
return x
现代模型如GPT-4和LLaMA已经用SwiGLU或GeGLU等门控激活函数替代了ReLU，效果有显著提升。SwiGLU的公式为：FFN(x) = (Swish(xW) * xV)，其中Swish(x) = x * sigmoid(x)。下面是SwiGLU的实现：
2
3
4
5
6
7
8
9
10
11
12
def __init__(self, d_model, d_ff, dropout=0.1):
super().__init__()
# SwiGLU需要两个投影矩阵，所以d_ff通常设为 (2/3) * 4 * d_model
self.w_gate = nn.Linear(d_model, d_ff, bias=False)
self.w_up = nn.Linear(d_model, d_ff, bias=False)
self.w_down = nn.Linear(d_ff, d_model, bias=False)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
gate = F.silu(self.w_gate(x)) # SiLU = Swish
return self.w_down(self.dropout(gate * self.w_up(x)))
5.1 组装完整的编码器层
现在我们可以将多头注意力、前馈网络和层归一化组装成一个完整的编码器层：
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
super().__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
def forward(self, x, mask=None):
# 子层1: 多头自注意力 + 残差 + LayerNorm
attn_output, attn_weights = self.self_attn(x, x, x, mask)
# 子层2: 前馈网络 + 残差 + LayerNorm
ff_output = self.feed_forward(attn_output)
return ff_output, attn_weights
class TransformerEncoder(nn.Module):
def __init__(self, vocab_size, d_model, num_heads, d_ff,
num_layers, max_len=5000, dropout=0.1):
super().__init__()
self.d_model = d_model
# 词嵌入层
self.token_embedding = nn.Embedding(vocab_size, d_model)
self.positional_encoding = PositionalEncoding(d_model, max_len, dropout)
# 编码器层堆叠
self.layers = nn.ModuleList([
EncoderLayer(d_model, num_heads, d_ff, dropout)
for _ in range(num_layers)
])
self.layer_norm = nn.LayerNorm(d_model)
def forward(self, x, mask=None):
# 词嵌入 + 缩放 + 位置编码
x = self.token_embedding(x) * math.sqrt(self.d_model)
x = self.positional_encoding(x)
# 逐层传递
attn_weights_all = []
for layer in self.layers:
x, attn_weights = layer(x, mask)
attn_weights_all.append(attn_weights)
x = self.layer_norm(x)
return x, attn_weights_all
注意词嵌入后面乘以sqrt(d_model)的操作：这是原始论文中的技巧，目的是让嵌入向量的尺度与位置编码的尺度相匹配，防止位置编码在相加时主导信号。
六、掩码机制详解
Transformer中使用两种掩码：Padding Mask和Sequence Mask（因果掩码）。Padding Mask用于处理变长序列，将padding位置的注意力权重设为负无穷，使模型不去关注padding token。Sequence Mask用于解码器的自注意力层，确保在生成第i个token时只能看到位置0到i-1的token，防止信息泄露。
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
"""
创建padding mask
Args:
seq: (batch_size, seq_len) 输入token id序列
Returns:
mask: (batch_size, 1, 1, seq_len) 1表示有效位置，0表示padding
"""
mask = (seq != pad_idx).unsqueeze(1).unsqueeze(2)
return mask # (batch, 1, 1, seq_len)
def create_causal_mask(seq_len):
"""
创建因果（上三角）掩码
Returns:
mask: (seq_len, seq_len) 上三角为0，下三角和对角线为1
"""
mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).unsqueeze(0)
return mask # (1, 1, seq_len, seq_len)
def create_combined_mask(seq, pad_idx=0):
"""组合padding mask和causal mask"""
pad_mask = create_padding_mask(seq, pad_idx) # (batch, 1, 1, seq_len)
causal_mask = create_causal_mask(seq.size(1)) # (1, 1, seq_len, seq_len)
combined = pad_mask & causal_mask
return combined
因果掩码使用torch.tril生成下三角矩阵——对角线及以下的元素为1，上方为0。这与上三角填充负无穷配合，确保当前位置只能注意到过去和自身的位置。结合padding mask时，用按位与操作合并两个掩码。
七、完整模型训练示例
下面我们用一个简单的文本分类任务来演示如何使用上面实现的Transformer编码器。我们将构建一个情感分析模型，输入一段文本，输出正面或负面的分类：
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(self, vocab_size, d_model, num_heads, d_ff,
num_layers, num_classes, max_len=512, dropout=0.1):
super().__init__()
self.encoder = TransformerEncoder(
vocab_size, d_model, num_heads, d_ff,
num_layers, max_len, dropout
)
# 分类头
self.classifier = nn.Linear(d_model, num_classes)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
# 编码器输出: (batch, seq_len, d_model)
encoded, attn_weights = self.encoder(x, mask)
# 取序列的第一个token（类似[CLS]）作为分类表示
cls_output = encoded[:, 0] # (batch, d_model)
cls_output = self.dropout(cls_output)
# 分类
logits = self.classifier(cls_output)
return logits, attn_weights
# 模型实例化与训练循环
model = TransformerClassifier(
vocab_size=30000,
d_model=256,
num_heads=8,
d_ff=1024,
num_layers=6,
num_classes=2,
max_len=512,
dropout=0.1
).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10000)
criterion = nn.CrossEntropyLoss()
# 训练循环
for epoch in range(num_epochs):
model.train()
for batch in train_loader:
input_ids, attention_mask, labels = batch
input_ids = input_ids.cuda()
attention_mask = attention_mask.cuda()
labels = labels.cuda()
# 创建padding mask
pad_mask = create_padding_mask(input_ids, pad_idx=0)
# 前向传播
logits, _ = model(input_ids, pad_mask)
# 计算损失
loss = criterion(logits, labels)
# 反向传播
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
scheduler.step()
print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")
在训练循环中，我们使用了梯度裁剪（clip_grad_norm_）来防止梯度爆炸，这在训练深层Transformer时尤为重要。AdamW优化器配合余弦退火学习率调度是训练Transformer的标准选择，权重衰减设为0.01有助于正则化。
八、Pre-LN vs Post-LN架构选择
原始Transformer论文使用Post-LN结构，即先计算子层输出再做残差连接和层归一化。然而Post-LN在深层网络中训练不稳定，需要学习率预热（warmup）策略。Pre-LN结构将层归一化移到子层之前，训练更稳定，被GPT-2、LLaMA等现代模型广泛采用。两种结构对比如下：
2
3
4
5
output = LayerNorm(x + Sublayer(x))
# Pre-LN (现代模型常用)
output = x + Sublayer(LayerNorm(x))
Pre-LN的关键优势在于残差路径上没有层归一化，梯度可以更通畅地流回输入，使得深层网络训练更稳定。但Post-LN在最终输出上有层归一化，理论上能产生更规范的表示。实践中，当你需要训练深层Transformer（超过12层）时，建议使用Pre-LN。下面是Pre-LN版本的编码器层：
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def __init__(self, d_model, num_heads, dropout=0.1):
super().__init__()
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
self.attention = ScaledDotProductAttention()
self.dropout = nn.Dropout(dropout)
self.layer_norm = nn.LayerNorm(d_model)
def forward(self, query, key, value, mask=None):
batch_size = query.size(0)
residual = query
# Pre-LN: 先归一化再计算
query = self.layer_norm(query)
key = self.layer_norm(key)
value = self.layer_norm(value)
Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
K = self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
V = self.W_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
if mask is not None:
mask = mask.unsqueeze(1)
context, attn_weights = self.attention(Q, K, V, mask)
context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
output = self.W_o(context)
output = self.dropout(output)
# 残差连接（无额外归一化）
output = output + residual
return output, attn_weights
九、训练优化技巧总结
训练Transformer模型时，除了架构选择，还有许多优化技巧可以显著影响最终效果：
- 学习率预热：在训练初期逐步增加学习率，避免随机初始化导致的梯度震荡。典型的预热步数为4000步，预热后可以切换到余弦退火或线性衰减。
- 梯度裁剪：将梯度范数限制在1.0以内，防止训练发散。
- Dropout策略：在注意力权重、FFN内部和残差连接上分别使用Dropout，通常0.1到0.3之间。
- 权重初始化：使用Xavier初始化或Kaiming初始化，避免初始梯度过大或过小。
- 标签平滑：将one-hot标签平滑为软标签，防止模型过度自信，通常平滑因子设为0.1。
- 混合精度训练：使用torch.cuda.amp进行半精度训练，节省显存并加速计算。
混合精度训练的示例代码：
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
scaler = GradScaler()
for epoch in range(num_epochs):
for batch in train_loader:
input_ids, attention_mask, labels = batch
input_ids, attention_mask, labels = input_ids.cuda(), attention_mask.cuda(), labels.cuda()
optimizer.zero_grad()
with autocast():
pad_mask = create_padding_mask(input_ids, pad_idx=0)
logits, _ = model(input_ids, pad_mask)
loss = criterion(logits, labels)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
十、从原理到实践的总结
通过本文的逐层实现，我们从Self-Attention的基础运算出发，构建了完整的Multi-Head Attention、位置编码、前馈网络、编码器层，并最终组装了一个可训练的Transformer分类模型。关键要点总结如下：
- 自注意力机制通过Q、K、V的点积运算捕捉序列内部的依赖关系，是Transformer的核心创新。
- 多头注意力将输入投影到多个子空间并行计算注意力，让模型能同时关注不同层面的信息。
- 位置编码补偿了注意力机制的排列不变性，正弦编码具有良好的外推性，而RoPE和ALiBi等现代方案更适合长序列。
- Pre-LN结构比Post-LN更适合训练深层Transformer，是现代大模型的默认选择。
- 掩码机制分为padding mask和causal mask，分别用于处理变长序列和因果生成约束。
- 训练优化方面，学习率预热、梯度裁剪、混合精度等技巧对于稳定训练至关重要。
理解Transformer的内部实现对于使用和改进现代深度学习模型至关重要。当你掌握了这些基础知识后，再阅读BERT、GPT、LLaMA等模型的源码就会变得水到渠成。建议读者在本文代码基础上，尝试实现完整的解码器层、加入交叉注意力、并构建一个完整的序列到序列翻译模型，以巩固对这些概念的理解。