第三章 反向传播的完整推导
先明确上游到底传入什么
反向接口中的 dout 是损失对最终输出 $y$ 的梯度 $g$。若融合了 SiLU,不能直接将它当作卷积线性输出 $z$ 的梯度。完整链条是:
$$ g=\frac{\partial\mathcal L}{\partial y} \longrightarrow\delta=\frac{\partial\mathcal L}{\partial z} \longrightarrow(dx,dw,d\beta). $$
无激活时 $\delta=g$。SiLU 定义为 $\phi(z)=z\sigma(z)$,其中 $\sigma(z)=1/(1+e^{-z})$。利用乘积法则:
$$ \phi'(z)=\sigma(z)+z\sigma(z)(1-\sigma(z)), $$
$$ \boxed{\delta=g\,\sigma(z)[1+z(1-\sigma(z))]}. $$
这里必须使用激活前的 $z$。把 $y$ 代入上述导数,通常会得到错误结果;例如 $z=0$ 时导数为 $1/2$,不是 1。
偏置梯度
由于每个位置的 $z_{b,c,t}$ 对同通道 $\beta_c$ 的导数都是 1:
$$ \boxed{\frac{\partial\mathcal L}{\partial\beta_c} =\sum_{b=0}^{B-1}\sum_{t=0}^{L-1}\delta_{b,c,t}}. $$
偏置沿 batch 和时间共享,所以梯度沿这两维归约。padding 对应的零输入不会消除有效输出位置的 bias 梯度。
权重梯度
固定某个权重 $w_{c,k}$:
$$ \frac{\partial z_{b,c,t}}{\partial w_{c,k}} =\widetilde x_{b,c,t-K+1+k}. $$
套用链式法则:
$$ \boxed{\frac{\partial\mathcal L}{\partial w_{c,k}} =\sum_b\sum_t\delta_{b,c,t}\widetilde x_{b,c,t-K+1+k}}. $$
对零初始状态,只有 $t\ge K-1-k$ 的项有效,因此也可写为:
$$ dw_{c,k}=\sum_b\sum_{t=K-1-k}^{L-1} \delta_{b,c,t}x_{b,c,t-K+1+k}. $$
当下界超过上界时和为零。这一点对 $L<K$ 很重要。权重最左项在序列开头参与的有效乘积更少,但梯度无需除以有效项数:loss 的求和或平均方式已经体现在上游 $g$ 中。
输入梯度与反向方向
固定输入位置 $s$,找出哪些输出使用了它。由
$$ s=t-K+1+k $$
得到 $t=s+K-1-k$。因此:
$$ \boxed{dx_{b,c,s}=\sum_{k=0}^{K-1} w_{c,k}\widetilde\delta_{b,c,s+K-1-k}}. $$
这里越过序列右端的 $\delta$ 取零。换成滞后权重 $a_j=w_{K-1-j}$ 后:
$$ dx_s=\sum_{j=0}^{K-1}a_j\widetilde\delta_{s+j}. $$
前向读取左边的输入,输入反向读取右边的输出梯度。这不违反自回归因果性:训练时已知整段损失,要汇总当前输入对后续多个损失项的影响;这些梯度不会作为前向输入泄漏给模型。
矩阵表达最直接:若 $z=T(w)x+\beta$,则 $dx=T(w)^\top\delta$。前向的下三角矩阵转置后变成上三角,梯度方向自然反过来。
完整数值例子
沿用第二章 $x=[1,2,3,4]$、$w=[2,3,5]$、$\beta=1$,无激活。取上游 $g=[1,2,3,4]$,即标量损失 $\mathcal L=\sum_tg_ty_t=242$。
输入梯度为:
$$ \begin{aligned} dx_0&=5\cdot1+3\cdot2+2\cdot3=17,\ dx_1&=5\cdot2+3\cdot3+2\cdot4=27,\ dx_2&=5\cdot3+3\cdot4=27,\ dx_3&=5\cdot4=20. \end{aligned} $$
权重梯度为:
$$ \begin{aligned} dw_0&=3\cdot1+4\cdot2=11,\ dw_1&=2\cdot1+3\cdot2+4\cdot3=20,\ dw_2&=1\cdot1+2\cdot2+3\cdot3+4\cdot4=30. \end{aligned} $$
偏置梯度为 $d\beta=1+2+3+4=10$。这个例子同时检查核方向、左 padding、右侧梯度边界和参数共享归约。
从数学公式到并行所有权
计算 $dx_s$ 有两种思路。scatter 思路逐个输出 $t$ 把贡献加到它使用的输入位置,会产生多个任务对同一 $dx$ 的竞争。gather 思路让一个任务负责某段 $dx$,主动读取对应的右侧 $\delta$,可以让每个 $dx$ 只写一次。
$dw$ 和 $d\beta$ 则不同:所有 batch 和时间块都贡献到同一个通道的小型参数梯度。必须归约。可先在线程或任务内累加,再在块内归约,最后原子加到全局;也可以写入各块 partial buffer 后启动第二阶段归约。
这不是同一个问题的两种命名。$dx$ 的核心是邻域读取,$dw$ 的核心是多位置共享参数带来的求和。MLU 上不应为了复用代码而强行让两者采用相同任务划分。
训练存储与重计算
SiLU 反向需要 $z$。可以前向保存完整 $z$,也可以反向读取 $x,w,\beta$ 重新计算。后者增加约 $2BCLK$ FLOPs,却省去一个大激活张量的写入、保存和读取。小 $K$ 下这个交换可能划算,但最终取决于带宽、计算与融合程度。
Dao 的 Python autograd 包装保存输入、权重、bias、序列标识与初始状态,反向 CUDA 路径在需要时重建预激活。具体归约和 halo 组织见第六章。Dao autograd 接口
分段数据的公式扩展
若采用 packed sequence,令 $m_b(t,s)$ 在两位置属于同一序列时为 1,否则为 0,则:
$$ z_{b,c,t}=\beta_c+\sum_kw_{c,k}x_{b,c,s}m_b(t,s),\quad s=t-K+1+k. $$
$dw$ 的每项乘积也乘同一个 $m_b(t,s)$,$dx_s$ 中的每项贡献乘 $m_b(s+K-1-k,s)$。mask 必须前反向一致。$d\beta$ 仍汇总有效输出的 $\delta$;若还有 padding token 不参与 loss,应由相应的 loss mask 或算子契约处理。
下一章:Dao 接口与源码地图。