# Lecture 3 · 架构与超参数

> **CS336: Language Modeling from Scratch** · Stanford · Spring 2025
> 📅 Apr 8 · 💻 [不可执行讲义（PDF）](https://github.com/stanford-cs336/spring2025-lectures/blob/main/nonexecutable/2025%20Lecture%203%20-%20architecture.pdf) · 讲者 Tatsunori Hashimoto

---

## 承上启下

**上一讲（[[Lecture 2 · PyTorch 与资源核算]]）**：

- 训练语言模型之前要先把资源账算清楚：参数量、激活内存、optimizer state 与 FLOPs 共同决定「这个方案是否现实」。
- 记住了训练总计算量的经验公式 $C \approx 6ND$（$N$ 为参数量，$D$ 为训练 token 数）：前向约 $2ND$，反向约 $4ND$。
- Tensor 的 dtype、stride、view/copy 这些底层机制看似琐碎，但直接决定能在给定显存里训多大的模型。

**本讲（架构与超参数）**：

- 回到 Transformer 本身：现代大模型与 2017 年的原始 Transformer 相比，到底改了哪些地方？为什么改？
- 区分两类设计：哪些几乎已经形成全行业共识（如 Pre-Norm），哪些只是模型家族之间的偏好（如并行 block、soft-capping）。
- 超参数如何选：FFN 维度、head 数与 head 维度、深宽比、词表大小、正则化，以及一批针对 softmax 的稳定性技巧。

> [!note] 本讲一句话总览
> 现代 LLM 并不是「原始 Transformer 放大」这么简单，而是逐步收敛到一套稳定配方：**Pre-Norm + RMSNorm + SwiGLU + RoPE + 无 bias + GQA/MQA + 若干 softmax 稳定技巧**。但别把经验法则当物理定律——很多设计只是被大规模实践筛选出来的局部共识，讲者反复强调的方法论是：**看多个模型家族的「共同选择」，比看单一论文的消融实验更可靠**。

---

## 1. 从原始 Transformer 到现代变体

### 1.1 原始 Transformer 的选择

Vaswani et al. (2017) 的标准 Transformer 是为机器翻译设计的 encoder-decoder 结构，它的默认配置在今天看来处处「过时」：

- 位置编码用正弦绝对位置嵌入（sinusoidal absolute position embedding），直接加在词嵌入上；
- FFN 用 ReLU 激活，且带 bias；
- 归一化用 Post-LayerNorm，即 LayerNorm 放在残差相加**之后**；
- 注意力是标准的多头注意力（multi-head attention, MHA），每个 query head 配一组独立的 key/value；
- 在 embedding、attention、FFN 多处使用 dropout。

```mermaid
graph TB
    X["输入"] --> ATT["Multi-Head Attention"]
    ATT --> ADD1["Add"]
    ADD1 --> LN1["LayerNorm"]
    LN1 --> FFN["ReLU FFN"]
    FFN --> ADD2["Add"]
    ADD2 --> LN2["LayerNorm"]
```

这套结构在机器翻译（几千万到几亿参数、有监督数据）上工作得很好，但换到**大规模自回归语言模型**（几十亿参数、几万亿 token、混合精度、上千块 GPU）的场景后，稳定性与效率问题被急剧放大——这正是后面每一节改动的动机。

### 1.2 作业中实现的现代简化版

CS336 Assignment 1 实现的是更接近现代 LLM（LLaMA 风格）的版本：

| 组件 | 原始 Transformer | 现代 LLM 常见选择 |
| :--- | :--- | :--- |
| Norm 位置 | Post-Norm | **Pre-Norm** |
| Norm 类型 | LayerNorm | **RMSNorm** |
| 位置编码 | Sinusoidal / learned absolute | **RoPE** |
| FFN 激活 | ReLU | **SwiGLU / GeGLU** |
| Linear bias | 常有 | **常去掉 bias** |
| Attention KV | MHA | GQA / MQA（推理优化） |

> [!tip] 直觉
> 这张表不是「审美升级」，每一行背后都有一个具体的痛点：Pre-Norm 与 RMSNorm 解决**训练稳定性**（深层网络的 loss spike）；去 bias 与 RMSNorm 顺带减少**数据搬运**；SwiGLU 在同等 FLOPs 下换取**小而稳定的质量收益**；RoPE 解决**长上下文的相对位置泛化**；GQA/MQA 则是为**推理时的 KV cache** 买单。读现代架构时，先问「这个改动在优化哪一个维度」，比背结论更有用。

---

## 2. Norm：Pre-Norm 是最大共识

### 2.1 Post-Norm 的问题

Post-Norm 把 LayerNorm 放在残差相加之后：

$$y = \mathrm{LayerNorm}\big(x + \mathrm{Sublayer}(x)\big)$$

其中 $x$ 是进入这个子层的残差流（residual stream），$\mathrm{Sublayer}$ 是 attention 或 FFN。注意 LayerNorm 直接作用在**主残差路径**上——从输入到输出的每一条路径都要穿过每层的 LayerNorm。

这在大模型里带来一串问题：

- 梯度必须穿过几十上百个串联的 LayerNorm 才能传回底层，传播不稳定。Xiong et al. (2020) 的分析表明，Post-Norm 在初始化时靠近输出层的梯度幅度过大；
- 因此 Post-Norm 严重依赖学习率 warmup，warmup 设置稍有不慎训练就会发散；
- 网络越深，训练中越容易出现 loss spike（损失突然飙升甚至 NaN）。

### 2.2 Pre-Norm

Pre-Norm 把归一化移到子层**输入**处，残差相加保持「干净」：

$$y = x + \mathrm{Sublayer}\big(\mathrm{Norm}(x)\big)$$

优点：

- 残差流本身不再被 norm 打断，从输入到输出始终存在一条**恒等路径**（identity path）；
- 梯度可以沿这条恒等路径无衰减地传回底层，与 ResNet 的经验一脉相承（残差网络与反向传播的基础见 CS221 的 [[Lecture 4 · 深度学习]]）；
- 大模型训练明显更稳定，基本不再需要精细调 warmup；
- 实践中可以用更大的学习率。

```mermaid
graph LR
    X["x"] --> ADD["+"]
    X --> N["Norm"]
    N --> SUB["Attention / FFN"]
    SUB --> ADD
    ADD --> Y["y"]
```

> [!tip] 直觉
> 把残差流想象成一条「信息高速公路」：Pre-Norm 只在匝道口（进入子层前）做归一化，主路畅通无阻；Post-Norm 则相当于在主路上每隔一层设一个收费站，几十层叠下来，梯度早已面目全非。这也是为什么 2020 年之后几乎所有新模型（GPT-3、LLaMA、PaLM、Chinchilla……）都用 Pre-Norm。

> [!note] 讲义观点
> Pre-Norm 是现代 LLM 架构里少数真正接近「所有人都同意」的设计——讲义把它列为共识度最高的一项。近年唯一的例外方向不是回到 Post-Norm，而是在 Pre-Norm 基础上**加更多 norm**（见下节）。

### 2.3 Double Norm / 残差流外的 Norm

一些近期模型在 Pre-Norm 之外做了加法：

- **Gemma 2**：每个子层前后各放一个 RMSNorm（「三明治」式双 Norm），即 $y = x + \mathrm{Norm}\big(\mathrm{Sublayer}(\mathrm{Norm}(x))\big)$；Grok 也用了类似结构；
- **OLMo 2**：把 norm 移到子层输出侧（$y = x + \mathrm{Norm}(\mathrm{Sublayer}(x))$）——形式上像 Post-Norm，但 norm 仍在残差流之外，恒等路径没有被打断。

动机是一致的：保住 Pre-Norm 的训练稳定性（不动残差流），同时**约束每个 block 输出的尺度**，抑制训练后期激活值缓慢增长导致的异常。这个方向尚未形成统一标准，但「norm 不打断残差流」这条底线已经是共识。

---

## 3. LayerNorm vs RMSNorm

### 3.1 两者差异

LayerNorm 先减均值再除以标准差，带可学习的缩放 $\gamma$ 和偏移 $\beta$：

$$\mathrm{LN}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \gamma + \beta, \qquad \mu = \frac{1}{d}\sum_{i=1}^{d} x_i, \quad \sigma^2 = \frac{1}{d}\sum_{i=1}^{d} (x_i - \mu)^2$$

RMSNorm（Zhang & Sennrich, 2019）去掉减均值与偏移，只按均方根（root mean square）缩放：

$$\mathrm{RMSNorm}(x) = \frac{x}{\mathrm{RMS}(x)} \odot \gamma, \qquad \mathrm{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}$$

符号说明：$x \in \mathbb{R}^d$ 是单个 token 的激活向量，$d$ 是模型维度，$\gamma, \beta \in \mathbb{R}^d$ 是可学习的逐维参数，$\epsilon$ 是防止除零的小常数，$\odot$ 表示逐元素相乘。RMSNorm 的直觉：归一化真正起作用的是**控制向量的尺度**，减不减均值其实无关紧要。

| 对比 | LayerNorm | RMSNorm |
| :--- | :--- | :--- |
| 减均值 | 是 | 否 |
| 归一化方差 | 是 | 用 RMS |
| bias 参数 | 有 | 通常无 |
| 计算量 | 稍高 | 稍低 |
| 常见模型 | GPT-2/3、OPT、BLOOM | LLaMA、PaLM、T5、Chinchilla |

### 3.2 为什么 RMSNorm 有用？

表面解释：更少 FLOPs、更少参数。但这个解释站不住——norm 层的 FLOPs 在整个模型里占比不到 0.5%，参数占比更是可以忽略。

真正的要点：**FLOPs 不等于 runtime**。

- Norm 类算子的算术强度极低：每个激活值只做几次加乘，却要把整块激活从 HBM 读进来再写回去；
- 它们是典型的 **memory-bound（访存受限）** 算子——瓶颈在内存带宽而不是算力。虽然 FLOPs 占比不到 0.5%，这类小算子在未做融合的实现里可能吃掉两位数百分比的 wall-clock 时间；
- 少算一次均值、少加一个 bias，意味着少一轮数据搬运，在 memory-bound 区域就能真实提速。

> [!tip] 直觉
> 评价一个架构改动「省不省」，要问的不是「省了多少乘加」，而是「省了多少 HBM 往返」。RMSNorm 是这条原则的最小案例：效果与 LayerNorm 打平（大量模型验证），但访存更少、参数更少，于是成为默认选择。[[Lecture 5 · GPU]] 与 [[Lecture 6 · Kernel 与 Triton]] 会系统展开「算术强度」与 kernel 融合。

### 3.3 去掉 bias

现代 Transformer 普遍去掉所有线性层的 bias：

```python
nn.Linear(d_model, d_ff, bias=False)
```

原因可以串成一句话：bias 的收益本来就小，而它的成本不止参数本身。

- bias 参数量相对权重矩阵可以忽略，对最终效果的贡献在大模型上测不出来；
- Norm 层已经提供了可学习的逐维缩放（$\gamma$），表达能力上有冗余；
- 去掉 bias 少一类参数、少一次广播加法的访存，还能简化 kernel；
- 经验上去 bias 的模型训练**更稳定**——LLaMA、PaLM 等主流模型均已验证可行。

---

## 4. FFN 与激活：从 ReLU 到 SwiGLU

### 4.1 标准 FFN

原始 FFN 是「升维—非线性—降维」两层结构：

$$\mathrm{FFN}(x) = \max(0,\, xW_1)\, W_2$$

其中 $W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}$ 负责升维，$W_2 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}$ 负责降维，$\max(0,\cdot)$ 即 ReLU。经典维度选择是

