写自定义算子的时候，广播是最容易"前向能跑、反向静默错"的地方。前向时框架帮你把形状 $(n,)$ 的 bias 铺成 $(m,n)$ 参与加法，一切都自然；反向时如果你把形状 $(m,n)$ 的上游梯度直接赋给 bias，形状对不上会报错还算幸运，怕的是某些路径下 numpy 再帮你广播一次，梯度数值就全错了。这篇把广播的雅可比写出来，看 sum 归约到底从哪来。
前向：广播就是乘一个"堆叠单位阵"
以线性层的 bias 加法为例：$x\in\mathbb{R}^{m\times n}$ 是 batch，$b\in\mathbb{R}^n$ 是 bias，前向
$$ y_{ij}=x_{ij}+b_j,\qquad Y=X+\mathbf{1}_m b^\top. $$
把 $b$ 的视角抽出来：$b\mapsto \mathbf{1}_m b^\top$ 是一个线性映射 $\mathbb{R}^n\to\mathbb{R}^{m\times n}$。向量化之后它的雅可比一目了然：
$$ \mathrm{vec}(\mathbf{1}_m b^\top)=(I_n\otimes \mathbf{1}_m)\,b, $$
也就是说广播的雅可比是 $mn\times n$ 的矩阵 $I_n\otimes\mathbf{1}_m$——把 $n\times n$ 单位阵竖着叠 $m$ 份（每份之间以全 1 向量的方式错行排列）。这个矩阵每一列恰好有 $m$ 个 1，对应 $b_j$ 被复制到的 $m$ 个位置。
前向是复制，雅可比是"一列多个 1"的稀疏矩阵。 反向只是把它转置。
反向：转置之后就是按列求和
设上游梯度 $G=\partial\ell/\partial Y\in\mathbb{R}^{m\times n}$，用雅可比左乘的 VJP 形式（前面雅可比那篇的约定）：
$$ \bar b=(I_n\otimes\mathbf{1}_m)^\top\,\mathrm{vec}(G)=(I_n\otimes\mathbf{1}_m^\top)\,\mathrm{vec}(G). $$
$(I_n\otimes\mathbf{1}_m^\top)$ 每一行是 $m$ 个 1 落在对应列的位置上，乘出来的效果：
$$ \boxed{\;\bar b_j=\sum_{i=1}^{m} G_{ij}\;} $$
即 $\bar b=\mathrm{rowsum}(G)$，或者写成矩阵形式 $\bar b=G^\top\mathbf{1}_m$。沿被广播的轴求和，这就是"sum 回原始形状"的代数出处：$b_j$ 在前向被用了 $m$ 次，链式法则要求把 $m$ 条路径上的梯度全部加起来。一次都不能少，也一次都不能多。
用微分法验证一遍更短：$d\ell=\mathrm{tr}(G^\top dY)=\mathrm{tr}(G^\top \mathbf{1}_m\, db^\top)=\mathbf{1}_m^\top G\,db^\top$，整理成内积形式 $\mathrm{tr}(\bar b\, db^\top)$ 读出 $\bar b^\top=\mathbf{1}_m^\top G$，一样。
一般规则与 sum/mean 的尺度差
把结论推广到任意广播：前向把张量 $A$ 沿某些轴复制到形状 $S$，反向时对 $G$ 沿这些轴 sum，并 squeeze 掉新增的维度，得到与 $A$ 完全相同的形状。口诀只有一句：梯度形状必须等于参数形状，差出来的轴一律 sum 掉。
这里有个尺度问题经常被忽略。对比两种 loss 写法对同一个 bias 的梯度：
$$ \ell_{\text{sum}}=\sum_{i,j}Y_{ij}^2\ \Rightarrow\ \bar b_j=2\sum_i Y_{ij};\qquad \ell_{\text{mean}}=\frac{1}{mn}\sum_{i,j}Y_{ij}^2\ \Rightarrow\ \bar b_j=\frac{2}{mn}\sum_i Y_{ij}. $$
广播归约本身是 sum，不含任何平均因子；$1/m$ 或 $1/(mn)$ 全部来自前向 loss 里有没有除。batch 从 32 改到 256 时如果你用 mean loss，梯度尺度不变、学习率不用动；用 sum loss，梯度放大 8 倍，学习率要跟着调。很多"换 batch size 就训崩"的事故，根源是把这两种约定混着用。
框架里的实现与常见踩坑
PyTorch 的 autograd 在所有二元算子的反向里都内置了这条归约：torch/csrc/autograd/functions/utils.cpp 附近的路径最终走到 sum_to_size(grad, self.sizes())，把梯度 sum/squeeze 回输入形状。TensorFlow 里对应的函数叫 BroadcastGradientArgs，它先算两个输入各自该在哪些轴上归约，再分别 reduce_sum。你自己写自定义 autograd.Function 时框架不会替你做这一步——backward 返回的张量形状必须自己 sum 到位，这是自定义算子反向 bug 的高发区。
三个我实际踩过或见过的坑：
- keepdim 遗忘：手动归约时
G.sum(dim=0)得到 $(n,)$，但如果参数本身是 $(1,n)$ 注册的，要用G.sum(dim=0, keepdim=True)。形状差一个维度，后续param.grad累加直接报 RuntimeError，还算显眼。 - numpy 的静默再广播：纯 numpy 实现里
b_grad = G之后如果写b -= lr * b_grad，$b$ 是 $(n,)$、$b_{\text{grad}}$ 是 $(m,n)$，numpy 会把 $b$ 广播成 $(m,n)$ 做减法然后赋值失败或形状膨胀，错误表现为训练无声地烂掉而不是崩溃。 - 分布式下的重复归约：数据并行时每个 rank 已经对本卡 batch 做了 sum，all-reduce 默认又是 sum，如果你用 mean loss，全局限度上梯度多乘了一个 world size。DeepSpeed/Megatron 的梯度归约模式要和 loss 约定对齐。
任意广播的 einsum 视角
逐轴复制之外，实战中更常见的是多维混合广播，比如注意力里 scores 加 mask：$\text{scores}\in\mathbb{R}^{(b,h,n,n)}$ 加 $\text{mask}\in\mathbb{R}^{(1,1,n,n)}$，或 RoPE 里 $(b,n,d)$ 乘 $(1,n,d)$ 的 cos 表。这类算子的反向用 einsum 记号写最不容易错：把前向写成
$$ Y_{\alpha\beta}=\sum_{\gamma}\,T_{\alpha\beta\gamma}\,B_{\gamma}, $$
其中 $T$ 是广播复制操作的稀疏张量（元素只有 0/1），那么反向永远是同一个动作：$\bar B_\gamma=\sum_{\alpha\beta}G_{\alpha\beta}T_{\alpha\beta\gamma}$。用大白话说：前向把 $B$ 的每个元素撒到哪些输出位置，反向就把那些位置的梯度捡回来加在一起。形状对齐、轴对齐这些操作细节都是这句话的机械执行。
顺带一提，$X$ 那边的梯度没有归约问题：$Y=X+\mathbf{1}_mb^\top$ 对 $X$ 的雅可比是单位阵，$\bar X=G$ 原样透传。广播加法前向复制多少份，只影响被复制那一方的反向。
expand 与 repeat：两种"复制"的反向是同一件事吗
PyTorch 里把小张量变大有两条路：expand 做零拷贝的视图广播，repeat 真正复制内存。前向数值完全一样，反向也殊途同归——expand 的反向是 sum_to_size，repeat 的反向是分段求和，数学上都是"沿复制轴 sum"。差别只在性能与内存：repeat 前向花内存、反向走连续内存的归约；expand 前向免费，反向要对 stride 为 0 的视图做归约。写性能敏感代码时这个选择才值得纠结，数学正确性上两者无条件等价。
另一个容易混淆的邻居是 gather/scatter 一族：广播是"一个源到多个确定位置"，scatter 是"多个源按索引到多个位置"。后者的反向同样把同目标位置的梯度相加（scatter_add 的反向是 gather，gather 的反向是 scatter_add），可以看成广播归约在任意索引下的推广。embedding 查表的反向——nn.Embedding 把梯度 scatter_add 回词表行——就是这条规则的日常实例：同一个词在 batch 里出现几次，它的梯度就累加几次。
数值验证
再补一个多维混合广播的对照：$(m,n)$ 与 $(n,)$、$(m,1)$ 同时参与的加法，两边的归约轴不同，是最容易写错的场景，一并验证：
30 行验证广播归约，顺带对比 mean 约定的尺度：
import torch
torch.manual_seed(0)
m, n = 5, 4
x = torch.randn(m, n, dtype=torch.float64)
b = torch.randn(n, dtype=torch.float64, requires_grad=True)
w = torch.randn(m, n, dtype=torch.float64)
y = x + b # 广播前向
loss = (w * y).sum() # 上游梯度即 w
loss.backward()
manual = w.sum(dim=0) # 沿广播轴 sum
assert torch.allclose(b.grad, manual, atol=1e-12)
loss_mean = ((w * (x + b)).sum()) / m
b.grad = None
loss_mean.backward()
assert torch.allclose(b.grad, manual / m, atol=1e-12) # mean 约定差一个 1/m
# 双方向广播：(m,n) + (n,) + (m,1)
r = torch.randn(m, 1, dtype=torch.float64, requires_grad=True)
y2 = x + b + r
(w * y2).sum().backward() # b.grad 已被累加过，先清零重来
b.grad = None
(w * y2).sum().backward()
assert torch.allclose(b.grad, w.sum(dim=0)) # 沿行归约
assert torch.allclose(r.grad, w.sum(dim=1, keepdim=True)) # 沿列归约且 keepdim
注意 r 的形状是 $(m,1)$，归约后必须 keepdim=True 才能对上——这就是上一节坑 1 的最小复现。
小结
广播的雅可比是堆叠的单位阵 $I_n\otimes\mathbf{1}_m$，转置之后自然变成沿广播轴求和，梯度形状因此恰好回到参数形状。sum 还是 mean 的尺度差不在广播本身，而在前向 loss 的约定里。这条归约规则看着琐碎，却是 bias、LayerNorm 的 $\gamma/\beta$、注意力 mask 加法这些 everyday 算子反向正确性的共同地基。下一篇看矩阵乘法的梯度：维度对齐法则如何把一切形状检查归结为一句话。