查看全课大纲（第 23 / 40 节）

从零实现一个最小 Transformer block

一个 Transformer block 最少要几个零件？

小象实战讲义 · 深度学习基础：从神经网络到 Transformer

前四节我们拆出了注意力、多头、位置编码、残差与 LayerNorm 这些零件。本节把它们装成一个能跑的最小单元：一个 Transformer block（编码器层）。学完你能回答：两个子层各自负责什么？残差加在哪、归一化放在前还是后，为什么会有两种流派？以及一个 block 到底有多少参数——用代码算，不靠估。

先修：缩放点积注意力的四步（本章 5.1 节）、nn.Linear 与 nn.LayerNorm 的用法、残差连接的基本形式。后续：下一节讲训练稳定性与复杂度，之后进入预训练范式。

💡 核心导读

两个子层：多头自注意力负责「位置之间交换信息」，前馈网络负责「逐位置做非线性变换」。
两条残差：每个子层外面都套一层 $x + \mathrm{Sublayer}(x)$ ，让梯度有一条直达通路。
两种归一化位置：论文原式是 post-LN，现代大模型普遍用 pre-LN；差别在深层训练是否容易。
一个 block 的参数量可以手算：四个投影矩阵 + 两个前馈矩阵 + 两个 LayerNorm，按 $d_{\text{model}}$ 与 $d_{ff}$ 就能对出来。

🧭 前置知识

缩放点积注意力： $\mathrm{softmax}(QK^{\top}/\sqrt{d_k})V$ 与多头拼接后的输出投影。
形状约定：(batch, seq_len, d_model)，下文记 $d_{\text{model}}$ 为模型宽度。
nn.LayerNorm(n) 对最后一维归一化，eps 默认为 1e-5（PyTorch 官方文档）。

整体结构：两个子层，两条残差

论文对编码器层的描述很具体（arXiv:1706.03762 §3.1）：每一层有两个子层，第一个是多头自注意力，第二个是逐位置的全连接前馈网络；每个子层外面都套一层残差连接，然后做层归一化，即

$\mathrm{output} = \mathrm{LayerNorm}\bigl(x + \mathrm{Sublayer}(x)\bigr)$

为了让残差能直接相加，所有子层的输入输出宽度都必须等于 $d_{\text{model}}$ （论文取 512）。这条约束决定了整个 block 的「等宽」设计：不管你中间怎么变宽变窄，出口必须回到 $d_{\text{model}}$ 。

x (B, L, d_model)
   │
   ├──────────────► 多头自注意力 ──► dropout ──┐
   │                                          + ──► LayerNorm ──► h
   └──────────────────────────────────────────┘
   │
   ├──────────────► 前馈网络 FFN ───► dropout ──┐
   │                                          + ──► LayerNorm ──► y
   └──────────────────────────────────────────┘

注意图里的两个细节，它们都有一手出处：dropout 加在子层输出上、加回残差之前，基础模型取 0.1（arXiv:1706.03762 §5.4）；归一化在残差相加之后，这是论文的原式（同 §3.1），也就是后来被称为 post-LN 的版本。

论文把这样的层堆叠 $N=6$ 次（§3.1）。堆叠的意义不在于「更深更强」，而在于分工逐层抽象：底层更多在做局部混合与逐位置变换，上层再在这之上继续组合。参数并不随层数共享——每一层有自己的一套投影与前馈矩阵，所以 $N$ 越大，参数量线性增长。

子层一：多头自注意力在 block 里做什么

单看注意力，它做的是「每个位置从其他位置取信息」。放进 block 之后，它承担的是位置之间的信息交换：把上一层各位置已经加工过的表示重新混合，让一个位置能看到与它相关的其他位置。

实现上有两个容易忽略的点：

投影后要拆头：(B, L, d_model) 先 view 成 (B, L, h, d_k) 再 transpose 成 (B, h, L, d_k)，注意力在最后两维上算，最后一维是 $d_k$ 。
拼接后要再投影： $h$ 个头的结果拼回 $d_{\text{model}}$ 后还要过一次输出投影 $W^O$ （arXiv:1706.03762 §3.2.2），否则各头之间没有可学习的混合方式。

这两步的细节已在 5.2 节（多头注意力）展开，这里只需要知道：它的输入输出形状都是 (B, L, d_model)，这是能塞进残差的前提。

子层二：前馈网络为什么是「逐位置」的

