CS336 · 第 3 课

Transformer 架构与超参数

从 Vaswani et al. 2017 的"Attention Is All You Need"到今天的前沿模型，Transformer 架构经历了一系列关键演化。这一课拆解现代 Transformer 的每一个组件，理解每个设计选择背后的 trade-off。

本讲问题：有了 token 之后，还需要决定一个 Transformer block 如何组织归一化、attention、FFN 和位置表示；这些选择会同时影响训练稳定性、参数量和计算成本。

x:[B,T,d]
├─ norm → Q,K,V → causal attention → [B,T,d] ─┐
│                                             ├─ residual add
└─────────────────────────────────────────────┘
      ↓ norm → MLP d→d_ff→d → residual add → next block

architecture choice → tensor shape → parameters / FLOPs / memory
hyperparameter value → capacity / optimization / hardware efficiency

架构演化的核心问题

原版 Transformer 提出了一个优雅的设计。但在实际训练大规模模型时，研究者发现了很多可以改进的地方。这一课围绕几个核心问题展开：

Transformer 回顾

原版架构（Vaswani et al., 2017）

每个 Transformer block 包含：

Input: x  (batch, seq_len, d_model)

# 1. Multi-Head Self-Attention
attention_output = MultiHeadAttention(x, x, x)  # Q=K=V=x
x = LayerNorm(x + attention_output)              # Post-norm + residual

# 2. Feed-Forward Network
ffn_output = FFN(x)                              # Linear → ReLU → Linear
x = LayerNorm(x + ffn_output)                    # Post-norm + residual

# 3. 重复 L 层

# 4. 最终输出
output = Linear(x)  # d_model → vocab_size

关键维度

Pre-Norm vs Post-Norm

原版：Post-Norm

# Post-norm (Vaswani 2017)
x = LayerNorm(x + MultiHeadAttention(x))
x = LayerNorm(x + FFN(x))

现代：Pre-Norm

# Pre-norm (GPT-2 及以后主流)
x = x + MultiHeadAttention(LayerNorm(x))
x = x + FFN(LayerNorm(x))

为什么 Pre-Norm 更好？

这不是“Pre-norm 在所有指标上都更好”。Post-norm 或混合方案在合适初始化与缩放下也能稳定训练，并可能有不同的表示与收敛特性。架构选择应结合深度、初始化和训练 recipe 判断。

Norm 位置的现代变体

官方 Lecture 3 特别讨论了把额外的 norm 放在残差分支输出、但不放进恒等残差主干的做法；Grok、Gemma 2、OLMo 2 等模型采用了不同组合。它们不应笼统称为“DeepNet 的 Sandwich Norm”，也不是 DeepSeek-V3 的统一配置。

# 一类示意：保留干净的 residual stream，规范化分支输入/输出
x = x + PostNorm(Attention(PreNorm(x)))

激活函数的演化

ReLU → GELU

原版 Transformer 用 ReLU。GPT-2 开始用 GELU：

ReLU(x) = max(0, x)

GELU(x) = x × Φ(x)   （Φ 是标准正态的 CDF）
        ≈ 0.5 × x × (1 + tanh(√(2/π) × (x + 0.044715 × x³)))

为什么使用 GELU？ GELU 是平滑的门控式非线性，按输入大小连续缩放信号，而 ReLU 对负值直接截断。它在 GPT/BERT 系列中表现良好，但不能仅凭“0 点不可导”解释优劣——深度学习可以正常优化 ReLU，最终选择主要来自经验结果与硬件代价。

SwiGLU：现代 LLM 的标准

Llama 系列和 Mistral 使用 SwiGLU：

# 标准 FFN
FFN(x) = W₂ × ReLU(W₁ × x)

# SwiGLU FFN（Noam Shazeer 2020 系统研究）
FFN(x) = W_down × (Swish(W_gate × x) ⊙ (W_up × x))

# 其中 Swish(x) = x × sigmoid(x)

SwiGLU 的优势：

• 门控机制让模型可以"选择性地通过"信息
• 实验表明在相同 FLOPs 下表现更好
• 若保持原来的中间维度，三个矩阵会增加参数量；实践中通常缩小 d_ff 以保持参数/FLOPs 近似可比

位置编码的演化

Sinusoidal（原版 Transformer）

PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

优势：可以外推到更长序列（理论上）。

劣势：实际表现不如可学习嵌入和 RoPE；推理时需要预先计算。

Learned Embedding

# GPT-2 风格
position_embedding = nn.Embedding(max_seq_len, d_model)

优势：简单，可学习。

劣势：不能外推到训练时未见过的位置；需要固定 max_seq_len。

RoPE（Rotary Position Embedding）—— 常见现代选择

RoPE 的核心思想：不直接编码位置，而是让 attention 分数自然包含相对位置信息。

# RoPE 数学
对于维度 i (0, 2, 4, ...):
θ_i = 1 / (10000^(2i/d))

位置 m 的旋转:
R(m) = diag([e^{imθ_0}, e^{imθ_1}, ...])

# 应用到 query 和 key:
q_rotated = R(m) @ q
k_rotated = R(n) @ k

# attention score = q_rotated @ k_rotated^T
# 只依赖 (m - n) —— 相对位置！

RoPE 的优势：

• 编码相对位置（attention 天然需要这个）
• 可通过位置插值、频率缩放等方法扩展上下文，但外推质量不是自动保证
• Llama、Mistral、Qwen、DeepSeek 等大量模型采用；也仍有模型使用相对偏置、ALiBi、NoPE 或混合方案

ALiBi（Attention with Linear Biases）

