最近重新补大模型基础时,我发现自己虽然记得 Q、K、V 和注意力公式,却很难把每一步的形状变化、计算目的与完整的数据流真正串起来。因此我从输入表示开始,把标准 Transformer 中的注意力、多头拼接、残差连接、LayerNorm 和前馈网络重新推了一遍。

1. 输入与符号

设一句话经过分词后共有 $n$ 个 token,每个 token 用 $d_{model}$ 维向量表示。把这些向量逐行排列:

$$X\in\mathbb{R}^{n\times d_{model}}.$$

这里的 $X$ 已经包含词嵌入和位置信息。实际上,在原始 Transformer 中,

$$X=\sqrt{d_{model}}\,X_{embedding}+PE,$$

其中 $PE$ 是位置编码。词嵌入的分量在常见初始化下通常处于 $1/\sqrt{d_{model}}$ 的量级,而正弦、余弦位置编码的分量处于 $1$ 的量级。将词嵌入乘以 $\sqrt{d_{model}}$,可以使两者在训练初期具有相近的数值尺度,避免位置信息压过词义信息。这个缩放发生在输入端,与注意力中除以 $\sqrt{d_k}$ 的作用不同。$X$ 作为 Transformer Block 的输入。

注意力可以理解为:当前位置通过 Query 表达自己想找什么,用 Query 和其他位置的 Key 做匹配,再从匹配到的位置读取 Value 中的信息。

2. Q、K、V 与单头注意力

定义三个可学习的投影矩阵:

$$W^Q,W^K\in\mathbb{R}^{d_{model}\times d_k}, \qquad W^V\in\mathbb{R}^{d_{model}\times d_v}.$$

由输入 $X$ 得到

$$Q=XW^Q\in\mathbb{R}^{n\times d_k},$$

$$K=XW^K\in\mathbb{R}^{n\times d_k},$$

$$V=XW^V\in\mathbb{R}^{n\times d_v}.$$

有核心关系:

$$\mathrm{Score}=QK^\top.$$

因为

$$Q\in\mathbb{R}^{n\times d_k}, \qquad K^\top\in\mathbb{R}^{d_k\times n},$$

所以

$$QK^\top\in\mathbb{R}^{n\times n}.$$

其中第 $i$ 行第 $j$ 列为

$$\mathrm{Score}_{ij}=q_i^\top k_j.$$

它表示第 $i$ 个 token 的 Query 与第 $j$ 个 token 的 Key 的匹配程度。

将分数缩放并对每一行做 Softmax:

$$A= \operatorname{Softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right) \in\mathbb{R}^{n\times n}.$$

每一行满足

$$\sum_{j=1}^{n}A_{ij}=1.$$

最后使用注意力权重对 Value 加权求和:

$$Z=AV\in\mathbb{R}^{n\times d_v}.$$

因此单头缩放点积注意力的完整公式为

