本文正在参加 人工智能创作者扶持计划

一些缘由和动机

  机缘巧合之下，最近读了关于 Sharpness-Aware Minimization 的两篇文章。第一篇奠基的文章是在 ICLR 2021 上发表的，虽然这篇工作里还有一些坑没有填并不能说是非常完备，但是也确实产生了很大的影响。借着刚读完的热劲，在此简单记录一下 (之后的论文笔记我会朝着更精简的道路一路狂奔的)。

0.SAM 的论文信息

1.锐度与泛化能力的感性认知

  如果损失函数的锐度很大，那么它会更加敏感地捕捉模型在训练集上的表现，使得模型需要更好地适应训练数据才能获得较低的损失值。而泛化能力则要求模型在训练集外也有很好的表现，锐度过大的损失函数很可能导致模型在训练集外表现不佳。因此模型要兼顾锐度与平滑。

图 1.(左子图) 表示通过切换到 SAM 获得的错误率降低。每个点都是不同的数据集/模型/数据扩充。(中间子图) 表示用 SGD 训练的 ResNet 收敛到的急剧最小值。(右子图) 表示用SAM训练的相同ResNet收敛到的宽最小值。

2.Sharpness-Aware Minimization (SAM)

  受损失函数图像的锐度和泛化之间的联系的启发，我们提出了一种不同的方法：不是寻找仅具有低训练损失值 $L_{\mathbf{S}}(w)$ 的参数值 $w$，而是寻找整个邻域具有一致低的参数值训练损失值（等效于具有低损失和低曲率的邻域）。以下定理通过在邻域训练损失方面限制泛化能力来说明这种方法的动机

定理 1 (非正式地表述). 对于任何 $\rho >0$ 且在从分布 $\mathcal{D}$ 生成的训练集 $\mathcal{S}$ 上有大概率满足

$$L_{\mathcal{D}}(\mathbf{w})\le \underset{\Vert ϵ{\Vert }_2\le \rho }{\max }{L}_{\mathcal{S}}(\mathbf{w}+\mathbf{ϵ})+h\left(\Vert \mathbf{w}{\Vert }_2^2/{\rho }^2\right)$$

  其中 $h:{\mathbb{R}}_{+}\to {\mathbb{R}}_{+}$ 是严格递增的函数 ( 在某些 $L_{\mathcal{D}}(w)$ 的技术条件下 )。为了明确锐度项，我们可以将上述不等式的右侧重写为

$$\left[\underset{\Vert ϵ{\Vert }_2\le \rho }{\max }{L}_{\mathcal{S}}(\mathbf{w}+\mathbf{ϵ})-L_{\mathcal{S}}(\mathbf{w})\right]+L_{\mathcal{S}}(\mathbf{w})+h\left(\Vert \mathbf{w}{\Vert }_2^2/{\rho }^2\right)$$

  方括号中的通过测量通过将 $w$ 移动到附近的参数值来提高训练损失的速度来捕获 $w$ 处 L_boldsymbol{S} 的锐度；然后将这个锐度项与训练损失值本身和 $w$ 大小的正则化器相加。鉴于特定函数 $h$ 受到证明细节的严重影响，我们将第二项用超参数 $\lambda$ 替换为 $\lambda \Vert w{\Vert }_2^2$ ，产生标准 L2-正则化项。因此，受边界项的启发，我们建议通过解决以下锐化感知最小化 (SAM) 问题来选择参数值 :

$$\underset{\mathbf{w}}{\min }{L}_{\mathcal{S}}^{SAM}(\mathbf{w})+\lambda \Vert \mathbf{w}{\Vert }_2^2\text{ where }{L}_{\mathcal{S}}^{SAM}(\mathbf{w})\triangleq \underset{\Vert \mathbf{ϵ}{\Vert }_p\le \rho }{\max }{L}_S(\mathbf{w}+\mathbf{ϵ})\text{\hspace{1em}(1)}$$

  其中 $\rho >0$ 是超参数且范数维度 $p\in [1,\infty ]$，为了最小化 $L_{\mathcal{S}}^{SAM}(\mathbf{w})$，通过对 ${\nabla }_w{L}_{\mathcal{S}}^{SAM}(\mathbf{w})$ 进行内极值微分，我们得到了一个高效的近似，这反过来使我们能够将随机梯度下降直接应用于 SAM 目标。以此继续，我们首先通过 $L_{\mathcal{S}}(w+ϵ)$ 在 0 附近的一阶泰勒展开来近似内部最大化问题，得到 :

