随机梯度下降（SGD）原理与实现：从无偏梯度估计到动态学习率调度

随机梯度下降（Stochastic Gradient Descent, SGD）是深度学习中训练神经网络的事实标准优化算法。本文以《动手学深度学习》（D2L）随机梯度下降章节为骨架，系统讲解 SGD 的数学原理、与批量梯度下降的复杂度对比、动态学习率的必要性、凸目标下的收敛性分析，并结合本仓库 d2l/ 工具库中的 train_2d、show_trace_2d、SGD 等实现，给出可直接运行的二维优化实验代码，帮助读者理解"为什么深度学习训练离不开 SGD 与学习率调度"。

从梯度下降到随机梯度下降：动机与定义

在前面的章节中，我们已经反复使用随机梯度下降来训练模型，但没有解释它为什么有效。梯度下降章节介绍了基本梯度下降原理，本节则深入探讨随机梯度下降本身。

在深度学习中，目标函数通常是训练数据集中每个样本损失函数的平均值。给定包含 $n$ 个样本的训练数据集，设 $f_i(\mathbf{x})$ 是第 $i$ 个样本对应的损失函数（$\mathbf{x}$ 为参数向量），则目标函数为

$$f(\mathbf{x}) = \frac{1}{n} \sum_{i = 1}^n f_i(\mathbf{x}).$$

其在 $\mathbf{x}$ 处的梯度为

$$\nabla f(\mathbf{x}) = \frac{1}{n} \sum_{i = 1}^n \nabla f_i(\mathbf{x}).$$

如果使用（批量）梯度下降，每次迭代的计算代价为 $\mathcal{O}(n)$，随样本数 $n$ 线性增长。因此，当训练数据集越大，梯度下降每轮迭代的代价就越高。

随机梯度下降（SGD）降低了每次迭代的计算成本。在 SGD 的每次迭代中，我们从 $\{1,\ldots, n\}$ 中均匀随机采样一个样本索引 $i$，用 $\nabla f_i(\mathbf{x})$ 来更新 $\mathbf{x}$：

$$\mathbf{x} \leftarrow \mathbf{x} - \eta \nabla f_i(\mathbf{x}),$$

其中 $\eta$ 是学习率。可以看到，每次迭代的计算代价从梯度下降的 $\mathcal{O}(n)$ 降到了常数 $\mathcal{O}(1)$。更重要的是，随机梯度 $\nabla f_i(\mathbf{x})$ 是完整梯度 $\nabla f(\mathbf{x})$ 的无偏估计：

$$\mathbb{E}_i \nabla f_i(\mathbf{x}) = \frac{1}{n} \sum_{i = 1}^n \nabla f_i(\mathbf{x}) = \nabla f(\mathbf{x}).$$

这意味着平均而言，随机梯度是梯度的良好估计——这是 SGD 一切收敛性结论的基石。

用噪声模拟随机梯度：二维实验环境

为了直观比较梯度下降与 SGD，本节在目标函数 $f(x_1, x_2) = x_1^2 + 2x_2^2$ 上，通过向梯度添加均值为 0、方差为 1 的随机噪声来模拟随机梯度下降（与真实 SGD 的等价性见文末练习 2）。

#@tab all
def f(x1, x2):  # Objective function
    return x1 ** 2 + 2 * x2 ** 2

def f_grad(x1, x2):  # Gradient of the objective function
    return 2 * x1, 4 * x2

以 PyTorch 为例，带噪声的 SGD 更新函数如下（MXNet 与 TensorFlow 版本结构一致，仅随机数 API 不同）：

#@tab pytorch
def sgd(x1, x2, s1, s2, f_grad):
    g1, g2 = f_grad(x1, x2)
    # Simulate noisy gradient
    g1 += torch.normal(0.0, 1, (1,)).item()
    g2 += torch.normal(0.0, 1, (1,)).item()
    eta_t = eta * lr()
    return (x1 - eta_t * g1, x2 - eta_t * g2, 0, 0)

注意两点设计：

• 参数 (s1, s2) 是优化器的内部状态变量，本节的朴素 SGD 不使用它们（恒返回 0），但后续的 Momentum、Adagrad、RMSProp 等算法会用到，因此接口保持统一；
• 更新步长写作 eta_t = eta * lr()，其中 lr() 是一个学习率函数——这正是为下文"动态学习率"预留的钩子，当前先让它返回常数 1。

