Transformer 解剖：从 Attention 到推理系统

第 5 章 Transformer Block：FFN、LayerNorm、Residual 的角色

作者杨艺韬 · 5,453 字

第 5 章 Transformer Block：FFN、LayerNorm、Residual 的角色

第 2、3、4 章我们把 Multi-Head Self-Attention 拆得很透了。但单靠 Attention 没法搭起一个 Transformer——Llama 70B 的 80 层、GPT-3 的 96 层，每一层都不是纯 attention，而是一个完整的 Transformer Block：Attention + 前馈网络（FFN）+ 残差连接（Residual）+ 层归一化（LayerNorm）。

为什么是这四件套？为什么少了任何一个 Transformer 都训不动？这一章我们把 Block 的每个零件单独打开看，再合在一起看它们是怎么互锁的。读完这章你能：

在白板上画出一个完整的 Transformer Block 的数据流（含张量形状）；
解释为什么 FFN 必须存在、为什么是 4× $d_{\text{model}}$ 、为什么 ReLU 后来被 SwiGLU 替代；
解释 Pre-LN 和 Post-LN 的差别，以及为什么今天的大模型几乎一边倒选 Pre-LN；
解释 RMSNorm 为什么取代了 LayerNorm；
估算一个模型的参数量、计算量在 Attention 和 FFN 之间是怎么分配的。

5.1 为什么 Attention 不够：缺非线性

回想一下 Self-Attention 的输出是什么：

\text{output}_i = \sum_{j=1}^{T} A_{ij} v_j = \sum_{j=1}^{T} A_{ij} (x_j W_V)

每个输出位置是 value 向量的加权和。其中 $v_j = x_j W_V$ 是输入的线性投影，权重 $A_{ij}$ 来自 softmax，但加权求和本身仍然是输入的线性组合——把 $A_{ij}$ 视作系数，整个 attention 输出就是 $V$ 的线性组合。

唯一的非线性来自 softmax。但 softmax 只作用在 attention 权重上，不作用在 value 本身的特征空间里。结果就是：经过若干层 Self-Attention 之后，输出仍然停留在「输入向量空间」附近的某种线性变换——这远远不够用来表达自然语言里的复杂结构。

这是一个深度学习里很基础的原则：纯线性的网络（无论多深）等价于一层线性变换。要让深度网络真正「深」起来——也就是每一层都能引入新的表达能力——必须每一层都加上一个非线性变换。

Transformer Block 里这件事由 Feed-Forward Network（前馈网络，FFN）来做。

5.2 FFN：每个 token 独立的非线性变换

FFN 的设计简单到出奇：

\text{FFN}(x) = \text{Activation}(x W_1 + b_1) W_2 + b_2

也就是一个两层全连接 + 一个非线性激活函数。 $W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ffn}}}$ 把维度从 $d_{\text{model}}$ 升到 $d_{\text{ffn}}$ （典型 $4 d_{\text{model}}$ ），激活后再用 $W_2 \in \mathbb{R}^{d_{\text{ffn}} \times d_{\text{model}}}$ 降回 $d_{\text{model}}$ 。

关键点：FFN 是逐 token 独立的。每个位置的输入向量分别通过 FFN，不同位置之间没有交互。这个性质是 Self-Attention 的镜像——Self-Attention 让位置之间交互、FFN 让特征在每个位置上做非线性变换，两者互补。

flowchart TB
  X["输入 (T, d_model)"] --> SA[Self-Attention<br/>位置之间交互]
  SA --> FFN[FFN<br/>每个位置独立做<br/>非线性变换]
  FFN --> OUT["输出 (T, d_model)"]
  SUB1[所有位置 互相打通] -.对应.-> SA
  SUB2[每个位置 单独加工] -.对应.-> FFN

可以把它理解成「人开会」的隐喻：Self-Attention 是大家聚在一起讨论（位置间信息交换），FFN 是会后每个人回去自己消化加工（位置内非线性变换）。Transformer Block 是「讨论 → 消化」的循环——每一层都做一次。

FFN 的尺寸为什么是 4×