前馈网络（FFN）是两层线性变换夹一个 ReLU（论文原式）：

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

「逐位置」（position-wise）的意思是：对序列里每个位置的那个 $d_{\text{model}}$ 维向量分别做同一套变换，位置之间不交互。用代码写就是 nn.Linear 作用在最后一维上，前面所有维度自动并行——所以它在实现上和「对每个位置跑同一个 MLP」完全等价。

论文给的宽度是 $d_{ff} = 2048$ ，即 $d_{\text{model}}$ 的 4 倍（§3.3）。这个 4 倍不是随手取的，它让 FFN 成为 block 里参数最多的一块：两个矩阵分别是 $d_{\text{model}} \times d_{ff}$ 与 $d_{ff} \times d_{\text{model}}$ ，参数量约 $2 d_{\text{model}} d_{ff}$ ，是注意力四个投影（约 $4 d_{\text{model}}^2$ ）的两倍。

残差与 LayerNorm：为什么缺一不可

残差解决的是「深层网络梯度传不回去」。 $y = x + f(x)$ 里， $\partial y / \partial x = I + \partial f/\partial x$ ，恒等项保证梯度至少能原样往回走一层；这一思想来自 ResNet（He et al., 2015，arXiv:1512.03385），Transformer 直接沿用。

LayerNorm 解决的是「每个位置的数值尺度漂移」。它沿特征维（最后一维）计算均值与方差，再做可学习的仿射变换：

$\mathrm{LN}(x) = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} \cdot \gamma + \beta$

这里的细节值得记牢：归一化维度是每个位置自己的特征维，与 batch 无关、与其他位置无关——这也是它与 BatchNorm 的根本区别（BatchNorm 跨样本统计，见第 2 章正则化一节）。PyTorch 的 nn.LayerNorm(n) 里 n 就是被归一化的最后若干维，eps 默认 1e-5，elementwise_affine 默认开启。

post-LN 与 pre-LN：差一个加号的位置

把归一化放在残差之后（论文原式）还是之前，看起来只是顺序问题，实际影响训练稳定性：

post-LN： $x \leftarrow \mathrm{LN}(x + \mathrm{Sublayer}(x))$ 。论文 §3.1 的写法，也是原始 Transformer 的实现。它的输出被强制归一化，但深层堆叠时靠近输出的梯度尺度偏大，需要仔细的学习率 warmup。
pre-LN： $x \leftarrow x + \mathrm{Sublayer}(\mathrm{LN}(x))$ 。归一化放在子层入口，残差主干是一条不被归一化打断的直通路。Xiong 等人从初始化时的梯度尺度分析证明了 pre-LN 的梯度更平稳，因此可以去掉 warmup 而达到可比效果（arXiv:2002.04745，ICML 2020）。

今天的开源大模型绝大多数用 pre-LN 或其变体。本节的实现按论文原式写 post-LN，练习里再让你改成 pre-LN 对比——两种都写过，读别人的代码才不会卡在「为什么这里多一个 norm」。

顺带把这条差别和工程现象连起来：post-LN 对学习率 warmup 更敏感——Xiong 等人正是从「warmup 为什么必要」这个问题切入，才证明了归一化位置会改变初始化时的梯度尺度。这也是不少现代模型改用 pre-LN 的动机之一。

从零实现一个最小 block

下面把三个类拼起来：多头自注意力、前馈网络、以及包住它们的 block。为了 CPU 上跑得快，宽度压到 32、4 个头、 $d_{ff}=64$ ，形状关系与论文完全一致。

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"
        self.h = num_heads
        self.d_k = d_model // num_heads
        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)   # 拼接后的输出投影

    def forward(self, x):
        B, L, _ = x.shape
        q = self.w_q(x).view(B, L, self.h, self.d_k).transpose(1, 2)
        k = self.w_k(x).view(B, L, self.h, self.d_k).transpose(1, 2)
        v = self.w_v(x).view(B, L, self.h, self.d_k).transpose(1, 2)
        attn = F.scaled_dot_product_attention(q, k, v)       # (B, h, L, d_k)
        attn = attn.transpose(1, 2).contiguous().view(B, L, self.h * self.d_k)
        return self.w_o(attn)

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.fc2(F.relu(self.fc1(x)))                 # 逐位置，形状不变