$$\boxed{ \operatorname{Attention}(Q,K,V) = \operatorname{Softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V }$$

第 $i$ 个位置的输出可以展开为

$$z_i=\sum_{j=1}^{n}A_{ij}v_j.$$

当前位置生成 Query,与各位置的 Key 计算匹配程度,再按照注意力权重从对应的 Value 中提取并汇总信息。

3. 为什么除以根号 $d_k$

假设查询和键的每一维彼此独立,并且满足

$$\mathbb{E}[q_r]=\mathbb{E}[k_r]=0, \qquad \operatorname{Var}(q_r)=\operatorname{Var}(k_r)=1.$$

点积为

$$q^\top k=\sum_{r=1}^{d_k}q_rk_r.$$

每一项的均值为 $0$,方差近似为 $1$,所以

$$\operatorname{Var}(q^\top k) = \sum_{r=1}^{d_k}\operatorname{Var}(q_rk_r) \approx d_k.$$

点积的标准差因此约为

$$\sqrt{d_k}.$$

当 $d_k$ 很大时,$QK^\top$ 的绝对值会变大,Softmax 容易接近 one-hot,非最大位置的梯度变小。除以 $\sqrt{d_k}$ 后,

$$\operatorname{Var} \left( \frac{q^\top k}{\sqrt{d_k}} \right) \approx 1,$$

分数的尺度更稳定。这里的 $d_k$ 是一个注意力头中 Query 和 Key 的维度,并不是序列长度。

4. 多头注意力

设注意力头数为 $h$,通常取

$$d_k=d_v=\frac{d_{model}}{h}.$$

第 $i$ 个头拥有独立的投影矩阵:

$$W_i^Q,W_i^K\in \mathbb{R}^{d_{model}\times d_k}, \qquad W_i^V\in \mathbb{R}^{d_{model}\times d_v}.$$

于是

$$Q_i=XW_i^Q,\qquad K_i=XW_i^K,\qquad V_i=XW_i^V.$$

并且

$$head_i = \operatorname{Attention}(Q_i,K_i,V_i) \in\mathbb{R}^{n\times d_v}.$$

将所有头沿特征维拼接:

$$H= \operatorname{Concat} (head_1,head_2,\ldots,head_h) \in\mathbb{R}^{n\times hd_v}.$$

当 $d_v=d_{model}/h$ 时,$hd_v=d_{model}$。最后乘输出投影矩阵

$$W^O\in\mathbb{R}^{hd_v\times d_{model}},$$

得到

$$\boxed{ \operatorname{MHA}(X) = \operatorname{Concat} (head_1,\ldots,head_h)W^O }$$

输出形状重新变成

$$\operatorname{MHA}(X) \in\mathbb{R}^{n\times d_{model}},$$

因此可以与输入 $X$ 做残差相加。多头的意义在于同时得到 $h$ 张不同的注意力矩阵,让不同投影空间分别学习不同关系。

例如

$$d_{model}=512,\qquad h=8.$$

$$d_k=d_v=64.$$

缩放时使用的是 $\sqrt{64}$,不是 $\sqrt{512}$。

4.1 Concat 到底是什么

Concat 是 concatenate(拼接)的缩写。它不做加法,也不做平均,而是把各个注意力头在同一个 token 上得到的特征首尾接起来。

例如,序列中有两个 token,每个头为每个 token 输出两个特征。若

$$head_1= \begin{bmatrix} a_1&b_1\\ a_2&b_2 \end{bmatrix}, \qquad head_2= \begin{bmatrix} c_1&d_1\\ c_2&d_2 \end{bmatrix},$$

那么沿特征维拼接后

$$\operatorname{Concat}(head_1,head_2)= \begin{bmatrix} a_1&b_1&c_1&d_1\\ a_2&b_2&c_2&d_2 \end{bmatrix}.$$

行数仍然是 token 数,变化的是每行的特征数。因此,$h$ 个形状为 $n\times d_v$ 的头拼接后,形状为 $n\times(hd_v)$。当 $d_v=d_{model}/h$ 时,拼接结果正好回到 $n\times d_{model}$。随后乘 $W^O$,不是为了再次改变总维数,而是把各个头原本并排放置的特征进行一次可学习的混合。

多头注意力从 QKV 到 Concat 与输出投影的完整过程

5. 残差连接与 LayerNorm

设某个 token 的特征为

$$x=(x_1,x_2,\ldots,x_{d_{model}}).$$

LayerNorm 在这一个 token 的全部特征维上计算均值和方差:

$$\mu = \frac{1}{d_{model}} \sum_{j=1}^{d_{model}}x_j,$$

$$\sigma^2 = \frac{1}{d_{model}} \sum_{j=1}^{d_{model}}(x_j-\mu)^2.$$

标准化后再进行可学习的仿射变换:

$$\operatorname{LayerNorm}(x)_{j}= \gamma_{j}\frac{x_{j}-\mu}{\sqrt{\sigma^{2}+\epsilon}} + \beta_{j}.$$

其中,可学习参数 $\gamma\in\mathbb{R}^{d_{model}}$ 与 $\beta\in\mathbb{R}^{d_{model}}$。每个 token 独立计算自己的均值和方差,但所有 token 共享同一组 $\gamma$ 和 $\beta$。

残差连接要求子层输出和输入形状相同。原始 Transformer 的注意力子层使用 post-LN:

$$\boxed{ Y_{attn} = \operatorname{LayerNorm} \left( X+\operatorname{Dropout}(\operatorname{MHA}(X)) \right) }$$

残差中的 $X$ 提供恒等路径,使信息和梯度可以绕过子层直接传播;LayerNorm 负责稳定每个 token 的特征尺度。

5.1 LayerNorm 到底在做什么

LayerNorm 的作用是为每个 token 的特征建立稳定的数值尺度。以 $x=(1,2,3)$ 为例,其均值为 $2$、方差为 $2/3$;减去均值并除以标准差后,各维特征被调整到均值约为 $0$、方差约为 $1$ 的尺度。随后再使用可学习参数 $\gamma$ 和 $\beta$ 恢复模型所需的缩放与平移,因此归一化不会限制网络只能输出固定分布。

LayerNorm 对每个 token 分别处理,只在该 token 的 $d_{model}$ 个特征之间计算均值和方差,不负责 token 之间的信息传递。它可以抑制不同层之间激活尺度的持续漂移,改善优化条件,并使梯度和注意力分数更加稳定。若直接去掉 LayerNorm,浅层模型仍可能训练,但深层 Transformer 的残差累积容易使特征尺度不断变化,进而导致 Softmax 过早饱和、梯度消失或爆炸、训练速度变慢,甚至出现损失发散。无归一化模型并非不可实现,但通常需要专门的初始化、残差缩放或其他稳定化设计。

LayerNorm 发生在注意力子层和 FFN 子层的附近。原始 Transformer 使用 Post-LN,即先让子层计算,再和输入做残差相加,最后归一化:

$$y=\operatorname{LayerNorm}\bigl(x+\operatorname{Sublayer}(x)\bigr).$$

现代 GPT 常使用 Pre-LN,即先归一化,再把归一化后的结果送进子层,最后与未经归一化的残差主路相加:

$$y=x+\operatorname{Sublayer}\bigl(\operatorname{LayerNorm}(x)\bigr).$$

6. 前馈神经网络 FFN

前馈神经网络是信息只从输入层流向输出层、内部不存在循环或反馈连接的神经网络。在 Transformer 中,FFN 是对每个 token 独立应用、但在所有位置共享参数的两层多层感知机。注意力负责在 token 之间交换信息,FFN 则沿特征维进行非线性变换。其表达式为

$$\operatorname{FFN}(x) = \phi(xW_1+b_1)W_2+b_2,$$

其中

$$W_1\in \mathbb{R}^{d_{model}\times d_{ff}}, \qquad W_2\in \mathbb{R}^{d_{ff}\times d_{model}}.$$

原始 Transformer 使用

$$\phi=\operatorname{ReLU},$$

并通常取

$$d_{ff}=4d_{model}.$$

FFN 先把特征从 $d_{model}$ 扩展到 $d_{ff}$,经过非线性函数后再压回 $d_{model}$。标准 post-LN 写法为

$$\boxed{ Y_{next} = \operatorname{LayerNorm} \left( Y_{attn} + \operatorname{Dropout} (\operatorname{FFN}(Y_{attn})) \right) }$$

这样,一个标准 Transformer Block 的核心就是

$$X \quad\text{经过 MHA、残差和 LayerNorm 得到}\quad Y_{attn},$$

再经过 FFN、残差和 LayerNorm 得到

$$Y_{next}.$$

7. 整个 Transformer 的数据流

原始 Transformer 是编码器—解码器结构。编码器读取源序列并形成上下文表示;解码器一边查看已经生成的目标 token,一边读取编码器输出,最后预测下一个 token。

原始 Transformer 编码器—解码器结构与 Add & Normalize

每个编码器 Block 内部的数据顺序是:自注意力 $→$ 残差相加与 LayerNorm $→$ FFN $→$ 残差相加与 LayerNorm。自注意力让每个 token 汇总整条源序列的信息,FFN 再分别加工每个 token 的特征。这个 Block 重复 $N$ 次,输出记作 $H$。

每个解码器 Block 多一个交叉注意力:先做带因果掩码的自注意力,使当前位置只能看见已经生成的内容;再做交叉注意力,其中 Query 来自解码器当前状态,而 Key、Value 来自编码器输出 $H$;最后经过 FFN。线性层把最后的 $d_{model}$ 维向量映射成词表大小的 logits,Softmax 再把它变成下一个 token 的概率分布。

GPT 是 decoder-only Transformer,没有单独的编码器和交叉注意力。输入 token 经过词嵌入与位置编码后,依次通过若干个“因果自注意力 + FFN”Block,最后归一化并映射到词表:

GPT-2 的 decoder-only 堆叠结构

8. Post-LN 与 Pre-LN:GPT 为什么改变位置

现代模型常见的 pre-LN:

$$Y_{attn} = X+ \operatorname{MHA} (\operatorname{LayerNorm}(X)),$$

$$Y_{next} = Y_{attn} + \operatorname{FFN} (\operatorname{LayerNorm}(Y_{attn})).$$

Pre-LN 将 LayerNorm 放在子层之前,深层训练通常更稳定。现代 Transformer 的 FFN 常使用 GELU 或 SwiGLU;原始 Transformer 则采用 Post-LN 与 ReLU。

Post-LN 与 Pre-LN 都包含残差连接和 LayerNorm。

两种结构的差别不在于有没有 LayerNorm,而在于 LayerNorm 位于残差加法的哪一侧。

Post-LN 与 Pre-LN Transformer 层对照

原始 Transformer 的一个子层写成

$$y=\operatorname{LayerNorm}\bigl(x+\operatorname{Sublayer}(x)\bigr).$$

GPT-2 所采用的 Pre-LN 写成

$$y=x+\operatorname{Sublayer}\bigl(\operatorname{LayerNorm}(x)\bigr).$$

将它代入完整的 GPT Block,第 $l$ 层的数据流为

$$u_l=x_l+\operatorname{CausalMHA}\bigl(\operatorname{LayerNorm}(x_l)\bigr),$$

$$x_{l+1}=u_l+\operatorname{MLP}\bigl(\operatorname{LayerNorm}(u_l)\bigr).$$

堆叠完全部 Block 后,再执行一次 Final LayerNorm,然后送入词表投影层。实现时,Pre-LN 的关键是保留一条从 $x_l$ 到 $x_{l+1}$ 的直接恒等主路:LayerNorm 和子层都放在旁边的残差分支上,主路本身不经过归一化。这使深层网络反向传播时,梯度拥有一条更直接的路径。

从训练角度看,研究表明原始 Post-LN 在初始化时,靠近输出层的参数可能出现较大的期望梯度,因此通常依赖学习率 warm-up 来减轻早期训练的不稳定;Pre-LN 在初始化时的梯度更规整,深层模型往往更容易优化。这是训练性质的差异,并不表示 Pre-LN 在所有指标上必然更好。

8.1 原始资料

本节结构与结论主要对照 Attention Is All You NeedLanguage Models are Unsupervised Multitask LearnersOn Layer Normalization in the Transformer Architecture。其中 Post-LN/Pre-LN 对照图取自 Xiong 等人的论文,其余结构图取自 Jay Alammar 的 The Illustrated TransformerThe Illustrated GPT-2,原图遵循 CC BY-NC-SA 4.0 许可。

9. 标准 Transformer 中的其余必要部分

自注意力本身不包含顺序,所以输入需要加入位置编码。原始正弦位置编码为

$$PE_{pos,2i} = \sin \left( \frac{pos}{10000^{2i/d_{model}}} \right),$$

$$PE_{pos,2i+1} = \cos \left( \frac{pos}{10000^{2i/d_{model}}} \right).$$

解码器的因果注意力还要在 Score 上加入掩码 $M$:

$$A= \operatorname{Softmax} \left( \frac{QK^\top}{\sqrt{d_k}}+M \right),$$

其中未来位置对应 $M_{ij}=-\infty$,Softmax 后权重为 $0$。交叉注意力仍使用同一个注意力公式,只是 Q 来自解码器,K 和 V 来自编码器输出。

需要优先手写并掌握的核心公式只有四组:

$$Q=XW^Q,\qquad K=XW^K,\qquad V=XW^V.$$

$$\operatorname{Attention}(Q,K,V) = \operatorname{Softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V,$$

$$\operatorname{MHA}(X) = \operatorname{Concat} (head_1,\ldots,head_h)W^O,$$

以及标准 Transformer Block 的两次 Add and Norm:

$$Y_{attn} = \operatorname{LayerNorm} \left( X+\operatorname{MHA}(X) \right),$$

$$Y_{next} = \operatorname{LayerNorm} \left( Y_{attn}+\operatorname{FFN}(Y_{attn}) \right).$$

掌握时应同时说出每个矩阵的形状,并解释为什么缩放使用 $\sqrt{d_k}$、为什么多头输出必须投影回 $d_{model}$、为什么 FFN 最终也要回到 $d_{model}$。

最后修改:2026 年 08 月 03 日
如果觉得我的文章对你有用,请随意赞赏