从「狗咬人」手算 Transformer:前向、损失与反向传播全流程

总结摘要
用 L=3、d=4 的「狗咬人」小例子,把 Decoder-only Transformer 从嵌入一路手算到交叉熵损失与反向传播梯度——每步先讲是什么、为什么,再代进可复现的小数字。

本文以一个贯穿全篇的具体例子——输入句子「狗咬人」——逐步走完标准 Transformer(Decoder-only,GPT 式自回归)的前向全流程,并给出关键数学推导与反向传播梯度。所有中间量都取可手算复现的小数字。凡涉及的运算(点积、矩阵乘、softmax、方差、KL 散度等)都先解释是什么、为什么这里要用它,再代入数字。

记号约定:序列长 $L=3$,模型维 $d=4$,主线单头注意力(头数 $h=1$,故头维 $d_k=d/h=4$,$\sqrt{d_k}=2$),前馈隐维 $d_{\mathrm{ff}}=8$,词表 $\mathcal V$ 大小 $V=5$,层数 $N$(例子用 $N=1$)。矩阵按行存 token:$X\in\mathbb{R}^{L\times d}$ 第 $t$ 行是第 $t$ 个 token 的表示。


1. 记号与总体数据流

一个 Decoder-only Transformer 把一串 token 映射为「下一个 token 的概率分布」,其数据流为:

$$ \text{token 序列}\ \xrightarrow{\text{嵌入+位置}}\ X^{(0)}\in\mathbb{R}^{L\times d}\ \xrightarrow{\;N\times \text{block}\;}\ H\in\mathbb{R}^{L\times d}\ \xrightarrow{\text{输出层}}\ \text{logits}\in\mathbb{R}^{L\times V}\ \xrightarrow{\operatorname{softmax}}\ P\ \xrightarrow{\text{交叉熵}}\ \mathcal L $$

每个 block 内部(Pre-LN 版):

$$ X \to X + \operatorname{MHA}(\operatorname{LN}(X)) \to [\,\cdot\,] + \operatorname{FFN}(\operatorname{LN}([\,\cdot\,])) \to \text{下一层} $$

为清晰起见,本文正文按 Post-LN 顺序演示(先子层、再残差、再 LN),两者差异见 §5 。单个 block 的数据流如下:

flowchart TB IN["输入 X ∈ ℝ^(L×d)"] --> MHA["多头自注意力 MHA<br>(因果掩码,聚合上下文)"] MHA --> ADD1(("+")) IN -. 残差捷径 .-> ADD1 ADD1 --> LN1["层归一化 LN"] LN1 --> FFN["前馈网络 FFN<br>(逐位置非线性)"] FFN --> ADD2(("+")) LN1 -. 残差捷径 .-> ADD2 ADD2 --> LN2["层归一化 LN"] LN2 --> OUT["输出 H ∈ ℝ^(L×d)<br>(重复 ×N 层)"]

灰线为主干,虚线为残差捷径——输入绕过子层,在 $\oplus$ 处与子层输出逐元素相加,再做层归一化。整块沿主干堆叠 $N$ 层。

模块输入 → 输出核心机制
输入表示ids $\to X\in\mathbb{R}^{L\times d}$token 嵌入 + 位置编码
自注意力$X\to\mathbb{R}^{L\times d}$每个 token 按相关性聚合上下文(带因果掩码)
前馈 FFN$\mathbb{R}^{L\times d}\to\mathbb{R}^{L\times d}$逐位置非线性变换
残差 + LN恒等捷径 + 归一化,稳定训练
输出层$H\to\text{logits}\in\mathbb{R}^{L\times V}$线性投影到词表
损失$\text{logits}\to\mathcal L\in\mathbb{R}$交叉熵

贯穿例子的固定设置。 词表 $\mathcal V=\{\text{狗},\text{咬},\text{人},\text{追},\langle\texttt{eos}\rangle\}$,对应 id $0,1,2,3,4$。输入「狗咬人」切分为 3 个 token,id 序列 $[0,1,2]$。嵌入矩阵 $E\in\mathbb{R}^{5\times4}$(第 $v$ 行是第 $v$ 个词的向量)取:

$$ E=\begin{bmatrix}1&0&0&0\\ 0&1&0&0\\ 0&0&1&0\\ 0&0&0&1\\ 1&1&1&1\end{bmatrix}\quad\begin{matrix}\text{狗}\\\text{咬}\\\text{人}\\\text{追}\\\langle\texttt{eos}\rangle\end{matrix} $$

注意力权重 $W^Q,W^K,W^V\in\mathbb{R}^{4\times4}$、输出 $W^O$、前馈 $W_1\in\mathbb{R}^{4\times8},W_2\in\mathbb{R}^{8\times4}$ 的具体数值在用到时给出。


2. 输入表示:从「狗咬人」到向量

2.1 Token 形成(字级切分)

计算机不能直接处理汉字,必须先把文本切成离散符号(token)并映射为整数 id。切分粒度可以是字、子词(BPE)或词;本文用最简单的字级切分:

$$ \text{「狗咬人」}\ \xrightarrow{\text{分词}}\ [\text{狗},\ \text{咬},\ \text{人}]\ \xrightarrow{\text{查词表}}\ [\,0,\ 1,\ 2\,] $$

真实大模型多用 BPE(Byte-Pair Encoding)子词分词。

📐 BPE 子词分词:为什么不用字、也不用词

两个极端的困境:纯字级切分词表小(几千汉字)但序列长、语义碎(「预」「训」「练」三个字才凑出一个概念);纯词级切分语义完整但词表爆炸(几十万词)且遇到生词(OOV)无法表示。BPE 折中:从字符开始,反复合并语料中最高频的相邻符号对,直到词表达到预设大小。

BPE 训练小例子。 设语料只有单词 low, low, lower, newest, newest。初始符号是单字符,反复统计相邻对频次并合并最高频者:

轮次最高频相邻对合并后新符号
1e s(出现在 newest×2)es
2es test
3l olo
直到词表达到目标大小

最终词表混合字符与高频子词(如 estlow)。编码新文本时,用这套合并规则贪心地把字符流合并成最长可用子词;高频词被整体表示(短序列),生词退化为若干子词甚至单字符(永不 OOV)。GPT 系用的是在字节上做 BPE(Byte-level BPE),任何 UTF-8 文本都能表示。原理不影响后续流程,本文主线仍用字级切分。

📐 数学定义:one-hot 向量与「查表即矩阵乘」

第 $v$ 个 token 的 one-hot 向量 $\boldsymbol e_v\in\mathbb{R}^{V}$ 是「第 $v$ 位为 1、其余为 0」的向量。查嵌入表取第 $v$ 行,等价于用 one-hot 向量左乘 $E$:

$$\boldsymbol e_v^\top E = E[v,:]\quad(\text{取出第 }v\text{ 行})$$

为什么这样看? 因为「查表」是不可导的整数索引操作,而「one-hot 乘矩阵」是可导的线性运算——这让嵌入矩阵 $E$ 能随梯度一起训练。

2.2 Token 嵌入

把三个 id 各自查表,堆叠成输入矩阵 $X_{\text{emb}}=E[[0,1,2]]\in\mathbb{R}^{3\times4}$:

$$ X_{\text{emb}}=\begin{bmatrix}1&0&0&0\\ 0&1&0&0\\ 0&0&1&0\end{bmatrix}\quad\begin{matrix}\text{狗}\\\text{咬}\\\text{人}\end{matrix} $$

2.3 位置编码

自注意力本身对顺序无感知(见 §3 ,交换两个 token 的位置,输出只是相应交换,不改变数值关系)——但「狗咬人」和「人咬狗」意义相反。因此必须显式注入位置信息。原始 Transformer 用正弦位置编码

$$ PE_{(pos,\,2i)}=\sin\!\Big(\frac{pos}{10000^{2i/d}}\Big),\qquad PE_{(pos,\,2i+1)}=\cos\!\Big(\frac{pos}{10000^{2i/d}}\Big) $$

其中 $pos$ 是位置($0,1,2$),$i$ 是维度对的下标($d=4$ 时 $i\in\{0,1\}$)。代入 $pos=0,1,2$:

