Reproduce “Sharpness-Aware Training for Free” on OpenResearch

About this paper

Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent Y. F. Tan, Joey Tianyi Zhou

Sharpness-Aware Training for Free (SAF) and its memory-efficient variant (MESA) improve deep neural network generalization by encouraging flatter loss minima. These methods achieve performance comparable to or exceeding Sharpness-Aware Minimization (SAM) while significantly reducing computational overhead, with SAF training at nearly the same speed as a standard optimizer.

Key insights

• Minimizing the difference in training loss or output distributions between consecutive optimization steps implicitly guides the model towards flatter loss minima, similar to explicit sharpness regularization.
• The change in training loss after a standard gradient update is approximately proportional to the square of the gradient norm, which is directly related to the sharpness measure of the loss landscape.
• Incorporating a trajectory loss, computed either by recalling past network outputs (SAF) or by using an Exponential Moving Average model (MESA), effectively regularizes sharpness with minimal computational overhead.

Summary by alphaXiv.