# 不加位置编码，直接在 attention score 上加偏置
scores = Q @ K^T / √d
scores[i, j] -= |i - j| × slope_head_i  # 距离越远，惩罚越大

ALiBi 不需要位置 embedding 表，并在原论文设置中展示了长度外推能力；RoPE 与 ALiBi 的优劣取决于训练长度、缩放方法、架构和任务，不能概括为一个方法在“大多数任务”上必然更好。

归一化：LayerNorm → RMSNorm

LayerNorm

# 原版 LayerNorm
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
x_norm = (x - mean) / sqrt(var + ε)
output = γ × x_norm + β

参数量：2 × d_model（γ 和 β）

RMSNorm

Llama 系列使用 RMSNorm——去掉 mean centering：

# RMSNorm
rms = sqrt(mean(x²) + ε)
x_norm = x / rms
output = γ × x_norm  # 没有 β！

归一化的位置

# 常见现代配置: Pre-RMSNorm
x = x + Attention(RMSNorm(x))
x = x + FFN(RMSNorm(x))

# 有些模型在 attention 内部也做 norm (QK-norm)
# 特别是长上下文模型，防止 attention logits 爆炸

FFN 架构选择

标准 FFN

FFN(x) = W₂ × activation(W₁ × x)
d_ff = 4 × d_model  （标准比例）

门控 FFN（Gated Linear Unit, GLU）

# 标准 GLU
GLU(x) = σ(W_gate × x) ⊙ (W_up × x)   # ⊙ 是逐元素乘
FFN(x) = W_down × GLU(x)

# SwiGLU（用 Swish 替代 sigmoid）
SwiGLU(x) = Swish(W_gate × x) ⊙ (W_up × x)

# GeGLU（用 GELU 替代 sigmoid）
GeGLU(x) = GELU(W_gate × x) ⊙ (W_up × x)

为什么门控更好？

• 门控让模型学会"过滤"——某些维度可以完全关闭
• 类似于 LSTM/GRU 的门控机制
• 实验上在相同 FLOPs 下表现更好

注意力机制的细节

Multi-Head Attention（MHA）

# 输入: Q, K, V  (batch, seq, d_model)
# 拆分成 h 个头
Q = rearrange(Q, 'b s (h d) -> b h s d', h=num_heads)
K = rearrange(K, 'b s (h d) -> b h s d', h=num_heads)
V = rearrange(V, 'b s (h d) -> b h s d', h=num_heads)

# Scaled dot-product attention
scores = Q @ K^T / √d_k
scores = causal_mask(scores)  # 因果掩码
attn = softmax(scores)
output = attn @ V

# 合并头
output = rearrange(output, 'b h s d -> b s (h d)')

Grouped-Query Attention（GQA）

下一课会详细讲，这里先预览：

FlashAttention

不是架构创新，而是计算方法的革命：

超参数选择

核心超参数

d_model vs n_layers 的权衡

对于固定的参数量预算，应该选宽模型（大 d_model, 少层）还是深模型（小 d_model, 多层）？

• 深模型：每层 FLOPs 少（d_model 小），但层数多 → 顺序依赖性更强，pipeline 并行效率低。
• 宽模型：每层 FLOPs 多（d_model 大），层数少 → 更容易并行，但单卡内存压力大。
• 经验：前沿模型倾向于适中的深度（80-120 层），用 d_model 控制总参数量。

参数量估算

# Transformer 参数量（忽略 embedding 和 final head）
# Attention（MHA 假设）: 4 × d_model²
# GQA/MQA 的 K、V 投影更小，不能直接套 4d²
# FFN (SwiGLU): 3 × d_model × d_ff  (W_gate, W_up, W_down)
# RMSNorm: 2 × d_model  (每层两个 norm)

# 使用 SwiGLU (d_ff = 8/3 × d_model):
per_layer = 4d² + 3d × (8/3)d + 2d = 4d² + 8d² + 2d = 12d² + 2d
total_params ≈ n_layers × 12 × d_model²

Chinchilla 比例

Hoffmann et al. (2022) 的实验常被概括为下面这个经验预算：

D ≈ 20 × N

其中:
D = 训练 token 数
N = 模型参数量

例如: 70B 模型 → ~1.4T tokens (compute-optimal)

边界：20 tokens/parameter 是 Chinchilla 70B 配置及其量级的常用速记，不是跨架构、数据质量和推理成本都不变的物理常数。现代模型常为了摊薄推理成本而显著 overtrain。

现代模型配置对比

参数初始化策略

Xavier/Glorot 初始化

# 目标: 让输出的方差 ≈ 输入的方差
# 对于 (input_dim, output_dim) 的权重矩阵:
std = sqrt(2 / (input_dim + output_dim))  # Xavier uniform
std = sqrt(1 / input_dim)                  # 简化版本（前向）

w = torch.randn(input_dim, output_dim) * std

Transformer 特定的初始化

很多现代模型对残差分支做特殊初始化：

# GPT-2 风格
# 残差投影初始化为 0.02（很小），让训练初期接近 identity
w = torch.randn(d_model, d_model) * 0.02

# T5 / PaLM 风格
# 按层深度缩放: std = base_std / sqrt(2 * n_layers)
# 深层层的残差贡献更小，保持梯度稳定

In-Context-Learning 友好的初始化

PaLM 和后续模型使用特殊的初始化来增强 ICL 能力：

• Output projection 初始化为 0（让训练初期输出接近 0）
• Attention 的 W_o 初始化为较小的值
• 这些细节在实践中影响显著

知识检查

如果只能记住 5 件事

延伸阅读