[论文解读] Why Does Sharpness-Aware Minimization Generalize Better Than SGD?

Zixiang Chen, Junkai Zhang+4 more|arXiv (Cornell University)|Oct 11, 2023

Machine Learning and Data Classification被引用 4

一句话总结

本文为两层ReLU卷积网络中Sharpness-Aware Minimization（SAM）为何比SGD泛化性能更优提供了理论解释，尤其是在标签噪声条件下。研究表明，SAM通过偏好更平坦的极小值，防止了早期阶段的噪声记忆，从而实现有效的特征学习，而SGD在信号强度微弱时则会失效。关键结果是基于信噪比的泛化性能正式相变。

ABSTRACT

The challenge of overfitting, in which the model memorizes the training data and fails to generalize to test data, has become increasingly significant in the training of large neural networks. To tackle this challenge, Sharpness-Aware Minimization (SAM) has emerged as a promising training method, which can improve the generalization of neural networks even in the presence of label noise. However, a deep understanding of how SAM works, especially in the setting of nonlinear neural networks and classification tasks, remains largely missing. This paper fills this gap by demonstrating why SAM generalizes better than Stochastic Gradient Descent (SGD) for a certain data model and two-layer convolutional ReLU networks. The loss landscape of our studied problem is nonsmooth, thus current explanations for the success of SAM based on the Hessian information are insufficient. Our result explains the benefits of SAM, particularly its ability to prevent noise learning in the early stages, thereby facilitating more effective learning of features. Experiments on both synthetic and real data corroborate our theory.

研究动机与目标

• 理解SAM为何在非线性、非光滑神经网络中，特别是在标签噪声条件下，泛化性能优于SGD。
• 刻画在两层ReLU卷积网络中，SGD导致有害过拟合与良性过拟合的条件。
• 正式证明SAM即使在SGD失效时也能实现良性过拟合，其机制在于防止训练初期的噪声学习。
• 在非光滑损失景观下，建立SAM与SGD在泛化误差方面的理论区分。

提出的方法

• 分析两层ReLU卷积网络的损失景观，表明其非光滑性，从而否定基于Hessian的解释。
• 引入一个由信号强度 $\|\bm{\mu}\|_2$ 和输入维度 $d$ 定义的数据模型，以研究标签噪声下的泛化性能。
• 使用概率浓度不等式来界定SAM和SGD的测试误差，重点关注预测的置信度边界。
• 通过分析网络权重范数与激活输出范数，推导出预测置信度保持较大的条件，从而确保低测试误差。
• 采用相变分析方法，表明泛化性能存在基于 $\|\bm{\mu}\|_2$ 相对于 $d^{1/4}$ 的清晰边界。
• 证明当 $\|\bm{\mu}\|_2 \geq \widetilde{\Omega}(1)$ 时，SAM的测试误差保持在 $\leq p + \epsilon$，而SGD在 $\|\bm{\mu}\|_2 \leq O(d^{1/4})$ 时失效。

实验结果

研究问题

• RQ1 在何种条件下，SGD在两层ReLU网络中因有害过拟合而无法泛化？
• RQ2 为何SAM在非光滑、非线性神经网络中，特别是在存在标签噪声时，泛化性能优于SGD？
• RQ3 信号强度 $\|\bm{\mu}\|_2$ 在决定SAM与SGD泛化性能方面起什么作用？
• RQ4 SAM如何在训练初期防止噪声学习，从而促进更优的特征学习？
• RQ5 基于信噪比，SAM与SGD之间的泛化误差是否存在相变？

主要发现

• 对于SGD，当 $\|\bm{\mu}\|_2 \leq O(d^{1/4})$ 时，测试误差至少为 $p + 0.1$，表明存在有害过拟合。
• 对于SGD，当 $\|\bm{\mu}\|_2 \geq \Omega(d^{1/4})$ 时，测试误差被限制在 $p + \epsilon$ 以内，表明为良性过拟合。
• 对于SAM，当 $\|\bm{\mu}\|_2 \geq \widetilde{\Omega}(1)$ 时，测试误差被限制在 $p + \epsilon$ 以内，确保即使SGD失效，也能实现良性过拟合。
• SAM在训练初期有效防止了噪声学习，从而相比SGD实现了更高效的特征学习。
• 本文建立了泛化性能的正式相变，表明在SGD失效的低信号区域，SAM优于SGD。
• 理论分析证实，SAM的成功并非源于基于Hessian的光滑性假设，因此适用于非光滑ReLU网络。