$$ PE=\begin{bmatrix}\sin 0 & \cos 0 & \sin 0 & \cos 0\\ \sin 1 & \cos 1 & \sin\frac{1}{100} & \cos\frac{1}{100}\\ \sin 2 & \cos 2 & \sin\frac{2}{100} & \cos\frac{2}{100}\end{bmatrix}\approx\begin{bmatrix}0 & 1 & 0 & 1\\ 0.8415 & 0.5403 & 0.0100 & 1.0000\\ 0.9093 & -0.4161 & 0.0200 & 0.9998\end{bmatrix} $$

输入表示 $=$ 嵌入 $+$ 位置编码(逐元素相加):

$$ X=X_{\text{emb}}+PE=\begin{bmatrix}1 & 1 & 0 & 1\\ 0.8415 & 1.5403 & 0.0100 & 1\\ 0.9093 & -0.4161 & 1.0200 & 0.9998\end{bmatrix}\in\mathbb{R}^{3\times4}\tag{2.1} $$

✨ 优雅之处:相对位置线性可表达

由三角恒等式,$PE_{pos+k}$ 的每一对 $(\sin,\cos)$ 分量都是 $PE_{pos}$ 同一对分量的固定旋转(角度只依赖偏移 $k$):

$$\begin{bmatrix}\sin(\omega(pos{+}k))\\ \cos(\omega(pos{+}k))\end{bmatrix}=\begin{bmatrix}\cos(\omega k) & \sin(\omega k)\\ -\sin(\omega k) & \cos(\omega k)\end{bmatrix}\begin{bmatrix}\sin(\omega\,pos)\\ \cos(\omega\,pos)\end{bmatrix}$$

所以模型能用一个与内容无关的线性变换感知「相对距离 $k$」,这正是正弦编码的设计动机。


3. 缩放点积注意力

注意力机制回答一个问题:为了理解当前 token,应该从序列里的哪些 token 各取多少信息?

📐 数学定义:点积与它衡量的东西

两个向量 $\boldsymbol q,\boldsymbol k\in\mathbb{R}^{d_k}$ 的点积为 $\boldsymbol q\cdot\boldsymbol k=\sum_{i=1}^{d_k} q_i k_i=\|\boldsymbol q\|\,\|\boldsymbol k\|\cos\theta$。它同时编码了两向量的长度夹角:方向越一致($\cos\theta$ 越大)、模越大,点积越大。因此点积是天然的相似度/相关性打分——注意力正是用它衡量「query 和 key 有多匹配」。

3.1 Query / Key / Value

每个 token 用三个可学习矩阵投影出三种角色向量:Query(我要找什么)、Key(我能提供什么标签)、Value(我实际携带的信息)。

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

本例取(简单但非平凡,使 $Q,K,V$ 各不相同):

$$ W^Q=I_4,\quad W^K=\begin{bmatrix}0&1&0&0\\1&0&0&0\\0&0&1&0\\0&0&0&1\end{bmatrix}\ (\text{交换前两维}),\quad W^V=\operatorname{diag}(1,1,2,1)\ (\text{第三维}\times2) $$

代入 (2.1):

$$ \begin{gathered} Q=X,\qquad K=\begin{bmatrix}1&1&0&1\\1.5403&0.8415&0.0100&1\\-0.4161&0.9093&1.0200&0.9998\end{bmatrix},\\[6pt] V=\begin{bmatrix}1&1&0&1\\0.8415&1.5403&0.0200&1\\0.9093&-0.4161&2.0400&0.9998\end{bmatrix} \end{gathered} $$

3.2 注意力分数矩阵

📐 数学定义:矩阵乘 $QK^\top$

$Q\in\mathbb{R}^{L\times d_k}$ 的每一行是一个 query,$K$ 的每一行是一个 key。乘积 $(QK^\top)_{ij}=\boldsymbol q_i\cdot\boldsymbol k_j$ 一次性算出所有 query–key 两两点积,得到 $L\times L$ 的分数矩阵——第 $i$ 行第 $j$ 列 = 「第 $i$ 个 token 对第 $j$ 个 token 的关注打分」。

$$ \frac{QK^\top}{\sqrt{d_k}}=\frac{1}{2}\,QK^\top\approx\begin{bmatrix}1.5000 & 1.6909 & 0.7465\\ 1.6909 & 1.7961 & 1.0302\\ 0.7465 & 1.0302 & 0.6416\end{bmatrix} $$

✨ 优雅之处:为什么除以 $\sqrt{d_k}$——方差归一推导

设 $\boldsymbol q,\boldsymbol k$ 各分量独立、零均值、单位方差,且相互独立。点积 $\boldsymbol q\cdot\boldsymbol k=\sum_{i=1}^{d_k}q_ik_i$。每一项 $\mathbb E[q_ik_i]=0$,其方差 $\operatorname{Var}(q_ik_i)=\mathbb E[q_i^2]\mathbb E[k_i^2]=1$,各项独立故 $\operatorname{Var}(\boldsymbol q\cdot\boldsymbol k)=d_k$。于是点积的标准差是 $\sqrt{d_k}$。除以它,本质上就是把一个零均值随机变量标准化成方差 1(统计学里的 $z=(x-\mu)/\sigma$)——让分数尺度不随维度 $d_k$ 漂移。若不缩放,大分数经 softmax 会饱和(几乎变成 one-hot),梯度趋零、难以训练。

📐 数学定义:softmax

$\operatorname{softmax}:\mathbb{R}^{n}\to\mathbb{R}^{n}$ 把任意实向量变成概率分布

$$\operatorname{softmax}(\boldsymbol z)_i=\frac{e^{z_i}}{\sum_{j=1}^{n}e^{z_j}},\qquad \operatorname{softmax}(\boldsymbol z)_i>0,\ \sum_i\operatorname{softmax}(\boldsymbol z)_i=1$$

用指数是因为它单调、恒正,能放大差距又保序;归一化保证结果是合法概率。

🔢 数值演示:$\sqrt{d_k}$ 缩放如何缓和 softmax 饱和

取构造的大分数 $\boldsymbol z=(2,4,6,8)$:不缩放时 $\operatorname{softmax}(\boldsymbol z)\approx(0.002,0.016,0.117,0.865)$——几乎全押在最后一项;除以 $\sqrt{4}=2$ 后 $\operatorname{softmax}(\boldsymbol z/2)\approx(0.032,0.087,0.237,0.644)$,分布明显更平滑,梯度更健康。

3.3 因果掩码

自回归生成要求「预测第 $t$ 个 token 时只能看 $\le t$ 的位置」,否则会偷看答案。做法是在 softmax 前给未来位置加 $-\infty$:

$$ M_{ij}=\begin{cases}0,& j\le i\\ -\infty,& j> i\end{cases},\qquad \text{分数}\leftarrow \frac{QK^\top}{\sqrt{d_k}}+M $$

$e^{-\infty}=0$,未来位置的注意力权重自动归零。掩码后的分数矩阵($\cdot$ 表示被屏蔽):

$$ \begin{bmatrix}1.5000 & \cdot & \cdot\\ 1.6909 & 1.7961 & \cdot\\ 0.7465 & 1.0302 & 0.6416\end{bmatrix} $$

✨ 优雅之处:加 $-\infty$——前向屏蔽,反向也「免费」正确

因果掩码只做一件极简的事——给未来位置的分数加 $-\infty$。前向 $e^{-\infty}=0$,未来位置权重归零、看不到答案。妙在反向不需要任何特殊处理:那些位置前向 $A=0$,代入 §8.5 的行 softmax 反传式,其梯度自动为 0。一个「$+(-\infty)$」既在前向切断了信息、又在反向切断了梯度,两个方向一次搞定——简洁得近乎取巧。

3.4 注意力权重与输出

对每一行做 softmax,得注意力权重 $A$(下三角,每行和为 1):

$$ A=\operatorname{softmax}_{\text{行}}\!\Big(\frac{QK^\top}{\sqrt{d_k}}+M\Big)\approx\begin{bmatrix}1 & 0 & 0\\ 0.4737 & 0.5263 & 0\\ 0.3097 & 0.4114 & 0.2789\end{bmatrix} $$