$$d_{\text{ff}} = 4 \times d_{\text{model}}$$

后来 GPT 系列与 BERT 把 ReLU 换成更平滑的 GeLU，其余结构不变。

### 4.2 GLU 家族

GLU（Gated Linear Unit, Shazeer 2020）给 FFN 加一个**门控分支**：一条分支过激活函数当「门」，另一条分支保持线性，两者逐元素相乘：

$$\mathrm{GLU\text{-}FFN}(x) = \big(\sigma_{\text{act}}(xW_{\text{gate}}) \odot xW_{\text{up}}\big)\, W_{\text{down}}$$

SwiGLU 取激活函数为 SiLU（也叫 Swish）：$\mathrm{SiLU}(z) = z \cdot \sigma(z)$，$\sigma$ 是 sigmoid：

$$\mathrm{SwiGLU}(x) = \big(\mathrm{SiLU}(xW_{\text{gate}}) \odot xW_{\text{up}}\big)\, W_{\text{down}}$$

符号说明：$W_{\text{gate}}, W_{\text{up}} \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}$ 是两个并行的升维矩阵（门控分支与主分支），$W_{\text{down}} \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}$ 降维回残差流。

| 激活 | 常见模型 |
| :--- | :--- |
| ReLU | 原始 Transformer、T5 早期 |
| GeLU | GPT 系列、BERT |
| GeGLU | T5 v1.1、PaLM 一些变体 |
| SwiGLU | LLaMA、Gemma、DeepSeek 等 |