$${ϵ}^{\ast }(\mathbf{w})\triangleq \underset{\Vert ϵ{\Vert }_p\le \rho }{\mathrm{argmax}}{L}_{\mathcal{S}}(\mathbf{w}+ϵ)\approx \underset{\Vert ϵ{\Vert }_p\le \rho }{\mathrm{argmax}}{L}_{\mathcal{S}}(\mathbf{w})+{ϵ}^T{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})=\underset{\Vert ϵ{\Vert }_p\le \rho }{\mathrm{argmax}}{ϵ}^T{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})$$

  反过来，求解该近似值 $\widehat{ϵ}(w)$ 由经典对偶范数问题的解给出 ( $|\cdot {|}^{q-1}$表示元素绝对值和幂 ) :

$$\widehat{\mathbf{ϵ}}(\mathbf{w})=\rho \mathrm{sign}\left({\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})\right){\left|{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})\right|}^{q-1}/{\left({\Vert {\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})\Vert }_q^q\right)}^{1/p}\text{\hspace{1em}(2)}$$

  其中 $\frac{1}{p}+\frac{1}{q}=1$，代回式 (1) 并求导有 :

$$\begin{array}{rl}{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}^{SAM}(\mathbf{w})& \approx {\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w}+\widehat{ϵ}(\mathbf{w}))={\frac{d(\mathbf{w}+\widehat{ϵ}(\mathbf{w}))}{d\mathbf{w}}{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})|}_{\mathbf{w}+\widehat{ϵ}(w)}\\ & ={{\nabla }_w{L}_{\mathcal{S}}(\mathbf{w})|}_{\mathbf{w}+\widehat{ϵ}(\mathbf{w})}+{\frac{d\widehat{ϵ}(\mathbf{w})}{d\mathbf{w}}{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})|}_{\mathbf{w}+\widehat{ϵ}(\mathbf{w})}\end{array}$$

  对于 ${\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}^{SAM}(\mathbf{w})$ 的近似可以通过自动微分直接计算。虽然这个计算隐含地依赖于 $L_{\mathcal{S}}$ 的 Hessian 矩阵，因为 $\widehat{ϵ}(w)$ 本身是 ${\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})$ 的函数，Hessian 只能通过 Hessian 向量积计算，这使得可以在不具体化 Hessian 矩阵的情况下计算。尽管如此，为了进一步加速计算，我们去掉了二阶项。得到我们的最终梯度近似 :

$${{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}^{SAM}(\mathbf{w})\approx {\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(w)|}_{\mathbf{w}+\widehat{ϵ}(\mathbf{w})}\text{\hspace{1em}(3)}$$

  以下给出对应算法和示意图

算法 1. (左子图) Sharpness-Aware Minimization 算法对应流程

图 2. (右子图) SAM参数更新示意图

0.LookSAM 的论文信息

1.这篇工作的核心点 : 降低计算量

  SAM 的更新规则需要在每一步进行两次连续的（不可并行化的）梯度计算，这会使计算开销增加一倍。在本文中，我们提出了一种新颖的算法 LookSAM，它只定期计算内部梯度上升，以显着降低 SAM 的额外训练成本。

2.LookSAM 的具体技术细节

  在 SAM 中梯度为 $g_{\mathbf{s}}={{\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w})|}_{\mathbf{w}+\widehat{ϵ}}$，在下图中用蓝色箭头进行表示。几种梯度分别如下图所示。