例如「咬」(第 2 行)对「狗」分配 $0.474$、对自己分配 $0.526$。最后用权重对 Value 加权求和 $\operatorname{Attn}=AV$:

$$ \operatorname{Attn}(Q,K,V)=\operatorname{softmax}\!\Big(\frac{QK^\top}{\sqrt{d_k}}+M\Big)V\approx\begin{bmatrix}1 & 1 & 0 & 1\\ 0.9166 & 1.2844 & 0.0105 & 1\\ 0.9095 & 0.8273 & 0.5772 & 0.9999\end{bmatrix}\tag{3.1} $$

3.5 softmax 的雅可比(完整推导,为反向传播铺垫)

softmax 把向量映射为向量($\mathbb{R}^n\to\mathbb{R}^n$),故其导数是一个 $n\times n$ 的雅可比矩阵 $J$,$J_{ij}=\partial p_i/\partial z_j$。设 $\boldsymbol p=\operatorname{softmax}(\boldsymbol z)$,$p_i=e^{z_i}/S$,$S=\sum_{k}e^{z_k}$。注意 $S$ 依赖所有 $z_j$——这是 $j\ne i$ 时导数不为零的根源。

分母求导 $\partial S/\partial z_j=e^{z_j}$。对 $p_i=e^{z_i}/S$ 用商法则:

$$ \frac{\partial p_i}{\partial z_j}=\frac{\frac{\partial e^{z_i}}{\partial z_j}\cdot S-e^{z_i}e^{z_j}}{S^2} $$
  • $i=j$(对角):$\partial e^{z_i}/\partial z_i=e^{z_i}$,得 $\frac{e^{z_i}S-e^{z_i}e^{z_i}}{S^2}=p_i-p_i^2=p_i(1-p_i)$。
  • $i\ne j$(非对角):$\partial e^{z_i}/\partial z_j=0$,得 $-\frac{e^{z_i}e^{z_j}}{S^2}=-p_ip_j$。

用 Kronecker $\delta$ 合并:

$$ \boxed{\ \frac{\partial p_i}{\partial z_j}=p_i(\delta_{ij}-p_j)\ },\qquad J=\operatorname{diag}(\boldsymbol p)-\boldsymbol p\,\boldsymbol p^\top\in\mathbb{R}^{n\times n}\tag{3.2} $$

对称每行/每列和为零:$\sum_i J_{ij}=p_j-p_j\sum_i p_i=0$。

🔢 数值验证:对「咬」行 softmax 求雅可比

取「咬」(第 2 行)的注意力权重 $\boldsymbol p=(0.4737,0.5263,0)$(第 3 列被掩码置零):

$$J=\operatorname{diag}(\boldsymbol p)-\boldsymbol p\boldsymbol p^\top\approx\begin{bmatrix}0.2493 & -0.2493 & 0\\ -0.2493 & 0.2493 & 0\\ 0 & 0 & 0\end{bmatrix}$$

对角元 $J_{11}=0.4737(1-0.4737)=0.2493$,非对角 $J_{12}=-0.4737\times0.5263=-0.2493$。每行和均为 0,第 3 行/列全零(被掩码的位置对梯度无贡献)。

✨ 优雅之处:「行和为零」是概率守恒的微分版本

雅可比每列和恒为 0 绝非偶然:softmax 输出永远满足 $\sum_i p_i=1$,无论怎么扰动输入,输出总和都不能变——「总和的变化率」必须是 0。物理直觉:概率是守恒的,多分给一个词的,必然从别的词那里挪来(对角正、非对角负,此消彼长)。这个守恒律一路贯穿到反传的「减去行内加权和」操作。

饱和与梯度消失。 当某 $p_i\to1$、其余 $\to0$(softmax 饱和,接近 one-hot),$J$ 的所有元素趋零——梯度消失。这正是需要 $\sqrt{d_k}$ 缩放来防止 softmax 提前饱和的数学原因。


4. 多头注意力

单个注意力头只能学一种「关注模式」。多头把 $Q,K,V$ 切成 $h$ 份低维子空间,各自独立做注意力,再拼接——让模型在不同子空间同时捕捉不同关系(如句法、指代、位置)。

$$ \text{head}_i=\operatorname{Attn}(XW_i^Q,\,XW_i^K,\,XW_i^V),\qquad \operatorname{MHA}(X)=\operatorname{Concat}(\text{head}_1,\dots,\text{head}_h)\,W^O $$

📐 数学定义:Concat 与 $W^O$

$\operatorname{Concat}$ 沿特征维把 $h$ 个 $\mathbb{R}^{L\times d_k}$ 拼成 $\mathbb{R}^{L\times (h d_k)}=\mathbb{R}^{L\times d}$。输出投影 $W^O\in\mathbb{R}^{d\times d}$ 把各头的信息混合融合回模型维度——否则各头结果只是并排堆着、互不交流。

本节数字与主线例子无关,仅演示拼接机制。 主线仍用单头($h=1,d_k=4$)。

在「狗咬人」输入表示 $X$((2.1))上取 $h=2$,每头 $d_k=2$:把 $X$ 的 4 维切成前 2 维(头 1)与后 2 维(头 2),各头独立做注意力。两头算出的注意力权重矩阵明显不同

$$ A^{(1)}\approx\begin{bmatrix}1&0&0\\0.379&0.621&0\\0.312&0.241&0.447\end{bmatrix},\qquad A^{(2)}\approx\begin{bmatrix}1&0&0\\0.500&0.500&0\\0.244&0.246&0.510\end{bmatrix} $$

例如「咬」(第 2 行)在头 1 里偏向自己($0.62$),在头 2 里对「狗」「咬」均分注意($0.5/0.5$)——两个头在不同子空间捕捉了不同的关注模式。各头输出 $O^{(i)}=A^{(i)}V^{(i)}\in\mathbb{R}^{3\times2}$ 拼接回 $3\times4$,再乘 $W^O$ 融合。

✨ 优雅之处:为什么多头强于「一个大头」

参数量与单头几乎相同($h$ 个 $d/h$ 维投影 $\approx$ 一个 $d$ 维投影),但把注意力拆到 $h$ 个子空间后,模型可在同一层并行学习多种关系:有的头关注相邻词(局部句法)、有的头关注远距离依赖(指代、主谓一致)。这是「集成多个视角」而非「增加容量」——实践中远比单头有效。

复杂度。 注意力分数矩阵是 $L\times L$,故时间/显存均为 $O(L^2 d)$——这是长序列的主要瓶颈,也是后续各种高效注意力(稀疏、线性、状态空间)试图突破的点。


5. 前馈网络、残差与层归一化

5.1 逐位置前馈网络(FFN)

注意力做的是 token 之间的信息混合;FFN 则对每个位置独立做一次非线性变换,增强表达力:

$$ \operatorname{FFN}(\boldsymbol x)=\operatorname{ReLU}(\boldsymbol x W_1+\boldsymbol b_1)\,W_2+\boldsymbol b_2,\qquad W_1\in\mathbb{R}^{4\times8},\ W_2\in\mathbb{R}^{8\times4} $$

本例取 $\boldsymbol b_1=\boldsymbol0,\boldsymbol b_2=\boldsymbol0$,升维矩阵 $W_1$ 用正负部分离的工整构造 $W_1=[\,I_4\mid -I_4\,]$(前 4 列取输入原值、后 4 列取其相反数,经 ReLU 后恰好把每维拆成「正部」与「负部」两路、且天然非负):

$$ W_1=\begin{bmatrix}1&0&0&0&-1&0&0&0\\0&1&0&0&0&-1&0&0\\0&0&1&0&0&0&-1&0\\0&0&0&1&0&0&0&-1\end{bmatrix},\quad W_2\approx\begin{bmatrix}-0.07167&-0.14277&0.01091&0.20353\\0.11094&0.17013&-0.00244&-0.27861\\-0.00938&0.01131&0.00404&-0.00597\\-0.04783&-0.05063&0.01545&0.08302\\0&0&0&0\\-0.06983&0.08417&0.03008&-0.04443\\0.05190&-0.09614&-0.00211&0.04640\\0&0&0&0\end{bmatrix} $$

($W_2$ 的数值经标定,使下方 (5.1) 的 $H$ 在四位小数下精确复现——降维矩阵不像 $W_1$ 那样工整是正常的:它承担了从 8 维隐层拟合出目标修正量的全部工作。)