> [!note]
> Shazeer 在 GLU 论文结尾写道：这些结构为何有效「我们不给出解释，只能归功于神的眷顾（divine benevolence）」。这句半开玩笑的话很能代表本讲的方法论现实：SwiGLU 的优势是**实证筛选**出来的，而非理论推导的必然。

### 4.3 GLU 的维度修正

GLU 有**三个**矩阵（gate、up、down）而标准 FFN 只有两个。若仍取 $d_{\text{ff}} = 4d_{\text{model}}$，参数量与 FLOPs 会变成标准 FFN 的 1.5 倍。为了公平对齐，令参数量相等：

$$3 \cdot d_{\text{model}} \cdot d_{\text{ff}}' = 2 \cdot d_{\text{model}} \cdot (4 d_{\text{model}}) \;\Longrightarrow\; d_{\text{ff}}' = \frac{8}{3} d_{\text{model}}$$

所以 GLU 系模型常用 $d_{\text{ff}} \approx \frac{8}{3} d_{\text{model}} \approx 2.67 d_{\text{model}}$，使总计算量与传统 $4\times$ FFN 相当。

> [!example] 数值演算
> LLaMA-2 7B：$d_{\text{model}} = 4096$，按 $\frac{8}{3} \times 4096 = 10922.67$，实际取 $d_{\text{ff}} = 11008$——向上取整到 256 的倍数，让矩阵维度对 GPU 的 tile 尺寸友好。这类「凑整」在真实模型配置里随处可见。