原始论文里 $d_{\text{ffn}} = 4 d_{\text{model}}$ ，这成了一个事实标准。今天大多数模型仍然遵循这个比例（少数现代变体在 SwiGLU 下用 $\frac{8}{3} d_{\text{model}}$ ，下文会讲）。

为什么是 4？这没有一个严格证明，但有几条经验观察：

第一，FFN 是 Transformer 的「记忆容量」。研究表明，FFN 的权重很大程度上扮演着「key-value 字典」的角色——Geva et al.（2021）的论文 Transformer Feed-Forward Layers Are Key-Value Memories 证明 FFN 第一层把输入向量「映射成 key」，第二层把「value 取回来」。 $d_{\text{ffn}}$ 越大，能存的 key-value 对越多，模型记忆容量越大。

第二， $4d$ 大致是表达力和参数量之间的甜点。再小（比如 $2d$ ）模型表达力不够；再大（比如 $8d$ ）参数量爆炸但收益边际递减。 $4d$ 是工业界反复实验找出的平衡点。

第三，FFN 占 Transformer 大约 2/3 的参数。一个 Transformer Block 的参数量（不含归一化和 bias）：

Attention： $4 d_{\text{model}}^2$ （ $W_Q, W_K, W_V, W_O$ 各一个 $d \times d$ ）
FFN： $d_{\text{model}} \cdot 4 d_{\text{model}} \cdot 2 = 8 d_{\text{model}}^2$ （ $W_1, W_2$ 各一个 $d \times 4d$ ）

合起来 $12 d^2$ ，FFN 占 $8/12 = 2/3$ 。也就是说Transformer 大部分参数其实在 FFN 里——这个直觉对很多读者是反常识的，因为 attention 名声更大。

激活函数：从 ReLU 到 GELU 到 SwiGLU

原始 Transformer 用的激活函数是 ReLU： $\text{ReLU}(x) = \max(0, x)$ 。简单粗暴有效。

GPT-2 改成了 GELU（Gaussian Error Linear Unit）：

\text{GELU}(x) = x \cdot \Phi(x)

其中 $\Phi(x)$ 是标准正态分布的累积分布函数。GELU 是 ReLU 的「平滑版」——在 0 附近不再有硬切，而是平滑地过渡。BERT、GPT-2/3、原始 Llama 1 都用 GELU。

到 Llama 1 之后，激活函数又一次升级到 SwiGLU（Swish-Gated Linear Unit），由 Shazeer 在论文 GLU Variants Improve Transformer（2020）提出。SwiGLU 是「门控版 FFN」：

\text{SwiGLU}(x, W, V, W_2) = (\text{Swish}(xW) \odot xV) W_2

其中 $\text{Swish}(x) = x \cdot \sigma(x)$ （sigmoid 加权）， $\odot$ 是逐元素相乘。和原始 FFN 比，SwiGLU 多了一个门控分支：用 $xV$ 作为门，控制 $\text{Swish}(xW)$ 的激活强度——本质是「让模型学会哪个特征该开、哪个该关」。

flowchart LR
  X["输入 x"] --> W["× W (扩展维度)"]
  X --> V["× V (门控分支)"]
  W --> SW["Swish 激活"]
  V --> GATE
  SW --> GATE["逐元素相乘 ⊙"]
  GATE --> W2["× W_2 (降回 d_model)"]
  W2 --> OUT[输出]

SwiGLU 多了一个矩阵 $V$ ，参数量大约是原始 FFN 的 1.5 倍。为了控制总参数量不爆炸，实际工程把 $d_{\text{ffn}}$ 从 $4d$ 改成 $\frac{8}{3} d$ ——参数量大致与原 FFN 相当但表达力更强。这就是 Llama 系列模型卡片上 intermediate_size 通常约等于 $\frac{8}{3} \cdot \text{hidden\_size}$ 的来历（实际还会向上对齐到 256 的倍数）。

为什么 SwiGLU 比 GELU 好？ 经验观察 + 一些理论：门控让 FFN 的表达力近似翻倍，因为它能「对每个特征独立打开/关闭」。在 Llama-1 / 7B 的实验里，相同参数量下 SwiGLU 比 GELU 困惑度低 0.2~0.3——足以让今天的大模型几乎全部切到 SwiGLU。

