让炼丹更科学一些（七）：步长调度与权重平均

精通炼丹的同学都知道，步长调度，或者说学习率调度（LR Schedule），对模型最终效果来说至关重要。我们在前几篇文章中已经推导过，即便只考虑SGD，终点收敛的最优学习率函数也具有Warmup-Decay这一实践常用的形态。不过，近来也有一些工作，比如Schedule-Free，尝试通过模型权重的某种平均来替代学习率调度，也取得了一些进展。

如何从理论上分析学习率调度与权重平均的联系呢？它们在多大程度上可以相互替换呢？这篇文章我们来尝试回答一下这些问题。

权重平均[#](https://spaces.ac.cn/archives/11804#%E6%9D%83%E9%87%8D%E5%B9%B3%E5%9D%87)

对于权重平均，事实上我们在《让炼丹更科学一些（三）：SGD的终点损失收敛》就已经给出过一个基本结果

$$ (\text{1}) \mathbb{E} \left[\right. L \left(\right. \bar{\mathbf{\mathit{\theta}}}_{T} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. \leq \frac{\parallel \mathbf{\mathit{\theta}}_{1} - \mathbf{\mathit{\theta}}^{*} \parallel^{2}}{2 T \eta_{T}} + \frac{G^{2}}{2 T} \sum_{t = 1}^{T} \frac{\eta_{t}^{2}}{\eta_{T}} $$

其中$\bar{\mathbf{\mathit{\theta}}}_{T} = \frac{1}{T} \sum_{t = 1}^{T} \mathbf{\mathit{\theta}}_{t}$。这个结果意味着只要我们选择常数学习率$\eta \propto 1 / \sqrt{T}$，右端就可以达到$\mathcal{O} \left(\right. 1 / \sqrt{T} \left.\right)$的理想收敛速度，左端则意味着$\bar{\mathbf{\mathit{\theta}}}_{T}$收敛于$\mathbf{\mathit{\theta}}^{*}$，于是我们设计如下优化器

$$ (\text{2}) \begin{matrix}\mathbf{\mathit{\theta}}_{t + 1} = & \mathbf{\mathit{\theta}}_{t} - \eta \mathbf{\mathit{g}} \left(\right. \mathbf{\mathit{x}}_{t} , \mathbf{\mathit{\theta}}_{t} \left.\right) \\ \mathbf{\mathit{\mu}}_{t + 1} = & \left(\right. 1 - c_{t + 1} \left.\right) \mathbf{\mathit{\mu}}_{t} + c_{t + 1} \mathbf{\mathit{\theta}}_{t + 1}\end{matrix} $$

其中$c_{t} = 1 / t$。也就是说，新增一步完全不干扰训练轨迹的迭代，把训练轨迹等权平均起来，最终用来做预测的模型权重并不是终点值$\mathbf{\mathit{\theta}}_{T}$，而是平均值$\mathbf{\mathit{\mu}}_{T}$。然而，该做法实践表现并不佳，这是理论与实践不符的经典例子之一。 不过，如果我们将$c_{t}$改为某个与$t$无关的适当的常数，比如$10^{- 3}$，那么实践中往往有着不俗的表现——这正是笔者之前打比赛常用的指数滑动平均（EMA）技巧。相比$c_{t} = 1 / t$，这种滑动平均的做法更倾向于平均最近的优化轨迹，而非整个历史轨迹，它的表现优异也表明我们不需要关注太久远的历史结果。

如何从理论上解释这种滑动平均技巧，或者更一般地，解释任意权重平均的可行性，便是接下来的核心任务。

再次推广[#](https://spaces.ac.cn/archives/11804#%E5%86%8D%E6%AC%A1%E6%8E%A8%E5%B9%BF)

在文章《让炼丹更科学一些（四）：新恒等式，新学习率》中，为了将平均损失的结果推广到终点损失，我们引入了恒等式

$$ (\text{3}) q_{T} = \frac{1}{w_{1 : T}} \sum_{t = 1}^{T} w_{t} q_{t} + \sum_{k = 1}^{T - 1} \left(\right. \frac{1}{w_{k + 1 : T}} - \frac{1}{w_{k : T}} \left.\right) \sum_{t = k + 1}^{T} w_{t} \left(\right. q_{t} - q_{k} \left.\right) $$

其中$w_{k : T} \triangleq \sum_{t = k}^{T} w_{t}$。这一节我们再次将它推广为

$$ (\text{4}) \frac{1}{v_{1 : T}} \sum_{t = 1}^{T} v_{t} q_{t} = \frac{1}{w_{1 : T}} \sum_{t = 1}^{T} w_{t} q_{t} + \frac{1}{v_{1 : T}} \sum_{k = 1}^{T - 1} \left(\right. \frac{v_{k + 1 : T}}{w_{k + 1 : T}} - \frac{v_{k : T}}{w_{k : T}} \left.\right) \sum_{t = k + 1}^{T} w_{t} \left(\right. q_{t} - q_{k} \left.\right) $$

直观来看，它提供了将一种加权求和转换成另一种加权求和的思路，当$v_{T} = 1$且剩下的$v_{t}$等于0时，它便退化为终点形式。证明实际上也是之前证明的推广，记$S_{k} = \frac{1}{w_{k : T}} \sum_{t = k}^{T} w_{t} q_{t}$，可以写出

$$ (\text{5}) \begin{matrix}v_{k : T} S_{k} - v_{k + 1 : T} S_{k + 1} = & v_{k : T} \left(\right. \frac{w_{k + 1 : T} S_{k + 1} + w_{k} q_{k}}{w_{k : T}} \left.\right) - v_{k + 1 : T} S_{k + 1} \\ = & v_{k} q_{k} + \left(\right. \frac{v_{k : T}}{w_{k : T}} w_{k} - v_{k} \left.\right) q_{k} + \left(\right. \frac{v_{k : T}}{w_{k : T}} w_{k + 1 : T} - v_{k + 1 : T} \left.\right) S_{k + 1} \\ = & v_{k} q_{k} + \left(\right. \frac{v_{k : T}}{w_{k : T}} w_{k + 1 : T} - v_{k + 1 : T} \left.\right) \left(\right. S_{k + 1} - q_{k} \left.\right) \\ = & v_{k} q_{k} + \left(\right. \frac{v_{k : T}}{w_{k : T}} - \frac{v_{k + 1 : T}}{w_{k + 1 : T}} \left.\right) \left(\right. w_{k + 1 : T} S_{k + 1} - w_{k + 1 : T} q_{k} \left.\right) \\ = & v_{k} q_{k} + \left(\right. \frac{v_{k : T}}{w_{k : T}} - \frac{v_{k + 1 : T}}{w_{k + 1 : T}} \left.\right) \sum_{t = k + 1}^{T} w_{t} \left(\right. q_{t} - q_{k} \left.\right)\end{matrix} $$

第三个等号是因为直接将第二个等号的两个括号相加等于0，所以这两个括号互为相反数。现在左右两端对$k = 1 sim T - 1$求和，左端等于$v_{1 : T} S_{1} - v_{T} q_{T}$，将$v_{T} q_{T}$挪到右端，两端除以$v_{1 : T}$并稍加整理，即得待证恒等式$(\text{4})$。

放缩变换[#](https://spaces.ac.cn/archives/11804#%E6%94%BE%E7%BC%A9%E5%8F%98%E6%8D%A2)

接着我们沿着上一篇《让炼丹更科学一些（六）：自上而下的精妙构造》的思路，去推导更一般的收敛结论。设$q_{t} = \mathbb{E} \left[\right. L \left(\right. \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right.$，代入到式$(\text{4})$得

$$ (\text{6}) \frac{1}{v_{1 : T}} \sum_{t = 1}^{T} v_{t} \mathbb{E} \left[\right. L \left(\right. \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. = \frac{1}{w_{1 : T}} \sum_{t = 1}^{T} w_{t} \mathbb{E} \left[\right. L \left(\right. \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. + \frac{1}{v_{1 : T}} \sum_{k = 1}^{T - 1} \left(\right. \frac{v_{k + 1 : T}}{w_{k + 1 : T}} - \frac{v_{k : T}}{w_{k : T}} \left.\right) \sum_{t = k + 1}^{T} w_{t} \mathbb{E} \left[\right. L \left(\right. \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}_{k} \left.\right) \left]\right. $$

利用$L$的凸性，对$\mathbb{E} \left[\right. L \left(\right. \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right.$和$\mathbb{E} \left[\right. L \left(\right. \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}_{k} \left.\right) \left]\right.$进行放缩，然后像上一篇文章的“恒等变换”一节那样仔细整理，可得

$$ (\text{7}) \frac{1}{v_{1 : T}} \sum_{t = 1}^{T} v_{t} \mathbb{E} \left[\right. L \left(\right. \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. \leq \frac{1}{w_{1 : T}} \sum_{t = 1}^{T} w_{t} \mathbb{E} \left[\right. \mathbf{\mathit{g}} \left(\right. \mathbf{\mathit{x}}_{t} , \mathbf{\mathit{\theta}}_{t} \left.\right) \cdot \left(\right. \mathbf{\mathit{\psi}}_{t} - \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. $$

其中

$$ (\text{8}) \mathbf{\mathit{\psi}}_{t} = \frac{w_{1 : T}}{v_{1 : T}} \left[\right. \frac{v_{t : T}}{w_{t : T}} \mathbf{\mathit{\theta}}_{t} - \sum_{k = 1}^{t - 1} \left(\right. \frac{v_{k + 1 : T}}{w_{k + 1 : T}} - \frac{v_{k : T}}{w_{k : T}} \left.\right) \mathbf{\mathit{\theta}}_{k} \left]\right. $$

并且直接验证可得

$$ (\text{9}) \mathbf{\mathit{\psi}}_{t + 1} - \mathbf{\mathit{\psi}}_{t} = \frac{v_{t + 1 : T}}{v_{1 : T}} \frac{w_{1 : T}}{w_{t + 1 : T}} \left(\right. \mathbf{\mathit{\theta}}_{t + 1} - \mathbf{\mathit{\theta}}_{t} \left.\right) $$

因此，如果我们让$\mathbf{\mathit{\psi}}_{t}$按照$\mathbf{\mathit{\psi}}_{t + 1} = \mathbf{\mathit{\psi}}_{t} - w_{t} \mathbf{\mathit{g}} \left(\right. \mathbf{\mathit{x}}_{t} , \mathbf{\mathit{\theta}}_{t} \left.\right)$更新，那么对应的$\mathbf{\mathit{\theta}}_{t}$的更新规则是

$$ (\text{10}) \mathbf{\mathit{\theta}}_{t + 1} = \mathbf{\mathit{\theta}}_{t} - \underset{\eta_{t}}{\underbrace{\frac{v_{1 : T}}{v_{t + 1 : T}} \frac{w_{t} w_{t + 1 : T}}{w_{1 : T}}}} \mathbf{\mathit{g}} \left(\right. \mathbf{\mathit{x}}_{t} , \mathbf{\mathit{\theta}}_{t} \left.\right) $$

一般结论[#](https://spaces.ac.cn/archives/11804#%E4%B8%80%E8%88%AC%E7%BB%93%E8%AE%BA)

接下来的操作就比较朴素了，对式$(\text{7})$右端应用《让炼丹更科学一些（五）：基于梯度精调学习率》等文章的结论得

$$ (\text{11}) \frac{1}{w_{1 : T}} \sum_{t = 1}^{T} w_{t} \mathbb{E} \left[\right. \mathbf{\mathit{g}} \left(\right. \mathbf{\mathit{x}}_{t} , \mathbf{\mathit{\theta}}_{t} \left.\right) \cdot \left(\right. \mathbf{\mathit{\psi}}_{t} - \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. \leq \frac{1}{2 w_{1 : T}} \left(\right. R^{2} + \sum_{t = 1}^{T} w_{t}^{2} G_{t}^{2} \left.\right) $$

其中$R = \parallel \mathbf{\mathit{\theta}}_{1} - \mathbf{\mathit{\theta}}^{*} \parallel , G_{t}^{2} = \mathbb{E} \left[\right. \parallel \mathbf{\mathit{g}} \left(\right. \mathbf{\mathit{x}}_{t} , \mathbf{\mathit{\theta}}_{t} \left.\right) \parallel^{2} \left]\right.$，约定$\mathbf{\mathit{\psi}}_{1} = \mathbf{\mathit{\theta}}_{1}$。至于左端，我们应用凸性得

$$ (\text{12}) \mathbb{E} \left[\right. L \left(\right. \frac{1}{v_{1 : T}} \sum_{t = 1}^{T} v_{t} \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. \leq \frac{1}{v_{1 : T}} \sum_{t = 1}^{T} v_{t} \mathbb{E} \left[\right. L \left(\right. \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. $$

综合起来是

$$ (\text{13}) \mathbb{E} \left[\right. L \left(\right. \frac{1}{v_{1 : T}} \sum_{t = 1}^{T} v_{t} \mathbf{\mathit{\theta}}_{t} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. \leq \frac{1}{2 w_{1 : T}} \left(\right. R^{2} + \sum_{t = 1}^{T} w_{t}^{2} G_{t}^{2} \left.\right) $$

右端的最小值在

$$ (\text{14}) w_{t} = \frac{R G_{t}^{- 2}}{\sqrt{Q_{T}}} , Q_{T} = \sum_{k = 1}^{T} G_{k}^{- 2} $$

取到，所以意味着最优学习率具有形式

$$ (\text{15}) \eta_{t} = \frac{v_{1 : T}}{v_{t + 1 : T}} \frac{w_{t} w_{t + 1 : T}}{w_{1 : T}} = \frac{v_{1 : T}}{v_{t + 1 : T}} \frac{R G_{t}^{- 2}}{\sqrt{Q_{T}}} \left(\right. 1 - \frac{Q_{t}}{Q_{T}} \left.\right) $$

结果解读[#](https://spaces.ac.cn/archives/11804#%E7%BB%93%E6%9E%9C%E8%A7%A3%E8%AF%BB)

上一篇文章《让炼丹更科学一些（六）：自上而下的精妙构造》给出的终点收敛结果是

$$ (\text{16}) \mathbb{E} \left[\right. L \left(\right. \mathbf{\mathit{\theta}}_{T} \left.\right) - L \left(\right. \mathbf{\mathit{\theta}}^{*} \left.\right) \left]\right. \leq \frac{1}{2 w_{1 : T}} \left(\right. R^{2} + \sum_{t = 1}^{T} w_{t}^{2} G_{t}^{2} \left.\right) $$

右端跟式$(\text{13})$是一致的，但学习率变为$\eta_{t} = \frac{w_{t} w_{t + 1 : T}}{w_{1 : T}}$。由此可见，平均权重如果想要跟终点权重效果相近，学习率需要多乘以一个$\frac{v_{1 : T}}{v_{t + 1 : T}}$，或者反过来说，在相同的学习率下，权重平均的效果相当于学习率乘以$\frac{v_{t + 1 : T}}{v_{1 : T}}$后的终点效果。很明显$\frac{v_{t + 1 : T}}{v_{1 : T}}$小于等于1且单调递减，乘上这一项就起到了类似于LR Decay的作用。 让我们考虑一些具体例子。首先是最简单的$v_{t} \equiv 1$，对应开头的式$(\text{2})$，此时$\frac{v_{t + 1 : T}}{v_{1 : T}} = 1 - t / T$，正好是线性衰减，这意味着“常数学习率 + 等权平均”理论上跟线性衰减的学习率有着相近的效果。然而，实测效果往往是，直接将学习率做线性衰减的终点效果不错，但平均权重也就是式$(\text{2})$的效果欠佳。

回顾整个推导过程，可以猜测问题大概出在“凸函数”的假设上。我们这个系列文章，其实都依赖于“目标函数是凸的”这一假设，但实际情况往往是非凸的，凸优化的结论无法精确复现。折中方案是假设局部近似为凸，那么可以在有限区间内对等权重平均，又或者直接考虑它的平滑版本——指数滑动平均（EMA），此时$v_{t} = \gamma^{- t}$（其中$0 < \gamma < 1$），那么

$$ (\text{17}) \frac{v_{t + 1 : T}}{v_{1 : T}} = \frac{1 - \gamma^{T - t}}{1 - \gamma^{T}} $$

它维持了一个$\mathcal{O} \left(\right. \left(\right. 1 - \gamma \left.\right)^{- 1} \left.\right)$的平均区间。由于指数函数的无记忆性，如果我们在常数学习率训练的同时运行EMA，那么相当于每一步同时得到了经过$\mathcal{O} \left(\right. \left(\right. 1 - \gamma \left.\right)^{- 1} \left.\right)$步线性衰减后的权重，可以提前窥见线性衰减后的效果。

相关工作[#](https://spaces.ac.cn/archives/11804#%E7%9B%B8%E5%85%B3%E5%B7%A5%E4%BD%9C)

对于式$(\text{2})$效果欠佳的问题，论文《The Road Less Scheduled》提出了“Schedule-Free”，它将求梯度的点改为了$\mathbf{\mathit{\theta}}_{t}$与$\mathbf{\mathit{\mu}}_{t}$的插值：

$$ (\text{18}) \begin{matrix}\hat{\mathbf{\mathit{\theta}}}_{t} = & \left(\right. 1 - \beta \left.\right) \mathbf{\mathit{\theta}}_{t} + \beta \mathbf{\mathit{\mu}}_{t} \\ \mathbf{\mathit{\theta}}_{t + 1} = & \mathbf{\mathit{\theta}}_{t} - \eta \mathbf{\mathit{g}} \left(\right. \mathbf{\mathit{x}}_{t} , \hat{\mathbf{\mathit{\theta}}}_{t} \left.\right) \\ \mathbf{\mathit{\mu}}_{t + 1} = & \left(\right. 1 - c_{t + 1} \left.\right) \mathbf{\mathit{\mu}}_{t} + c_{t + 1} \mathbf{\mathit{\theta}}_{t + 1}\end{matrix} $$

直观来想，就是原本的平均操作完全独立于训练过程，模型从未感知过$\mathbf{\mathit{\mu}}_{t}$究竟好不好，现在我们将求梯度的点换成$\mathbf{\mathit{\theta}}_{t}$与$\mathbf{\mathit{\mu}}_{t}$的插值，模型能感知到$\mathbf{\mathit{\mu}}_{t}$的好坏，从而作出调整。这个改动理论上是有保证的，同时确实能显著改进平均权重$\mathbf{\mathit{\mu}}_{T}$的效果。 不过，随着训练规模的扩大，Schedule-Free逐渐表现出了颓势，作者近来又提出了ScheduleFree+，给原本的Schedule-Free打上了不少“补丁”，使它重新成为了SOTA。不过，Plus版额外添加的操作比较多，整套算法显得复杂起来了，理解难度也增加不少，所以这里就不详细展开了，后面我们再回来讨论它。

值得一提的是，Schedule-Free及其Plus版的作者，正是前几篇文章我们介绍过的经典论文《Optimal Linear Decay Learning Rate Schedules and Further Refinements》的作者，这几篇工作都成功将经典凸优化理论用到了实践训练之上，而不是停留在玩具模型的“纸上谈兵”之上，值得我们致敬！

延伸思考[#](https://spaces.ac.cn/archives/11804#%E5%BB%B6%E4%BC%B8%E6%80%9D%E8%80%83)

真正值得思考的是，Schedule-Free的设想真的能成立吗？所谓Schedule-Free，它有两个目标，一是希望以常数学习率从头训到尾，二是希望各种超参数不依赖于总训练步数$T$。如果满足这两个条件，那么训练过程就可以随时停止，也可以随时续训，并且随便一步停下来，它都是截止到当前步数的最优解，确实非常理想。

然而，这两个目标本身是值得商榷的。首先，对于实际训练来说，常数学习率还是动态的Cosine Decay、Linear Decay等，并无本质区别，反而都是调参对比选最优，所以追求常数学习率实践意义不大；而且即便是Schedule-Free及其Plus版，都没摆脱Warmup的依赖，所以实际上它们都还某种程度上依赖着人为定制的LR Schedule。

其次，希望各种超参数不依赖于总训练步数$T$，确实是一个好目标，但只从理论上看它都很难成立。还是以式$(\text{2})$为例，即便在满足所有假设下，式$(\text{2})$能表现良好，但它的最优学习率$\eta$依然是依赖于$T$；同样，Schedule-Free无法摆脱这一事实，它是通过引入新的假设，来获得$\eta$不依赖于$T$的结果。

我们也可以从式$(\text{15})$来观察，对于等权平均有$\frac{v_{1 : T}}{v_{t + 1 : T}} = 1 / \left(\right. 1 - t / T \left.\right)$，进一步假设$G_{t}$是常数$G$，那么$1 - Q_{t} / Q_{T} = 1 - t / T$，刚好抵消，因此最优学习率是$\frac{R G^{- 2}}{\sqrt{Q_{T}}}$，$Q_{T}$依赖于$T$。另一方面，$G_{t}$通常是先大后小，所以根据式$(\text{15})$反比于$G_{t}^{2}$的特点，早期学习率应该更小，这反而可以解释Schedule-Free仍无法彻底摆脱Schedule、仍需要Warmup的现象。

为此，ScheduleFree+以及一些同期工作，已经逐渐从“Schedule-Free”往“LR-Free”考虑，但总的来说，这一块仍处于“方兴未艾”的阶段，有待更本质的探索。

文章小结[#](https://spaces.ac.cn/archives/11804#%E6%96%87%E7%AB%A0%E5%B0%8F%E7%BB%93)

本文引入了新的恒等式，继续推广上一篇文章的结果，这使得我们可以从理论上定量地揭示权重平均与学习率衰减的联系。此外，我们还讨论了等权平均效果不如滑动平均、Schedule-Free仍无法彻底摆脱Schedule等问题。

_转载到请包括本文地址：[https://spaces.ac.cn/archives/11804](https://spaces.ac.cn/archives/11804 "让炼丹更科学一些（七）：步长调度与权重平均")_

_更详细的转载事宜请参考：_[《科学空间FAQ》](https://spaces.ac.cn/archives/6508#%E6%96%87%E7%AB%A0%E5%A6%82%E4%BD%95%E8%BD%AC%E8%BD%BD/%E5%BC%95%E7%94%A8 "《科学空间FAQ》")