> [!note] 当前经验
> SwiGLU / GeGLU 通常带来小而稳定的收益（perplexity 提升约 1–2%），但不是没有它就训不好模型——GPT-3 用 GeLU 也很强。把它理解为「几乎免费的小改进」，而非决定性因素。

---

## 5. Transformer Block：串行还是并行？

传统 block 串行执行 attention 与 FFN，两次进入残差流：

$$x_1 = x + \mathrm{Attn}\big(\mathrm{Norm}_1(x)\big), \qquad y = x_1 + \mathrm{FFN}\big(\mathrm{Norm}_2(x_1)\big)$$

并行 block 让两者同时从同一个输入出发，一次相加：

$$y = x + \mathrm{Attn}\big(\mathrm{Norm}(x)\big) + \mathrm{FFN}\big(\mathrm{Norm}(x)\big)$$

优点：

- attention 与 FFN 共享同一个 norm，少算一次；
- attention 的输入投影与 FFN 的升维投影可以**融合成一次更大的矩阵乘**，提高 GPU 利用率；
- 两个分支没有串行依赖，对张量并行等系统优化更友好。

采用过并行结构的模型包括 GPT-J、PaLM、GPT-NeoX、Falcon、部分 Cohere 模型等。

缺点与现状：

- FFN 看不到本层 attention 的输出，表达上是真实的结构差异，质量影响需要实证（PaLM 报告在大规模下退化很小）；
- 它始终没有成为主流——LLaMA 系等近年开源模型基本都回到了串行 block。

> [!tip] 直觉
> 并行 block 是「用一点表达能力换系统效率」的典型交易。当年 PaLM 在 TPU 上追求极致吞吐时值得做；当矩阵乘融合等收益能靠更好的 kernel 实现时，这笔交易就不再划算——架构选择永远要放在当时的系统背景下理解。

---

## 6. 位置编码：RoPE 成为主流

### 6.1 绝对位置编码

原始 Transformer 使用正弦位置嵌入或可学习的绝对位置嵌入（learned absolute embedding），做法都是把一个只依赖位置 $i$ 的向量 $p_i$ **加**到词嵌入上。

问题：

- 长度外推差：learned embedding 在超出训练长度的位置上根本没有参数；
- 语言建模真正需要的是**相对位置**信息——「这个词在那个词前面 3 个位置」比「这个词在第 1024 位」重要得多，绝对编码让模型必须自己学会做这个换算。

### 6.2 RoPE 的目标

RoPE（Rotary Position Embedding, Su et al. 2021）直接把目标写成一个函数方程：希望 query 与 key 的内积只依赖内容与**相对位置差**：

$$\langle f_q(x, i),\; f_k(y, j)\rangle = g(x, y, i - j)$$

符号说明：$x, y$ 是位置 $i, j$ 上的 query/key 内容向量，$f_q, f_k$ 是带位置信息的变换，$g$ 是某个只通过 $i-j$ 依赖位置的函数。满足这个性质，attention score 就天然具有平移不变性——整段文本平移后分数不变。

### 6.3 RoPE 怎么做？

解就是**旋转**：把 $d$ 维向量按维度两两配对成 $d/2$ 个二维平面，位置 $m$ 上的向量在第 $i$ 个平面内旋转角度 $m\theta_i$：

$$f(x, m) = R_{\Theta, m}\, x, \qquad R^{(i)}_{\Theta, m} = \begin{pmatrix} \cos m\theta_i & -\sin m\theta_i \\ \sin m\theta_i & \cos m\theta_i \end{pmatrix}, \qquad \theta_i = 10000^{-2i/d}$$

两个都被旋转过的向量做内积时，旋转角相减，只剩 $\;(m - n)\theta_i$——这正是要的相对位置性质。不同平面的 $\theta_i$ 从高频到低频排布（和正弦编码同源）：高频维度分辨近距离，低频维度分辨远距离。

```python
def apply_rope(x, cos, sin):
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    return torch.stack([
        x1 * cos - x2 * sin,
        x1 * sin + x2 * cos,
    ], dim=-1).flatten(-2)
```

与 sinusoidal embedding 的区别，逐条说清：