主流模型激活函数对照：

模型	激活	FFN 类型	$d_{\text{ffn}}$
原始 Transformer	ReLU	标准 FFN	$4d$
BERT / GPT-2/3	GELU	标准 FFN	$4d$
Llama 1/2/3	SwiGLU	GLU 变种	$\sim \frac{8}{3} d$
Mistral / Qwen / DeepSeek	SwiGLU	GLU 变种	$\sim \frac{8}{3} d$
PaLM	SwiGLU	GLU 变种	$\sim \frac{8}{3} d$

直觉总结：FFN 是「每个 token 单独经过一个深度小、参数大的非线性层」，给整个模型提供主要的非线性表达力和大部分参数容量。激活函数从 ReLU 到 SwiGLU 是工程上「在不改变结构的前提下挤出更多表达力」的演化。

5.3 Residual：让深度可训

设想一下：堆叠 80 层 Transformer，每一层做 Attention + FFN。从第 80 层往前，loss 对第 1 层的梯度要穿过 80 次复合运算——每次都包含矩阵乘法和非线性。链式法则里这相当于 80 个雅可比矩阵相乘，方差累积、梯度消失或爆炸的风险极高。

这是深度神经网络从 2010 年代初一直被困扰的问题。直到 2015 年 ResNet（He et al.）提出残差连接（residual connection / skip connection）：

y = x + F(x)

每一层的输入 $x$ 直接「跳过」函数 $F$ ，加到输出上。这意味着梯度可以从 $y$ 直接传到 $x$ ，完全绕过 $F$ ：

\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial y} \cdot \left(1 + \frac{\partial F}{\partial x}\right)

那个「+1」就是关键——它保证梯度有一条始终为 1 的快速通道，无论 $F$ 的雅可比是多少，梯度至少能完整地传到底层。

Transformer 把这个想法用在了每个子层（Attention 子层和 FFN 子层）上：

y = x + \text{SubLayer}(\text{LN}(x))

具体到 Block：

x_in
 ├─→ LN → MHA → out1
 │                |
 ├────────────────+→ x_mid
 │
 ├─→ LN → FFN → out2
 │                |
 ├────────────────+→ x_out

每个子层（Attention、FFN）都有一条 residual 旁路。这样从输入到输出至少存在一条「线性恒等」的路径——梯度可以无损地传到任何一层。

flowchart TB
  XIN[x_in] --> LN1[LayerNorm]
  XIN --> ADD1["+ residual"]
  LN1 --> MHA[Multi-Head<br/>Attention]
  MHA --> ADD1
  ADD1 --> XMID[x_mid]
  XMID --> LN2[LayerNorm]
  XMID --> ADD2["+ residual"]
  LN2 --> FFN[FFN]
  FFN --> ADD2
  ADD2 --> XOUT[x_out]

Residual 的几何理解

另一个非常有用的视角：把每一层的 $F(x)$ 看成是「对 $x$ 的微调」。Residual 让模型默认「保留 $x$ 的大部分信息，在它上面做一点微调」。深度网络的每一层不需要重新发明轮子——它继承上一层的状态，只在必要的方向上加工。

这个性质和 Transformer 训练动力学非常契合。当 80 层堆叠时，每层只需要做「小改动」，最终累计成大变化。如果没有 residual，每层都要从零构建一个完整的输出——80 次完整构建几乎不可能稳定收敛。

实践中的一个观察：当你在训练一个 Transformer 时把 residual 砍掉（强制 $y = F(\text{LN}(x))$ ），模型会在 100~1000 步内立刻发散——loss 飙升到 NaN。residual 不是锦上添花，是必要条件。

5.4 LayerNorm：稳定每个位置的输入分布

加入 residual 之后，新问题来了：每一层 residual 把 $x + F(x)$ 输出，下一层的输入就是「原始输入 + 多次累加的修改量」。如果不约束，这个输出的方差会逐层放大（每加一次 $F(x)$ 都引入新方差），到第 80 层已经爆炸。