图 3. LookSAM 更新的可视化。蓝色箭头 $g_{\mathbf{s}}$ 是 SAM 针对较平坦区域的梯度。黄色箭头 $-\eta {\nabla }_w{\mathcal{L}}_{\mathbf{S}}(w)$ 表示 SGD 梯度。$g_{\mathbf{h}}$ (棕色箭头) 和 $g_{\mathbf{v}}$ (红色箭头) 是 $g_{\mathbf{s}}$ 的正交梯度分量，分别与 SGD 梯度平行和垂直。

  文章中提出了一种新颖的 LookSAM 算法来应对这一挑战。主要思想是研究如何重用信息以防止每次都计算 SAM 的梯度。如图 3 所示，与 SGD 梯度 (黄色箭头) 相比，SAM 的梯度 $g_{\mathbf{s}}={{\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w})|}_{\mathbf{w}+\widehat{ϵ}}$ 提升到一个更平坦的区域（蓝色箭头）。为了更直观地了解这个平坦区域，我们基于泰勒展开重写了 SAM 的更新：

$$\begin{array}{rl}{\nabla }_{\mathbf{w}}& {{\mathcal{L}}_S(\mathbf{w})|}_{\mathbf{w}+\widehat{ϵ}}={\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w}+\widehat{\mathbf{ϵ}})\\ & \approx {\nabla }_{\mathbf{w}}\left[{\mathcal{L}}_S(\mathbf{w})+\widehat{ϵ}\cdot {\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w})\right]\\ & ={\nabla }_{\mathbf{w}}\left[{\mathcal{L}}_S(\mathbf{w})+\frac{\rho }{\Vert {\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w})\Vert }{\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w}{)}^T{\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w})\right]\\ & ={\nabla }_{\mathbf{w}}\left[{\mathcal{L}}_S(\mathbf{w})+\rho \Vert {\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w})\Vert \right]\end{array}\text{\hspace{1em}(4)}$$

  我们发现 SAM 的梯度由两部分组成：原始梯度 ${\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w})$ 和原始梯度 $\Vert {\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w})\Vert$ 的 L2-范数的梯度。我们认为优化梯度的 L2-范数可以提示模型收敛到平坦区域，因为平坦区域通常意味着低梯度范数值。因此，SAM 的更新可分为两部分 : 第一部分 (记为 $g_{\mathbf{h}}$) 是降低损失值，第二部分 (记为 $g_{\mathbf{v}}$) 是将更新偏向平坦区域。更具体地说，$g_{\mathbf{h}}$是普通的 SGD 的梯度方向，即使没有 SAM，也需要在每一步计算。因此，SAM 的额外计算代价主要是由第二部分 $g_{\mathbf{v}}$ 引起的。已知 SAM 的梯度(蓝色箭头) 和 SGD 的梯度方向 ($g_{\mathbf{h}}$)，我们可以进行投影得到 $g_{\mathbf{v}}$ :

$$g_{\mathbf{v}}={{\nabla }_{\mathbf{w}}{\mathcal{L}}_S(\mathbf{w})|}_{\mathbf{w}+\widehat{\mathbf{ϵ}}}\cdot \sin (\theta )\text{\hspace{1em}(5)}$$

图 4. $g_{\mathbf{s}}$, $g_{\mathbf{h}}$ 和 $g_{\mathbf{v}}$ 的每 5 步之间的梯度差异 (即 $\Vert g_s^t-g_s^{t+k}\Vert$)， $g_{\mathbf{v}}$ 的变化比 $g_{\mathbf{h}}$ 和 $g_{\mathbf{s}}$ 平滑得多

  一个关键的观察结果是 $g_{\mathbf{v}}$ 的变化比 $g_{\mathbf{h}}$ 和 $g_{\mathbf{s}}$ 慢得多。在图 4 中，我们通过 SAM 的整个训练过程绘制了迭代 $t$ 和迭代 $t+5$ 之间这三个组件的变化，结果表明 $g_{\mathbf{v}}$ (绿线) 的差异显示出比 $g_{\mathbf{h}}$ (橙色线) 和 $g_{\mathbf{s}}$ (蓝线) 更稳定的模式。直观地说，这意味着指向平坦区域的方向在几次迭代中不会显着变化。所以优化呼之欲出 : 减少 $g_{\mathbf{v}}$ 的采样和计算次数，在小批量中只计算一次，即可大幅加快速度。

算法 2. LookSAM 算法对应流程

  Layer-Wise 的优化策略在此不做过多赘述。