Transformer 里九成的 FLOPs 花在矩阵乘法上，矩阵乘法的梯度也因此是整个反向传播里被调用次数最多的公式。好消息是它只有一条核心结论配一句形状检查口诀，坏消息是记错转置方向的代价是形状静默错位、数值全错。这篇用微分法把 $\nabla_A$ 和 $\nabla_B$ 各推一遍，然后提炼出一个不背公式也能写对的检查方法。
微分法推导：一遍就够
设 $Y=AB$，$A\in\mathbb{R}^{m\times k}$，$B\in\mathbb{R}^{k\times n}$，上游梯度 $G=\partial\ell/\partial Y\in\mathbb{R}^{m\times n}$。乘积的微分：
$$ dY=dA\,B+A\,dB. $$
套上 trace 内积 $d\ell=\mathrm{tr}(G^\top dY)$（整个系列的约定，见 trace 技巧那篇）：
$$ d\ell=\mathrm{tr}(G^\top dA\,B)+\mathrm{tr}(G^\top A\,dB). $$
两项分别用 trace 的循环置换整理成"梯度转置乘以微分"的标准形。第一项：$\mathrm{tr}(G^\top dA\,B)=\mathrm{tr}(B\,G^\top dA)$，要配成 $\mathrm{tr}((\nabla_A\ell)^\top dA)$，所以
$$ (\nabla_A\ell)^\top=BG^\top\ \Rightarrow\ \boxed{\ \nabla_A\ell=GB^\top\ } $$
第二项：$\mathrm{tr}(G^\top A\,dB)=\mathrm{tr}((A^\top G)^\top dB)$，直接读出
$$ \boxed{\ \nabla_B\ell=A^\top G\ } $$
推导全程没有逐分量展开，trace 循环置换做了全部重活。注意两个结果的"分工"：$\nabla_A$ 用 $G$ 左乘、$B^\top$ 右乘；$\nabla_B$ 用 $A^\top$ 左乘、$G$ 右乘。转置永远落在不求导的那个因子上。
维度对齐法则：一句话的形状检查
记不住转置在哪边？不用记。有一条机械法则可以从零重建公式：
梯度形状必须等于参数形状，且公式里每个维度都要能正确收缩。
以 $\nabla_A$ 为例重建一遍：$\nabla_A$ 必须是 $(m,k)$。手头可用的材料只有 $G\,(m,n)$ 和 $B\,(k,n)$。$m$ 只能由 $G$ 提供，$k$ 只能由 $B$ 提供，剩下 $n$ 是两个矩阵都有、且恰好各出现一次的维度——把它收缩掉：$G$ 的列对 $B$ 的列，即 $GB^\top$，$(m,n)(n,k)\to(m,k)$。唯一解，没有第二种摆法。
同理 $\nabla_B$ 要是 $(k,n)$：$k$ 来自 $A\,(m,k)$、$n$ 来自 $G\,(m,n)$，收缩共同的 $m$：$A^\top G$。这套推理对任何矩阵乘法梯度都成立，我把它叫做材料唯一性：形状约束强到公式被唯一确定。
三个高频实例的对号入座
先说一个用法上的提醒：本节所有公式默认上游梯度 $G$ 已经"原样到达"，即中间的激活函数、归一化等操作都已被链式法则处理过。实际回传时，矩阵乘收到的是经过逐点激活雅可比（对角阵，$\odot\,\sigma'$）调制过的梯度——形状不变，所以本节法则与激活函数的逐点缩放正交叠加，互不干扰。
线性层 $Y=XW$（$X$ 是 $(m,k)$ 的 batch，$W$ 是 $(k,n)$）：$\nabla_W=X^\top G$，$\nabla_X=GW^\top$。这就是 PyTorch nn.Linear 反向的两个 addmm（aten/src/ATen/native/Linear.cpp 里 linear_backward 的形状）。顺带看出一个工程事实：算 $\nabla_W$ 需要一个 $(k,m)(m,n)$ 的矩阵乘，FLOPs 和前向 $(m,k)(k,n)$ 完全相同——矩阵乘的反向代价是两个前向，这是"反向约 2 倍前向算力"经验法则的主要来源。
向量情形 $y=Wx$：把 $x$ 看成 $(k,1)$ 的矩阵，法则给出 $\nabla_W=Gx^\top$（外积）、$\nabla_x=W^\top G$。外积形式解释了 Hebb 式直觉：权重的梯度正比于"上游梯度 × 输入"的逐对乘积。
batch 矩阵乘：$(b,m,k)@(b,k,n)$ 逐 batch 独立，梯度逐 batch 套公式，形状照旧。但注意 $b$ 维如果被广播（比如共享的 $B$），就叠加了上一篇的广播归约——$B$ 的梯度要对 $b$ 维再 sum 一次。框架里 torch.matmul 的反向把这两种归约叠在一起处理，手写时最容易漏的就是这一层。
反向算力账本与布局细节
把"反向是两个前向"这句话算细一点。前向 $Y=AB$：$mkn$ 次乘加，约 $2mkn$ FLOPs。反向两个矩阵乘 $GB^\top$ 和 $A^\top G$ 各 $2mkn$，合计 $4mkn$——正好两倍。注意这里有个顺序自由度：$\nabla_A=GB^\top$ 没有别的算法，但先算哪个梯度影响显存峰值，框架的执行顺序不是随意的。
布局（layout）是另一层工程现实。$A^\top G$ 里的转置通常不落盘成新矩阵，而是以 stride 标记传给 BLAS——cuBLAS 的 cublasGemmEx 直接接受 CUBLAS_OP_T，转置融化在 kernel 的访存模式里。但转置访存的合并度差，性能上 $(k,m)(m,n)$ 与 $(k,m)$ 按行读 $G$ 的列并不等价，所以高性能实现里反向 kernel 对输入布局有偏好，这也是 cuDNN/cuBLAS 为前向反向分别选不同算法的原因。手写 CUDA 时如果想"顺便"省掉转置，先确认你换来的访存模式不亏。
形状检查在调试里的实际用法
这个法则不只是推导工具，更是调试工具。我review反向代码时的固定动作：
- 对每个参数写下目标形状，挨个核对梯度张量的 shape——不一致必错，没有例外；
- 一致之后再抽查一个分量的数值（用数值梯度，下一篇的主题），因为形状对不代表转置摆对——$(k,n)$ 的梯度拿错成 $(n,k)$ 再转置回来这类错误，形状检查是放过的；
- batch 维的归约方向单独核，它是唯一"梯度形状小于上游梯度形状"的维度。
链式堆叠时的形状推演
单层对了不代表叠起来对，多层形状推演有一个固定节奏：从 loss 往回走，每过一层矩阵乘，梯度的形状经历一次"$G_{\text{out}}\,(m,n)\to G_{\text{in}}=G_{\text{out}}W^\top\,(m,k)$"的变换——上游梯度的行数永远跟 batch 走，列数从下一层的输出维换成输入维。三层 MLP $Y=XW_1W_2W_3$ 一路回传，$\bar X=((GW_3^\top)W_2^\top)W_1^\top$，链式右乘转置，顺序与权重梯度互不干扰。这个推演在纸上十秒就能完成，却能在写代码前抓掉大半形状错误。
注意力里的形状是这套推演的综合测验：$Q=XW^Q$（$(n,d)@(d,d_k)$）、scores $=QK^\top$（$(n,d_k)@(d_k,n)$）、输出 $=PV$（$(n,n)@(n,d_v)$）。scores 的梯度 $(n,n)$ 回传到 $Q$ 时用维度对齐法则：$\nabla_Q=(\nabla_S)K$，$(n,n)(n,d_k)\to(n,d_k)$——材料唯一性再次锁定答案，不需要记任何注意力专用公式。
数值验证
import torch
torch.manual_seed(0)
m, k, n = 5, 3, 4
A = torch.randn(m, k, dtype=torch.float64, requires_grad=True)
B = torch.randn(k, n, dtype=torch.float64, requires_grad=True)
G = torch.randn(m, n, dtype=torch.float64)
Y = A @ B
(Y * G).sum().backward()
assert A.grad.shape == A.shape and B.grad.shape == B.shape
assert torch.allclose(A.grad, G @ B.T, atol=1e-12)
assert torch.allclose(B.grad, A.detach().T @ G, atol=1e-12)
两行断言把形状法则和转置方向同时钉死。如果公式里 $B$ 忘了转置，G @ B 在 $k\neq n$ 时直接形状报错，$k=n$ 时才危险——方阵碰巧不报错，数值却是错的。这也是我为什么建议单元测试里故意用非方阵。
batch 矩阵乘的情形顺手也验了，顺带确认广播叠加时的归约方向：
Ab = torch.randn(2, m, k, dtype=torch.float64, requires_grad=True)
Gb = torch.randn(2, m, n, dtype=torch.float64)
(Ab @ B * Gb).sum().backward()
# B 被两个 batch 共用：梯度要对 batch 维再 sum（上一篇的规则叠加）
manual_B = torch.einsum("bmk,bmn->kn", Ab.detach(), Gb)
assert torch.allclose(B.grad, manual_B, atol=1e-10)
einsum 的记号在这里显示出价值：把指标写出来，哪些轴收缩、哪些轴保留一目了然，与"公共维收缩"法则逐字对应。
小结
$Y=AB$ 的梯度是 $\nabla_A=GB^\top$、$\nabla_B=A^\top G$，转置永远落在不求导的因子上；即使忘掉结论，"梯度形状等于参数形状 + 公共维收缩"也能唯一重建公式。矩阵乘反向的代价是两倍前向 FLOPs，这条账在估算训练算力时天天要用。有了形状法则兜底，下一篇进入验证环节：手写的反向到底对不对，有限差分给你一个不依赖直觉的裁判。