深度神经网络（dnns）的泛化能力与极值点的平坦程度密切相关，因此出现了 sharpness-aware minimization (sam) 算法来寻找更平坦的极值点以提高泛化能力。本文重新审视 sam 的损失函数，提出了一种更通用、有效的方法 wsam，通过将平坦程度作为正则化项来改善训练极值点的平坦度。通过在各种公开数据集上的实验表明，与原始优化器、sam 及其变体相比，wsam 在绝大多数情形都实现了更好的泛化性能。wsam 在蚂蚁内部数字支付、数字金融等多个场景也被普遍采用并取得了显著效果。该文被 kdd '23 接收为 oral paper。
论文地址：https://arxiv.org/pdf/2305.15817.pdf代码地址：https://github.com/intelligent-machine-learning/dlrover/tree/master/atorch/atorch/optimizers随着深度学习技术的发展，高度过参数化的 dnns 在 cv 和 nlp 等各种机器学习场景下取得了巨大的成功。虽然过度参数化的模型容易过拟合训练数据，但它们通常具有良好的泛化能力。泛化的奥秘受到越来越多的关注，已成为深度学习领域的热门研究课题。
最新的研究显示，泛化能力与极值点的平坦程度密切相关。换句话说，损失函数的“地貌”中存在平坦的极值点可以实现更小的泛化误差。sharpness-aware minimization (sam) [1] 是一种用于寻找更平坦极值点的技术，被认为是当前最有前途的技术方向之一。sam技术被广泛应用于计算机视觉、自然语言处理和双层学习等多个领域，并在这些领域中明显优于之前的最先进方法
为了探索更平坦的最小值，sam 定义损失函数 l 在 w 处的平坦程度如下：
gsam [2] 证明了  是局部极值点 hessian 矩阵最大特征值的近似，表明  确实是平坦（陡峭）程度的有效度量。然而  只能用于寻找更平坦的区域而不是最小值点，这可能导致损失函数收敛到损失值依然很大的点（虽然周围区域很平坦）。因此，sam 采用 ，即   作为损失函数。它可以视为在  和 之间寻找更平坦的表面和更小损失值的折衷方案，在这里两者被赋予了同等的权重。
本文重新思考了  的构建，将  视为正则化项。我们开发了一个更通用、有效的算法，称为 wsam（weighted sharpness-aware minimization），其损失函数加入了一个加权平坦度项  作为正则项，其中超参数控制了平坦度的权重。在方法介绍章节，我们演示了如何通过来指导损失函数找到更平坦或更小的极值点。我们的关键贡献可以总结如下。
我们提出 wsam，将平坦度视为正则化项，并在不同任务之间给予不同的权重。我们提出一个“权重解耦”技术来处理更新公式中的正则化项，旨在精确反映当前步骤的平坦度。当基础优化器不是 sgd 时，如 sgdm 和 adam，wsam 在形式上与 sam 有显著差异。消融实验表明，这种技术在大多数情况下可以提升效果。我们在公开数据集上验证了 wsam 在常见任务中的有效性。实验结果表明，与 sam 及其变体相比，wsam 在绝大多数情形都有着更好的泛化性能。预备知识sam 是解决由公式（1）定义的  的极小极大最优化问题的一种技术。
首先，sam 使用围绕 w 的一阶泰勒展开来近似内层的最大化问题，即、
其次，sam 通过采用  的近似梯度来更新 w ，即
其中第二个近似是为了加速计算。其他基于梯度的优化器（称为基础优化器）可以纳入 sam 的通用框架中，具体见algorithm 1。通过改变 algorithm 1 中的 和，我们可以获得不同的基础优化器，例如 sgd、sgdm 和 adam，参见 tab. 1。请注意，当基础优化器为 sgd 时，algorithm 1 回退到 sam 论文 [1] 中的原始 sam。
方法介绍wsam 的设计细节在此，我们给出的正式定义，它由一个常规损失和一个平坦度项组成。由公式（1），我们有
其中  。当=0 时， 退化为常规损失；当 =1/2 时， 等价于  ；当 >1/2 时， 更注重平坦度，因此与 sam 相比更容易找到具有较小曲率而非较小损失值的点；反之亦然。
包含不同基础优化器的 wsam 的通用框架可以通过选择不同的  和  来实现，见 algorithm 2。例如，当  和  时，我们得到基础优化器为 sgd 的 wsam，见 algorithm 3。在此，我们采用了一种“权重解耦”技术，即  平坦度项不是与基础优化器集成用于计算梯度和更新权重，而是独立计算（algorithm 2 第 7 行的最后一项）。这样，正则化的效果只反映了当前步骤的平坦度，而没有额外的信息。为了进行比较，algorithm 4 给出了没有“权重解耦”（称为 coupled-wsam）的 wsam。例如，如果基础优化器是 sgdm，则 coupled-wsam 的正则化项是平坦度的指数移动平均值。如实验章节所示，“权重解耦”可以在大多数情况下改善泛化表现。
fig. 1 展示了不同取值下的 wsam 更新过程。当 时，介于  和  之间，并随着增大逐渐偏离 。
简单示例为了更好地说明 wsam 中 γ 的效果和优势，我们设置了一个二维简单示例。如 fig. 2 所示，损失函数在左下角有一个相对不平坦的极值点（位置：(-16.8, 12.8)，损失值：0.28），在右上角有一个平坦的极值点（位置：(19.8, 29.9)，损失值：0.36）。损失函数定义为: ，这里  是单变量高斯模型与两个正态分布之间的 kl 散度，即  ，其中  和  。
我们使用动量为 0.9 的 sgdm 作为基础优化器，并对 sam 和 wsam 设置=2 。从初始点 (-6, 10) 开始，使用学习率为 5 在 150 步内优化损失函数。sam 收敛到损失值更低但更不平坦的极值点，=0.6的 wsam 也类似。然而，=0.95 使得损失函数收敛到平坦的极值点，说明更强的平坦度正则化发挥了作用。
实验我们在各种任务上进行了实验，以验证 wsam 的有效性。
图像分类我们首先研究了 wsam 在 cifar10 和 cifar100 数据集上从零开始训练模型的效果。我们选择的模型包括 resnet18 和wideresnet-28-10。我们使用预定义的批大小在 cifar10 和 cifar100 上训练模型，resnet18 和 wideresnet-28-10 分别为 128，256。这里使用的基础优化器是动量为 0.9 的 sgdm。按照 sam [1] 的设置，每个基础优化器跑的 epoch 数是 sam 类优化器的两倍。我们对两种模型都进行了 400 个 epoch 的训练（sam 类优化器为 200 个 epoch），并使用 cosine scheduler 来衰减学习率。这里我们没有使用其他高级数据增强方法，例如 cutout 和 autoaugment。
对于两种模型，我们使用联合网格搜索确定基础优化器的学习率和权重衰减系数，并将它们保持不变用于接下来的 sam 类优化器实验。学习率和权重衰减系数的搜索范围分别为 {0.05, 0.1} 和 {1e-4, 5e-4, 1e-3}。由于所有 sam 类优化器都有一个超参数(邻域大小)，我们接下来在 sam 优化器上搜索最佳的并将相同的值用于其他 sam 类优化器。的搜索范围为 {0.01, 0.02, 0.05, 0.1, 0.2, 0.5}。最后，我们对其他 sam 类优化器各自独有的超参进行搜索，搜索范围来自各自原始文章的推荐范围。对于 gsam [2]，我们在 {0.01, 0.02, 0.03, 0.1, 0.2, 0.3} 范围内搜索。对于 esam [3]，我们在 {0.4, 0.5, 0.6} 范围内搜索 ，在 {0.4, 0.5, 0.6} 范围内搜索 ，在 {0.4, 0.5, 0.6} 范围内搜索。对于 wsam，我们在 {0.5, 0.6, 0.7, 0.8, 0.82, 0.84, 0.86, 0.88, 0.9, 0.92, 0.94, 0.96} 范围内搜索。我们使用不同的随机种子重复实验 5 次，计算了平均误差和标准差。我们在单卡 nvidia a100 gpu 上进行实验。每个模型的优化器超参总结在 tab. 3 中。
tab. 2 给出了在不同优化器下，resnet18、wrn-28-10 在 cifar10 和 cifar100 上测试集的 top-1 错误率。相比基础优化器，sam 类优化器显著提升了效果，同时，wsam 又显著优于其他 sam 类优化器。
imagenet 上的额外训练我们在imagenet数据集上进一步使用data-efficient image transformers网络结构进行实验。我们恢复了一个预训练的deit-base检查点，然后继续训练三个epoch。模型使用批大小256进行训练，基础优化器为动量0.9的sgdm，权重衰减系数为1e-4，学习率为1e-5。我们在四卡nvidia a100 gpu上重复运行5次，并计算平均误差和标准差
我们在 {0.05, 0.1, 0.5, 1.0,⋯ , 6.0} 中搜索 sam 的最佳。最佳的=5.5 被直接用于其他 sam 类优化器。之后，我们在{0.01, 0.02, 0.03, 0.1, 0.2, 0.3}中搜索 gsam 的最佳 ，并在 0.80 到 0.98 之间以 0.02 的步长搜索wsam 的最佳。
模型的初始 top-1 错误率为 18.2%，在进行了三个额外的 epoch 之后，错误率如 tab. 4 所示。我们没有发现三个 sam-like 优化器之间有明显的差异，但它们都优于基础优化器，表明它们可以找到更平坦的极值点并具有更好的泛化能力。
标签噪声的鲁棒性如先前的研究 [1, 4, 5] 所示，sam 类优化器在训练集存在标签噪声时表现出良好的鲁棒性。在这里，我们将 wsam 的鲁棒性与 sam、esam 和 gsam 进行了比较。我们在 cifar10 数据集上训练 resnet18 200 个 epoch，并注入对称标签噪声，噪声水平为 20%、40%、60% 和 80%。我们使用具有 0.9 动量的 sgdm 作为基础优化器，批大小为 128，学习率为 0.05，权重衰减系数为 1e-3，并使用 cosine scheduler 衰减学习率。针对每个标签噪声水平，我们在 {0.01, 0.02, 0.05, 0.1, 0.2, 0.5} 范围内对 sam 进行网格搜索，确定通用的值。然后，我们单独搜索其他优化器特定的超参数，以找到最优泛化性能。我们在 tab. 5 中列出了复现我们结果所需的超参数。我们在 tab. 6 中给出了鲁棒性测试的结果，wsam 通常比 sam、esam 和 gsam 都具有更好的鲁棒性。
探索几何结构的影响sam 类优化器可以与 asam [4] 和 fisher sam [5] 等技术相结合，以自适应地调整探索邻域的形状。我们在 cifar10 上对 wrn-28-10 进行实验，比较 sam 和 wsam 在分别使用自适应和 fisher 信息方法时的表现，以了解探索区域的几何结构如何影响 sam 类优化器的泛化性能。
除了和之外的参数，我们复用了图像分类中的配置。根据先前的研究 [4, 5]，asam 和 fisher sam 的通常较大。我们在 {0.1, 0.5, 1.0,…, 6.0} 中搜索最佳的，asam 和 fisher sam 最佳的均为 5.0。之后，我们在 0.80 到 0.94 之间以 0.02 的步长搜索 wsam 的最佳，两种方法最佳均为 0.88。
令人惊讶的是，如 tab. 7 所示，即使在多个候选项中，基准的 wsam 也表现出更好的泛化性。因此，我们建议直接使用具有固定的基准 wsam 即可。
消融实验在本节中，我们进行消融实验，以深入理解 wsam 中“权重解耦”技术的重要性。如wsam 的设计细节所述，我们将不带“权重解耦”的 wsam 变体（算法 4）coupled-wsam 与原始方法进行比较。
结果如 tab. 8 所示。coupled-wsam 在大多数情况下比 sam 产生更好的结果，wsam 在大多数情况下进一步提升了效果，证明“权重解耦”技术的有效性。
极值点分析在这里，我们通过比较 wsam 和 sam 优化器找到的极值点之间的差异，进一步加深对 wsam 优化器的理解。极值点处的平坦（陡峭）度可通过 hessian 矩阵的最大特征值来描述。特征值越大，越不平坦。我们使用 power iteration 算法来计算这个最大特征值。
tab. 9 显示了 sam 和 wsam 优化器找到的极值点之间的差异。我们发现，vanilla 优化器找到的极值点具有更小的损失值但更不平坦，而 sam 找到的极值点具有更大的损失值但更平坦，从而改善了泛化性能。有趣的是，wsam 找到的极值点不仅损失值比 sam 小得多，而且平坦度十分接近 sam。这表明，在寻找极值点的过程中，wsam 优先确保更小的损失值，同时尽量搜寻到更平坦的区域。
超参敏感性与 sam 相比，wsam 具有一个额外的超参数，用于缩放平坦（陡峭）度项的大小。在这里，我们测试 wsam 的泛化性能对该超参的敏感性。我们在 cifar10 和 cifar100 上使用 wsam 对 resnet18 和 wrn-28-10 模型进行了训练，使用了广泛的取值。如 fig. 3 所示，结果表明 wsam 对超参的选择不敏感。我们还发现，wsam 的最优泛化性能几乎总是在 0.8 到 0.95 之间。
以上就是更通用、有效，蚂蚁自研优化器wsam入选kdd oral的详细内容。