class TransformerBlock(nn.Module):
    """post-LN 版本，严格对应论文 §3.1 的写法。"""
    def __init__(self, d_model=32, num_heads=4, d_ff=64, dropout=0.0):
        super().__init__()
        self.attn = MultiHeadSelfAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        x = self.norm1(x + self.drop(self.attn(x)))           # 子层一 + 残差 + LN
        x = self.norm2(x + self.drop(self.ffn(x)))            # 子层二 + 残差 + LN
        return x

block = TransformerBlock()
x = torch.randn(2, 6, 32)          # batch=2, seq=6, d_model=32
y = block(x)
print("输入形状 =", tuple(x.shape))
print("输出形状 =", tuple(y.shape))
print("输出均值 =", round(float(y.mean()), 4))
print("输出标准差 =", round(float(y.std()), 4))
counts = {name: sum(p.numel() for p in child.parameters()) for name, child in block.named_children()}
print("多头注意力参数量 =", counts["attn"])
print("前馈网络参数量 =", counts["ffn"])
print("两个 LayerNorm 参数量 =", counts["norm1"] + counts["norm2"])
print("block 总参数量 =", sum(p.numel() for p in block.parameters()))

输入形状 = (2, 6, 32)
输出形状 = (2, 6, 32)
输出均值 = 0.0
输出标准差 = 1.0013
多头注意力参数量 = 4096
前馈网络参数量 = 4192
两个 LayerNorm 参数量 = 128
block 总参数量 = 8416

输出里三件事值得对照：形状没变（(2, 6, 32) 进、(2, 6, 32) 出，才能一层层堆叠）；参数量分布（玩具配置取 $d_{ff}=2d_{\text{model}}$ ，FFN 与注意力基本相当；论文配置 $d_{ff}=4d_{\text{model}}$ 时 FFN 才约为注意力的两倍，下文对账会看到）；输出均值接近 0（post-LN 的最后一个操作就是归一化，均值被压到 0 附近）。

参数量对账：从玩具尺寸推到论文尺寸

同一个结构，把宽度换成论文取值（ $d_{\text{model}}=512$ 、 $d_{ff}=2048$ ）就能算出真实规模。下面这段只做整数运算，用来核对「一个 block 到底多大」：

d_model, d_ff = 512, 2048

attn_params = 4 * d_model * d_model          # W_q, W_k, W_v, W_o（无 bias）
ffn_params = 2 * d_model * d_ff + d_ff + d_model
ln_params = 2 * 2 * d_model                  # 两个 LayerNorm，各 gamma + beta

print("注意力投影参数量 =", attn_params)
print("前馈网络参数量 =", ffn_params)
print("两个 LayerNorm 参数量 =", ln_params)
print("单层合计 =", attn_params + ffn_params + ln_params)
print("FFN 是注意力的", round(ffn_params / attn_params, 4), "倍")

注意力投影参数量 = 1048576
前馈网络参数量 = 2099712
两个 LayerNorm 参数量 = 2048
单层合计 = 3150336
FFN 是注意力的 2.0024 倍

把「单层合计」乘 6（论文的 $N=6$ ）就是编码器主体的参数量级。这个数不是背下来的，而是按矩阵形状数出来的——你以后看到任何「XX 模型有 X 亿参数」的说法，都可以用同样方式拆开核对。

需要区分两个常被混为一谈的量：参数量是模型权重占用的规模，与输入多长无关；显存占用还包含前向留下的中间结果（激活）。训练时每个 block 都要保存注意力权重、子层输入等中间量供反向使用，序列越长这部分越大——所以同一个模型，推理能跑的长序列，训练未必塞得下。这也解释了为什么长上下文训练要靠梯度检查点、分块计算这类技巧，而不是单纯加显存。

还有一点：残差和 LayerNorm 都不改变宽度。残差是逐元素相加，LayerNorm 保持形状只做归一化，所以它们贡献的参数量很少（每层两个，各只有 $\gamma$ 与 $\beta$ ）。真正决定 block 规模的是四个注意力投影和两个前馈矩阵——这也意味着压缩模型时，动手改 FFN 的宽度收益最大。

顺带说明一个常被忽略的事实：block 本身不含任何位置信息。注意力对输入做的是「按内容加权」，如果你把序列顺序打乱，输出只是跟着重排，模型无法察觉顺序变了。位置编码因此必须加在 block 栈的入口（嵌入层之后），而不是塞进 block 内部——这正是位置编码一节（5.3）里那些正弦/可学习向量的落点。