对应的 train_2d 与 show_trace_2d 工具函数定义在 d2l/torch.py（MXNet 与 TensorFlow 版本分别在 d2l/mxnet.py 和 d2l/tensorflow.py 中）。train_2d 从初始点 $(-5, -2)$ 出发迭代指定步数，逐轮记录 $(x_1, x_2)$ 轨迹并打印最终坐标；show_trace_2d 则在二维等高线图上叠加优化轨迹。

在常数学习率 $\eta = 0.1$ 下运行 50 步：

#@tab all
def constant_lr():
    return 1

eta = 0.1
lr = constant_lr  # Constant learning rate
d2l.show_trace_2d(f, d2l.train_2d(sgd, steps=50, f_grad=f_grad))

实验结果表明：相比 梯度下降章节 中平滑的轨迹，SGD 的变量轨迹噪声大得多，这是梯度的随机性所致——即使已经逼近最小值，仍会受到瞬时梯度经 $\eta \nabla f_i(\mathbf{x})$ 注入的不确定性影响。50 步后解的质量依然不佳，而且继续增加步数也不会改善（读者可自行加大步数验证）。这迫使我们只能改变学习率 $\eta$：选得太小，初期无法取得有意义的进展；选得太大，又得不到好解。解决这一矛盾的唯一办法是随着优化进行动态地降低学习率——这也是在 sgd 步函数中预留学习率函数 lr 的原因所在。

动态学习率：三种基础衰减策略

将固定的 $\eta$ 替换为随时间变化的学习率 $\eta(t)$，会增加控制优化算法收敛的复杂度。关键在于确定 $\eta$ 应该衰减得多快：太快会过早停止优化，太慢则浪费大量优化时间。下面是几种基础调整策略（更高级的调度策略在学习率调度章节讨论）：

$$\begin{aligned} \eta(t) & = \eta_i \textrm{ if } t_i \leq t \leq t_{i+1} && \textrm{piecewise constant} \\ \eta(t) & = \eta_0 \cdot e^{-\lambda t} && \textrm{exponential decay} \\ \eta(t) & = \eta_0 \cdot (\beta t + 1)^{-\alpha} && \textrm{polynomial decay} \end{aligned}$$

• 分段常数（piecewise constant）：每当优化进展停滞时降低学习率，例如从 $\eta_0$ 降到 $\eta_0/10$。这是训练深度网络的常见策略。
• 指数衰减（exponential decay）：衰减过于激进，常常在算法收敛之前就导致过早停止。
• 多项式衰减（polynomial decay）：$\alpha = 0.5$ 是受欢迎的选择。在凸优化情形下，已有大量证明表明该速率表现良好。

指数衰减的失败案例

#@tab all
def exponential_lr():
    # Global variable that is defined outside this function and updated inside
    global t
    t += 1
    return math.exp(-0.1 * t)

t = 1
lr = exponential_lr
d2l.show_trace_2d(f, d2l.train_2d(sgd, steps=1000, f_grad=f_grad))

exponential_lr 通过 global t 在函数外部维护迭代计数器。正如预期，参数方差显著减小了，但代价是无法收敛到最优解 $\mathbf{x} = (0, 0)$——即使迭代 1000 步，距离最优解依然很远，算法根本没有收敛。

多项式衰减的正确打开方式

#@tab all
def polynomial_lr():
    # Global variable that is defined outside this function and updated inside
    global t
    t += 1
    return (1 + 0.1 * t) ** (-0.5)

t = 1
lr = polynomial_lr
d2l.show_trace_2d(f, d2l.train_2d(sgd, steps=50, f_grad=f_grad))

当学习率按步数的平方根倒数（即 $\alpha = 0.5$）衰减时，仅 50 步收敛效果就明显变好。

学习率调度还有更多选择：可以从小学习率起步、快速攀升再缓慢下降（即 warmup 预热策略），也可以在大小学习率之间交替。本节将重点放在可以进行完整理论分析的场景——凸优化设定下的学习率上。对于一般的非凸问题，由于最小化非线性非凸问题通常是 NP 困难的，很难获得有意义的收敛保证。

凸目标函数的收敛性分析