这就需要归一化（normalization）来约束每一层的输入分布。

LayerNorm 的定义

LayerNorm（Ba et al., 2016）是为 RNN/序列模型设计的归一化方法：对每个 token 的特征向量，沿着特征维做归一化：

\text{LN}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

其中：

$\mu = \frac{1}{d} \sum_{i=1}^{d} x_i$ 是特征的均值
$\sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2$ 是特征的方差
$\gamma, \beta \in \mathbb{R}^d$ 是可学习的缩放和偏移参数
$\epsilon$ 是数值稳定项（典型 $10^{-5}$ ）

注意 LayerNorm 是对每个 token 单独做的——位置 $i$ 的归一化只用位置 $i$ 自己的特征，与其他位置无关。这是它和 BatchNorm 最大的区别：

特性	BatchNorm	LayerNorm
归一化沿哪个维度	batch 维	特征维
是否跨样本	是	否
训练/推理是否一致	不一致（推理用滑动平均）	一致
序列长度变化	敏感	不敏感
对 batch_size 敏感	是（小 batch 不稳定）	否

LayerNorm 的「不跨样本」和「不依赖 batch_size」是它在序列模型里大获成功的关键。语言任务里 batch 之间序列长度可以差很多、batch_size 经常被显存约束（特别是大模型），BatchNorm 的统计量会非常不稳定。

Pre-LN vs Post-LN：一个分水岭

原始 Transformer 论文里 LayerNorm 放在 residual 之后——叫 Post-LN：

y = \text{LN}(x + \text{SubLayer}(x))

但很快人们发现，Post-LN 训练 Transformer 很不稳定，特别是在层数多、模型大的情况下。GPT-2 论文专门提到这个问题——他们把 LayerNorm 移到 residual 之前，得到 Pre-LN：

y = x + \text{SubLayer}(\text{LN}(x))

flowchart TB
  subgraph postln ["Post-LN（原始论文，不稳定）"]
    P_X[x] --> P_SUB[SubLayer]
    P_X --> P_ADD["+"]
    P_SUB --> P_ADD
    P_ADD --> P_LN[LayerNorm]
    P_LN --> P_OUT[y]
  end
  subgraph preln ["Pre-LN（GPT-2 之后主流）"]
    PR_X[x] --> PR_LN[LayerNorm]
    PR_X --> PR_ADD["+"]
    PR_LN --> PR_SUB[SubLayer]
    PR_SUB --> PR_ADD
    PR_ADD --> PR_OUT[y]
  end

为什么 Pre-LN 更稳定？关键在 residual 的「恒等通道」。Pre-LN 的残差连接把 $x$ 直接加到 $\text{SubLayer}(\text{LN}(x))$ 上，不经过任何归一化——梯度可以原路返回，链式法则上是 $1 + \frac{\partial F}{\partial x}$ 。

Post-LN 把 LayerNorm 放在 residual 之后，残差连接中间多了一次归一化。LayerNorm 的雅可比矩阵是非平凡的（涉及标准差的倒数等），把它接到 residual 路径上，梯度的「恒等通道」被破坏了——梯度不再保证至少为 1，会随层数衰减。

Xiong et al.（2020）在论文 On Layer Normalization in the Transformer Architecture 里给了严格的数学分析：Pre-LN 下损失对参数的梯度上界与层数无关，Post-LN 下梯度上界随层数指数增长——意味着大模型用 Post-LN 几乎一定会爆。

实践中：

Pre-LN：今天的 GPT、Llama、Mistral、DeepSeek 几乎全部用 Pre-LN。优点是训练稳定，可以堆很深。
Post-LN：BERT 仍用 Post-LN（24 层以下勉强能稳定），但 BERT-large（24 层）的训练就需要复杂的 warmup。
Sandwich-LN：极少数模型在子层前后都加 LN——更稳定但开销大，几乎被淘汰。
DeepNorm：微软提出的 Post-LN 改进版，让 Post-LN 也能训上千层。技术上可行，但实际工业基本用 Pre-LN。

总结：你在写一个新 Transformer，无论模型多大，默认选 Pre-LN。

RMSNorm：LayerNorm 的简化版

