Reproduce “Sharpness-Aware Minimization Enhances Feature Quality via Balanced Learning” on OpenResearch

About this paper

Jacob Mitchell Springer, Vaishnavh Nagarajan, Aditi Raghunathan

Researchers from Carnegie Mellon University and Google Research demonstrate that Sharpness-Aware Minimization (SAM) improves out-of-distribution generalization by implicitly balancing the quality of diverse features, counteracting the simplicity bias of standard optimizers. The work identifies two core mechanisms: an importance-weighting effect that re-weights training examples and a learning-rate-scaling effect that adaptively adjusts effective learning rates for different features.

Key insights

• SAM implicitly re-weights training examples, giving more importance to those poorly fit by hard features but well-fit by easy features, thus ensuring strong learning signals for all features (importance-weighting effect).
• SAM adaptively scales effective learning rates, suppressing them for well-learned features while increasing them for hard-to-learn features, promoting balanced progress across feature learning (learning-rate-scaling effect).
• These two identified mechanisms collectively allow SAM to overcome the simplicity bias of standard optimizers, fostering the learning of a more diverse set of high-quality features.

Summary by alphaXiv.