第2课：Transformer架构详解
引言
欢迎来到《从零构建大型语言模型：Python实现20亿参数LLM的完整指南》的第二课。在上一课中，我们回顾了大型语言模型的发展历程，比较了当前市场上的主流模型，并勾勒出了我们的课程项目：构建一个20亿参数的LLM。
今天，我们将深入探讨现代LLM的核心基础——Transformer架构。自2017年Google研究团队在论文《Attention Is All You Need》中提出以来，Transformer架构已经成为NLP领域的主导范式，彻底改变了机器学习对序列数据的处理方式。无论是OpenAI的GPT系列，还是我们将要构建的模型，都以这一架构为基础。
在本课中，我们将分解Transformer的各个核心组件，解析其工作原理，理解其数学基础，为我们从零构建LLM奠定坚实的理论基础。让我们开始这段探索之旅！
1. Transformer架构概览
传统序列模型的局限性
在Transformer出现之前，处理序列数据（如文本）的主流模型是循环神经网络（RNN）及其变种LSTM和GRU。这些模型通过维持一个随时间步更新的隐藏状态来捕捉序列信息，但它们存在几个关键限制：
- 长距离依赖问题：随着序列长度增加，早期信息容易"丢失"
- 并行计算受限：由于顺序处理的本质，难以有效利用现代GPU并行计算能力
- 梯度消失/爆炸：尽管LSTM/GRU在一定程度上缓解了这个问题，但长序列处理时仍然存在挑战
Transformer的革命性创新
Transformer架构通过几个创新性设计解决了上述问题：
- 完全抛弃递归结构：不依赖顺序处理，使得并行计算成为可能
- 引入注意力机制：允许模型直接关注输入序列中的任何位置，有效捕捉长距离依赖
- 结合位置编码：虽然没有固有的序列顺序感，但通过位置编码保留了位置信息
Transformer整体架构
标准的Transformer架构包含编码器（Encoder）和解码器（Decoder）两部分，形成了一个经典的序列到序列（Seq2Seq）结构：
- 编码器：将输入序列转换为连续表示（潜在空间表示）
- 解码器：根据编码器的输出和之前生成的标记预测下一个标记
然而，在现代LLM中，我们通常只使用架构的一部分：
- 编码器模型（如BERT）：专注于理解输入文本，适用于分类、情感分析等任务
- 解码器模型（如GPT）：专注于生成文本，适用于文本生成、对话等任务
- 编码器-解码器模型（如T5）：适用于翻译、摘要等转换任务
在我们的LLM项目中，我们将重点关注解码器架构，这是现代文本生成模型如GPT系列的基础。
2. 注意力机制：Transformer的核心
为什么需要注意力机制？
注意力机制的核心思想是允许模型在处理当前标记时，有选择地"关注"输入序列中的相关部分。这种机制模拟了人类认知中的选择性注意，即我们能够从复杂信息中筛选出最相关的部分。
自注意力机制详解
自注意力（Self-attention）是Transformer中最关键的组件，也称为缩放点积注意力（Scaled Dot-Product Attention）。其计算过程可分为以下步骤：
-
线性投影：将输入向量转换为查询（Query）、键（Key）和值（Value）三种表示：
- Query (Q): 表示"我想要寻找什么信息"
- Key (K): 表示"我包含什么信息"
- Value (V): 表示"我的实际内容是什么"
-
计算注意力分数：通过Query和Key的点积衡量相似度：
Attention Scores = (Q × K^T) / √d_k其中，√d_k是缩放因子，防止点积在高维空间中变得过大。
-
应用Softmax：将分数转换为概率分布：
Attention Weights = Softmax(Attention Scores) -
加权聚合：根据注意力权重聚合Value向量：
Output = Attention Weights × V
整个自注意力机制可以用一个公式表示：
Attention(Q, K, V) = Softmax((Q × K^T) / √d_k) × V
代码实现示例
下面是自注意力机制的PyTorch简化实现：
import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module):
def __init__(self, embed_size, heads):
super(SelfAttention, self).__init__()
self.embed_size = embed_size
self.heads = heads
self.head_dim = embed_size // heads
assert (self.head_dim * heads == embed_size), "Embedding size needs to be divisible by heads"
# 创建查询、键、值和输出的线性层
self.q = nn.Linear(embed_size, embed_size)
self.k = nn.Linear(embed_size, embed_size)
self.v = nn.Linear(embed_size, embed_size)
self.fc_out = nn.Linear(embed_size, embed_size)
def forward(self, values, keys, query, mask=None):
N = query.shape[0] # 批次大小
value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]
# 线性投影
values = self.v(values) # (N, value_len, embed_size)
keys = self.k(keys) # (N, key_len, embed_size)
queries = self.q(query) # (N, query_len, embed_size)
# 分割成多头
values = values.reshape(N, value_len, self.heads, self.head_dim)
keys = keys.reshape(N, key_len, self.heads, self.head_dim)
queries = queries.reshape(N, query_len, self.heads, self.head_dim)
# 计算注意力分数
energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
# energy 维度: (N, heads, query_len, key_len)
# 应用掩码（如果提供）
if mask is not None:
energy = energy.masked_fill(mask == 0, float("-1e20"))
# 缩放注意力分数
energy = energy / (self.head_dim ** (1/2))
# 应用softmax
attention = torch.softmax(energy, dim=3) # (N, heads, query_len, key_len)
# 加权聚合
out = torch.einsum("nhql,nlhd->nqhd", [attention, values])
# out 维度: (N, query_len, heads, head_dim)
# 合并多头结果
out = out.reshape(N, query_len, self.heads * self.head_dim)
# out 维度: (N, query_len, embed_size)
# 最后的线性层
out = self.fc_out(out) # (N, query_len, embed_size)
return out
多头注意力机制
为了增强模型的表示能力，Transformer使用了多头注意力（Multi-Head Attention）。这一机制允许模型同时关注不同位置的不同表示子空间，捕捉更丰富的信息：
- 将输入分割成多个"头"（head）
- 每个头独立进行自注意力计算
- 合并各头的输出结果并通过线性变换
多头注意力可以表示为：
MultiHeadAttention(Q, K, V) = Concat(head_1, head_2, ..., head_h) × W^O
其中：
head_i = Attention(Q × W_i^Q, K × W_i^K, V × W_i^V)
实践中，这种机制能够显著提升模型性能，GPT-3使用96个注意力头，而我们的20亿参数模型可能会使用20-32个注意力头。
3. 编码器与解码器的工作原理
编码器结构
Transformer编码器由多个相同层堆叠而成，每层包含两个子层：
- 多头自注意力层：允许编码器关注输入序列中的不同部分
- 前馈神经网络：对每个位置独立应用相同的全连接网络
每个子层都使用残差连接（Residual Connection）和层归一化（Layer Normalization）：
LayerNorm(x + Sublayer(x))
编码器的主要特点是每个位置可以看到完整的输入序列（双向注意力），这使其特别适合理解任务。
解码器结构
Transformer解码器同样由多个相同层堆叠而成，但每层包含三个子层：
- 掩码多头自注意力层：只允许关注当前位置及其之前的位置（防止信息泄露）
- 编码器-解码器注意力层：关注编码器的输出
- 前馈神经网络：与编码器中的相同
解码器中的掩码机制至关重要，它确保模型在预测第i个标记时，只能看到第i-1个及之前的标记，这与自回归生成的本质一致。
解码器模型（如GPT）
在现代LLM中，尤其是像GPT这样的自回归生成模型，我们主要使用Transformer解码器（或其变体）。这些模型去除了编码器-解码器注意力层，只保留了掩码自注意力和前馈网络。
GPT架构的关键特点：
- 只使用解码器组件（掩码自注意力）
- 使用因果掩码（causal mask）确保自回归属性
- 通常堆叠更多层以增加模型容量
我们的20亿参数模型将遵循类似架构，使用24-36个Transformer解码器层。
4. 前馈神经网络
在每个注意力层之后，Transformer使用一个简单但强大的前馈神经网络（Feed-Forward Network，FFN）处理每个位置的表示。
结构与功能
FFN在每个位置上独立应用，由两个线性变换和一个非线性激活函数组成：
FFN(x) = max(0, x × W_1 + b_1) × W_2 + b_2
这里使用ReLU作为激活函数，但在现代实现中，我们通常使用GELU或SiLU（Swish）等更平滑的激活函数。
FFN的作用是引入非线性并增强模型的表示能力。虽然结构简单，但其参数量通常是模型总参数的大部分，因为内部维度通常设置为嵌入维度的4倍。
代码实现示例
class FeedForward(nn.Module):
def __init__(self, embed_size, ff_hidden_size, dropout=0.1):
super(FeedForward, self).__init__()
self.fc1 = nn.Linear(embed_size, ff_hidden_size)
self.fc2 = nn.Linear(ff_hidden_size, embed_size)
self.gelu = nn.GELU()
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# x shape: (batch_size, seq_len, embed_size)
x = self.fc1(x) # (batch_size, seq_len, ff_hidden_size)
x = self.gelu(x)
x = self.dropout(x)
x = self.fc2(x) # (batch_size, seq_len, embed_size)
return x
5. 位置编码与上下文理解
位置信息的重要性
Transformer架构本身不包含序列顺序的内置感知，因为自注意力对所有位置一视同仁。然而，序列中标记的位置对于理解语言至关重要（如"狗咬人"和"人咬狗"的意思完全不同）。
正弦位置编码
原始Transformer论文使用了正弦和余弦函数构造的位置编码：
PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
其中：
- pos是标记在序列中的位置
- i是维度索引
- d_model是模型的嵌入维度
这种编码有几个有趣的属性：
- 每个位置有唯一的编码
- 编码之间的相对关系可以通过线性投影学习
- 理论上可以外推到未见过的序列长度
代码实现示例
import numpy as np
def get_positional_encoding(max_seq_len, embed_size):
positional_encoding = np.zeros((max_seq_len, embed_size))
for pos in range(max_seq_len):
for i in range(0, embed_size, 2):
# 为偶数索引使用正弦
positional_encoding[pos, i] = np.sin(pos / (10000 ** (i / embed_size)))
# 为奇数索引使用余弦（如果在维度范围内）
if i + 1 < embed_size:
positional_encoding[pos, i + 1] = np.cos(pos / (10000 ** (i / embed_size)))
# 转换为PyTorch张量
return torch.FloatTensor(positional_encoding)
可学习位置编码
虽然固定的正弦位置编码在原始Transformer中表现良好，但许多现代实现（包括GPT系列）使用可学习的位置嵌入，允许模型自己学习最优的位置表示：
class LearnablePositionalEncoding(nn.Module):
def __init__(self, max_seq_len, embed_size):
super(LearnablePositionalEncoding, self).__init__()
self.positional_embedding = nn.Parameter(torch.zeros(max_seq_len, embed_size))
nn.init.normal_(self.positional_embedding, std=0.02)
def forward(self, x):
# x shape: (batch_size, seq_len, embed_size)
seq_len = x.size(1)
return x + self.positional_embedding[:seq_len, :]
相对位置编码
最近的研究表明，相对位置编码（RPE）可能比绝对位置编码更有效。RPE直接在注意力计算中编码位置关系，而不是在输入嵌入中添加位置信息。
T5、DeBERTa和PaLM等模型采用了各种相对位置编码方案，它们通常能够捕捉更长距离的依赖关系，并在长文本处理上表现更佳。
位置编码与上下文窗口
位置编码方案直接影响模型的上下文窗口大小（即模型能"看到"多远的上下文）。GPT-2有1024标记的上下文窗口，GPT-3扩展到2048，而最新的模型如Claude和GPT-4甚至支持超过10万标记的上下文窗口。
在我们的20亿参数模型中，我们将使用2048的上下文窗口，并实现可学习的位置编码，以平衡性能和计算需求。
6. 组合完整的Transformer层
现在，我们已经理解了Transformer的核心组件，让我们看看如何将它们组合成一个完整的Transformer层，特别是针对解码器模型。
解码器层结构
一个完整的GPT风格的解码器层包括：
- 掩码多头自注意力
- 残差连接和层归一化
- 前馈神经网络
- 另一个残差连接和层归一化
代码实现
class DecoderLayer(nn.Module):
def __init__(self, embed_size, num_heads, ff_hidden_size, dropout=0.1):
super(DecoderLayer, self).__init__()
self.attention = SelfAttention(embed_size, num_heads)
self.norm1 = nn.LayerNorm(embed_size)
self.ff = FeedForward(embed_size, ff_hidden_size, dropout)
self.norm2 = nn.LayerNorm(embed_size)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
# 自注意力块
attention_output = self.attention(x, x, x, mask)
x = self.norm1(x + self.dropout(attention_output))
# 前馈网络块
ff_output = self.ff(x)
x = self.norm2(x + self.dropout(ff_output))
return x
完整的GPT模型框架
最后，我们可以组合多个解码器层构建一个完整的GPT风格模型：
class GPTModel(nn.Module):
def __init__(
self,
vocab_size,
embed_size,
num_layers,
num_heads,
ff_hidden_size,
max_seq_len,
dropout=0.1
):
super(GPTModel, self).__init__()
# 标记嵌入
self.token_embedding = nn.Embedding(vocab_size, embed_size)
# 位置编码
self.position_embedding = nn.Parameter(torch.zeros(max_seq_len, embed_size))
# 解码器层
self.layers = nn.ModuleList(
[DecoderLayer(embed_size, num_heads, ff_hidden_size, dropout)
for _ in range(num_layers)]
)
# 输出层
self.norm = nn.LayerNorm(embed_size)
self.fc_out = nn.Linear(embed_size, vocab_size, bias=False)
# 权重绑定（可选，但GPT模型通常使用）
self.fc_out.weight = self.token_embedding.weight
# 初始化
self.apply(self._init_weights)
def _init_weights(self, module):
if isinstance(module, (nn.Linear, nn.Embedding)):
module.weight.data.normal_(mean=0.0, std=0.02)
if isinstance(module, nn.Linear) and module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.LayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
def get_causal_mask(self, seq_len):
# 创建因果掩码（下三角矩阵）
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
return ~mask # 反转使得1表示可见，0表示掩码
def forward(self, x):
batch_size, seq_len = x.shape
# 获取嵌入
token_embeds = self.token_embedding(x) # (batch_size, seq_len, embed_size)
# 添加位置嵌入
position_embeds = self.position_embedding[:seq_len, :]
x = token_embeds + position_embeds
# 创建因果掩码
mask = self.get_causal_mask(seq_len).to(x.device)
# 应用解码器层
for layer in self.layers:
x = layer(x, mask)
# 应用最终层归一化
x = self.norm(x)
# 预测下一个标记
logits = self.fc_out(x)
return logits
这个模型实现了GPT风格架构的核心组件，为我们的20亿参数模型提供了基础框架。在实际实现中，我们还需要添加更多优化和功能，如混合精度训练、梯度累积、检查点保存等。
7. 扩展到20亿参数规模
在理解了基础架构后，让我们简要讨论如何将这个模型扩展到20亿参数规模。
参数数量估计
Transformer模型的参数主要集中在以下几个部分：
- 标记嵌入：vocab_size × embed_size
- 位置嵌入：max_seq_len × embed_size
- 每个注意力层：4 × embed_size × embed_size（Q,K,V投影矩阵和输出投影矩阵）
- 每个前馈网络：embed_size × ff_hidden_size + ff_hidden_size × embed_size
对于我们的20亿参数模型，一个可能的配置是：
- 词汇表大小(vocab_size): 50,257（与GPT-2相同）
- 嵌入维度(embed_size): 2,048
- 层数(num_layers): 24
- 注意力头数(num_heads): 16
- 前馈网络隐藏维度(ff_hidden_size): 8,192
- 最大序列长度(max_seq_len): 2,048
这样的配置会产生约20.5亿参数，其中大部分来自多头注意力层和前馈网络。
内存与计算考虑
训练如此大规模的模型需要考虑内存需求和计算效率：
- 梯度检查点：只保存部分中间激活以节省内存
- 混合精度训练：使用FP16/BF16减少内存占用和加速计算
- 模型并行：跨多个GPU分割模型
- 梯度累积：多步累积梯度以模拟更大的批量
在实现我们的20亿参数模型时，这些优化技术将是必不可少的。
总结与展望
在本课中，我们深入探讨了Transformer架构的核心组件，包括自注意力机制、前馈神经网络、位置编码以及编码器和解码器的工作原理。我们还提供了关键组件的代码实现，并讨论了如何将这些组件组合成完整的GPT风格模型。
通过理解这些基础组件，我们现在已经为构建我们的20亿参数LLM奠定了理论基础。Transformer的设计精巧而强大，通过自注意力机制有效捕捉长距离依赖，通过并行计算提高训练效率，为现代大型语言模型的成功提供了关键支撑。
在下一课中，我们将探讨语言模型的数学基础，包括概率语言模型、交叉熵损失函数、梯度下降和优化器选择，为我们的实际模型训练做进一步准备。
延伸阅读
- Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS.
- Radford, A., et al. (2019). Language Models are Unsupervised Multitask Learners. OpenAI.
- Brown, T. B., et al. (2020). Language Models are Few-Shot Learners. NeurIPS.
- Child, R., et al. (2019). Generating Long Sequences with Sparse Transformers. arXiv.
- Dai, Z., et al. (2019). Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. ACL.
思考问题:
- 自注意力机制与传统RNN相比，在处理长序列时有哪些优势和潜在劣势？
- 在扩展模型规模时，应该优先增加哪些架构参数（层数、宽度、注意力头数等）？为什么？
- 固定的正弦位置编码和可学习的位置嵌入各有什么优缺点？在什么情况下一种可能优于另一种？
- 如果你有有限的计算资源，如何在不牺牲太多性能的情况下简化Transformer架构？
在下一课中，我们将深入探讨语言模型的数学基础，为我们构建的20亿参数LLM的训练和优化奠定坚实的理论基础！