到了 2020 年之后，LayerNorm 又被进一步简化为 RMSNorm（Root Mean Square Normalization，Zhang & Sennrich, 2019）：

\text{RMSNorm}(x) = \gamma \cdot \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}}

和 LayerNorm 比，RMSNorm 去掉了均值的减法——只用 RMS（root mean square，根均方）做缩放，没有偏移项 $\beta$ 。

直觉：LayerNorm 的「减均值」在数据分布相对均匀时贡献很小，工程上可以省掉。少一个减法、少一个 $\beta$ 参数，每层节省一点点开销，整个 80 层模型积累下来效果可观。

实验：Llama 论文和后续 ablation 显示，RMSNorm 替换 LayerNorm 几乎不影响模型质量，但训练速度提升 7%~10%。

主流大模型几乎全部已切到 RMSNorm：Llama 1/2/3、Mistral、Qwen、DeepSeek、Yi——RMSNorm 是新一代标配。

归一化变种	公式核心	使用模型
BatchNorm	沿 batch 维归一化	不用于序列模型
LayerNorm	沿特征维归一化（含均值减法）	原始 Transformer / BERT / GPT-2/3
RMSNorm	沿特征维归一化（只 RMS 缩放）	Llama / Mistral / Qwen / DeepSeek
GroupNorm	沿特征分组归一化	CV 模型偶用，NLP 不用

5.5 完整 Block 的数据流

把 5.2、5.3、5.4 节合在一起，今天主流 Transformer Block 长这个样子（Pre-LN + RMSNorm + GQA + SwiGLU）：

flowchart TB
  XIN["x_in (T, d)"] --> RES1["+ residual"]
  XIN --> RMS1[RMSNorm]
  RMS1 --> RMSout1[normalized x]
  RMSout1 --> ROPE[RoPE 旋转 Q,K]
  ROPE --> MHA[Multi-Head<br/>Self-Attention<br/>GQA]
  MHA --> RES1
  RES1 --> XMID["x_mid (T, d)"]
  XMID --> RES2["+ residual"]
  XMID --> RMS2[RMSNorm]
  RMS2 --> SWI[SwiGLU FFN]
  SWI --> RES2
  RES2 --> XOUT["x_out (T, d)"]

形式化的数据流：

1. x ← x_in                              shape (T, d)
2. h ← x + Attention(RMSNorm(x))         shape (T, d)
3. y ← h + FFN(RMSNorm(h))               shape (T, d)
4. x_out ← y                             shape (T, d)

这就是 Llama 一个 Block 的全部内容。整个模型：

embedding → [Block × N] → final RMSNorm → LM head

每个 Block 的输入输出形状完全一致 $(T, d_{\text{model}})$ ——这是 Transformer 能堆叠任意深度的关键性质。Llama-3 70B 的 80 层就是 80 个这样的 Block 串起来的。

5.6 参数量与计算量分解

来一笔账。设 $d = d_{\text{model}}$ 、头数为 $h$ 、 $d_{\text{ffn}} = 4d$ （标准 FFN），单 Block 的参数：

子模块	参数量	占比
MHA: $W_Q, W_K, W_V, W_O$	$4d^2$	33%
FFN: $W_1, W_2$ （标准）	$8d^2$	67%
RMSNorm × 2（仅 $\gamma$ ）	$2d$	<1%
合计	$\approx 12d^2$	100%

如果是 SwiGLU（ $d_{\text{ffn}} = \frac{8}{3}d$ ），FFN 的参数：

\text{SwiGLU FFN} = d \cdot \frac{8}{3}d \cdot 3 = 8d^2

（三个矩阵： $W$ 、 $V$ 、 $W_2$ ）。所以总参数仍然是 $\approx 12d^2$ ——SwiGLU 在保持参数量不变的前提下提升了表达力，这就是 5.2 节我们说「Llama 的 intermediate_size 是 $\frac{8}{3}$ 倍」的工程合理性。

对应的整模型参数： $L$ 层 Block，每层 $12d^2$ ，再加 embedding $V \cdot d$ （ $V$ 是词表大小）和 LM head（如果不共享权重，又是一个 $V \cdot d$ ）。Llama-3 70B 的实际参数：