以下针对凸目标函数的 SGD 收敛性分析是可选的，主要用于传达直觉，只涉及最简单的证明之一（Nesterov 与 Vial 的经典结果，见 d2l.bib 中的 Nesterov.Vial.2000 条目）。当目标函数性质特别良好时，存在更先进的证明技术。

设定：假设对所有 $\boldsymbol{\xi}$，目标函数 $f(\boldsymbol{\xi}, \mathbf{x})$ 关于 $\mathbf{x}$ 是凸的。考虑 SGD 更新：

$$\mathbf{x}_{t+1} = \mathbf{x}_{t} - \eta_t \partial_\mathbf{x} f(\boldsymbol{\xi}_t, \mathbf{x}),$$

其中 $f(\boldsymbol{\xi}_t, \mathbf{x})$ 是第 $t$ 步从某分布中抽取的训练样本 $\boldsymbol{\xi}_t$ 对应的目标函数。定义期望风险

$$R(\mathbf{x}) = E_{\boldsymbol{\xi}}[f(\boldsymbol{\xi}, \mathbf{x})],$$

其关于 $\mathbf{x}$ 的最小值为 $R^*$，并设 $\mathbf{x}^*$ 为最小化器（假设存在于 $\mathbf{x}$ 的定义域内）。此时可以追踪当前参数 $\mathbf{x}_t$ 与风险最小化器 $\mathbf{x}^*$ 之间的距离随时间的变化：

$$\begin{aligned} &\|\mathbf{x}_{t+1} - \mathbf{x}^*\|^2 \\ =& \|\mathbf{x}_{t} - \eta_t \partial_\mathbf{x} f(\boldsymbol{\xi}_t, \mathbf{x}) - \mathbf{x}^*\|^2 \\ =& \|\mathbf{x}_{t} - \mathbf{x}^*\|^2 + \eta_t^2 \|\partial_\mathbf{x} f(\boldsymbol{\xi}_t, \mathbf{x})\|^2 - 2 \eta_t \left\langle \mathbf{x}_t - \mathbf{x}^*, \partial_\mathbf{x} f(\boldsymbol{\xi}_t, \mathbf{x})\right\rangle. \end{aligned}$$

:eqlabel:eq_sgd-xt+1-xstar

有界梯度假设：假设随机梯度 $\partial_\mathbf{x} f(\boldsymbol{\xi}_t, \mathbf{x})$ 的 $\ell_2$ 范数被常数 $L$ 界定，则

$$\eta_t^2 \|\partial_\mathbf{x} f(\boldsymbol{\xi}_t, \mathbf{x})\|^2 \leq \eta_t^2 L^2.$$

:eqlabel:eq_sgd-L

利用凸性控制内积：我们对距离的期望变化更感兴趣——对任何特定样本序列，距离完全可能增加（取决于遇到的 $\boldsymbol{\xi}_t$），因此必须约束上式中的内积项。由凸函数性质，对所有 $\mathbf{x}$ 和 $\mathbf{y}$ 有 $f(\mathbf{y}) \geq f(\mathbf{x}) + \langle f'(\mathbf{x}), \mathbf{y} - \mathbf{x} \rangle$，于是

$$f(\boldsymbol{\xi}_t, \mathbf{x}^*) \geq f(\boldsymbol{\xi}_t, \mathbf{x}_t) + \left\langle \mathbf{x}^* - \mathbf{x}_t, \partial_{\mathbf{x}} f(\boldsymbol{\xi}_t, \mathbf{x}_t) \right\rangle.$$

:eqlabel:eq_sgd-f-xi-xstar

将 :eqref:eq_sgd-L 与 :eqref:eq_sgd-f-xi-xstar 代入 :eqref:eq_sgd-xt+1-xstar，得到第 $t+1$ 步参数距离的界：

$$\|\mathbf{x}_{t} - \mathbf{x}^*\|^2 - \|\mathbf{x}_{t+1} - \mathbf{x}^*\|^2 \geq 2 \eta_t (f(\boldsymbol{\xi}_t, \mathbf{x}_t) - f(\boldsymbol{\xi}_t, \mathbf{x}^*)) - \eta_t^2 L^2.$$

:eqlabel:eqref_sgd-xt-diff

这意味着只要当前损失与最优损失之差超过 $\eta_t L^2/2$，我们就在取得进展。由于该差值必然趋于零，学习率 $\eta_t$ 也必须趋于零（即必须衰减）。