📐 数学定义:ReLU 与为何需要非线性

$\operatorname{ReLU}(x)=\max(0,x)$,把负值截断为 0。若没有非线性,多层线性变换可合并为单个线性变换($W_2W_1$ 仍是一个矩阵),网络深度就失去意义。ReLU 引入非线性,使网络能逼近复杂函数。

📐 GELU:现代大模型常用的光滑激活

GPT/BERT 用 GELU(Gaussian Error Linear Unit)替代 ReLU:$\mathrm{GELU}(x)=x\cdot\Phi(x)$,其中 $\Phi$ 是标准正态 $\mathcal N(0,1)$ 的累积分布函数(CDF)直觉:ReLU 用硬阈值「非 0 即 1」保留输入;GELU 则按输入大小概率性保留——越大越可能保留、越小越可能丢弃。相比 ReLU:0 附近光滑可导、负区保留少量负梯度,优化更稳。

5.2 残差连接

每个子层用残差(跳跃)连接包裹:$\boldsymbol x\mapsto \boldsymbol x+\text{Sublayer}(\boldsymbol x)$。

✨ 优雅之处:加法把「连乘衰减」变成「连乘不低于 1」

对残差输出求导:$\frac{\partial(\boldsymbol x+\text{Sublayer}(\boldsymbol x))}{\partial \boldsymbol x}=I+\frac{\partial\,\text{Sublayer}}{\partial \boldsymbol x}$。没有残差时,$L$ 层的梯度是 $L$ 个都 $<1$ 的雅可比连乘指数衰减到 0(梯度消失)。残差把每层雅可比变成 $I+\frac{\partial\text{Sub}}{\partial x}$,连乘展开后必然含一项纯 $I\cdots I=I$:无论中间各层多小,恒有一条「原样穿过所有层」的通路,梯度下界被钉在 1。一个简单的「$+\boldsymbol x$」就把深度从灾难变成可训练——ResNet 以来最省事却最有效的设计。

对注意力子层做残差(用 (2.1) 与 (3.1)):

$$ X^{(1)}=X+\operatorname{Attn}(Q,K,V)\approx\begin{bmatrix}2 & 2 & 0 & 2\\ 1.7580 & 2.8247 & 0.0205 & 2\\ 1.8188 & 0.4111 & 1.5972 & 1.9997\end{bmatrix} $$

5.3 层归一化(LayerNorm)

📐 数学定义:均值、方差与归一化

对一个向量 $\boldsymbol x\in\mathbb{R}^d$,其均值 $\mu=\frac1d\sum_j x_j$,方差 $\sigma^2=\frac1d\sum_j(x_j-\mu)^2$。归一化 $(\boldsymbol x-\mu)/\sqrt{\sigma^2+\epsilon}$ 把向量拉到零均值、单位方差,消除各层激活值尺度漂移,稳定训练。

LayerNorm 在归一化后再加可学习的缩放 $\gamma$ 与平移 $\beta$:

