Lookahead Optimizer/zh

Article
Topic area	Optimization
Prerequisites	Stochastic Gradient Descent, Adam Optimizer

概述

Lookahead 优化器是一种用于深度学习中随机优化的包装算法,由 Zhang、Lucas、Hinton 和 Ba 于 2019 年提出。 Lookahead 并不取代 SGD 或 Adam 等现有优化器,而是维护两组权重:一组快速权重由内部优化器更新若干步,另一组慢速权重周期性地向快速权重靠拢。每次同步之后,快速权重重置为慢速权重,然后过程重复。这种双循环结构降低了内部优化器轨迹的方差,并倾向于在几乎不增加成本的情况下改善泛化能力。

Lookahead 被广泛用作对现有训练流水线的即插即用增强,因为它保留了内部优化器的行为,只额外引入两个超参数:内部步数 $ k $ 与慢速权重的插值系数 $ \alpha $。

直觉

随机梯度方法本质上是有噪声的:每一步都基于一个小批量计算,因此参数空间中的轨迹会围绕下降方向振荡而非精确地沿其前进。诸如重球法或自适应学习率等标准的解决方法会对当前的梯度信号进行调整。Lookahead 则在更长的时间尺度上工作:将内部优化器的一连串更新视为一次探索阶段,然后只采纳所得位移的一部分。

该算法常被描述为口号"k 步前进,1 步后退"。内循环侦察损失曲面的某个区域;外循环则沿着从原慢速点到新快速点的连线保守地移动慢速权重。由于慢速权重只移动位移的一个分数 $ \alpha $,内层轨迹的瞬时振荡会被抑制,而一致的下降方向会逐渐积累。

算法

设 $ \phi $ 表示慢速权重,$ \theta $ 表示快速权重。设 $ A $ 为任意内部优化器(例如带动量的 SGD 或 Adam),并设 $ L $ 表示损失函数。Lookahead 在一个内循环和一次外部更新之间交替。

1. 初始化 $ \phi_0 $,并令 $ \theta_{0,0} \leftarrow \phi_0 $。
2. 对每一个外部步骤 $ t = 0, 1, 2, \dots $:
   1. 对于 $ i = 1, \dots, k $:抽取一个小批量 $ d_i $,并更新 $ \theta_{t,i} \leftarrow A(\theta_{t,i-1}, \nabla L(\theta_{t,i-1}, d_i)) $。
   2. 更新慢速权重:$ \phi_{t+1} \leftarrow \phi_t + \alpha\,(\theta_{t,k} - \phi_t) $。
   3. 重置快速权重:$ \theta_{t+1,0} \leftarrow \phi_{t+1} $。

慢速权重 $ \phi $ 即训练结束时返回并在推理时使用的权重。快速权重只是临时状态。

更新规则

外循环的更新可以紧凑地写成线性插值的形式:

$ {\displaystyle \phi_{t+1} = (1 - \alpha)\,\phi_t + \alpha\,\theta_{t,k}.} $

当 $ \alpha = 1 $ 时,慢速权重在每次同步时都跳到快速权重处,Lookahead 退化为只运行内部优化器。当 $ \alpha = 0 $ 时,慢速权重永远不动。实际取值位于 $ (0, 1] $ 中,常用的默认值为 $ \alpha = 0.5 $。

一种等价的表述是对每 $ k $ 步采样一次的快速检查点序列做指数移动平均。不过与标准的移动平均不同,Lookahead 还会在每次外部更新后将快速权重重置为慢速权重,从而将两条轨迹耦合起来,而不是让它们各自漂离。

超参数

在内部优化器自身的超参数之外,Lookahead 还额外引入了两个超参数。

• 同步周期 $ k $:外部更新之间所执行的内部步数。典型取值为 $ 5 $、$ 10 $ 或 $ 20 $。较大的 $ k $ 可以摊销外部步骤的(微小)开销,但会让内部优化器在两次同步之间漂移更多。
• 慢速权重步长 $ \alpha $:被应用于慢速权重的内部位移的比例。典型值为 $ 0.5 $ 或 $ 0.8 $。较小的 $ \alpha $ 提供更强的方差缩减,但进度也更慢。

原文献报告在 $ k \in \{5, 10\} $ 与 $ \alpha \in \{0.5, 0.8\} $ 的取值下,该方法在图像分类与语言建模任务上都表现稳健,表明很少需要精细调参。

方差缩减

在二次损失近似下,Zhang 等人证明:在内部优化器动态固定的情况下,慢速权重的方差满足

$ {\displaystyle \mathrm{Var}(\phi_{t+1}) = (1-\alpha)^2 \,\mathrm{Var}(\phi_t) + \alpha^2 \,\mathrm{Var}(\theta_{t,k}),} $

因此慢速权重在稳态时的方差是内部优化器方差的 $ \alpha / (2 - \alpha) $。当 $ \alpha = 0.5 $ 时,慢速权重的方差为内部方差的三分之一,这可能转化为更平坦的极小值与更好的泛化能力。该分析假设内部轨迹不相关且损失局部为二次形式,因此对深度网络来说只是参考而非保证,但与实践中观察到的经验规律相吻合。

实践考虑

内存与计算开销都很有限。Lookahead 需要额外存储一份参数副本(即慢速权重),并每隔 $ k $ 步进行一次按元素插值,因此对 $ k \geq 5 $ 而言,相较于内部优化器的实际墙钟开销通常低于百分之一。

当 Lookahead 与批归一化或其他带有运行时统计量的模块结合使用时,这些层内部维护的缓冲(运行均值与方差)不是优化器的参数,因此不参与插值。多数实现将其继续附着在模型上,由内循环的前向传播来更新,这也是推荐的约定。

关于学习率调度,通常保留内部优化器的调度方案而不做修改。Lookahead 本身并不需要预热,但预热并无害处,且在 Transformer 训练中对内部优化器仍然有用。

变体与相关方法

有若干后续工作扩展了 Lookahead 或与其密切相关。

• Stochastic Weight Averaging(SWA)在训练结束时对权重取平均,而不是在训练过程中进行。SWA 产生一个平均后的检查点,而 Lookahead 在整个优化过程中维持一个滑动的插值。
• Ranger 将 Lookahead 与 Adam 的 RAdam 变体结合,在计算机视觉基准上颇受欢迎。
• Polyak 平均对过去迭代进行移动平均,可视为不带快慢权重重置的、重复 $ \alpha $ 插值的极限形式。
• Reptile 是一种与之具有相同外部更新形式的元学习算法,只是其内循环在一个采样的任务上训练,而非一个采样的小批量。

比较

与原始的 SGD 或 Adam 相比,Lookahead 用少量额外内存换取更低的轨迹方差,并且常常带来测试准确率上的小幅改善。与 SWA 相比,Lookahead 不需要在训练末期专门进行平均阶段,且可以在任意检查点上进行评估。与动量相比,这两个时间尺度是正交的:动量在每一步层面上平滑梯度信号,而 Lookahead 在每 $ k $ 步层面上平滑参数轨迹,二者通常会一起使用。

局限性

Lookahead 的经验收益较为有限,且在不同任务上并不一致。在调优良好且使用强学习率调度的基线上,改善幅度可能落入运行间波动的噪声范围。该方法还引入了两个超参数,尽管默认值 $ k=5 $ 与 $ \alpha=0.5 $ 通常已经够用。最后,方差缩减的分析假设损失局部为二次形式,且内部动态近似平稳,但在非凸的深度网络训练中,这两个假设都不严格成立,因此理论保证不能直接适用。