Why Does Sharpness-Aware Minimization Generalize Better Than SGD?
Zixiang Chen Junkai Zhang Yiwen Kou Xiangning Chen Cho-Jui Hsieh Quanquan Gu
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.