[论文解读] Surrogate Gap Minimization Improves Sharpness-Aware Training

Juntang Zhuang, Boqing Gong+7 more|arXiv (Cornell University)|Mar 15, 2022

Image Enhancement Techniques被引用 13

一句话总结

该论文提出 GSAM，一种新型训练方法，通过引入代理间隙（surrogate gap）——一种基于 Hessian 的曲率的计算高效代理——改进了 Sharpness-Aware Minimization (SAM)。该方法实现了损失与曲率的直接联合最小化。GSAM 每轮更新包含两个步骤：最小化扰动损失（如 SAM 所做），以及在正交方向上上升以减小代理间隙，从而获得更平坦的极小值并提升泛化性能。在 ImageNet 上使用 ViT-B/32 时，GSAM 相较于 AdamW 实现了 5.4% 的 top-1 准确率提升。

ABSTRACT

The recently proposed Sharpness-Aware Minimization (SAM) improves generalization by minimizing a extit{perturbed loss} defined as the maximum loss within a neighborhood in the parameter space. However, we show that both sharp and flat minima can have a low perturbed loss, implying that SAM does not always prefer flat minima. Instead, we define a extit{surrogate gap}, a measure equivalent to the dominant eigenvalue of Hessian at a local minimum when the radius of the neighborhood (to derive the perturbed loss) is small. The surrogate gap is easy to compute and feasible for direct minimization during training. Based on the above observations, we propose Surrogate extbf{G}ap Guided extbf{S}harpness- extbf{A}ware extbf{M}inimization (GSAM), a novel improvement over SAM with negligible computation overhead. Conceptually, GSAM consists of two steps: 1) a gradient descent like SAM to minimize the perturbed loss, and 2) an extit{ascent} step in the extit{orthogonal} direction (after gradient decomposition) to minimize the surrogate gap and yet not affect the perturbed loss. GSAM seeks a region with both small loss (by step 1) and low sharpness (by step 2), giving rise to a model with high generalization capabilities. Theoretically, we show the convergence of GSAM and provably better generalization than SAM. Empirically, GSAM consistently improves generalization (e.g., +3.2\% over SAM and +5.4\% over AdamW on ImageNet top-1 accuracy for ViT-B/32). Code is released at \url{ https://sites.google.com/view/gsam-iclr22/home}.

研究动机与目标

• 为解决 SAM 的局限性，即在尖锐和平坦极小值中均可能产生低扰动损失，通过识别更可靠的曲率度量来改进。
• 定义一种计算上可行的代理间隙，以近似局部极小值处主导 Hessian 特征值，尤其适用于小邻域半径。
• 开发一种联合最小化扰动损失与代理间隙的训练方法，以实现更平坦的极小值和更优的泛化性能。
• 证明所提出的 GSAM 算法的收敛性，并在标准假设下建立其在泛化性能上优于 SAM 的理论优势。
• 通过实证验证 GSAM 在多种架构（ResNets、Vision Transformers、MLP-Mixer）上的表现，结果一致优于 SAM 和基线优化器。

提出的方法

• 将代理间隙定义为 $ h(w) = f_p(w) - f(w) $，其中 $ f_p(w) $ 为扰动损失，$ f(w) $ 为标准损失，其在局部极小值处近似主导 Hessian 特征值。
• 将梯度 $ \nabla f(w) $ 分解为平行于 $ \nabla f_p(w) $ 和正交于 $ \nabla f_p(w) $ 的分量，从而实现在正交方向上的定向优化。
• 对 $ f_p(w) $ 执行梯度下降步骤，以最小化扰动损失，保持与 SAM 相同的更新规则。
• 沿正交分量 $ \nabla_\perp f(w) $ 执行上升步骤，以最小化代理间隙 $ h(w) $，由于正交性，不会改变 $ f_p(w) $。
• 引入超参数 $ \alpha $ 控制上升步长的大小，实现损失与曲率最小化之间的灵活权衡。
• 通过假设邻域半径 $ \rho_t $ 递减，确保理论收敛性，支持 GSAM 的收敛性分析。

实验结果

研究问题

• RQ1 SAM 中的扰动损失能否可靠地区分尖锐与平坦极小值，还是可能在两者中均产生低值？
• RQ2 是否存在一种计算高效的 Hessian 曲率度量替代方法，可在训练过程中使用？
• RQ3 通过最小化代理间隙（即扰动损失与标准损失之差）是否能有效引导优化过程趋向更平坦的极小值？
• RQ4 结合扰动损失下降与正交方向上升以减小代理间隙的两步更新策略，是否能带来优于 SAM 的泛化性能？
• RQ5 在标准假设下，能否理论证明 GSAM 收敛，并且其泛化性能可被严格证明优于 SAM？

主要发现

• 当邻域半径 $ \rho $ 较小时，代理间隙 $ h(w) = f_p(w) - f(w) $ 在理论上等价于局部极小值处的主导 Hessian 特征值，使其成为一种有效且高效的曲率代理。
• 在 ImageNet 上使用 ViT-B/32 时，GSAM 相较于 SAM 实现 3.2% 的 top-1 准确率提升，相较 AdamW 实现 5.4% 的提升，证明了其一致的泛化增益。
• 消融实验表明，GSAM 中用于减小代理间隙的上升步骤是性能提升的主要驱动力，无论采用恒定还是衰减的 $ \rho_t $ 策略均成立。
• 实证验证显示，$ \cos\theta_t $（即 $ \nabla f(w) $ 与 $ \nabla f_p(w) $ 之间夹角的余弦值）在整个训练过程中始终高于 0.9，支持了高维参数空间中梯度近乎对齐的理论假设。
• 随着 $ \alpha $ 增大，代理间隙减小；随着训练步数增加，代理间隙增大，表明 GSAM 成功地随时间将模型引导至更平坦的极小值。
• GSAM 具有广泛适用性，且相较于 SAM 的计算开销可忽略不计，使其适用于大规模训练。