- sinusoidal 是把位置向量**加**到 embedding 上，位置信息与内容信息混在同一个向量里，随层数加深逐渐被稀释；
- RoPE 是在**每一层的 attention 里旋转 query/key**，位置信息在每层注意力计算时都新鲜注入；
- RoPE 不改变 value——位置只应该影响「关注谁」（score），不应该污染「取回什么内容」（value）；
- 旋转是正交变换，不改变向量模长，对数值稳定性友好。

> [!tip] 直觉
> 用复数看最简洁：把每对维度看成一个复数，RoPE 就是乘以 $e^{\mathrm{i} m\theta}$。两个复数做（共轭）内积时相位相减，得到 $e^{\mathrm{i}(m-n)\theta}$——绝对位置自动消掉，只剩相对位置。「加法编码」做不到这一点，因为加法在内积里会产生内容×位置的交叉项。

### 6.4 长上下文相关变体

RoPE 之后出现了一整条「上下文长度扩展」的技术线，各自一句话定位：

- **Position Interpolation（线性插值）**：把超出训练长度的位置线性压缩回训练区间，加少量微调即可外推；
- **NTK-aware scaling**：只改 RoPE 的 base（如 10000 → 500000），让低频维度多拉伸、高频维度少动，保住近距离分辨率；
- **YaRN**：按频率分段采用不同插值强度，并对 attention 加温度校正，是目前开源界常用方案；
- **LongRoPE**：对每个维度搜索独立的缩放因子，把上下文推到百万 token 级；
- **NoPE**：干脆不加显式位置编码——causal mask 本身泄露了位置信息，小模型上可行，大规模验证不足；
- **滑动窗口 + 少量全局层**：多数层只看局部窗口（不需要位置外推），少数全局层配合 RoPE scaling 负责远距离信息（如 Gemma 3）。

本讲重点不是这些细节，而是理解一个结论：**位置编码的选择直接决定长上下文的泛化方式**，这也是 RoPE（相对位置、可插值）胜过绝对编码的根本原因。

---

## 7. 超参数：哪些选择已经很「无聊」？

### 7.1 FFN 维度

经验规则：

| FFN 类型 | 常见比例 |
| :--- | :--- |
| ReLU / GeLU FFN | $d_{\text{ff}} \approx 4 \times d_{\text{model}}$ |
| SwiGLU / GeGLU | $d_{\text{ff}} \approx \frac{8}{3} \times d_{\text{model}}$ |

Kaplan et al. 的扫描实验显示，$d_{\text{ff}} / d_{\text{model}}$ 在 1~10 的宽区间内 loss 几乎不变（一个平坦的盆地），所以大家干脆沿用保守选择。极端反例是 T5 11B（比例高达 64）——能训出来，但没有成为惯例。缩放实验的方法论见 [[Lecture 9 · 缩放定律 I：基础]]。

### 7.2 Attention head 与 model dim

几乎所有模型遵循约定

$$d_{\text{model}} = n_{\text{heads}} \times d_{\text{head}}$$

即所有 head 的维度拼起来恰好等于模型维度。这**不是数学必需**——head 总维度完全可以与 $d_{\text{model}}$ 不同，只是 1:1 最省心，投影矩阵恰好是方阵。

常见 $d_{\text{head}}$ 取 64、80、96、128，现代大模型多用 128。例如 GPT-3 175B：$d_{\text{model}} = 12288 = 96 \text{ heads} \times 128$。head 维度太小会限制每个 head 内的表达能力（attention score 的秩），太大则减少 head 数、降低多视角并行关注的能力——64~128 是长期实践稳定下来的折中区间。

### 7.3 深宽比

问题：同样参数量，是更深（多层）还是更宽（大 $d_{\text{model}}$）？

- Kaplan et al. 发现宽深比（aspect ratio）$d_{\text{model}} / n_{\text{layers}}$ 在约 100 附近有一个很平坦的最优区间——性能对这个比例不敏感；
- 现代模型大多落在 100~200：GPT-3 是 $12288/96 = 128$，LLaMA-3 70B 是 $8192/80 \approx 102$；
- 极深的模型串行依赖长，流水线并行难做、推理延迟高；极宽的模型单层矩阵大，更吃单卡显存与张量并行通信。

> [!tip] 直觉
> 深宽比不是纯算法问题，而是**算法与系统的交界**：既然 loss 对深宽比不敏感（平坦盆地），实际选择就由系统约束主导——深度受流水线并行与延迟限制，宽度受张量并行与显存限制。「性能平坦区内挑系统最舒服的点」是大模型超参选择的通用模式。

### 7.4 词表大小

| 场景 | 常见词表大小 |
| :--- | :--- |
| 英文/单语模型 | 30k–50k（GPT-2: 50257） |
| 多语言/产品模型 | 100k–250k（Llama 3: 128k，Qwen: ~152k，Gemma: 256k） |

