CVPR2023论文学习_Gradient Norm Aware Minimization Seeks First-Order Flatness and Improves Generalization

文章链接：paper

Starting Point

当给定扰动半径内存在单个极小值或多个极小值时，零阶平坦度不足以区分具有低泛化误差的极小值与具有高泛化误差的极小值。

提出了一阶平坦性，这是一种更强的平坦性度量，关注扰动半径内的最大梯度范数，该扰动半径既包含局部极小值处的Hessian最大特征值，也包含SAM的正则化函数。

提出了一种新的优化器GAM，相较于SAM (Sharpness Aware Minimization) 能够寻找更平坦的最小值

In Depth: first order flatness

zeroth-order flatness
$L^{\mathrm{sam}}(\mathbf{\theta })=\hat{L}(\mathbf{\theta })+{\max }_{{\mathbf{\theta }}^{\prime }\in B(\mathbf{\theta },\rho )}\left(\hat{L}\left({\mathbf{\theta }}^{\prime }\right)-\hat{L}(\mathbf{\theta })\right)$

ρ-zeroth-order flatness
$R_{\rho }^{(0)}(\mathbf{\theta })\triangleq {\max }_{{\mathbf{\theta }}^{\prime }\in B(\mathbf{\theta },\rho )}\left(\hat{L}\left({\mathbf{\theta }}^{\prime }\right)-\hat{L}(\mathbf{\theta })\right),\forall \mathbf{\theta }\in \Theta$

ρ-first-order flatness
$R_{\rho }^{(1)}(\mathbf{\theta })\triangleq \rho \cdot {\max }_{{\mathbf{\theta }}^{\prime }\in B(\mathbf{\theta },\rho )}\Vert \nabla \hat{L}\left({\mathbf{\theta }}^{\prime }\right)\Vert ,\forall \mathbf{\theta }\in \Theta$，这里$\rho$是控制邻域大小的扰动半径，其衡量了领域内的最大梯度范数。直观地说，ρ-first-order flatness要求损失函数在θ的邻域中不应发生剧烈变化，故由梯度范数衡量。

Relationship with Eigenvalue

作者随后证明了ρ-first-order flatness与Hessian阵的最大特征值${\lambda }_{\max }\left({\nabla }^2\hat{L}\left({\mathbf{\theta }}^{\ast }\right)\right))$间的关系，因为其被证明是泛化能力的有效度量之一，故有以下引理：

Gradient Norm Aware Minimization

随后，将上述ρ-first-order flatness定义引入优化器设计。将总损失函数设计为：$L^{\text{overall }}(\mathbf{\theta })=L^{\text{oracle }}(\mathbf{\theta })+\alpha R_{\rho }^{(1)}(\mathbf{\theta })$，其中$\alpha$控制正则强度，则总损失函数梯度为：$\nabla L^{\text{overall }}(\mathbf{\theta })=\nabla L^{\text{oracle }}(\mathbf{\theta })+\alpha \nabla R_{\rho }^{(1)}(\mathbf{\theta })$，对这两部分分别求导，$\nabla L^{\text{oracle }}(\mathbf{\theta })$直接求导即可，对于第二部分$\nabla R_{\rho }^{(1)}(\mathbf{\theta })$，利用对$\mathbf{\theta }$邻域泰勒展开并忽略高阶项，有
$$\begin{array}{rl}& \nabla R_{\rho }^{(1)}(\mathbf{\theta })\approx \rho \cdot \nabla \Vert \nabla \hat{L}\left({\mathbf{\theta }}^{\mathrm{adv}}\right)\Vert ,{\mathbf{\theta }}^{\mathrm{adv}}=\mathbf{\theta }+\rho \cdot \frac{\mathbf{f}}{\Vert \mathbf{f}\Vert },\\ & \mathbf{f}=\nabla \Vert \nabla \hat{L}(\mathbf{\theta })\Vert =\frac{{\nabla }^2\hat{L}(\mathbf{\theta })\cdot \nabla \hat{L}(\mathbf{\theta })}{\Vert \nabla \hat{L}(\mathbf{\theta })\Vert }.\end{array}$$
算法流程如下：

Summary

相较于SAM，所提出的GAM转换了优化搜索的目标，即梯度二范数，并证明所提出优化器一定程度上包含了SAM与eig value；同时对于损失函数的改进也体现了正则化的思想。