梯度累积原理：为什么我们需要它？

做深度学习训练，尤其是大模型，最头疼的问题之一就是显存不够。你想想看，一个 batch 稍微大一点，显卡就爆了。这时候，梯度累积就派上用场了。

我个人习惯把梯度累积比作「攒钱」。你一次赚不了太多，那就分多次赚，攒够了再花。在模型训练里，就是一次算不了太多样本的梯度，那就分几次算，累积起来再更新参数。

为什么需要梯度累积？

说白了，就是为了突破显存限制。我刚开始用 JAX 做 ViT 训练时，遇到过一个大坑：想把 batch size 设成 256，结果 24G 的显卡直接 OOM。后来用梯度累积，把 batch size 拆成 4 个 micro-batch，每个 64，问题就解决了。

具体来说，梯度累积解决了三个核心问题：

• 显存瓶颈：大 batch size 需要更多显存来存储中间激活值。梯度累积让你用小 batch size 训练，却能达到大 batch size 的效果。
• 数据分布不均：有些场景下，样本差异很大。小 batch 可能让梯度方向不稳定。累积多个 batch 的梯度，能平滑更新方向。
• 分布式训练同步开销：在多卡训练中，频繁同步梯度很耗时。梯度累积可以减少同步次数，提升吞吐量。

梯度累积的数学原理

我们来拆解一下数学过程。假设你原本想用 batch size 为 B 的数据更新一次参数。标准流程是：

1. 前向传播：计算 loss
2. 反向传播：计算梯度 ∇θ L
3. 参数更新：θ = θ - η * ∇θ L

现在，我们把 batch 拆成 N 个 micro-batch，每个大小为 B/N。梯度累积的流程变成：

1. 初始化累积梯度：grad_acc = 0
2. 对每个 micro-batch i:
   a. 前向传播：计算 loss_i
   b. 反向传播：计算 ∇θ L_i
   c. 累积：grad_acc += ∇θ L_i
3. 参数更新：θ = θ - η * (grad_acc / N)

注意最后一步要除以 N。为什么？因为每个 micro-batch 的梯度大小是 B/N 样本的平均梯度。累积 N 次后，总和相当于 B 个样本的梯度总和。为了保持和原始 batch size 一致的更新幅度，必须除以 N。

梯度累积与 Batch Size 的关系

这里有个容易混淆的点：梯度累积后的等效 batch size 到底是多少？

答案是：等效 batch size = micro-batch size × 累积步数。

举个例子：

micro-batch size	累积步数	等效 batch size
32	4	128
64	2	128
16	8	128

但要注意，等效 batch size 只是数学上的等价。实际训练中，梯度累积和真正的大 batch 训练还是有区别的。我曾经在训练 BERT 时对比过：用 micro-batch 32 累积 4 步 vs 直接用 batch 128。结果发现，梯度累积的收敛曲线稍微抖动一些，但最终精度几乎一样。

为什么会这样？因为梯度累积相当于在参数更新之间引入了延迟。每个 micro-batch 的梯度是基于旧参数计算的，累积后再更新。而真正的大 batch 是一次性基于当前参数计算所有梯度。这个延迟会导致梯度方向略有偏差，但实践中影响不大。

梯度累积的适用场景

不是所有场景都适合梯度累积。我总结了几种典型情况：

• 大模型训练：GPT、BERT、ViT 这类模型，单卡根本放不下大 batch。梯度累积是标配。
• 高分辨率图像：比如 1024×1024 的医学影像，一张图就占很多显存。用梯度累积可以凑够 batch。
• 长序列 NLP：序列长度 4096 甚至 8192 的 Transformer，每个样本都很大。梯度累积能帮你跑起来。
• 强化学习：PPO 等算法需要收集多个 trajectory 再更新。梯度累积可以模拟这种「收集-更新」模式。

但有些场景不太适合：

• 小模型 + 小数据：没必要，直接跑大 batch 更简单。
• 在线学习：数据流式到达，需要实时更新。累积会引入延迟。
• 某些 GAN 训练：生成器和判别器的更新频率需要精细控制，累积可能打乱节奏。

知识体系图

下面这张图展示了梯度累积在整个训练流程中的位置：

嗯，这张图很直观。数据被拆成多个 micro-batch，各自计算梯度，然后累积起来，最后统一更新参数。这就是梯度累积的核心流程。