GPT模型架构详解

qll766

370人浏览 · 2026-05-26 17:24:01

qll766 · 2026-05-26 17:24:01 发布

1. GPT整体架构图

输入 token IDs (batch_size, seq_len)
    ↓
┌─────────────────────────────────┐
│ 1. 嵌入层 (Embedding Layer)      │
│    - Token Embedding            │
│    - Position Embedding         │
│    - Dropout                    │
└─────────────────────────────────┘
    ↓ (batch_size, seq_len, emb_dim)
┌─────────────────────────────────┐
│ 2. Transformer Blocks (×n层)    │
│    - Multi-Head Attention       │
│    - Feed Forward Network       │
│    - LayerNorm + Residual       │
└─────────────────────────────────┘
    ↓ (batch_size, seq_len, emb_dim)
┌─────────────────────────────────┐
│ 3. Final LayerNorm              │ 
└─────────────────────────────────┘
    ↓ (batch_size, seq_len, emb_dim)
┌─────────────────────────────────┐
│ 4. Output Head (Linear Layer)   │ 
└─────────────────────────────────┘
    ↓ (batch_size, seq_len, vocab_size)
输出 logits (未归一化的概率分布)

2. 逐层详细解析

Step 1: 嵌入层

self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
self.drop_emb = nn.Dropout(cfg["drop_rate"])

作用：

Token Embedding: 将离散的 token ID 映射为连续向量
- 输入：in_idx 形状 (batch_size, seq_len)，每个元素是 0~50256 的整数
- 输出：(batch_size, seq_len, 768)，每个 token 变成 768 维向量
Position Embedding: 注入位置信息
- Transformer 没有内置的顺序概念，需要显式添加位置编码
- torch.arange(seq_len) 生成 [0, 1, 2, …, seq_len-1]
- 每个位置有一个可学习的 768 维向量
相加操作: x = tok_embeds + pos_embeds
- 将语义信息和位置信息融合
- 这是标准做法，而非拼接（concat）
Dropout: 防止过拟合，随机丢弃部分神经元
Step 2: Transformer Blocks

self.trf_blocks = nn.Sequential(
    *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

核心处理：
堆叠 12 个 Transformer Block（GPT-124M 配置）
每个 Block 包含：

 输入 x
    ↓
  LayerNorm → Multi-Head Attention → Dropout → 残差连接
    ↓
  LayerNorm → FeedForward → Dropout → 残差连接
    ↓
  输出 x（相同维度）

关键点：

输入输出维度保持一致：(batch_size, seq_len, 768)
通过自注意力机制捕捉 token 之间的依赖关系
前馈神经网络提供非线性变换能力
残差连接缓解梯度消失问题
Step 3: Final LayerNorm

self.final_norm = LayerNorm(cfg["emb_dim"])

❓ 为什么 Transformer 之后还需要 LayerNorm？
原因 1：稳定输出分布
经过 12 层 Transformer 后，隐藏状态的数值范围可能变得不稳定：
某些维度可能非常大（爆炸）
某些维度可能非常小（消失）
不同样本的输出分布差异大
LayerNorm 确保输出具有：
均值 ≈ 0
方差 ≈ 1
这使得后续的线性层更容易学习
原因 2：Pre-LN vs Post-LN 架构选择
代码中使用的是 Pre-LN 架构（在每个子层之前归一化）：
#TransformerBlock 中的 Pre-LN

def forward(self, x):
    shortcut = x
    x = self.norm1(x)        # ← 在 Attention 之前归一化
    x = self.att(x)
    x = x + shortcut
    
    shortcut = x
    x = self.norm2(x)        # ← 在 FFN 之前归一化
    x = self.ff(x)
    x = x + shortcut
    return x

Pre-LN 的特点：
✅ 训练更稳定（深层网络也能收敛）
✅ 梯度流动更好
❌ 但最后一个 Block 的输出没有经过归一化
因此需要 Final LayerNorm 来统一所有输出的分布。

Step 4: Output Head

self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

❓ 为什么需要这个线性层？
原因 1：维度转换

Transformer 输出: (batch_size, seq_len, 768)   ← 隐藏空间
需要预测:         每个位置下一个 token 的概率  ← 词汇表空间

Linear 层的作用：
  768 维 → 50257 维（GPT-2 词汇表大小）

这是一个投影操作，将隐藏表示映射到词汇表空间。
原因 2：计算 logits（未归一化的分数）

logits = self.out_head(x)  # 形状: (batch_size, seq_len, 50257)

Logits 是每个 token 的"得分"，尚未经过 softmax
后续通过 softmax(logits) 得到概率分布
训练时使用 CrossEntropyLoss，它内部会自动应用 log-softmax
原因 3: bias=False 的设计

nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

为什么不用偏置项？
LayerNorm 已经做了中心化
Final LayerNorm 确保输入均值为 0
偏置项的作用被削弱
减少参数量
有 bias: 768 × 50257 + 50257 = 38,647,573 参数
无 bias: 768 × 50257 = 38,597,376 参数
节省约 50K 参数（虽小但有意义）
遵循原始 GPT 设计
OpenAI 的 GPT-2/GPT-3 都使用无偏置的输出层

3. 完整数据流示例

假设输入 “Hello, I am”

# 1. Tokenization
encoded = [15496, 11, 314, 716]  # 4 个 token

# 2. 输入形状
in_idx.shape = (1, 4)  # batch_size=1, seq_len=4

# 3. 嵌入层
tok_emb.shape = (1, 4, 768)
pos_emb.shape = (1, 4, 768)
x = tok_emb + pos_emb  # (1, 4, 768)

# 4. 经过 12 层 Transformer
x = self.trf_blocks(x)  # (1, 4, 768)
# 此时 x 包含了丰富的上下文信息

# 5. Final LayerNorm
x = self.final_norm(x)  # (1, 4, 768)
# 标准化输出分布

# 6. 输出层
logits = self.out_head(x)  # (1, 4, 50257)
# 每个位置的每个 token 都有一个得分

# 7. 预测下一个 token（只取最后一个位置）
last_logits = logits[:, -1, :]  # (1, 50257)
probabilities = torch.softmax(last_logits, dim=-1)
next_token = torch.argmax(probabilities)  # 最可能的下一个 token

4. 为什么要这样的设计？

信息论视角：

 隐藏空间 (768 维)          词汇表空间 (50257 维)
┌──────────────┐           ┌────────────────┐
│ 稠密表示      │  ──────→  │ 稀疏预测        │
│ Dense Rep.   │  Linear   │ Sparse Pred.   │
│ 捕捉语义关系  │  Layer    │ 具体 token 概率 │
└──────────────┘           └────────────────┘

隐藏空间：紧凑、连续的语义表示
词汇表空间：离散、高维的分类任务
线性层充当两个空间之间的桥梁。
几何视角：
Transformer 输出空间: