SAM解析：Sharpness-Aware Minimization for Efficiently Improving Generalization

论文：Sharpness-Aware Minimization for Efficiently Improving Generalization( ICLR 2021)

一、理论

综合了另一篇论文：ASAM: Adaptive Sharpness-Aware Minimization
for Scale-Invariant Learning of Deep Neural Networks 对理论部分这边的解释，同时这篇论文自己也对SAM做出了改进。

1.要解决什么问题？

（1）目标函数

$L_{\mathcal{Q}}(\mathbf{w})$是总体损失值，而训练集$L_{\mathcal{S}}(\mathbf{w})$是样本损失值，即训练集S的损失值（S i.i.d. from distribution$\mathcal{D}$），我们使用$L_{\mathcal{S}}(\mathbf{w})$来估计$L_{\mathcal{Q}}(\mathbf{w})$。

根据PAC-Bayesian generalization bound定理的非正式形式（严格的定理以及其证明在论文的附录 A.1）：

对于任何$\rho >0$（$\rho$为领域大小），从分布来看，生成的训练集大概率满足：
$$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}}_{+}$是单调递增函数。附录 A.1结尾给出了具体这个$h$函数是什么。
$$\begin{array}{rl}{L}_{\mathcal{D}}(\mathbf{w})& \le \underset{\Vert ϵ{\Vert }_2\le \rho }{\max }{L}_{\mathcal{S}}(\mathbf{w}+\mathbf{ϵ})+\\ & +\sqrt{\frac{k\log \left(1+\frac{\Vert \mathbf{w}{\Vert }_2^2}{{\rho }^2}{\left(1+\sqrt{\frac{\log (n)}{k}}\right)}^2\right)+4\log \frac{n}{\delta }+8\log (6n+3k)}{n-1}}\end{array}$$
那么也就是说，为了得到损失函数$L_{\mathcal{Q}}(\mathbf{w})$的极小解，现在要求：不等式右边式子，即${\max }_{\Vert ϵ{\Vert }_2\le \rho }{L}_{\mathcal{S}}(\mathbf{w}+\mathbf{ϵ})+h\left(\Vert \mathbf{w}{\Vert }_2^2/{\rho }^2\right)$的最小值，也就是一个极小-极大化minimax问题。

（2）引出锐度sharpness概念

这里把不等式的右边展开为：
$$\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)$$
这个中括号里的就是$L_{\mathcal{S}}$的锐度sharpness，可以衡量从$\mathbf{w}$移动到相临近的参数值，training loss增加的速度。这个式子就变成sharpness + training loss＋$w$的正则项了。

（3）进一步对目标函数进行变换

考虑到具体的函数$h$受证明细节的影响严重，把$h\left(\Vert \mathbf{w}{\Vert }_2^2/{\rho }^2\right)$替换成$\lambda \Vert w{\Vert }_2^2$，这样就产生了一个标准的L2正则项了，$\lambda \Vert w{\Vert }_2^2$为$w$的正则项，$\lambda$即为weight decay。

综合以上，现在的目标函数可以变成（公式0）：
$$\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{ϵ})$$把它叫做Sharpness-Aware Minimization (SAM) problem。这里，$\rho \ge 0$是一个超参数，以及$p\in [1,\infty ]$，论文说$p=2$通常是最优的（附录C.5）。

2.怎么解决？

为了最小化$L_{\mathcal{S}}^{SAM}(\mathbf{w})$，论文提出可以通过对inner maximization求微分来得出${\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}^{SAM}(\mathbf{w})$的近似值。

可是这个the inner maximization是什么东西呢？

the inner problem，指的是在固定损失函数后，如何在训练集上更新模型超参数使得损失函数值最小，得到当前最优的模型参数的问题。对应的数学表达式就是[超参数] = arg min[损失函数]。

inner maximization其实就是上面的使得损失函数值最小改成最大啦，用数学语言表示就是（公式1）：
$${\mathbf{ϵ}}^{\ast }(\mathbf{w})\triangleq \underset{\Vert \mathbf{ϵ}{\Vert }_p\le \rho }{\mathrm{argmax}}{L}_{\mathcal{S}}(\mathbf{w}+\mathbf{ϵ})$$意思是：在$w$的$ϵ$邻域$[\mathrm{a}-ϵ,\mathrm{a}+ϵ]$中，使得目标函数$L_{\mathcal{S}}(\mathbf{w}+\mathbf{ϵ})$最大的参数$ϵ$。