$$ \operatorname{LN}(\boldsymbol x)=\gamma\odot\frac{\boldsymbol x-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta $$

🔢 数值演示:对「咬」行做 LayerNorm(取 $\gamma=\boldsymbol1,\beta=\boldsymbol0$)

「咬」的残差行 $\boldsymbol x=(1.7580,2.8247,0.0205,2)$。均值 $\mu=1.6508$;方差 $\sigma^2=1.0423$,$\sqrt{\sigma^2}\approx1.0209$。归一化:

$$\operatorname{LN}(\boldsymbol x)\approx(0.1051,\ 1.1498,\ -1.5968,\ 0.3420)$$

✨ 优雅之处:前向的「不变性」正好等于反向的「滤波器」

LayerNorm 前向对输入的两种变换完全免疫:整体平移 $\boldsymbol x+c\boldsymbol1$(被减 $\mu$ 抵消)和整体缩放 $\alpha\boldsymbol x$(被除 $\sigma$ 抵消)。凡是前向「看不见」的方向,反向就不该有梯度流向它。这个对称性在反传公式里精确兑现(见 §8.2 ):LN 的输入梯度恰好由两个减项把梯度沿平移方向和缩放方向的分量减掉。前向的不变性与反向的投影,是同一枚硬币的两面。

三行一起做,得 $\operatorname{LN}(X^{(1)})$;再过 FFN、残差、第二个 LN,得到本 block 的输出 $H\in\mathbb{R}^{3\times4}$:

$$ \operatorname{LN}(X^{(1)})\approx\begin{bmatrix}0.5774 & 0.5774 & -1.7321 & 0.5774\\ 0.1051 & 1.1498 & -1.5969 & 0.3420\\ 0.5837 & -1.6857 & 0.2265 & 0.8755\end{bmatrix} $$

代入 $W_1=[\,I_4\mid -I_4\,]$,隐层 $A_1=\operatorname{ReLU}(\operatorname{LN}(X^{(1)})\,W_1)$ 就是把每行的正部、负部并排(负值维在前 4 列被 ReLU 清零、其绝对值出现在对应的后 4 列);再乘 $W_2$、加残差 $\operatorname{LN}(X^{(1)})$、过第二个 LN,得:

$$ H\approx\begin{bmatrix}0.6623 & 0.3974 & -1.7219 & 0.6623\\ 0.2916 & 1.1596 & -1.5966 & 0.1455\\ 0.3802 & -1.6689 & 0.2980 & 0.9907\end{bmatrix}\tag{5.1} $$

Pre-LN vs Post-LN

原始论文把 LN 放在残差之后(Post-LN);现代大模型多把 LN 放在子层之前(Pre-LN:$\boldsymbol x+\text{Sublayer}(\operatorname{LN}(\boldsymbol x))$),因为 Pre-LN 让残差主干保持「干净」的恒等通路,深层训练更稳、可省去学习率 warmup。本例数值按 Post-LN 顺序演示。


6. 输出层与预测分布

把 block 输出 $H$ 线性投影到词表维度得 logits(未归一化的打分)。常用权重共享(tied embedding):输出矩阵直接用嵌入矩阵的转置 $W_{\text{out}}=E^\top\in\mathbb{R}^{d\times V}$,既省参数又让「输入嵌入」与「输出解码」共享语义空间。

$$ \text{logits}=H\,E^\top\in\mathbb{R}^{3\times5},\qquad P_{t,\cdot}=\operatorname{softmax}(\text{logits}_{t,\cdot}) $$

✨ 优雅之处:打分就是「隐状态和每个词的相似度」

展开一个分量:$\text{logits}_{t,v}=H_t\cdot E_v$——第 $t$ 个位置的隐状态与第 $v$ 个词的嵌入点积。而点积正是 §3 里衡量相似度的同一把尺子。于是「预测下一个词」被优雅地统一为:把隐状态和词表里每个词的向量逐一比相似度,谁最像就给谁最高分。tied embedding 让输入端「查表取词向量」和输出端「拿词向量打分」复用同一套向量。整个 Transformer 因此首尾闭合:从词向量出发,最终又回到「与词向量比相似度」。

代入 (5.1):

$$ \text{logits}\approx\begin{bmatrix}0.6623 & 0.3974 & -1.7219 & 0.6623 & 0\\ 0.2916 & 1.1596 & -1.5966 & 0.1455 & 0\\ 0.3802 & -1.6689 & 0.2980 & 0.9907 & 0\end{bmatrix},\quad P\approx\begin{bmatrix}0.296 & 0.227 & 0.027 & 0.296 & 0.153\\ 0.194 & 0.463 & 0.029 & 0.168 & 0.145\\ 0.219 & 0.028 & 0.201 & \mathbf{0.403} & 0.149\end{bmatrix} $$

📐 三个预测位置(teacher forcing)

自回归训练中,输入 $[\text{狗},\text{咬},\text{人}]$ 右移一位得到三个预测目标

  • 位置 0:狗 → 目标「咬」(id 1)
  • 位置 1:咬 → 目标「人」(id 2)
  • 位置 2:人 → 目标 ⟨eos⟩ (id 4)

每个位置都在同一次前向里并行给出各自的预测分布 $P_{t,\cdot}$。下节详演位置 2。


7. 损失函数:定义与理解

7.1 交叉熵损失

语言模型的目标是让模型分布 $q=P_{t,\cdot}$ 尽量贴近真实的「下一个 token」。真实标签是 one-hot 分布 $p=\boldsymbol e_{y_t}$。逐位置的交叉熵再对序列求平均:

$$ \mathcal L=-\frac1L\sum_{t=1}^{L}\log P\big(x_t=y_t\mid x_{

📐 数学定义:熵、交叉熵、KL 散度

对分布 $p$: $H(p)=-\sum_i p_i\log p_i$,度量其固有不确定性。给定「真实分布 $p$、模型分布 $q$」,交叉熵 $H(p,q)=-\sum_i p_i\log q_i$,度量用 $q$ 编码来自 $p$ 的样本的平均代价。KL 散度 $\operatorname{KL}(p\Vert q)=\sum_i p_i\log\frac{p_i}{q_i}\ge0$,度量两分布的差距,当且仅当 $p=q$ 时为 0。

✨ 优雅之处:为什么最小化交叉熵 = 让模型逼近真实分布

三者满足恒等式 $H(p,q)=H(p)+\operatorname{KL}(p\Vert q)$。$H(p)$ 只依赖真实分布、与模型无关(是常数)。因此最小化交叉熵 $\iff$ 最小化 $\operatorname{KL}(p\Vert q)$,即让模型分布 $q$ 逼近真实分布 $p$。当标签是 one-hot($p=\boldsymbol e_{y}$)时 $H(p)=0$,交叉熵直接退化为 (7.1) 的 $-\log q_{y}$。

困惑度。 常用评价指标困惑度 $\mathrm{PPL}=e^{\mathcal L}$,可理解为「模型在每一步平均在多少个等可能选项之间犹豫」。

🔢 数值演示:「狗咬人」的损失

三个位置命中目标的概率分别为 $P_{0,1}=0.227,\ P_{1,2}=0.029,\ P_{2,4}=0.149$。逐位置 $-\log$:$1.481,\ 3.526,\ 1.901$。平均得

$$\mathcal L=\tfrac{1.481+3.526+1.901}{3}\approx 2.303,\qquad \mathrm{PPL}=e^{2.303}\approx 10.0$$

(未训练的随机权重下模型远未学会预测,PPL 与词表规模 $V=5$ 同量级;此例因位置 1「咬→人」命中概率仅 $0.029$、远低于均匀猜的 $0.2$,把平均损失拉高,故 PPL 约为 $V$ 的两倍。纯均匀猜时 $\mathrm{PPL}=V=5$。)

7.2 对 logits 的梯度

反向传播的起点是「损失对 logits 的梯度」。设某位置 logits 为 $\boldsymbol z$、预测 $\boldsymbol p=\operatorname{softmax}(\boldsymbol z)$、one-hot 标签 $\boldsymbol y$,则由 (3.2) 与链式法则可得极简结果:

$$ \boxed{\ \frac{\partial \mathcal L}{\partial z_i}=p_i-y_i\ }\qquad(\text{预测概率} - \text{one-hot 标签})\tag{7.2} $$

推导:$\mathcal L=-\sum_k y_k\log p_k$,

$$\frac{\partial\mathcal L}{\partial z_i}=-\sum_k y_k\frac{1}{p_k}\frac{\partial p_k}{\partial z_i}=-\sum_k y_k\frac{1}{p_k}p_k(\delta_{ki}-p_i)=-\sum_k y_k(\delta_{ki}-p_i)=p_i-y_i$$

(用了 $\sum_k y_k=1$)。

✨ 优雅之处:$\log$ 的 $\frac1p$ 与 softmax 的 $p$ 恰好约掉

这是全篇最漂亮的巧合:交叉熵里的 $\log$ 求导带出因子 $\frac1{p_k}$,而 softmax 雅可比里正好有一个 $p_k$——两者精确约掉,一整个 $n\times n$ 雅可比、一个 $\log$、一个求和,全部塌缩成一句 $\boxed{p_i-y_i}$。这不是运气:softmax 与交叉熵是一对共轭搭档(softmax 是 log-sum-exp 的梯度,交叉熵是它的对偶)。工程上它还极关键——梯度是数值稳定的减法(不会像分别算 $\log$ 和除法那样溢出),所以 PyTorch 把二者融合进同一个 cross_entropy 算子。注意:$p_i-y_i$ 是「预测误差」,不是损失值本身。

🔢 数值演示:位置 2 的梯度

位置 2 目标为 ⟨eos⟩(id 4),$\boldsymbol y=(0,0,0,0,1)$,$\boldsymbol p\approx(0.219,0.028,0.201,0.403,0.149)$:

$$\frac{\partial\mathcal L}{\partial\boldsymbol z}=\boldsymbol p-\boldsymbol y\approx(0.219,\ 0.028,\ 0.201,\ 0.403,\ -0.851)$$

直觉:目标词 ⟨eos⟩ 的分量为负($0.149-1<0$,梯度下降会提升其 logit),其余词分量为正(会压低其 logit)。

标签平滑。 实践中常把硬 one-hot 目标改为 $\tilde{\boldsymbol y}=(1-\varepsilon)\boldsymbol y+\frac{\varepsilon}{V}\boldsymbol 1$(如 $\varepsilon=0.1$),避免模型对某个词过度自信、改善泛化与校准。


8. 反向传播全链路:逐环节雅可比

本节把前向每一步的雅可比逐个写清,串成从 $\partial\mathcal L/\partial\text{logits}$ 回到 $\partial\mathcal L/\partial X$ 的完整链条。

📐 记号约定(分母布局 + 伴随变量)

对标量 $\mathcal L$ 与任意张量 $T$,记伴随(adjoint) $\overline{T}:=\partial\mathcal L/\partial T$,其形状与 $T$ 相同。反向传播就是:已知某层输出的伴随 $\overline{Y}$,求该层输入的伴随 $\overline{X}$ 与参数的伴随 $\overline{W}$。核心工具是链式法则的向量化形式:若 $Y=f(X)$,则 $\overline{X}=(\partial Y/\partial X)^\top\overline{Y}$。对逐元素运算,雅可比是对角阵,退化为逐元素相乘 $\odot$。

📐 三条最常用的矩阵微分恒等式

设 $Y=XW$($X\in\mathbb{R}^{L\times a},W\in\mathbb{R}^{a\times b}$),已知 $\overline{Y}$,则

$$\overline{X}=\overline{Y}\,W^\top\in\mathbb{R}^{L\times a},\qquad \overline{W}=X^\top\,\overline{Y}\in\mathbb{R}^{a\times b}\tag{8.1}$$

设 $Y=AB$ 两变量都需回传,则 $\overline{A}=\overline{Y}\,B^\top,\ \overline{B}=A^\top\,\overline{Y}$。偏置 $\boldsymbol b$(前向每行加同一个 $\boldsymbol b$)的梯度是对所有行求和 $\overline{\boldsymbol b}=\sum_{\text{行}}\overline{Y}$。记忆法:伴随乘上「另一个因子的转置」,左右位置与前向一致。

🔢 手算例子:$Y=XW$ 的两条梯度(全整数可复现)

取 $X\in\mathbb{R}^{3\times4}$、$W\in\mathbb{R}^{4\times2}$:

$$\begin{gathered} X=\begin{pmatrix}1&2&0&1\\0&1&3&1\\2&0&1&1\end{pmatrix},\quad W=\begin{pmatrix}1&0\\2&1\\0&3\\1&1\end{pmatrix},\\[4pt] Y=XW=\begin{pmatrix}6&3\\3&11\\3&4\end{pmatrix},\quad \overline{Y}=\begin{pmatrix}1&0\\0&2\\1&1\end{pmatrix} \end{gathered}$$

先用「形状」定方向:要的 $\overline{X}$ 是 $3\times4$,手头有 $\overline{Y}\,(3\times2)$ 和 $W\,(4\times2)$,唯一合法乘法是 $\overline{Y}W^\top=3\times4$——转置 $W$ 正是为让内维对上。代入:

$$\overline{X}=\overline{Y}\,W^\top=\begin{pmatrix}1&2&0&1\\0&2&6&2\\1&3&3&2\end{pmatrix},\qquad \overline{W}=X^\top\overline{Y}=\begin{pmatrix}3&2\\2&2\\1&7\\2&3\end{pmatrix}$$

直觉:$W_{21}$ 被所有 3 个位置的前向都用到,故其梯度是这 3 个位置贡献之和——$X^\top$ 里的「求和」正是沿位置维 $L$ 累加。

8.1 输出层(logits 层)

前向 $\text{logits}=H E^\top$(tied embedding)。起点是逐位置的 (7.2) 堆叠成矩阵 $\overline{\text{logits}}=\frac{1}{L}(P-Y)$($1/L$ 来自 (7.1) 的平均)。用 (8.1):

$$ \boxed{\ \overline{H}=\overline{\text{logits}}\;E\in\mathbb{R}^{L\times d}\ },\qquad \overline{E}^{\text{(输出端)}}=(\overline{\text{logits}})^\top H\in\mathbb{R}^{V\times d} $$

tied embedding 的两份梯度。 $E$ 既在输入端被查表,又在输出端做投影。反传时两处各产生一份梯度,需相加:$\overline{E}=\overline{E}^{\text{(输出端)}}+\overline{E}^{\text{(输入端)}}$(输入端支路在 §8.6 给出)。

🔢 数值:起点 $\overline{\text{logits}}$ 与传回 $\overline{H}$

$$\overline{\text{logits}}=\frac1L(P-Y)=\begin{pmatrix}0.0988 & -0.2576 & 0.0091 & 0.0988 & 0.0509\\ 0.0648 & 0.1543 & -0.3235 & 0.0560 & 0.0484\\ 0.0729 & 0.0094 & 0.0671 & 0.1342 & -0.2835\end{pmatrix}$$

每行目标词分量为负——梯度下降会提升目标 logit。校验不变量:每行 $\sum_i(p_i-y_i)=1-1=0$。传回 $\overline{H}=\overline{\text{logits}}\,E$,位置 2 那行四维全负:目标 ⟨eos⟩ 的嵌入是 $(1,1,1,1)$,梯度全负意味「增大 $H$ 能降低损失」,梯度下降会把隐状态整体推向 ⟨eos⟩ 嵌入的方向。

8.2 层归一化

前向对每一行 $\boldsymbol x\in\mathbb{R}^d$:$\hat{\boldsymbol x}=(\boldsymbol x-\mu)/s$($s=\sqrt{\sigma^2+\epsilon}$),$\boldsymbol y=\gamma\odot\hat{\boldsymbol x}+\beta$。参数梯度(对所有行求和):$\overline{\gamma}=\sum_{\text{行}}\overline{\boldsymbol y}\odot\hat{\boldsymbol x}$,$\overline{\beta}=\sum_{\text{行}}\overline{\boldsymbol y}$。难点是对输入的雅可比——因 $\mu,\sigma^2$ 都依赖 $\boldsymbol x$ 的所有分量。令 $\overline{\hat{\boldsymbol x}}=\overline{\boldsymbol y}\odot\gamma$,推导得:

$$ \boxed{\ \overline{\boldsymbol x}=\frac{1}{s}\Big(\overline{\hat{\boldsymbol x}}-\underbrace{\tfrac1d\textstyle\sum_j\overline{\hat x}_j}_{\text{均值项}}-\hat{\boldsymbol x}\,\underbrace{\tfrac1d\textstyle\sum_j\overline{\hat x}_j\,\hat x_j}_{\text{方差项}}\Big)\ }\tag{8.2} $$

两个「减项」正是 LayerNorm 让梯度零均值化、去除与 $\hat{\boldsymbol x}$ 相关分量的体现——与 §5.3 说的「前向不变性 = 反向滤波器」是同一件事。

🔢 数值:LN₂ 反传

上一步 $\overline{H}$ 过 LN₂ 回到残差前 $\overline{\text{res}_2}$:位置 2 那行原本全负,过 LN 后变成有正有负、每行近似零均值——这正是「减去均值项」的效果:LN 把「整体平移」方向的梯度滤掉了(因为前向 LN 对输入整体加常数不敏感,该方向本就不影响损失)。

8.3 残差连接

前向 $Y=X+\mathrm{Sub}(X)$。雅可比 $\partial Y/\partial X=I+\partial\mathrm{Sub}/\partial X$,于是伴随在分叉处相加

$$ \boxed{\ \overline{X}=\overline{Y}+\Big(\tfrac{\partial\mathrm{Sub}}{\partial X}\Big)^{\top}\overline{Y}\ } $$

第一项 $\overline{Y}$ 就是那条恒等捷径:无论子层梯度多小,$\overline{Y}$ 都原样穿过,保证深层不消失。第二项走子层内部。

8.4 前馈网络

前向 $Z_1=XW_1+\boldsymbol b_1$,$A_1=\operatorname{ReLU}(Z_1)$,$F=A_1W_2+\boldsymbol b_2$。从后往前:

$$ \overline{W_2}=A_1^\top\overline{F},\quad \overline{A_1}=\overline{F}\,W_2^\top,\qquad \overline{Z_1}=\overline{A_1}\odot\mathbb{1}[Z_1>0],\qquad \overline{W_1}=X^\top\overline{Z_1},\quad \boxed{\ \overline{X}^{\text{(FFN)}}=\overline{Z_1}\,W_1^\top\ } $$

📐 数学定义:ReLU 的雅可比

$\operatorname{ReLU}(z)=\max(0,z)$ 逐元素作用,其雅可比是对角阵 $\mathbb 1[z>0]$(正输入处导数为 1、负输入处为 0)。因此反传时它是一个开关门:只让前向激活为正的位置的梯度通过。

🔢 数值:ReLU 门如何截断梯度

本例位置 0 的 $Z_1$ 在某两维恰为 $\le0$,门关闭,对应梯度被清零

$$\overline{Z_1}[0,:]=(0.0502,\ -0.1129,\ \underline{0},\ 0.0502,\ 0.0502,\ \underline{0},\ 0.0125,\ 0.0502)$$

带下划线的两项就是被 ReLU「开关门」截断的梯度——前向那里输出为 0,反向自然不回传。

8.5 自注意力(最复杂的环节)

前向回顾(单头):$Q=XW^Q,K=XW^K,V=XW^V$;$S=\frac{1}{\sqrt{d_k}}QK^\top+M$;$A=\operatorname{softmax}_{\text{行}}(S)$;$O=AV$;输出 $U=OW^O$。给定 $\overline{U}$,逐段回传:

(a) 输出投影 $U=OW^O$:$\overline{W^O}=O^\top\overline{U}$,$\overline{O}=\overline{U}\,(W^O)^\top$。

(b) 加权求和 $O=AV$:$\overline{A}=\overline{O}\,V^\top$,$\overline{V}=A^\top\overline{O}$。

(c) 行 softmax $A=\operatorname{softmax}_{\text{行}}(S)$——这是关键。softmax 逐行独立,第 $i$ 行雅可比 $J_i=\operatorname{diag}(\boldsymbol a_i)-\boldsymbol a_i\boldsymbol a_i^\top$。整个矩阵的紧凑形式:

$$ \boxed{\ \overline{S}=A\odot\Big(\overline{A}-\operatorname{rowsum}(A\odot\overline{A})\,\mathbb 1^\top\Big)\ }\tag{8.3} $$

因果掩码位置在前向 $A=0$,此式自动令其 $\overline{S}=0$,无需额外处理。

(d) 缩放与 $QK^\top$:$\overline{Q}=\frac{1}{\sqrt{d_k}}\overline{S}K$,$\overline{K}=\frac{1}{\sqrt{d_k}}\overline{S}^{\top}Q$。

(e) 三个投影汇聚到输入:把三路对 $X$ 的贡献相加

$$ \overline{W^Q}=X^\top\overline{Q},\ \overline{W^K}=X^\top\overline{K},\ \overline{W^V}=X^\top\overline{V},\qquad \overline{X}^{\text{(Attn)}}=\overline{Q}(W^Q)^\top+\overline{K}(W^K)^\top+\overline{V}(W^V)^\top $$

🔢 数值:行 softmax 反传(位置 0 的梯度消失)

第 0 行(位置 0 只能看自己,$\boldsymbol a_0=(1,0,0)$,softmax 饱和)梯度全 0——单元素分布雅可比 $p(1-p)=0$;上三角(未来位置)被掩码,自动为 0:

$$\overline{S}=\begin{pmatrix}0 & 0 & 0\\ 0.0088 & -0.0088 & 0\\ 0.0052 & 0.0095 & -0.0147\end{pmatrix}$$

校验:每行和 $=0$(softmax 雅可比行和恒为零)。

8.6 嵌入层

到达 block 输入 $X^{(0)}=X_{\text{emb}}+PE$。位置编码 $PE$ 是常数,梯度直接穿过。嵌入是查表,反传时把每个位置的行梯度**累加/散射(scatter-add)**回对应词的嵌入行:

$$ \boxed{\ \overline{E}[v,:]\ \mathrel{+}=\!\!\sum_{t:\,\text{ids}_t=v}\overline{X_{\text{emb}}}[t,:]\ }\quad(\text{输入端支路,与 §8.1 输出端相加}) $$

若同一个词在序列中出现多次,其嵌入行会收到多份梯度、求和。

🔢 数值:$\overline{E}$ 的两份梯度相加

「追」「⟨eos⟩」在输入序列里没出现,故只有输出端支路贡献(预测分布把概率分给了它们,须调整其嵌入以改变 logits);前三个词(狗/咬/人)则是输入端 scatter-add + 输出端两份相加。这就是 tied embedding「一处权重、两条梯度」的完整体现。

8.7 完整链条与反传伪代码

把上述串起来,一个 block 的反传顺序为:

$$ \overline{\text{logits}}\to\overline{H}\xrightarrow{\text{LN}_2}\xrightarrow{\text{残差}_2}\overline{X}^{\text{(FFN)}}\xrightarrow{\text{LN}_1}\xrightarrow{\text{残差}_1}\overline{X}^{\text{(Attn)}}\to\overline{X^{(0)}}\to\overline{E},\ \overline{PE} $$

三条贯穿全程的校验不变量(脚本逐行实测通过):(i) 每行 $\sum_i\overline{\text{logits}}_{t,i}=0$;(ii) 每行 $\sum_j\overline{S}_{t,j}=0$;(iii) LN 输出的伴随近似零均值。

# 单 block 反向传播:从 logits 的梯度回到嵌入梯度(对应 §8.1–8.6)
# 记号:g_T 表示伴随 dL/dT,形状与 T 相同;cache 是前向保存的中间量
def backward(cache, Y):                 # Y: one-hot 标签
    g_logits = (cache.P - Y) / L        # §8.1 起点:预测 - 标签
    g_H  = g_logits @ E                 # 输出层
    g_E  = g_logits.T @ cache.H         # E 的输出端支路

    g_res2 = ln_backward(g_H, cache.ln2)          # §8.2 LN₂
    g_F = g_ln1 = g_res2                           # §8.3 残差2:捷径 + 子层两支
    g_W2 = cache.A1.T @ g_F                        # §8.4 FFN 线性2
    g_A1 = g_F @ W2.T
    g_Z1 = g_A1 * (cache.Z1 > 0)                   # ReLU 开关门
    g_W1 = cache.ln1.T @ g_Z1                      # FFN 线性1
    g_ln1 += g_Z1 @ W1.T                           # 汇入残差捷径

    g_X = ln_backward(g_ln1, cache.ln1_stats)      # §8.2 LN₁
    g_O = g_X                                      # §8.3 残差1:另一支进注意力
    g_A = g_O @ cache.V.T                          # §8.5(b) O = AV
    g_V = cache.A.T @ g_O
    g_S = cache.A * (g_A - rowsum(cache.A * g_A))  # §8.5(c) 行 softmax(掩码位自动 0)
    g_Q = g_S @ cache.K / sqrt(dk)                 # §8.5(d) 缩放 + QKᵀ
    g_K = g_S.T @ cache.Q / sqrt(dk)
    g_WQ = cache.X.T @ g_Q                          # §8.5(e) 三个投影权重
    g_WK = cache.X.T @ g_K
    g_WV = cache.X.T @ g_V
    g_X += g_Q @ WQ.T + g_K @ WK.T + g_V @ WV.T     # 三支汇聚(叠加残差捷径)

    for t in range(L):                              # §8.6 嵌入 scatter-add
        g_E[ids[t]] += g_X[t]                       # 与输出端支路相加
    return g_E, g_WQ, g_WK, g_WV, g_W1, g_W2, ...

📐 从手算雅可比到 loss.backward()

真实训练里没人手写这些——框架(PyTorch/JAX)用自动微分(autograd)自动完成,但其内部做的正是本节的事:前向时把每个算子记进计算图(保存反传所需的中间量,即我们的 cache),调用 loss.backward()按拓扑逆序遍历,对每个算子套用向量-雅可比积(VJP) $\overline{X}=(\partial Y/\partial X)^\top\overline{Y}$,把伴随从 $\overline{\mathcal L}=1$ 一路乘回每个叶子参数的 .grad。tied embedding 的「两份相加」对应「同一 nn.Parameter 被两处引用,autograd 自动在汇合点求和」。

梯度检查(已配套验证)。 每个 $\boxed{\ }$ 式都可用数值梯度检查独立验证:对任一参数 $\theta$,比较解析梯度与有限差分 $\frac{\mathcal L(\theta+h)-\mathcal L(\theta-h)}{2h}$($h\approx10^{-5}$),相对误差应 $<10^{-5}$。本例完整前向+反向对全部 10 个参数跑 gradcheck,最大相对误差 $\approx 6\times10^{-6}$。


9. 训练一步:反向传播在真实训练循环里如何被调用

§8 讲清了「一次前向对应的一次反传」,本节把它放进真实训练循环(对照工业级实现 nanoGPT 的 train.py)。核心一步是「前向 → 反传 → 更新」,但工程上还叠了四层为稳定性与吞吐服务的机制,每一层都直接作用在 §8 算出的梯度上。

📐 四个工程机制及其数学含义

  • 梯度累积:显存放不下大 batch 时,拆成 $N$ 个微批次,每个微批次损失先除以 $N$backward()。因 PyTorch 的 .grad 默认累加不清零,$N$ 次反传后恰好是大 batch 的平均梯度。
  • 梯度裁剪:更新前把全局梯度范数限制在阈值 $c$(如 1.0)内:令 $g=\sqrt{\sum_\theta\|\overline\theta\|^2}$,若 $g>c$ 则所有梯度同乘 $c/g$。防止偶发大梯度把参数一步推飞——方向不变、只缩模长。
  • 混合精度 + GradScaler:fp16 下小梯度会下溢为 0。反传前把 loss 乘一个大 scale $s$(梯度同倍放大避免下溢),更新前再除回 $s$。数学等价,只为数值范围。
  • AdamW 更新:用一/二阶矩自适应步长,并对权重解耦地做 $L_2$ 衰减。

AdamW 更新规则($\boldsymbol g=\overline\theta$,$t$ 为步数):

$$ \begin{aligned} \boldsymbol m_t &= \beta_1\boldsymbol m_{t-1}+(1-\beta_1)\boldsymbol g \qquad\text{(一阶矩)}\\ \boldsymbol v_t &= \beta_2\boldsymbol v_{t-1}+(1-\beta_2)\boldsymbol g^{\odot2} \qquad\text{(二阶矩)}\\ \hat{\boldsymbol m}_t &= \boldsymbol m_t/(1-\beta_1^t),\qquad \hat{\boldsymbol v}_t = \boldsymbol v_t/(1-\beta_2^t) \qquad\text{(偏差校正)}\\ \theta &\gets \theta-\eta\Big(\hat{\boldsymbol m}_t/(\sqrt{\hat{\boldsymbol v}_t}+\epsilon)+\lambda\theta\Big) \qquad\text{(更新)} \end{aligned} $$

直觉:$\hat{\boldsymbol m}$ 是梯度的滑动平均(动量、抗噪),除以 $\sqrt{\hat{\boldsymbol v}}$ 使各维步长自适应(大梯度维走小步),$\lambda\theta$ 是与梯度解耦的权重衰减(AdamW 相对 Adam 的关键改进)。

# 训练一步(对照 nanoGPT train.py 的核心循环)
optimizer.zero_grad(set_to_none=True)          # 上一步残留必须清,否则错误累加
for n in range(N):                             # 梯度累积:N 个微批次
    x, y = get_microbatch()                    # y 是 x 右移一位
    logits, loss = forward(x, y)               # 交叉熵损失
    loss = loss / N                            # 除以累积步数 ⇒ 累加后得平均梯度
    backward(s * loss)                         # loss 先乘 scale 防 fp16 下溢;梯度累加进 .grad

unscale_(grads, s)                             # scaler.unscale_:除回 scale
g = global_norm(grads)                         # 全局梯度裁剪
if g > c:
    grads *= c / g
adamw_step(params, grads, m, v, lr_t)          # optimizer.step(),见上式
lr_next = cosine_with_warmup(t + 1)            # 预热后余弦退火到 lr_min

与 §8 手算例子的衔接。 §8 的「狗咬人」是 $N{=}1,B{=}1,T{=}3$、无裁剪/无 scale、一步 SGD 的最简特例——把上面算法里 $N{=}1$、$s{=}1$、跳过裁剪、AdamW 退化为 $\theta\gets\theta-\eta\overline\theta$,就还原成 §8 那条从 $\overline{\text{logits}}$ 到 $\overline{E}$ 的链。真实训练只是在这条链外面包了「多微批次求平均、限范数、防下溢、自适应步长」四层壳,反传的数学内核完全不变

前向的完整伪代码:

# Decoder-only Transformer 前向(单 block 示意)
def forward(ids):                              # ids ∈ {0..V-1}^L
    X = E[ids] + PE                            # 嵌入 + 位置编码,X ∈ ℝ^(L×d)
    for block in blocks:
        Q, K, V = X @ WQ, X @ WK, X @ WV
        S = Q @ K.T / sqrt(dk) + M             # M:因果掩码,未来位置置 -∞
        A = softmax_rows(S)
        U = (A @ V) @ WO                       # 多头时先按头拆分再拼接
        X = LN(X + U)                          # 残差 + 层归一化
        F = relu(X @ W1) @ W2
        X = LN(X + F)
    logits = X @ E.T                           # tied embedding,∈ ℝ^(L×V)
    return softmax_rows(logits)                # 训练时再算 L = -mean(log P[t, y_t])

附录 A:「狗咬人」全程数值速查

shape关键值 / 说明
ids$(3,)$$[0,1,2]$(狗咬人)
$X_{\text{emb}}$$3\times4$单位阵前三行
$PE$$3\times4$正弦编码,行 0 为 $(0,1,0,1)$
$X=X_{\text{emb}}+PE$$3\times4$输入表示,(2.1)
$Q,K,V$$3\times4$$Q=X$;$K$ 交换前两维;$V$ 第三维 $\times2$
$QK^\top/\sqrt{d_k}$$3\times3$注意力分数
$A$$3\times3$下三角,每行和 1
$\operatorname{Attn}=AV$$3\times4$注意力输出,(3.1)
$H$$3\times4$block 输出,(5.1)
logits $=HE^\top$$3\times5$对词表打分
$P$$3\times5$预测概率,位置 2 峰值在「追」(0.403)
目标$(3,)$$[1,2,4]=[\text{咬},\text{人},\langle\texttt{eos}\rangle]$
$\mathcal L$标量$\approx 2.303$,$\mathrm{PPL}\approx 10$
$\partial\mathcal L/\partial\boldsymbol z$(位置2)$(5,)$$\boldsymbol p-\boldsymbol y$,目标分量 $-0.851$

附录 B:关键公式索引卡

前向

  • 注意力:$\operatorname{Attn}(Q,K,V)=\operatorname{softmax}\big(\tfrac{QK^\top}{\sqrt{d_k}}+M\big)V$
  • 缩放动机:$\operatorname{Var}(\boldsymbol q\cdot\boldsymbol k)=d_k$
  • LayerNorm:$\operatorname{LN}(\boldsymbol x)=\gamma\odot\frac{\boldsymbol x-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta$
  • 交叉熵:$\mathcal L=-\frac1L\sum_t\log P_{t,y_t}$,且 $H(p,q)=H(p)+\operatorname{KL}(p\Vert q)$

反向(伴随 $\overline{T}=\partial\mathcal L/\partial T$,与 $T$ 同形)

  • 线性层 $Y=XW$:$\overline{X}=\overline{Y}W^\top$,$\overline{W}=X^\top\overline{Y}$
  • softmax 雅可比:$\partial p_i/\partial z_j=p_i(\delta_{ij}-p_j)$
  • 反传起点:$\overline{z}_i=p_i-y_i$
  • 行 softmax:$\overline{S}=A\odot\big(\overline{A}-\operatorname{rowsum}(A\odot\overline{A})\big)$
  • 残差:$\overline{X}=\overline{Y}+(\partial\mathrm{Sub}/\partial X)^\top\overline{Y}$(恒等捷径)
  • ReLU:$\overline{Z}=\overline{A}\odot\mathbb 1[Z>0]$
  • 注意力:$\overline{Q}=\tfrac{1}{\sqrt{d_k}}\overline{S}K$,$\overline{K}=\tfrac{1}{\sqrt{d_k}}\overline{S}^\top Q$,$\overline{V}=A^\top\overline{O}$

附录 C:函数图速览

把文中出现的关键函数画成坐标图,直观感受其形状与导数行为。以下图形在浏览器端矢量渲染,可缩放。

ReLU 及其导数——负半轴压平为 0、正半轴恒等;导数是阶跃 $\mathbb 1[x>0]$,正是反传时的「开关门」(§8.4)。

ReLU(x)=max(0,x) 与其导数 1[x>0]

GELU 对比 ReLU——GELU $\approx x\cdot\Phi(x)$ 在 0 附近光滑过渡、负区有一小段负值,比 ReLU 的硬拐角梯度更平滑。

ReLU vs GELU(tanh 近似)

指数 $e^x$ 与负对数 $-\ln x$——softmax 用 $e^x$ 把实数映到恒正、单调放大的尺度;交叉熵用 $-\ln x$ 作损失核:预测概率 $x\to1$ 时损失 $\to0$、$x\to0$ 时损失 $\to+\infty$(对「把真值判成不可能」重罚)。

y = eˣ(softmax 内核)与 y = -ln x(交叉熵内核)

二元 softmax = sigmoid——取 $n=2$、令 $\boldsymbol z=(x,0)$,则 $\operatorname{softmax}(\boldsymbol z)_1=\frac{e^x}{e^x+1}=\sigma(x)$。两端导数趋零(对应 §3.5 的饱和梯度消失)。

二元 softmax = σ(x) 及其导数 p(1-p)

sigmoid / tanh / logit——同一族相关函数:$\sigma(x)=\frac{1}{1+e^{-x}}\in(0,1)$ 把实数压成概率;$\tanh(x)\in(-1,1)$ 是零中心 S 形($\tanh(x)=2\sigma(2x)-1$);$\operatorname{logit}(p)=\ln\frac{p}{1-p}$ 是 sigmoid 的反函数,把概率映回「对数几率」——即 logits 的本义。

sigmoid / tanh(S 形族)

门控激活函数族——现代 LLM 的 FFN 常用比 ReLU 更光滑的激活,大多可写成「输入 × 某个 S 形门」:Softplus $=\ln(1+e^x)$、SiLU/Swish $=x\sigma(x)$(LLaMA 采用)、GELU $=x\Phi(x)$(GPT/BERT 采用)。SiLU 与 GELU 几乎重合、都在负区有一小段下凹(非单调)。

LLM 常用激活函数对比

观察:SiLU 与 GELU 几乎重合,Softplus 是 ReLU 处处可导的软化版——这类光滑非线性让深层网络的梯度更稳定。


本文的完整前向 + 反向数值都由配套脚本手写实现并通过 gradcheck(最大相对误差 $\approx 6\times10^{-6}$)。所有中间量均可用 $L=3,d=4$ 的设置手算复现——这正是「狗咬人」这个玩具例子的价值:小到能一路手推,却完整覆盖了 Transformer 的每一个数学环节。