对期望取平均：对 :eqref:eqref_sgd-xt-diff 取期望，得

$$E\left[\|\mathbf{x}_{t} - \mathbf{x}^*\|^2\right] - E\left[\|\mathbf{x}_{t+1} - \mathbf{x}^*\|^2\right] \geq 2 \eta_t [E[R(\mathbf{x}_t)] - R^*] - \eta_t^2 L^2.$$

对时间求和（望远镜求和）：对 $t \in \{1, \ldots, T\}$ 求和并丢弃低阶项，由于和式可望远镜式相消，且 $\mathbf{x}_1$ 已知、期望可去除，得到

$$\|\mathbf{x}_1 - \mathbf{x}^*\|^2 \geq 2 \left (\sum_{t=1}^T \eta_t \right) [E[R(\mathbf{x}_t)] - R^*] - L^2 \sum_{t=1}^T \eta_t^2.$$

:eqlabel:eq_sgd-x1-xstar

定义加权平均参数：

$$\bar{\mathbf{x}} \stackrel{\textrm{def}}{=} \frac{\sum_{t=1}^T \eta_t \mathbf{x}_t}{\sum_{t=1}^T \eta_t}.$$

由

$$E\left(\frac{\sum_{t=1}^T \eta_t R(\mathbf{x}_t)}{\sum_{t=1}^T \eta_t}\right) = \frac{\sum_{t=1}^T \eta_t E[R(\mathbf{x}_t)]}{\sum_{t=1}^T \eta_t} = E[R(\mathbf{x}_t)],$$

再结合 Jensen 不等式（取 $i=t$、$\alpha_i = \eta_t/\sum_{t=1}^T \eta_t$）以及 $R$ 的凸性，可得 $E[R(\mathbf{x}_t)] \geq E[R(\bar{\mathbf{x}})]$，从而

$$\sum_{t=1}^T \eta_t E[R(\mathbf{x}_t)] \geq \sum_{t=1}^T \eta_t E\left[R(\bar{\mathbf{x}})\right].$$

最终收敛界：将其代入 :eqref:eq_sgd-x1-xstar，得到

$$\left[E[\bar{\mathbf{x}}]\right] - R^* \leq \frac{r^2 + L^2 \sum_{t=1}^T \eta_t^2}{2 \sum_{t=1}^T \eta_t},$$

其中 $r^2 \stackrel{\textrm{def}}{=} \|\mathbf{x}_1 - \mathbf{x}^*\|^2$ 是初始参数与最终结果之间距离的界。简言之，收敛速度取决于随机梯度范数的界 $L$ 与初始参数离最优性的距离 $r$。注意该界是对 $\bar{\mathbf{x}}$ 而非 $\mathbf{x}_T$ 给出的——因为 $\bar{\mathbf{x}}$ 是优化路径的平滑版本。

当 $r$、$L$ 和 $T$ 已知时，可以选取学习率 $\eta = r/(L \sqrt{T})$，此时上界为 $rL/\sqrt{T}$，即以 $\mathcal{O}(1/\sqrt{T})$ 的速率收敛到最优解。

关键结论：为何学习率必须衰减

上述分析从理论上印证了实验观察：只要当前损失与最优损失的差距大于 $\eta_t L^2/2$，SGD 就在进步；而该差距必然收敛到零，因此 $\eta_t$ 也必须消失。这与第三节"常数学习率导致无法收敛、指数衰减导致过早停止、多项式衰减（$\alpha=0.5$）表现良好"的实验结果完全一致。

随机梯度与有限样本：有放回采样为什么不好

到目前为止，关于随机梯度下降的讨论有些"草率"：我们设想从分布 $p(x, y)$ 中抽取实例 $x_i$（通常带标签 $y_i$），并据此更新模型参数。对有限样本量，我们断言离散分布

$$p(x, y) = \frac{1}{n} \sum_{i=1}^n \delta_{x_i}(x) \delta_{y_i}(y)$$

（$\delta_{x_i}$ 与 $\delta_{y_i}$ 为某种函数）允许我们对其执行随机梯度下降。