权衡（分词与 BPE 的基础见 [[Lecture 1 · 概览与分词]]）：

- 词表大：同一段文本切出的 token 更少（序列更短、有效上下文更长），多语言字符不会被切得支离破碎；代价是 embedding 与 LM head 的参数量 $2 \times V \times d_{\text{model}}$ 变大，softmax 计算变贵；
- 词表小：参数省，但稀有词与非英语文本被切碎成大量 token，同样的文本要花更多的计算与上下文预算。

近年趋势明显向大词表走：产品化模型要服务多语言用户，序列长度就是钱。

---

## 8. 正则化：Dropout 越来越少

预训练大模型的处境与经典监督学习很不一样：

- 数据量巨大（数万亿 token），且**通常只训练一个 epoch**——每个样本只见一次；
- 因此几乎不存在传统意义上的 train/test 过拟合，dropout 防过拟合的初衷失去了对象；
- dropout 还会降低有效算力利用率，并给大规模优化引入额外噪声。

于是很多现代 LLM 的做法是：

- dropout 直接设为 0；
- 保留 weight decay（典型值 0.1）；
- 靠数据规模、学习率调度与 batch size 来控制训练行为。

> [!warning] 易错点
> 两个常见误解要避开。其一，「大模型不需要正则化」不等于「微调也不需要」——小数据 SFT/对齐阶段的过拟合是真实存在的，dropout 可能回归。其二，weight decay 在 LLM 里的作用**不是防过拟合**：近期研究表明它主要通过与学习率调度（尤其 cosine decay 末期）的交互改变优化动态，让训练末期的 loss 降得更低——把它理解为「优化器超参」比「正则化器」更贴近实际。

---

## 9. 稳定性技巧：小 softmax，大麻烦

Transformer 里有两处 softmax：attention 内部一处、输出层一处。大规模训练的数值崩溃几乎都发生在这两处——logits 尺度失控导致 softmax 饱和或上溢。以下三个技巧都是在「管尺度」。

### 9.1 z-loss

定义输出 softmax 的对数配分函数（log normalizer）：

$$z = \log \sum_{i=1}^{V} e^{\ell_i}$$

其中 $\ell_i$ 是第 $i$ 个词的输出 logit，$V$ 是词表大小。z-loss 在交叉熵之外惩罚 $z$ 偏离 0：

$$\mathcal{L} = \mathcal{L}_{\text{CE}} + \alpha \, z^2$$

```python
loss = cross_entropy + alpha * z.pow(2).mean()
```

作用机制：softmax 本身对 logits 整体平移不变，所以交叉熵**不约束** logits 的绝对尺度，放任其漂移；$z^2$ 惩罚把尺度拉回 $z \approx 0$（即 $\sum e^{\ell_i} \approx 1$），防止 logits 爆炸，在低精度（bf16）下尤其重要。PaLM 用 $\alpha = 10^{-4}$，后来被众多开源模型沿用。

### 9.2 QK Norm

Attention score 为

$$\mathrm{score}_{ij} = \frac{q_i \cdot k_j}{\sqrt{d_{\text{head}}}}$$

如果 $q$、$k$ 的模长在训练中增长，score 的尺度随之膨胀，softmax 进入饱和区——注意力分布变成 one-hot，梯度消失，训练出现 spike。

QK norm 的做法是在计算 score 之前对 query 和 key 各做一次 RMSNorm / LayerNorm，从源头钉死 attention logits 的尺度。这一技巧起源于视觉大模型（ViT-22B），后被 OLMo 2、Gemma 3、Chameleon 等语言模型采用，消融实验普遍显示它显著减少 loss spike。

### 9.3 Logit Soft-Capping

用 $\tanh$ 把 logits 压进有限区间 $(-c, c)$：

$$\ell' = c \cdot \tanh\!\left(\frac{\ell}{c}\right)$$

```python
logits = cap * torch.tanh(logits / cap)
```

好处：logits 无论如何不会超过 cap 值 $c$，softmax 不可能上溢；Gemma 2 在 attention logits 用 $c = 50$、最终输出 logits 用 $c = 30$，Grok-1 也用了类似做法。

代价：

- $\tanh$ 在接近饱和处压缩差异，可能损失表达能力；
- cap 值 $c$ 是需要调的超参数；
- 与 FlashAttention 等融合 kernel 的兼容性差（softmax 前多了一个非线性）。Gemma 3 已放弃 soft-capping，改用 QK norm——可以视为这场「管尺度」竞争的阶段性结论。

> [!tip] 直觉
> 三个技巧管的是同一件事的不同位置：z-loss 管**输出 softmax** 的尺度，QK norm 与 soft-capping 管 **attention softmax** 的尺度。前者用软惩罚（loss 项），后两者用硬约束（结构上限制）。共同哲学：softmax 是指数函数，指数放大一切失控——凡有 softmax 处，必须有人管住喂给它的数。