$d = 8192$ , $h = 64$ （GQA: $h_{kv} = 8$ ）, $L = 80$ , $V = 128256$
$L \cdot 12d^2 \approx 80 \cdot 12 \cdot 8192^2 \approx 64.4$ G
embedding（不共享）: $128256 \cdot 8192 \approx 1.05$ G
LM head（不共享）: $128256 \cdot 8192 \approx 1.05$ G
加上 GQA 调整、RMSNorm 等小项 + intermediate FFN 的 SwiGLU 多出 1 个矩阵，校正后总计 $\approx 70.5$ G

公开报告 Llama-3 70B 实际参数 70.6B——和我们的估算几乎一致。Llama 3 词表扩到 128K 后放弃了 weight tying（Llama 1/2 共享 embedding，Llama 3 不共享），是为了保留扩词表带来的额外表达自由度。这个估算量级正确，足以让你看到模型卡片就估出参数量。

计算量 / FLOPs 分解（前向，单个 token，标准 FFN）：

子模块	FLOPs
QKV 投影	$\sim 6 d^2$
Attention（短上下文）	$\sim 2 T d$
输出投影	$\sim 2 d^2$
FFN	$\sim 16 d^2$
合计	$\sim 24 d^2 + 2Td$

上下文短时（ $T \ll d$ ）FFN 主导计算量。上下文长时（ $T \gg d$ ）attention 的 $T^2$ 项主导。实际工程中 $T$ 在 4K~128K 之间变化，计算瓶颈会从 FFN 慢慢转向 attention——这是第 13 章「长上下文之战」的工程动因。

计算-访存比（重要，预示推理工程）：

FFN：每次都把权重从 HBM 读出来跑一次大矩阵乘——计算密集（compute-bound）
Attention：每次都要读写 KV Cache、做 softmax——访存密集（memory-bound）

这两种瓶颈对应不同的优化策略。FFN 用低比特量化（INT4 / FP8）把权重压小，让 HBM 读得快——第 16 章。Attention 用 Flash Attention 把 KV 访问局部化，让 SRAM 替 HBM 干活——第 18 章。这两条优化在第六部分会展开。

5.7 一些容易踩的细节

细节一：LayerNorm/RMSNorm 在哪里？

确认两遍：Pre-LN 下，LayerNorm 在子层（Attention 或 FFN）之前，residual 连接绕过 LayerNorm：

# 正确（Pre-LN）
h = x + attention(rmsnorm(x))     # 注意 residual 加的是 x，不是 rmsnorm(x)
y = h + ffn(rmsnorm(h))           # 同上

很多新手写错成 h = rmsnorm(x) + attention(rmsnorm(x))——把 residual 也加了归一化值，破坏了梯度直达通道。

细节二：最后一层归一化

整个模型最后还有一个 final RMSNorm（在 LM head 之前）。这个常被忽略，但它对 logits 的稳定有帮助。Llama 的 model.norm 就是这个最后的 RMSNorm。

细节三：bias 还是不要 bias

LayerNorm 通常带 $\gamma$ 和 $\beta$ 两个可学习参数；RMSNorm 只带 $\gamma$ 。FFN 的 $W_1, W_2$ 在原始 Transformer 里有 bias，但 Llama / Mistral 等现代模型把 bias 全部去掉——参数减少、训练略加速、对质量几乎无影响。

细节四：QKV 投影是否带 bias

同样，原始论文里 $W_Q, W_K, W_V, W_O$ 都带 bias，现代模型几乎都去掉了。

细节五：embedding 维度和 head 维度

$d_{\text{model}}$ 必须能被 $h$ 整除（这样每个头维度 $d_k = d_{\text{model}} / h$ 是整数）。 $d_{\text{model}}$ 也要能被某些硬件友好的数（如 128）整除以让 GPU 矩阵乘高效。

细节六：FFN 的 $d_{\text{ffn}}$ 实际值