📝 动手练一练

把上面的 TransformerBlock.forward 改成 pre-LN 写法（ $x \leftarrow x + \mathrm{Sublayer}(\mathrm{LN}(x))$ ），跑一遍，观察输出均值与标准差的变化，并说明为什么深层网络更偏好这种写法。

👉 点击查看参考答案

思路：pre-LN 把 norm 挪到子层入口，残差主干不再被归一化打断。直接后果是输出不再被强制拉回均值 0、方差 1，所以均值/标准差与 post-LN 不同；对深层训练的好处是主干上的梯度尺度更平稳（arXiv:2002.04745 的结论）。

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

class PreLNBlock(nn.Module):
    def __init__(self, d_model=32, num_heads=4, d_ff=64):
        super().__init__()
        self.h, self.d_k = num_heads, d_model // num_heads
        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)
        self.fc1, self.fc2 = nn.Linear(d_model, d_ff), nn.Linear(d_ff, d_model)
        self.norm1, self.norm2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)

    def attn(self, x):
        B, L, _ = x.shape
        q = self.w_q(x).view(B, L, self.h, self.d_k).transpose(1, 2)
        k = self.w_k(x).view(B, L, self.h, self.d_k).transpose(1, 2)
        v = self.w_v(x).view(B, L, self.h, self.d_k).transpose(1, 2)
        o = F.scaled_dot_product_attention(q, k, v)
        return self.w_o(o.transpose(1, 2).contiguous().view(B, L, -1))

    def forward(self, x):
        x = x + self.attn(self.norm1(x))                       # 归一化在子层入口
        x = x + self.fc2(F.relu(self.fc1(self.norm2(x))))
        return x

block = PreLNBlock()
y = block(torch.randn(2, 6, 32))
print("pre-LN 输出均值 =", round(float(y.mean()), 4))
print("pre-LN 输出标准差 =", round(float(y.std()), 4))

pre-LN 输出均值 = -0.0943
pre-LN 输出标准差 = 1.0987

常见错误

错误理解或写法	现象或报错	原因	正确做法
子层输出宽度不等于 $d_{\text{model}}$ 却直接相加	`RuntimeError: The size of tensor a (64) must match the size of tensor b (32)`	残差要求两边形状完全一致	让每个子层出口都投影回 $d_{\text{model}}$ ，输出投影 $W^O$ 与 FFN 第二层就是干这个的
拆头后忘了 `transpose` 回 `(B, L, h, d_k)` 再拼接	形状对得上但语义错乱，训练照样能跑	注意力在 `(B, h, L, d_k)` 上算，拼回序列维必须先转置	严格按 `view → transpose(1,2) → 注意力 → transpose(1,2) → view` 的顺序
把 LayerNorm 写成 BatchNorm	小 batch 时结果剧烈波动，与序列长度耦合	LayerNorm 沿特征维归一化、与 batch 无关；BatchNorm 跨样本统计	用 `nn.LayerNorm(d_model)`，参数是被归一化的特征维数
推理时忘记切 `eval()`	同一输入两次前向结果不同	训练态的 dropout 仍在生效	推理前 `model.eval()`；本讲义示例取 `dropout=0.0` 正是为了让输出可复现

本章小结

一个 Transformer block 由两个子层组成：多头自注意力（位置间交换信息）与逐位置前馈网络（逐位置非线性变换）。
每个子层外面套残差，出口统一回到 $d_{\text{model}}$ ，否则无法相加也无法堆叠。
论文原式是 post-LN： $\mathrm{LN}(x + \mathrm{Sublayer}(x))$ ；pre-LN 把归一化挪到子层入口，梯度更平稳（arXiv:2002.04745）。
FFN 是 block 里参数最多的部分（ $d_{ff}=4d_{\text{model}}$ 时约为注意力投影的两倍）。
dropout 加在子层输出上、残差相加之前，基础模型取 0.1（arXiv:1706.03762 §5.4）。

📋 行动清单

把 d_model 从 32 改成 48、num_heads 改成 4，先算一遍 d_k，再跑代码核对参数量变化。
故意把 num_heads 设为 5（不能整除 32），确认 assert 抛出的信息，再改回来。
把 block 的 forward 改成 pre-LN 写法，跑一遍，比较输出均值与标准差。