---

## 10. Attention 变体：为推理买单

### 10.1 MQA / GQA

自回归生成时，每生成一个 token 都要读一遍全部历史的 key/value——KV cache 的大小直接决定生成阶段的内存带宽压力：

$$\text{KV cache 大小} = 2 \times n_{\text{layers}} \times n_{\text{kv}} \times d_{\text{head}} \times L \times \text{每元素字节数}$$

其中 $n_{\text{kv}}$ 是 key/value head 的数量，$L$ 是序列长度，系数 2 对应 K 和 V 各一份。三种方案只差在 $n_{\text{kv}}$：

- **MHA**：每个 query head 都有自己的 K/V head，$n_{\text{kv}} = n_{\text{heads}}$；
- **MQA**（Multi-Query Attention, Shazeer 2019）：所有 query head 共享**一组** K/V，$n_{\text{kv}} = 1$；
- **GQA**（Grouped-Query Attention, Ainslie et al. 2023）：query head 分组，每组共享一组 K/V，$1 < n_{\text{kv}} < n_{\text{heads}}$。

```mermaid
graph LR
    Q["多个 Q heads"] --> GQA["分组"]
    GQA --> KV["较少 K/V heads"]
    KV --> CACHE["更小 KV cache"]
```

收益与代价：

- KV cache 缩小 $n_{\text{heads}} / n_{\text{kv}}$ 倍，生成阶段的内存带宽压力等比下降，推理吞吐显著提高（细节见 [[Lecture 10 · 推理]]）；
- MQA 太激进，可能带来可测量的质量损失；GQA 取中间值（如 8 组），质量几乎无损——Llama 2 70B 与 Llama 3 全系都用 $n_{\text{kv}} = 8$。

> [!example] 算一算
> Llama 2 70B（80 层、64 heads、$d_{\text{head}} = 128$、bf16）：MHA 下每 token 的 KV cache 为 $2 \times 80 \times 64 \times 128 \times 2\,\text{B} \approx 2.6\,\text{MB}$；GQA（8 组）降为 $2 \times 80 \times 8 \times 128 \times 2\,\text{B} \approx 0.33\,\text{MB}$，缩小 8 倍。4k 上下文、批量 64 的服务场景下，这是「85 GB vs 10.7 GB」的差别——决定了能否塞进一张卡。

### 10.2 Sliding Window / Sparse Attention

完整 attention 对序列长度是 $O(L^2)$。一种被广泛采用的工程折中：

- 大多数层只看固定大小的局部窗口（如 Mistral 7B 的 4096 窗口），单层计算量降为 $O(L \cdot w)$；
- 少数层保留 full attention，负责远距离信息；
- 远处的信息也能通过**逐层传递**到达——第 $k$ 层的感受野可达 $k \times w$。

Mistral、Gemma 2（局部/全局层 1:1 交替）、Gemma 3（5:1 交替，且只对全局层做长上下文 RoPE scaling）、Cohere Command A 等都采用了这类思想。

> [!tip] 直觉
> 本节两个技巧的共同背景：训练只发生一次，推理要发生无数次，**推理成本已经反过来塑造架构**。GQA 砍的是 KV cache 的「头数」维度，滑动窗口砍的是「序列长度」维度——下一讲的 MoE 则砍每 token 激活的参数量，思路一脉相承。

---

## 总结

```mermaid
mindmap
  root((架构与超参数))
    Norm
      Pre-Norm
      RMSNorm
      no bias
    FFN
      GeLU
      SwiGLU
      d_ff 4x or 8/3x
    Position
      RoPE
      relative position
      long context scaling
    Attention
      MHA
      MQA
      GQA
      sliding window
    Hyperparameters
      head dim
      depth width
      vocab size
      dropout
    Stability
      z-loss
      QK norm
      soft-capping
```

**关键要点**：

1. **Pre-Norm 是最大共识**
   - norm 移出残差流，保住恒等路径，深层训练稳定性显著改善，warmup 依赖大减。

2. **RMSNorm 和去 bias 是效率友好的默认选择**
   - 省的不是 FLOPs 而是数据搬运；norm 类算子是 memory-bound 的典型。

3. **SwiGLU 是现代 FFN 主流**
   - 门控结构带来小而稳的收益；配 $d_{\text{ff}} \approx \frac{8}{3} d_{\text{model}}$ 使三矩阵结构与传统 $4\times$ FFN 计算量对齐。

4. **RoPE 让 attention 天然依赖相对位置**
   - 不是加位置向量，而是按位置旋转 Q/K；内积时绝对位置消掉，只剩 $i-j$。这也是长上下文插值方法（YaRN 等）的基础。