SwiGLU 配 $\frac{8}{3}d$ ，但实际上工程会向上对齐到 256 的倍数让 GPU 友好。Llama 7B 的 $d = 4096$ ，理论 $d_{\text{ffn}} = 4096 \cdot 8/3 \approx 10923$ ，实际取 11008（256 的整数倍）。

5.8 把整个 Transformer 模型串起来

到这里，整个模型可以一目了然：

flowchart TB
  TOK[输入 token ids] --> EMB["Token Embedding"]
  EMB --> X0[x_0]
  X0 --> B1[Block 1]
  B1 --> B2[Block 2]
  B2 --> BD[...]
  BD --> BL[Block L]
  BL --> XL[x_L]
  XL --> FN[Final RMSNorm]
  FN --> HEAD[LM Head]
  HEAD --> LOGITS["logits 形状<br/>(T, V)"]
  LOGITS --> SM[softmax]
  SM --> P["词表上的<br/>下一 token<br/>概率分布"]

每个 Block 的内部就是 5.5 节的那套数据流。整个模型就是「embedding → N 层 Block → final RMSNorm → LM head → logits」。LM head 把 $d_{\text{model}}$ 维向量投到词表大小 $V$ 维（典型 50K~150K），softmax 之后是下一 token 的概率分布。

LM head 共享权重（weight tying）是常见优化：把 LM head 的权重矩阵设为 token embedding 的转置，可以省一份 $V \cdot d$ 的参数（通常占模型总参数的 5%~15%）。不是所有模型都共享：Llama 1/2 共享，Llama 3 不共享（因为词表扩到了 128K，效果上差异显著）。

本章小结

Self-Attention 是线性混合，必须配 FFN 引入非线性——FFN 是「逐 token 独立的两层 MLP」，提供模型主要的非线性表达力和大约 2/3 的参数容量。
激活函数演化：ReLU → GELU → SwiGLU。SwiGLU 用门控分支翻倍表达力，配合 $d_{\text{ffn}} = \frac{8}{3}d$ 保持参数量。
Residual 是深度网络可训的必要条件——残差连接给梯度提供恒等通道，让 80 层堆叠也能稳定收敛。
LayerNorm 把每层每个 token 的特征分布约束住——和 BatchNorm 不同，LayerNorm 不跨样本、不依赖 batch_size。
Pre-LN 远比 Post-LN 稳定——Post-LN 把 LN 接到 residual 路径上破坏了梯度的恒等性。今天的大模型几乎全部用 Pre-LN。
RMSNorm 是 LayerNorm 的简化版——去掉均值减法和偏移参数，速度快 7-10%，质量几乎不损失。
完整 Block 的数据流：x → RMSNorm → MHA → +residual → RMSNorm → FFN → +residual。每个 Block 输入输出形状一致，可任意堆叠。
参数量分解：MHA 占 1/3，FFN 占 2/3。计算量上短上下文 FFN 主导，长上下文 Attention 的 $T^2$ 项主导——这预示了第六部分推理优化的两条路线。

到这里第二部分结束。我们已经把 Transformer 的「数学层」拆完了——Self-Attention、Multi-Head、位置编码、Block 组装，所有零件都讲清楚了。

下一章进入第三部分「架构家族」。同样的 Transformer Block，可以搭出三种不同的架构：Encoder-only（BERT 系）、Encoder-Decoder（T5 系）、Decoder-only（GPT 系）。它们各自适合什么任务？为什么 Decoder-only 最终统治了大模型时代？这是第 6 章要回答的问题。

第 5 章 Transformer Block：FFN、LayerNorm、Residual 的角色

5.1 为什么 Attention 不够：缺非线性

5.2 FFN：每个 token 独立的非线性变换

FFN 的尺寸为什么是 4×

激活函数：从 ReLU 到 GELU 到 SwiGLU

5.3 Residual：让深度可训

Residual 的几何理解

5.4 LayerNorm：稳定每个位置的输入分布

LayerNorm 的定义

Pre-LN vs Post-LN：一个分水岭

RMSNorm：LayerNorm 的简化版

5.5 完整 Block 的数据流

5.6 参数量与计算量分解

5.7 一些容易踩的细节

5.8 把整个 Transformer 模型串起来

本章小结

延伸阅读