但本书实际做法并非如此。本节的玩具示例只是向非随机梯度添加了噪声，即假装有样本对 $(x_i, y_i)$（其合理性见练习 2 的讨论）。更值得注意的是，此前所有章节的讨论都没有这样做，而是恰好一次地遍历所有实例。为什么这样更好？考虑反面情形：从离散分布中有放回地采样 $n$ 次观测。随机选中元素 $i$ 的概率是 $1/n$，因此"至少选中一次"的概率为

$$P(\textrm{choose~} i) = 1 - P(\textrm{omit~} i) = 1 - (1-1/n)^n \approx 1-e^{-1} \approx 0.63.$$

类似推理给出"恰好选中某个样本一次"的概率：

$${n \choose 1} \frac{1}{n} \left(1-\frac{1}{n}\right)^{n-1} = \frac{n}{n-1} \left(1-\frac{1}{n}\right)^{n} \approx e^{-1} \approx 0.37.$$

有放回采样相比无放回采样会导致更大的方差和更低的数据利用效率。因此实践中采用无放回采样（这也是本书各处的默认做法）。另外注意，多次遍历训练数据集时，每次遍历的顺序是随机打乱的。这一观点与小批量随机梯度下降章节互为补充：那里论证了为什么在"逐样本 SGD"与"全批量梯度下降"之间取折中的小批量方案在计算（向量化、缓存利用）上更高效。

仓库中的 SGD 实现印证

本节概念在本仓库的工具代码中有直接对应物：

• 动手实现的朴素 SGD：d2l/torch.py 中的 sgd(params, lr, batch_size) 函数在 torch.no_grad() 下对所有参数执行 param -= lr * param.grad / batch_size 并清零梯度，即对批量梯度取平均后的标准小批量更新，其数学形式与本文的 $\mathbf{x} \leftarrow \mathbf{x} - \eta \nabla f_i(\mathbf{x})$ 一一对应。
• 面向对象封装：d2l/torch.py 中的 SGD(d2l.HyperParameters) 类提供 step() 与 zero_grad() 方法，step() 内部即 param -= self.lr * param.grad，zero_grad() 负责清空每个参数的梯度，供 configure_optimizers（见同文件 L380-L382）等高层接口复用。
• 二维实验基础设施：train_2d 与 show_trace_2d（d2l/torch.py）即为本文实验所用的轨迹跟踪与可视化工具，三者接口一致（MXNet 见 d2l/mxnet.py，TensorFlow 见 d2l/tensorflow.py），train_2d 注释明确说明 s1、s2 是为 Momentum、Adagrad、RMSProp 预留的状态变量。
• 调度器的工程化落地：本文讨论的分段常数、指数、多项式衰减等策略在 学习率调度章节 中借助框架内置的 lr_scheduler 自动化实现，是"动态学习率"思想在生产训练中的标准形态。

总结

• 对凸问题，可以证明在相当宽泛的学习率选择下，随机梯度下降都能收敛到最优解。
• 对深度学习而言一般并非如此，但凸问题的分析给出了如何开展优化的宝贵洞见：渐进式地降低学习率，但不要太快。
• 学习率太小或太大都会引发问题。实践中，合适的学习率往往需要多次实验才能找到。
• 训练数据集样本越多，梯度下降每次迭代的计算成本越高，因此这些场景下更倾向于随机梯度下降。
• 在非凸情形下，SGD 一般没有最优性保证，因为需要检查的局部极小值数量可能是指数级的。

练习

1. 尝试 SGD 的不同学习率调度与不同迭代次数，特别是绘制与最优解 $(0, 0)$ 的距离随迭代次数的变化曲线。
2. 证明：对函数 $f(x_1, x_2) = x_1^2 + 2 x_2^2$，向梯度添加正态噪声等价于最小化损失函数 $f(\mathbf{x}, \mathbf{w}) = (x_1 - w_1)^2 + 2 (x_2 - w_2)^2$，其中 $\mathbf{x}$ 服从正态分布。
3. 比较从 $\{(x_1, y_1), \ldots, (x_n, y_n)\}$ 中有放回采样与无放回采样时 SGD 的收敛表现。
4. 如果某些梯度（或其中某些坐标）持续远大于其他梯度，你会如何修改 SGD 求解器？
5. 假设 $f(x) = x^2 (1 + \sin x)$，$f$ 有多少个局部极小值？能否改造 $f$ 使得最小化它需要评估所有局部极小值？