然后，y用在$x=w+ϵ$处，$ϵ\to 0$的一阶泰勒展开式去逼近一个the inner maximization的目标函数：

a. 在$x=x_0$处的一阶泰勒展开式：
$$f(x)=f\left(x_0\right)+f^{\prime }\left(x_0\right)\left(x-x_0\right)+o\left(x-x_0\right)$$

b. 在$x=w+ϵ$处，$ϵ\to 0$的一阶泰勒展开式：
$$L_{\mathcal{S}}(\mathbf{w}+\mathbf{ϵ})=L_{\mathcal{S}}(\mathbf{w})+{ϵ}^T{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})+o\left(ϵ\right)$$$$L_{\mathcal{S}}(\mathbf{w}+\mathbf{ϵ})\approx L_{\mathcal{S}}(\mathbf{w})+{ϵ}^T{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})$$

c. 又因为$$\underset{\Vert ϵ{\Vert }_p\le \rho }{\mathrm{argmax}}{L}_{\mathcal{S}}(\mathbf{w})=0$$
把b、c公式代入公式1，得到：
$$\begin{array}{rl}{\mathbf{ϵ}}^{\ast }(\mathbf{w})& \triangleq \underset{\Vert \mathbf{ϵ}{\Vert }_p\le \rho }{\mathrm{argmax}}{L}_{\mathcal{S}}(\mathbf{w}+\mathbf{ϵ})\\ & \approx \underset{\Vert ϵ{\Vert }_p\le \rho }{\mathrm{argmax}}{L}_{\mathcal{S}}(\mathbf{w})+{\mathbf{ϵ}}^T{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})\\ & =\underset{\Vert \mathbf{ϵ}{\Vert }_p\le \rho }{\mathrm{argmax}}{\mathbf{ϵ}}^T{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})\end{array}$$

现在，用$\widehat{\mathbf{ϵ}}(\mathbf{w})$表示$\mathbf{ϵ}(\mathbf{w})$的近似值，$\widehat{\mathbf{ϵ}}(\mathbf{w})$可以用经典对偶范数问题（dual
norm problem）的解法来得出：
$$\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}$$

其中，$1/p+1/q=1$。

之前说过论文提到$p=2$通常效果最优，这里就把$p=2$代入到上式中，$1/p=1/2$, $q-1=1$，则得到公式2：
$$\begin{array}{rl}\widehat{\mathbf{ϵ}}(\mathbf{w})& =\rho \mathrm{sign}\left({\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w})\right)\frac{|{\nabla }_{\mathbf{w}}{L}_S\left(\mathbf{w}\right)|}{{\Vert {\nabla }_{\mathbf{w}}{L}_S\left(\mathbf{w}\right)\Vert }_2}\\ & =\rho \frac{\nabla L_S\left(\mathbf{w}\right)}{{\Vert \nabla L_S\left(\mathbf{w}\right)\Vert }_2}\end{array}$$
把上式代入到公式0，得到：
$$\begin{array}{rl}{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}^{SAM}(\mathbf{w})& \approx {\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(\mathbf{w}+\widehat{\mathbf{ϵ}}(\mathbf{w}))\\ & ={\frac{d(\mathbf{w}+\widehat{\mathbf{ϵ}}(\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}$$
上面第二行用到了复合微分：
$$\begin{array}{rl}& \frac{df(g(x))}{dx}={\frac{dg(x)}{dx}df(x)|}_{g(x)}\end{array}$$
公式3：
$${{\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}^{SAM}(\mathbf{w})\approx {\nabla }_{\mathbf{w}}{L}_{\mathcal{S}}(w)|}_{\mathbf{w}+\widehat{\mathbf{ϵ}}(\mathbf{w})}$$
如Section 3中的结果所示，如果没有二阶项，这个近似值会产生一个有效的算法。在论文附录C.4中具体研究了这种影响，在开始的实验中，加入二阶项会令人惊讶地降低性能，进一步研究这些条款的影响应该是未来工作的优先事项。

我们通过将标准的数值优化器(如随机梯度下降SGD)应用于SAM目标函数$L_{\mathcal{S}}^{SAM}(\mathbf{w})$，其中使用公式3计算目标函数梯度，从而获得最终的SAM算法。

论文的Algorithm 1给出了SAM算法的伪代码，它使用SGD作为基本优化器。右边的Figure 2示意性地说明了SAM参数的单次迭代更新。