5. **GQA/MQA 是推理时代的架构选择**
   - 压缩 KV cache 的 head 维度，降低生成阶段的内存带宽压力，质量几乎无损。

6. **稳定性问题往往藏在 softmax**
   - 输出 logits、attention logits 都可能失控；z-loss（软惩罚）、QK norm（源头归一）、soft-capping（硬上限）都是在管喂给指数函数的尺度。

**下一讲预告（[[Lecture 4 · 混合专家模型 MoE]]）**：

本讲优化的都是**稠密模型**——每个 token 都激活全部参数。下一讲换一个方向：能不能在相同 FLOPs 下塞入更多参数？MoE 把 FFN 拆成多个专家并稀疏激活，用路由机制把「模型容量」和「每 token 计算量」解耦。

---

## 复习自测

> [!question]- Q1：从梯度路径的角度解释，为什么 Pre-Norm 比 Post-Norm 更稳定？
> Post-Norm 的 LayerNorm 位于残差相加之后，从输出到输入的每条梯度路径都要穿过每一层的 norm，深层网络中梯度幅度失控（初始化时靠近输出层的梯度过大），必须依赖精细的 warmup。Pre-Norm 把 norm 移到子层输入处，残差流保留一条不经过任何 norm 的恒等路径，梯度可以沿它无衰减地传回底层——与 ResNet 的恒等捷径同理。

> [!question]- Q2：算一算——SwiGLU FFN 要与 $4\times$ 的 ReLU FFN 参数量持平，$d_{\text{ff}}$ 应取多少？
> ReLU FFN 有两个矩阵，参数量 $2 \cdot d \cdot 4d = 8d^2$；SwiGLU 有三个矩阵（gate、up、down），参数量 $3 \cdot d \cdot d_{\text{ff}}'$。令两者相等得 $d_{\text{ff}}' = \frac{8}{3}d \approx 2.67d$。实际配置还会向上取整到 128/256 的倍数（如 LLaMA-2 7B 的 11008）。

> [!question]- Q3：Norm 层的 FLOPs 占比不到 0.5%，为什么换 RMSNorm、去掉 mean 和 bias 还能带来可感的加速？
> 因为 FLOPs 不等于 runtime。Norm 是算术强度极低的算子：每个元素只做几次运算，却要把整块激活从 HBM 读入再写回，瓶颈在内存带宽（memory-bound）而非算力。少一次均值计算、少一个 bias，就少一轮数据搬运，在访存受限区域直接转化为 wall-clock 加速。

> [!question]- Q4：RoPE 与「把位置向量加到 embedding 上」的本质区别是什么？为什么 RoPE 只作用于 Q/K 而不作用于 V？
> 加法编码把位置混进内容向量，内积时会产生内容×位置的交叉项，无法保证分数只依赖相对位置；RoPE 按位置旋转 Q/K，内积时旋转角相减，绝对位置自动消去，严格满足 $\langle f(x,i), f(y,j)\rangle = g(x,y,i-j)$。不旋转 V 是因为位置只应影响「关注谁」（score），不应改变被取回的内容表示本身。

> [!question]- Q5：算一算——Llama 2 70B（80 层、64 heads、$d_{\text{head}}=128$、bf16）在 MHA 与 GQA（8 组）下，每 token 的 KV cache 各多大？
> KV cache $= 2 \times n_{\text{layers}} \times n_{\text{kv}} \times d_{\text{head}} \times 2$ 字节。MHA：$2 \times 80 \times 64 \times 128 \times 2 \approx 2.6\,\text{MB}$/token；GQA 8 组：$n_{\text{kv}} = 8$，得 $\approx 0.33\,\text{MB}$/token，正好缩小 $64/8 = 8$ 倍。长上下文、大批量服务时，这个 8 倍决定了显存能否装下。

---

## 参考资料

- 💻 [2025 Lecture 3 - architecture.pdf（官方讲义）](https://github.com/stanford-cs336/spring2025-lectures/blob/main/nonexecutable/2025%20Lecture%203%20-%20architecture.pdf)
- 📄 [Attention Is All You Need](https://arxiv.org/abs/1706.03762)
- 📄 [On Layer Normalization in the Transformer Architecture](https://arxiv.org/abs/2002.04745)
- 📄 [Root Mean Square Layer Normalization](https://arxiv.org/abs/1910.07467)
- 📄 [GLU Variants Improve Transformer](https://arxiv.org/abs/2002.05202)
- 📄 [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864)
- 📄 [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://arxiv.org/abs/2305.13245)
- 📄 [PaLM: Scaling Language Modeling with Pathways（z-loss 出处）](https://arxiv.org/abs/2204.02311)
- 📄 [Small-scale proxies for large-scale Transformer training instabilities（QK norm / z-loss 消融）](https://arxiv.org/abs/2309.14322)